from __future__ import annotations

import asyncio
import contextlib
import re
import uuid
from typing import Any

import httpx
import pytest

from langgraph_sdk._async.http import HttpClient
from langgraph_sdk._async.stream import AsyncThreadStream
from langgraph_sdk._async.threads import ThreadsClient
from langgraph_sdk.stream.transport import (
    ProtocolSseTransport,
    ProtocolWebSocketTransport,
)
from streaming._events import (
    checkpoints_event,
    custom_event,
    lifecycle_completed_event,
    lifecycle_event,
    lifecycle_started_event,
    message_finish_event,
    message_start_event,
    tasks_start_event,
    tool_finished_event,
    tool_started_event,
    updates_event,
    values_event,
)
from streaming._fake_server import FakeServer


async def test_thread_agent_get_tree_fetches_assistant_graph():
    fake = FakeServer()
    fake.set_graph(
        {
            "nodes": [{"id": "agent", "type": "runnable", "data": {"name": "agent"}}],
            "edges": [{"source": "agent", "target": "__end__"}],
        }
    )
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(
            thread_id="t-1",
            assistant_id="agent",
            headers={"X-Custom-Header": "my-value"},
        ) as thread:
            graph = await thread.agent.get_tree(xray=True)

    assert graph["nodes"][0]["id"] == "agent"
    assert graph["edges"] == [{"source": "agent", "target": "__end__"}]
    assert fake.graph_request_params == [{"xray": "true"}]
    assert fake.graph_request_headers[0].get("x-custom-header") == "my-value"


async def test_thread_agent_get_tree_raises_after_close():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        stream = threads.stream(thread_id="t-1", assistant_id="agent")
        await stream.close()
        with pytest.raises(RuntimeError, match="closed"):
            await stream.agent.get_tree()


async def test_extensions_projection_empty_name_raises():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        stream = threads.stream(thread_id="t-1", assistant_id="agent")
        with pytest.raises(ValueError, match="non-empty"):
            stream.extensions[""]


async def test_extensions_projection_closed_stream_yields_nothing():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        # Enter and immediately exit so _controller is set but _closed is True.
        async with threads.stream(thread_id="t-1", assistant_id="agent") as stream:
            pass
        payloads = [p async for p in stream.extensions["progress"]]
    assert payloads == []


async def test_thread_stream_stores_thread_id_and_assistant_id():
    async with httpx.AsyncClient(base_url="http://test") as client:
        stream = AsyncThreadStream(
            http=HttpClient(client),
            thread_id="t-1",
            assistant_id="agent",
        )
        assert stream.thread_id == "t-1"
        assert stream.assistant_id == "agent"


async def test_aenter_returns_self():
    async with httpx.AsyncClient(base_url="http://test") as client:
        stream = AsyncThreadStream(
            http=HttpClient(client), thread_id="t-1", assistant_id="agent"
        )
        async with stream as entered:
            assert entered is stream


async def test_aexit_marks_closed():
    async with httpx.AsyncClient(base_url="http://test") as client:
        stream = AsyncThreadStream(
            http=HttpClient(client), thread_id="t-1", assistant_id="agent"
        )
        async with stream:
            assert stream._closed is False
        assert stream._closed is True


async def test_close_is_idempotent():
    async with httpx.AsyncClient(base_url="http://test") as client:
        stream = AsyncThreadStream(
            http=HttpClient(client), thread_id="t-1", assistant_id="agent"
        )
        await stream.close()
        await stream.close()  # must not raise
        assert stream._closed is True


async def test_threads_stream_returns_async_thread_stream_with_explicit_id():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        stream = threads.stream(thread_id="my-thread", assistant_id="agent")
        assert stream.thread_id == "my-thread"
        assert stream.assistant_id == "agent"


async def test_threads_stream_mints_uuid4_when_thread_id_none():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        stream = threads.stream(assistant_id="agent")
        # uuid4 format: 8-4-4-4-12 hex
        assert re.fullmatch(
            r"[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}",
            stream.thread_id,
        )
        # And it's actually parseable as a v4 UUID.
        assert uuid.UUID(stream.thread_id).version == 4


async def test_threads_stream_requires_assistant_id():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        with pytest.raises(TypeError):
            threads.stream(thread_id="t-1")  # ty: ignore[missing-argument]


async def test_threads_stream_headers_forwarded_to_commands():
    """Headers passed to `threads.stream()` are forwarded to /commands requests."""
    fake = FakeServer()
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(
            thread_id="t-1",
            assistant_id="agent",
            headers={"X-Custom-Header": "my-value"},
        ) as thread:
            await thread.run.start(input={})
    assert fake.command_request_headers, "no command requests captured"
    assert fake.command_request_headers[0].get("x-custom-header") == "my-value"


async def test_threads_stream_headers_forwarded_to_stream_events():
    """Headers passed to `threads.stream()` are forwarded to /stream/events requests."""
    fake = FakeServer()
    fake.script([lifecycle_event(seq=0)])
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(
            thread_id="t-1",
            assistant_id="agent",
            headers={"X-Custom-Header": "my-value"},
        ) as thread:
            await thread.run.start(input={})
            _ = [e async for e in thread.subscribe(["lifecycle"])]
    assert fake.stream_request_headers_list, "no stream/events requests captured"
    assert fake.stream_request_headers_list[0].get("x-custom-header") == "my-value"


async def test_no_headers_by_default():
    """When `headers` is omitted, `_headers` is an empty dict and no custom
    headers appear in command or stream requests.
    """
    fake = FakeServer()
    fake.script([lifecycle_event(seq=0)])
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            assert thread._headers == {}
            await thread.run.start(input={})
            _ = [e async for e in thread.subscribe(["lifecycle"])]
    # No custom header keys beyond the protocol-required / transport-required ones.
    protocol_keys = {
        "content-type",
        "accept",
        "cache-control",
        "host",
        "user-agent",
        "accept-encoding",
        "connection",
        "transfer-encoding",
        "content-length",
    }
    extra_command = {
        k for k in fake.command_request_headers[0] if k.lower() not in protocol_keys
    }
    extra_stream = {
        k for k in fake.stream_request_headers_list[0] if k.lower() not in protocol_keys
    }
    assert extra_command == set(), f"unexpected command headers: {extra_command}"
    assert extra_stream == set(), f"unexpected stream headers: {extra_stream}"


async def test_aenter_constructs_transport_with_thread_id():
    fake = FakeServer()
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        from langgraph_sdk._async.http import HttpClient
        from langgraph_sdk._async.threads import ThreadsClient

        threads = ThreadsClient(HttpClient(raw))
        stream = threads.stream(thread_id="t-1", assistant_id="agent")
        async with stream:
            assert stream._transport is not None
            assert stream._transport.thread_id == "t-1"


async def test_aenter_selects_websocket_transport():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        from langgraph_sdk._async.http import HttpClient
        from langgraph_sdk._async.threads import ThreadsClient

        threads = ThreadsClient(HttpClient(raw))
        stream = threads.stream(
            thread_id="t-1", assistant_id="agent", transport="websocket"
        )
        async with stream:
            assert isinstance(stream._transport, ProtocolWebSocketTransport)


async def test_aexit_closes_transport():
    fake = FakeServer()
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        from langgraph_sdk._async.http import HttpClient
        from langgraph_sdk._async.threads import ThreadsClient

        threads = ThreadsClient(HttpClient(raw))
        stream = threads.stream(thread_id="t-1", assistant_id="agent")
        async with stream:
            inner_transport = stream._transport
        assert inner_transport is not None
        assert isinstance(inner_transport, ProtocolSseTransport)
        assert inner_transport._closed is True


async def test_run_start_sends_command_with_assistant_id():
    fake = FakeServer()
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        from langgraph_sdk._async.http import HttpClient
        from langgraph_sdk._async.threads import ThreadsClient

        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            result = await thread.run.start(input={"x": 1})
    assert result == {"run_id": "run-1"}
    command = fake.received_commands[0]
    assert command["method"] == "run.start"
    assert command["params"]["assistant_id"] == "agent"
    assert command["params"]["input"] == {"x": 1}
    assert command["id"] == 1


async def test_command_ids_are_monotonic():
    fake = FakeServer()
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        from langgraph_sdk._async.http import HttpClient
        from langgraph_sdk._async.threads import ThreadsClient

        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={"x": 1})
            await thread.run.start(input={"x": 2})
    assert [c["id"] for c in fake.received_commands] == [1, 2]


async def test_run_start_forwards_config_and_metadata():
    fake = FakeServer()
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        from langgraph_sdk._async.http import HttpClient
        from langgraph_sdk._async.threads import ThreadsClient

        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(
                input={"x": 1},
                config={"recursion_limit": 5},
                metadata={"trace": "abc"},
            )
    params = fake.received_commands[0]["params"]
    assert params["config"] == {"recursion_limit": 5}
    assert params["metadata"] == {"trace": "abc"}


async def test_run_start_raises_outside_context_manager():
    import pytest

    async with httpx.AsyncClient(base_url="http://test") as raw:
        stream = AsyncThreadStream(
            http=HttpClient(raw), thread_id="t-1", assistant_id="agent"
        )
        with pytest.raises(RuntimeError, match="async with"):
            await stream.run.start(input={"x": 1})


async def test_run_start_raises_on_error_envelope():
    from starlette.applications import Starlette
    from starlette.responses import JSONResponse
    from starlette.routing import Route

    async def commands(_request):
        return JSONResponse(
            {
                "type": "error",
                "id": 1,
                "error": "invalid_argument",
                "message": "run.start requires an assistant_id.",
            }
        )

    app = Starlette(
        routes=[Route("/threads/{thread_id}/commands", commands, methods=["POST"])]
    )
    transport = httpx.ASGITransport(app=app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        import pytest

        from langgraph_sdk._async.http import HttpClient
        from langgraph_sdk._async.threads import ThreadsClient

        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            with pytest.raises(RuntimeError, match="invalid_argument"):
                await thread.run.start(input={"x": 1})


async def test_events_yields_raw_events_after_run_start():
    fake = FakeServer()
    fake.script(
        [
            lifecycle_event(seq=0),
            values_event(seq=1),
        ]
    )
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        from langgraph_sdk._async.http import HttpClient
        from langgraph_sdk._async.threads import ThreadsClient

        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            received = [e async for e in thread.events]
    methods = [e["method"] for e in received]
    assert methods == ["lifecycle", "values"]


async def test_events_subscribes_to_all_channels():
    fake = FakeServer()
    fake.script([])
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        from langgraph_sdk._async.http import HttpClient
        from langgraph_sdk._async.threads import ThreadsClient

        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            _ = [e async for e in thread.events]
    body = fake.stream_request_bodies[0]
    assert set(body["channels"]) == {
        "values",
        "updates",
        "messages",
        "tools",
        "lifecycle",
        "input",
        "checkpoints",
        "tasks",
        "custom",
    }


async def test_events_terminates_on_aexit():
    import asyncio

    import pytest

    fake = FakeServer()
    fake.script([lifecycle_event(seq=i) for i in range(5)])
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        from langgraph_sdk._async.http import HttpClient
        from langgraph_sdk._async.threads import ThreadsClient

        threads = ThreadsClient(HttpClient(raw))
        stream = threads.stream(thread_id="t-1", assistant_id="agent")
        async with stream as thread:
            await thread.run.start(input={})
            handle_events = thread.events
        # After __aexit__, further iteration must terminate cleanly.
        with pytest.raises(StopAsyncIteration):
            await asyncio.wait_for(handle_events.__anext__(), timeout=1.0)


async def test_events_raises_outside_context_manager():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        stream = AsyncThreadStream(
            http=HttpClient(raw), thread_id="t-1", assistant_id="agent"
        )
        with pytest.raises(RuntimeError, match="async with"):
            _ = stream.events


async def test_aexit_preserves_original_exception_if_close_raises():
    """If the body of `async with` raises, AND close() also raises, the
    body's exception must propagate. close()'s error is suppressed (chained
    as context on close_err, but does not replace the original)."""
    async with httpx.AsyncClient(base_url="http://test") as raw:
        thread = AsyncThreadStream(
            http=HttpClient(raw), thread_id="t-1", assistant_id="agent"
        )

        async def failing_close():
            raise RuntimeError("close failed")

        thread.close = failing_close  # ty:ignore[invalid-assignment]

        with pytest.raises(ValueError, match="original"):
            async with thread:
                raise ValueError("original")


async def test_events_property_returns_fresh_iterator_each_access():
    """Two separate accesses of `thread.events` must return independent
    subscriptions — the second access should produce a fresh iterator,
    even if both are accessed before either is drained."""
    fake = FakeServer()
    fake.script([])
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        from langgraph_sdk._async.http import HttpClient
        from langgraph_sdk._async.threads import ThreadsClient

        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            first_iter = thread.events
            second_iter = thread.events
            # Each property access must return a distinct iterator object.
            assert first_iter is not second_iter


async def test_fresh_thread_happy_path_end_to_end():
    """User passes no thread_id; SDK mints one and uses it in all URLs.

    Validates core surface end-to-end:
      - uuid4 minted at client.threads.stream()
      - run.start posted to /threads/<minted-id>/commands
      - events SSE opened at /threads/<minted-id>/stream/events
      - scripted events delivered to the user iterator
    """
    fake = FakeServer()
    fake.script([lifecycle_event(seq=0), values_event(seq=1)])

    posted_paths: list[str] = []

    class _PathSpyTransport(httpx.ASGITransport):
        async def handle_async_request(self, request):
            posted_paths.append(str(request.url.path))
            return await super().handle_async_request(request)

    spy = _PathSpyTransport(app=fake.app)
    async with httpx.AsyncClient(transport=spy, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(assistant_id="agent") as thread:
            assert uuid.UUID(thread.thread_id).version == 4
            result = await thread.run.start(input={"x": 1})
            assert result == {"run_id": "run-1"}
            received = [e async for e in thread.events]
    assert [e["method"] for e in received] == ["lifecycle", "values"]
    # Both POSTs must include the minted thread_id in the path.
    minted_id_paths = [p for p in posted_paths if thread.thread_id in p]
    assert any(p.endswith("/commands") for p in minted_id_paths)
    assert any(p.endswith("/stream/events") for p in minted_id_paths)


async def test_aenter_raises_after_close():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        stream = AsyncThreadStream(
            http=HttpClient(raw), thread_id="t-1", assistant_id="agent"
        )
        async with stream:
            pass
        # After exit, the stream is closed; re-entering must raise rather than
        # silently constructing a new transport that would leak on the next exit.
        with pytest.raises(RuntimeError, match="closed and cannot be re-entered"):
            async with stream:
                pass


async def test_register_subscription_assigns_monotonic_ids():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        stream = AsyncThreadStream(
            http=HttpClient(raw), thread_id="t-1", assistant_id="agent"
        )
        async with stream:
            sub_a = stream._register_subscription({"channels": ["values"]})
            sub_b = stream._register_subscription({"channels": ["messages"]})
            assert sub_a.id == 1
            assert sub_b.id == 2
            assert stream._subscriptions[sub_a.id] is sub_a
            assert stream._subscriptions[sub_b.id] is sub_b


async def test_unregister_subscription_removes_from_registry():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        stream = AsyncThreadStream(
            http=HttpClient(raw), thread_id="t-1", assistant_id="agent"
        )
        async with stream:
            sub = stream._register_subscription({"channels": ["values"]})
            stream._unregister_subscription(sub.id)
            assert sub.id not in stream._subscriptions


async def test_await_run_start_gate_honors_timeout():
    """Gate must raise asyncio.TimeoutError if run.start never completes
    within the configured timeout."""
    import asyncio

    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            # Install a never-resolving gate to simulate an in-flight
            # run.start that will not complete within the timeout window.
            loop = asyncio.get_running_loop()
            thread._run_start_ready = loop.create_future()
            with pytest.raises(asyncio.TimeoutError):
                await thread._await_run_start_gate(timeout=0.1)
            # Gate must still be pending after the timeout (no side effects).
            assert thread._run_start_ready is not None
            assert not thread._run_start_ready.done()


async def test_await_run_start_gate_returns_when_gate_resolves_in_time():
    """With a generous timeout and a gate that resolves promptly, the
    gate returns without raising."""
    import asyncio

    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            loop = asyncio.get_running_loop()
            gate: asyncio.Future[None] = loop.create_future()
            thread._run_start_ready = gate
            loop.call_later(0.01, lambda: gate.set_result(None))
            await thread._await_run_start_gate(timeout=1.0)


async def test_run_start_timeout_constructor_kwarg_forwarded_to_gate():
    """`run_start_timeout` constructor kwarg is stored and consulted by
    `_reconcile_stream` via `_await_run_start_gate`."""
    import asyncio

    async with httpx.AsyncClient(base_url="http://test") as raw:
        stream = AsyncThreadStream(
            http=HttpClient(raw),
            thread_id="t-1",
            assistant_id="agent",
            run_start_timeout=0.1,
        )
        async with stream as thread:
            loop = asyncio.get_running_loop()
            # Install a never-resolving gate.
            thread._run_start_ready = loop.create_future()
            with pytest.raises(asyncio.TimeoutError):
                # Reconcile must surface the timeout from the gate.
                await thread._reconcile_stream({"channels": ["lifecycle"]})


async def test_subscribe_waits_for_run_start_to_commit():
    """Subscribing before run.start commits must not race the server.

    With the gate: subscribers wait for run.start to return before opening
    their SSE. Without it, a fast subscribe would 404 against a thread the
    server hasn't created yet.
    """
    import asyncio

    fake = FakeServer()
    fake.script([])
    transport = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            # Kick run.start without awaiting — concurrently subscribe.
            run_task = asyncio.create_task(thread.run.start(input={}))
            sub_iter = thread.subscribe(["lifecycle"])
            # Drain one event or hit EOF. The iterator's first __anext__
            # awaits _reconcile_stream which awaits the gate.
            async for _ in sub_iter:
                break
            # If the gate works, run.start completed before the subscription
            # opened its SSE (and thus before iteration finished).
            assert run_task.done()


async def test_run_respond_dispatches_input_respond_command():
    fake = FakeServer()
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            # Simulate one outstanding interrupt.
            thread.interrupts.append(
                {"interrupt_id": "i-1", "value": None, "namespace": []}
            )
            thread.interrupted = True
            await thread.run.respond("yes")
    command = fake.received_commands[-1]
    assert command["method"] == "input.respond"
    assert command["params"]["interrupt_id"] == "i-1"
    assert command["params"]["response"] == "yes"
    assert command["params"]["namespace"] == []


async def test_run_respond_with_explicit_interrupt_id():
    fake = FakeServer()
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            thread.interrupts.extend(
                [
                    {"interrupt_id": "a", "value": None, "namespace": []},
                    {"interrupt_id": "b", "value": None, "namespace": []},
                ]
            )
            thread.interrupted = True
            await thread.run.respond("pick", interrupt_id="b")
    assert fake.received_commands[-1]["params"]["interrupt_id"] == "b"
    assert fake.received_commands[-1]["params"]["namespace"] == []


async def test_run_respond_raises_when_no_outstanding_interrupts():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            with pytest.raises(RuntimeError, match="no outstanding interrupt"):
                await thread.run.respond("yes")


async def test_run_respond_raises_when_ambiguous_interrupt_id():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            thread.interrupts.extend(
                [
                    {"interrupt_id": "a", "value": None, "namespace": []},
                    {"interrupt_id": "b", "value": None, "namespace": []},
                ]
            )
            thread.interrupted = True
            with pytest.raises(RuntimeError, match=r"ambiguous|interrupt_id"):
                await thread.run.respond("yes")


async def test_run_respond_snapshots_interrupts_under_lock():
    """`respond()` must take a snapshot of `interrupts` under the
    `_interrupts_lock`, so a concurrent terminal-event clear cannot
    invalidate the in-flight dispatch.

    Verifies: if `_interrupts_lock` is held when `respond()` is called,
    `respond()` blocks until the lock is released — proving it serializes
    with the terminal-clear path that takes the same lock.
    """
    import asyncio

    fake = FakeServer()
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            thread.interrupts.append(
                {"interrupt_id": "i-1", "value": None, "namespace": []}
            )
            thread.interrupted = True
            # Take the interrupts lock externally to block `respond()`.
            assert hasattr(thread, "_interrupts_lock"), (
                "AsyncThreadStream must expose _interrupts_lock"
            )
            await thread._interrupts_lock.acquire()
            try:
                # `respond()` must NOT complete while we hold the lock.
                task = asyncio.create_task(thread.run.respond("yes"))
                # Give the task a chance to start and reach the lock.
                await asyncio.sleep(0.05)
                assert not task.done(), (
                    "respond() should be blocked waiting for _interrupts_lock"
                )
            finally:
                thread._interrupts_lock.release()
            # Now `respond()` should complete.
            await asyncio.wait_for(task, timeout=1.0)
    command = fake.received_commands[-1]
    assert command["method"] == "input.respond"
    assert command["params"]["interrupt_id"] == "i-1"


async def test_terminal_lifecycle_clear_acquires_interrupts_lock():
    """Terminal lifecycle event clears `interrupts` under the same lock
    that `respond()` uses, preventing TOCTOU between snapshot and
    dispatch."""
    import asyncio

    fake = FakeServer()
    # No scripted events; we exercise `_apply_lifecycle_event` directly.
    fake.script([])
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            thread.interrupts.append(
                {"interrupt_id": "i-1", "value": None, "namespace": []}
            )
            thread.interrupted = True
            # Hold the lock; a completion event must block on it before
            # clearing interrupts.
            await thread._interrupts_lock.acquire()
            try:
                from typing import cast

                from langchain_protocol import Event

                terminal_event = cast(
                    Event,
                    {
                        "type": "event",
                        "method": "lifecycle",
                        "params": {
                            "namespace": [],
                            "data": {"event": "completed"},
                        },
                        "seq": 99,
                        "event_id": "evt-99",
                    },
                )
                clear_task = asyncio.create_task(
                    thread._apply_lifecycle_event(terminal_event)
                )
                await asyncio.sleep(0.05)
                # Interrupts must still be present — clear is blocked.
                assert thread.interrupted is True
                assert len(thread.interrupts) == 1
                assert not clear_task.done()
            finally:
                thread._interrupts_lock.release()
            await asyncio.wait_for(clear_task, timeout=1.0)
            assert thread.interrupted is False
            assert thread.interrupts == []


async def test_run_respond_raises_when_explicit_interrupt_id_not_outstanding():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            thread.interrupts.append(
                {"interrupt_id": "a", "value": None, "namespace": []}
            )
            thread.interrupted = True
            with pytest.raises(RuntimeError, match="does not match"):
                await thread.run.respond("yes", interrupt_id="nonexistent")


async def test_output_cancellation_does_not_trigger_new_fetch():
    """When the in-flight fetch task for thread.output is cancelled, a
    subsequent call must NOT spawn a fresh task and issue a new REST GET;
    `_get_task` should return the same (cancelled) task so awaiters share
    the CancelledError outcome.
    """
    fake = FakeServer()
    # No lifecycle terminal event — the fetch task will park on _run_done.
    fake.script([])
    fake.set_state({"messages": ["hello"]})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            output_awaitable = thread.output

            # Materialize the underlying task and let it park on _run_done.
            shared_task = output_awaitable._get_task()
            for _ in range(20):
                await asyncio.sleep(0)
                if not shared_task.done():
                    break

            # Simulate an in-flight fetch being cancelled.
            shared_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await shared_task
            assert shared_task.done()
            assert shared_task.cancelled()

            # A subsequent _get_task() must return the SAME cancelled task —
            # no respawn, no fresh REST fetch.
            second_task = output_awaitable._get_task()
            assert second_task is shared_task, (
                "expected cancelled task to be reused; got a fresh task"
            )
            # No state GET should have been issued (lifecycle never completed).
            assert fake.state_request_count == 0


async def test_output_with_timeout_raises_timeout_error():
    """thread.output.with_timeout(s) raises TimeoutError when the lifecycle
    never resolves within the budget."""
    fake = FakeServer()
    # Hold the stream open with a long inter-event delay so the lifecycle
    # watcher parks on the iterator (mid-sleep before a non-terminal event)
    # and `_run_done` never resolves before the timeout fires. Without the
    # delay, a clean EOF would resolve `_run_done` with an errored
    # `_RunTerminal` (per PR 7821) and `with_timeout` would raise that error
    # instead of `asyncio.TimeoutError`.
    fake.script([lifecycle_started_event(seq=0)], delay=10.0)
    fake.set_state({"never": "reached"})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            with pytest.raises(asyncio.TimeoutError):
                await thread.output.with_timeout(0.1)


async def test_output_with_timeout_returns_value_when_lifecycle_completes_in_time():
    """thread.output.with_timeout(s) returns the values dict when the lifecycle
    resolves within the budget."""
    fake = FakeServer()
    fake.script([lifecycle_completed_event(seq=0)])
    fake.set_state({"messages": ["hello"]})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            result = await thread.output.with_timeout(2.0)
    assert result == {"messages": ["hello"]}


async def test_output_with_timeout_returns_new_awaitable_not_self():
    """with_timeout() returns a fresh awaitable, leaving the original untouched."""
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            bounded = thread.output.with_timeout(0.5)
            assert bounded is not thread.output
            assert bounded._timeout == 0.5
            assert thread.output._timeout is None


async def test_threads_stream_accepts_websocket_transport_option():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        stream = threads.stream(
            thread_id="t-1",
            assistant_id="agent",
            transport="websocket",
        )
    assert stream._transport_kind == "websocket"


async def test_threads_stream_rejects_unknown_transport_option():
    async with httpx.AsyncClient(base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        with pytest.raises(ValueError, match="transport"):
            threads.stream(
                thread_id="t-1",
                assistant_id="agent",
                transport="bogus",  # ty: ignore[invalid-argument-type]
            )


async def test_v3_streaming_async_surface_smoke():
    import asyncio

    from streaming._events import (
        custom_event,
        lifecycle_completed_event,
        message_finish_event,
        message_start_event,
        message_text_delta_event,
        message_text_finish_event,
        tool_finished_event,
        tool_started_event,
        values_event,
    )

    fake = FakeServer()
    fake.set_state({"final": True})
    fake.script(
        [
            values_event(seq=1, values={"step": 1}),
            message_start_event(seq=2, message_id="msg-1"),
            message_text_delta_event(seq=3, text="hi", message_id="msg-1"),
            message_text_finish_event(seq=4, text="hi", message_id="msg-1"),
            message_finish_event(seq=5, message_id="msg-1"),
            tool_started_event(seq=6, tool_call_id="call-1", tool_name="search"),
            tool_finished_event(seq=7, tool_call_id="call-1", output={"ok": True}),
            custom_event(seq=8, name="progress", step=1),
            lifecycle_completed_event(seq=9),
        ]
    )
    async with httpx.AsyncClient(
        transport=fake.transport, base_url="http://test"
    ) as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            start = await thread.run.start(
                input={"messages": [{"role": "user", "content": "hi"}]}
            )

            # Start all consumers concurrently so their subscriptions are all
            # registered before the first SSE reconciliation. This means ONE
            # shared SSE opens with the union filter, dedup never rejects events.
            async def _get_first_values() -> Any:
                async for v in thread.values:
                    return v

            async def _get_message_streams():
                return [s async for s in thread.messages]

            async def _get_tool_calls():
                return [call async for call in thread.tool_calls]

            async def _get_progress():
                return [p async for p in thread.extensions["progress"]]

            first_values, message_streams, tool_calls, progress = await asyncio.gather(
                _get_first_values(),
                _get_message_streams(),
                _get_tool_calls(),
                _get_progress(),
            )
            # stream.text is already fully accumulated after gather completes.
            message_texts = [await s.text for s in message_streams]
            final = await thread.output

    assert start == {"run_id": "run-1"}
    assert first_values == fake.state["values"]
    assert message_texts == ["hi"]
    assert tool_calls[0].name == "search"
    assert progress == [{"name": "progress", "step": 1}]
    assert final == {"final": True}


async def test_interleave_projections_single_channel_values():
    fake = FakeServer()
    fake.script(
        [
            lifecycle_started_event(seq=0),
            values_event(seq=1, counter=1),
            values_event(seq=2, counter=2),
            lifecycle_completed_event(seq=3),
        ]
    )
    fake.set_state({"counter": 0})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            items = []
            async for ch, item in thread.interleave_projections(["values"]):
                items.append((ch, item))
    assert ("values", {"counter": 1}) in items
    assert ("values", {"counter": 2}) in items
    assert all(ch == "values" for ch, _ in items)


async def test_interleave_projections_values_and_messages_arrival_order():
    fake = FakeServer()
    fake.script(
        [
            lifecycle_started_event(seq=0),
            values_event(seq=1, counter=1),
            message_start_event(seq=2, message_id="m-1"),
            values_event(seq=3, counter=2),
            message_finish_event(seq=4, message_id="m-1"),
            lifecycle_completed_event(seq=5),
        ]
    )
    fake.set_state({"counter": 0})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            order = []
            async for ch, _ in thread.interleave_projections(["values", "messages"]):
                order.append(ch)
                if len(order) >= 3:
                    break
    assert order[:3] == ["values", "messages", "values"]


async def test_interleave_projections_mixes_builtin_and_extension():
    fake = FakeServer()
    fake.script(
        [
            lifecycle_started_event(seq=0),
            values_event(seq=1, counter=1),
            custom_event(seq=2, name="foo", hello="world"),
            lifecycle_completed_event(seq=3),
        ]
    )
    fake.set_state({"counter": 0})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            items = []
            async for ch, item in thread.interleave_projections(["values", "foo"]):
                items.append((ch, item))
    assert ("values", {"counter": 1}) in items
    # extension payload is the whole params.data including "name"; tuple uses bare name "foo"
    assert ("foo", {"name": "foo", "hello": "world"}) in items


async def test_interleave_projections_tool_calls_uses_public_name():
    fake = FakeServer()
    fake.script(
        [
            lifecycle_started_event(seq=0),
            tool_started_event(seq=1, tool_call_id="call-1", tool_name="search"),
            tool_finished_event(seq=2, tool_call_id="call-1", output={"ok": True}),
            lifecycle_completed_event(seq=3),
        ]
    )
    fake.set_state({})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            names = []
            async for ch, item in thread.interleave_projections(["tool_calls"]):
                names.append(ch)
                assert item.tool_call_id == "call-1"  # real ToolCallHandle
    # tuple uses the PUBLIC name "tool_calls", never the wire name "tools"
    assert names == ["tool_calls"]


async def test_interleave_projections_subgraphs_discovers_child():
    fake = FakeServer()
    fake.script(
        [
            lifecycle_started_event(seq=0),
            lifecycle_started_event(seq=1, namespace=["child"]),
            lifecycle_completed_event(seq=2),
        ]
    )
    fake.set_state({})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            discovered = []
            async for ch, handle in thread.interleave_projections(["subgraphs"]):
                discovered.append((ch, handle.path))
    assert ("subgraphs", ("child",)) in discovered


async def test_interleave_projections_inflight_tool_call_failed_on_break():
    """A tool handle held past an early break is failed in teardown, never left hanging."""
    fake = FakeServer()
    fake.script(
        [
            lifecycle_started_event(seq=0),
            tool_started_event(seq=1, tool_call_id="call-1", tool_name="search"),
            # no tool-finished: the call is still in flight when the consumer breaks
            lifecycle_completed_event(seq=2),
        ]
    )
    fake.set_state({})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            handle = None
            async for _, item in thread.interleave_projections(["tool_calls"]):
                handle = item
                break
            assert handle is not None
            # Without teardown finalization this would hang forever; wait_for
            # turns a regression into a TimeoutError rather than a RuntimeError.
            with pytest.raises(RuntimeError):
                await asyncio.wait_for(handle.output, timeout=2)


async def test_interleave_projections_inflight_subgraph_finished_on_terminal():
    """A discovered subgraph child with no terminal tasks-result is force-completed."""
    fake = FakeServer()
    fake.script(
        [
            lifecycle_started_event(seq=0),
            lifecycle_started_event(seq=1, namespace=["child"]),
            # no tasks-result for the child: it is still "started" at run end
            lifecycle_completed_event(seq=2),
        ]
    )
    fake.set_state({})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            child = None
            async for _, handle in thread.interleave_projections(["subgraphs"]):
                child = handle
            assert child is not None
            assert child.status == "completed"


@pytest.mark.parametrize("channel", ["lifecycle", "tools", "input"])
async def test_interleave_projections_rejects_reserved_channel(channel):
    """Reserved protocol channel names raise instead of silently no-op'ing.

    `infer_channel` treats these as first-class methods, but they have no
    interleave decoder, so routing them to the extension/`custom:` fallback
    would subscribe to a channel that never matches and yield nothing. Fail
    closed. (`updates`/`checkpoints`/`tasks` are supported and tested below.)
    """
    fake = FakeServer()
    fake.script([lifecycle_started_event(seq=0), lifecycle_completed_event(seq=1)])
    fake.set_state({})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            with pytest.raises(ValueError, match=channel):
                async for _ in thread.interleave_projections([channel]):
                    pass


async def test_interleave_projections_data_channels_yield_payloads():
    """`updates`/`checkpoints`/`tasks` yield their raw `params.data` payloads."""
    fake = FakeServer()
    fake.script(
        [
            lifecycle_started_event(seq=0),
            updates_event(seq=1, node={"v": 1}),
            checkpoints_event(seq=2, ts="t-0", v=4),
            tasks_start_event(seq=3, task_id="task-9"),
            lifecycle_completed_event(seq=4),
        ]
    )
    fake.set_state({})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            items = []
            async for ch, item in thread.interleave_projections(
                ["updates", "checkpoints", "tasks"]
            ):
                items.append((ch, item))
    assert ("updates", {"node": {"v": 1}}) in items
    assert ("checkpoints", {"ts": "t-0", "v": 4}) in items
    assert any(ch == "tasks" and item.get("id") == "task-9" for ch, item in items)


async def test_interleave_projections_data_channel_scoped_to_root_namespace():
    """A child-namespace checkpoint must not leak into a root interleave.

    `values` subscribes unscoped, so `compute_union_filter` widens the merged
    subscription to all namespaces; the `DataDecoder` root filter is what keeps
    a subgraph checkpoint out of the root projection (mirrors local scope).
    """
    fake = FakeServer()
    fake.script(
        [
            lifecycle_started_event(seq=0),
            checkpoints_event(seq=1, namespace=["child"], scope="child"),
            checkpoints_event(seq=2, scope="root"),
            lifecycle_completed_event(seq=3),
        ]
    )
    fake.set_state({"counter": 0})
    asgi = httpx.ASGITransport(app=fake.app)
    async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw:
        threads = ThreadsClient(HttpClient(raw))
        async with threads.stream(thread_id="t-1", assistant_id="agent") as thread:
            await thread.run.start(input={})
            checkpoints = []
            async for ch, item in thread.interleave_projections(
                ["values", "checkpoints"]
            ):
                if ch == "checkpoints":
                    checkpoints.append(item)
    assert {"scope": "root"} in checkpoints
    assert {"scope": "child"} not in checkpoints
