"""Structured-output admission, correlation, and lifecycle tests."""

import asyncio

from datetime import UTC, datetime

from unittest.mock import AsyncMock, Mock

from uuid import uuid4

import pytest

from pydantic import BaseModel, ConfigDict, Field, ValidationError

from copilot._jsonrpc import JsonRpcError

from copilot.session import CopilotSession

from copilot.session_events import SessionEvent

class Inventory(BaseModel):

model_config = ConfigDict(extra="forbid")

count: int

color: str

def event(kind, data, agent_id=None):

return SessionEvent.from_dict(

{

"id": str(uuid4()),

"timestamp": datetime.now(UTC).isoformat(),

"type": kind,

"data": data,

**({"agentId": agent_id} if agent_id else {}),

}

)

def assistant(content='{"count":42,"color":"red"}', origin="user-1", **extra):

return event(

"assistant.message",

{"messageId": str(uuid4()), "content": content, "originatingMessageId": origin, **extra},

)

def user(origin="user-1"):

return event("user.message", {"messageId": origin, "content": "inventory"})

def idle(**data):

return event("session.idle", data)

def fake_session(events=(), admission_error=None):

client = Mock()

session = CopilotSession("session-1", client)

async def request(method, params):

if method == "session.sendAndWait":

raise JsonRpcError(-32601, "Method not found")

assert method == "session.send"

for item in events:

session._dispatch_event(item)

if admission_error:

raise admission_error

return {"messageId": "user-1"}

client.request = AsyncMock(side_effect=request)

return session, client

def assert_clean(session):

assert not session._structured_waits

assert not session._event_handlers

@pytest.mark.parametrize("structured", [False, True])

@pytest.mark.parametrize("before_ack", [False, True])

@pytest.mark.parametrize(

"source",

[

"user",

"assistant",

"autopilot",

"system",

"session-memory-rollover",

"marked",

"pre-message",

"steering",

],

)

@pytest.mark.asyncio

async def test_wait_correlates_errors_independently_of_message_events(

structured, before_ack, source

):

events = [

assistant(origin="other" if source == "pre-message" else "user-1"),

assistant("other", "other")

if source == "assistant"

else event(

"user.message",

{

"messageId": "other",

"content": "other",

"source": "sdk" if source in ("steering", "pre-message") else source,

"isAutopilotContinuation": source == "marked",

},

),

event(

"session.error",

{

"errorType": "query",

"message": "provider failed",

"originatingMessageId": "other" if source in ("user", "assistant") else "user-1",

},

),

idle(),

]

session, client = fake_session(events if before_ack else ())

sending = asyncio.Event()

original = client.request.side_effect

async def request(method, params):

result = await original(method, params)

sending.set()

return result

client.request.side_effect = request

waiting = asyncio.create_task(

session.send_and_wait(

"waited", response_schema={"type": "object"} if structured else None, timeout=1

)

)

await sending.wait()

if not before_ack:

for item in events:

session._dispatch_event(item)

if source in ("user", "assistant"):

result = await waiting

assert result is not None

assert result.data.content == '{"count":42,"color":"red"}'

else:

with pytest.raises(Exception, match="provider failed"):

await waiting

assert_clean(session)

@pytest.mark.parametrize("before_ack", [False, True])

@pytest.mark.parametrize("blocked", [False, True])

@pytest.mark.asyncio

async def test_plain_wait_ignores_prior_idle_and_correlates_its_start(before_ack, blocked):

session, client = fake_session()

admission = asyncio.get_running_loop().create_future()

sending = asyncio.Event()

async def request(method, params):

if method == "session.sendAndWait":

raise JsonRpcError(-32601, "Method not found")

assert method == "session.send"

sending.set()

return await admission

client.request.side_effect = request

pending = asyncio.create_task(session.send_and_wait("queued", timeout=1))

await sending.wait()

if not before_ack:

admission.set_result({"messageId": "user-1"})

await asyncio.sleep(0)

session._dispatch_event(assistant("previous", "other"))

session._dispatch_event(

event(

"session.warning",

{

"warningType": "policy_blocked",

"message": "Blocked",

"originatingMessageId": "other",

},

)

)

session._dispatch_event(

event(

"session.warning",

{

"warningType": "policy_blocked",

"message": "Blocked",

"originatingMessageId": "user-1",

},

agent_id="child",

)

)

session._dispatch_event(idle(aborted=True))

session._dispatch_event(idle())

if before_ack:

admission.set_result({"messageId": "user-1"})

await asyncio.sleep(0)

assert not pending.done()

final = assistant("current")

session._dispatch_event(

event(

"session.warning",

{

"warningType": "policy_blocked",

"message": "Blocked",

"originatingMessageId": "user-1",

},

)

if blocked

else final

)

session._dispatch_event(idle())

assert await pending is (None if blocked else final)

assert_clean(session)

@pytest.mark.parametrize("terminal", ["session.idle", "session.autopilot_stopped"])

@pytest.mark.asyncio

async def test_plain_wait_rejects_current_abort(terminal):

session, _ = fake_session(

[

user(),

event(terminal, {"aborted": True, "reason": "aborted", "requestId": "user-1"}),

]

)

with pytest.raises(RuntimeError, match="aborted"):

await session.send_and_wait("cancel", timeout=1)

assert_clean(session)

@pytest.mark.parametrize("terminal", ["idle", "autopilot", "error"])

@pytest.mark.asyncio

async def test_plain_wait_preserves_pre_ack_outcome_through_retirement(terminal):

session, client = fake_session()

final = assistant()

async def request(method, params):

if method == "session.sendAndWait":

raise JsonRpcError(-32601, "Method not found")

assert method == "session.send"

session._dispatch_event(final)

session._dispatch_event(

event("session.autopilot_stopped", {"reason": "task_complete", "requestId": "user-1"})

if terminal == "autopilot"

else event("session.error", {"errorType": "test", "message": "original failure"})

if terminal == "error"

else idle()

)

session._mark_disconnected()

return {"messageId": "user-1"}

client.request.side_effect = request

if terminal == "error":

with pytest.raises(Exception, match="original failure"):

await session.send_and_wait("finish", timeout=1)

else:

assert await session.send_and_wait("finish", timeout=1) is final

@pytest.mark.parametrize("structured", [False, True])

@pytest.mark.asyncio

async def test_autopilot_error_terminal_does_not_require_session_error(structured):

session, _ = fake_session(

[event("session.autopilot_stopped", {"reason": "error", "requestId": "user-1"})]

)

with pytest.raises(RuntimeError, match="Autopilot request failed"):

await session.send_and_wait(

"fail", response_schema={"type": "object"} if structured else None, timeout=1

)

assert_clean(session)

@pytest.mark.parametrize("structured", [False, True])

@pytest.mark.asyncio

async def test_autopilot_terminal_buffers_pre_ack_and_ignores_other_requests(structured):

final = assistant('{"count":7,"color":"blue"}')

session, _ = fake_session(

[

assistant(),

event(

"user.message",

{

"content": "Continue",

"messageId": "other",

"source": "autopilot",

"isAutopilotContinuation": True,

},

),

event("session.autopilot_stopped", {"reason": "task_complete", "requestId": "other"}),

event("session.autopilot_stopped", {"reason": "task_complete"}),

event(

"session.autopilot_stopped",

{"reason": "task_complete", "requestId": "user-1"},

agent_id="child",

),

idle(mode="autopilot"),

final,

event("session.autopilot_stopped", {"reason": "task_complete", "requestId": "user-1"}),

]

)

result = await session.send_and_wait(

"inventory", response_schema={"type": "object"} if structured else None, timeout=1

)

assert result is final

assert not session._response_waits

assert_clean(session)

@pytest.mark.parametrize("structured", [False, True])

@pytest.mark.asyncio

async def test_autopilot_terminal_isolates_concurrent_waits(structured):

session, client = fake_session()

admitted = asyncio.Queue()

async def request(method, params):

if method == "session.sendAndWait":

raise JsonRpcError(-32601, "Method not found")

assert method == "session.send"

admitted.put_nowait(params["prompt"])

return {"messageId": params["prompt"]}

client.request.side_effect = request

schema = {"type": "object"} if structured else None

first = asyncio.create_task(session.send_and_wait("first", response_schema=schema, timeout=1))

second = asyncio.create_task(session.send_and_wait("second", response_schema=schema, timeout=1))

await admitted.get()

await admitted.get()

first_message = assistant(origin="first")

second_message = assistant('{"count":7,"color":"blue"}', origin="second")

session._dispatch_event(first_message)

session._dispatch_event(

event("session.autopilot_stopped", {"reason": "task_complete", "requestId": "first"})

)

session._dispatch_event(second_message)

session._dispatch_event(

event("session.autopilot_stopped", {"reason": "task_complete", "requestId": "second"})

)

assert await first is first_message

assert await second is second_message

assert_clean(session)

@pytest.mark.asyncio

async def test_typed_output_buffers_pre_ack_events_and_infers_schema():

session, client = fake_session([user(), assistant(), idle()])

result = await session.send_and_wait_typed("inventory", Inventory, timeout=1)

assert result == Inventory(count=42, color="red")

params = client.request.call_args.args[1]

assert params["responseFormat"] == {

"type": "json_schema",

"jsonSchema": {"name": "response", "strict": True, "schema": Inventory.model_json_schema()},

}

assert_clean(session)

@pytest.mark.asyncio

async def test_typed_output_validates_schema_aliases_including_nested_models():

class Answer(BaseModel):

model_config = ConfigDict(extra="forbid", validate_by_alias=False, validate_by_name=True)

value: int = Field(alias="answer")

class Result(BaseModel):

model_config = ConfigDict(extra="forbid", validate_by_alias=False, validate_by_name=True)

value: Answer = Field(alias="result")

values: list[Answer] = Field(alias="results")

content = '{"result":{"answer":42},"results":[{"answer":99}]}'

session, client = fake_session([assistant(content), idle()])

result = await session.send_and_wait_typed("answer", Result, timeout=1)

assert result.value.value == 42

assert result.values[0].value == 99

schema = client.request.call_args.args[1]["responseFormat"]["jsonSchema"]["schema"]

assert schema == Result.model_json_schema()

assert set(schema["properties"]) == {"result", "results"}

assert set(schema["$defs"]["Answer"]["properties"]) == {"answer"}

assert not Result.model_config["validate_by_alias"]

assert not Answer.model_config["validate_by_alias"]

assert_clean(session)

@pytest.mark.asyncio

async def test_raw_schema_is_forwarded_unchanged_and_result_remains_an_event():

schema = {"type": "object", "properties": {"count": {"type": "integer", "minimum": 2}}}

final = assistant()

session, client = fake_session([final, idle()])

result = await session.send_and_wait("inventory", response_schema=schema, timeout=1)

assert result is final

assert client.request.call_args.args[1]["responseFormat"]["jsonSchema"]["schema"] == schema

assert "additionalProperties" not in schema

assert_clean(session)

@pytest.mark.asyncio

async def test_correlation_tool_commentary_stop_corrections_and_autopilot():

subagent = assistant('{"count":999,"color":"wrong"}')

subagent.agent_id = "child"

events = [

idle(),

event("session.error", {"errorType": "test", "message": "before this run"}),

user(),

assistant("working", toolRequests=[{"toolCallId": "tool-1", "name": "inventory"}]),

assistant(),

idle(mode="autopilot"),

assistant('{"count":99,"color":"blue"}'),

subagent,

assistant('{"count":123,"color":"wrong"}', origin="other-user"),

idle(),

]

session, _ = fake_session(events)

assert await session.send_and_wait_typed("inventory", Inventory, timeout=1) == Inventory(

count=99, color="blue"

)

assert_clean(session)

@pytest.mark.parametrize(

("events", "error"),

[

([user(), idle()], "without a structured"),

([user(), assistant(" "), idle()], "without a structured"),

(

[

assistant(),

assistant("working", toolRequests=[{"toolCallId": "t", "name": "tool"}]),

idle(),

],

"without a structured",

),

([assistant(), idle(aborted=True)], "aborted"),

(

[user(), event("session.error", {"errorType": "test", "message": "provider failed"})],

"provider failed",

),

],

)

@pytest.mark.asyncio

async def test_failed_runs_do_not_return_stale_or_missing_results(events, error):

session, _ = fake_session(events)

with pytest.raises(RuntimeError, match=error):

await session.send_and_wait_typed("inventory", Inventory, timeout=1)

assert_clean(session)

@pytest.mark.parametrize(

"content",

[

"null",

"not JSON",

'{"count":"no","color":"red"}',

'{"count":42}',

'{"count":42,"color":"red","extra":1}',

],

)

@pytest.mark.asyncio

async def test_typed_output_validates_json(content):

session, _ = fake_session([assistant(content), idle()])

with pytest.raises(ValidationError):

await session.send_and_wait_typed("inventory", Inventory, timeout=1)

assert_clean(session)

@pytest.mark.asyncio

async def test_admission_error_is_not_hidden_by_buffered_success():

session, _ = fake_session([assistant(), idle()], ValueError("invalid schema"))

with pytest.raises(ValueError, match="invalid schema"):

await session.send_and_wait_typed("inventory", Inventory, timeout=1)

assert_clean(session)

@pytest.mark.parametrize("kind", ["timeout", "cancel", "disconnect"])

@pytest.mark.asyncio

async def test_wait_cleanup(kind, monkeypatch):

session, client = fake_session()

admitted = asyncio.get_running_loop().create_future()

# Expire only after admission, not while generating the response schema.

timeout_context = asyncio.timeout(None)

monkeypatch.setattr(asyncio, "timeout_at", lambda _: timeout_context)

async def request(method, params):

if method == "session.sendAndWait":

raise JsonRpcError(-32601, "Method not found")

admitted.set_result(None)

return {"messageId": "user-1"}

client.request.side_effect = request

waiting = asyncio.create_task(

session.send_and_wait_typed("inventory", Inventory, timeout=float("inf"))

)

try:

done, _ = await asyncio.wait((admitted, waiting), return_when=asyncio.FIRST_COMPLETED)

if waiting in done:

await waiting

pytest.fail("Structured wait completed before cleanup was exercised")

assert session._structured_waits

assert session._event_handlers

if kind == "cancel":

waiting.cancel()

error = asyncio.CancelledError

elif kind == "disconnect":

session._mark_disconnected()

error = RuntimeError

else:

timeout_context.reschedule(asyncio.get_running_loop().time())

error = TimeoutError

with pytest.raises(error):

await waiting

assert timeout_context.expired() == (kind == "timeout")

finally:

waiting.cancel()

await asyncio.gather(waiting, return_exceptions=True)

assert_clean(session)

@pytest.mark.asyncio

async def test_concurrent_typed_waits_keep_their_own_origin():

session, client = fake_session()

admitted = asyncio.Queue()

async def request(method, params):

if method == "session.sendAndWait":

raise JsonRpcError(-32601, "Method not found")

origin = params["prompt"]

admitted.put_nowait(origin)

return {"messageId": origin}

client.request.side_effect = request

first = asyncio.create_task(session.send_and_wait_typed("first", Inventory, timeout=1))

second = asyncio.create_task(session.send_and_wait_typed("second", Inventory, timeout=1))

await admitted.get()

await admitted.get()

session._dispatch_event(assistant(origin="first"))

session._dispatch_event(assistant('{"count":7,"color":"blue"}', origin="second"))

session._dispatch_event(idle())

assert (await first).count == 42

assert (await second).count == 7

assert_clean(session)

@pytest.mark.asyncio

async def test_typed_immediate_is_rejected_before_admission():

session, client = fake_session()

with pytest.raises(ValueError, match="immediate"):

await session.send_and_wait_typed("inventory", Inventory, mode="immediate")

client.request.assert_not_called()

assert_clean(session)

@pytest.mark.asyncio

async def test_disconnect_while_admission_is_pending_cancels_rpc_wait():

session, client = fake_session()

entered = asyncio.Event()

cancelled = asyncio.Event()

async def request(method, params):

entered.set()

try:

await asyncio.Future()

finally:

cancelled.set()

client.request.side_effect = request

waiting = asyncio.create_task(session.send_and_wait_typed("inventory", Inventory))

await entered.wait()

session._mark_disconnected()

with pytest.raises(RuntimeError, match="Session closed"):

await asyncio.wait_for(waiting, 1)

assert cancelled.is_set()

assert_clean(session)