"""Managed-worker failure, affinity, admission, and 70-client regression tests."""

import asyncio
from types import SimpleNamespace

import httpx
import pytest
from fastapi import HTTPException

from app.gateway.worker_pool.admission import Admission, bounded_worker_run
from app.gateway.worker_pool.controller import Controller, PoolUnavailable
from app.gateway.worker_pool.docker import DockerError, worker_spec
from app.gateway.worker_pool.settings import PoolSettings, SettingsStore


def source():
    return {
        "Image": "sha256:fixed",
        "Config": {"Env": ["AUTH_JWT_SECRET=fixed", "WEB_CONCURRENCY=8"], "WorkingDir": "/app"},
        "NetworkSettings": {"Networks": {"deerflow": {}}},
        "Mounts": [{"Type": "volume", "Name": "runtime", "Destination": "/data", "RW": True}, {"Type": "bind", "Source": "/var/run/docker.sock", "Destination": "/var/run/docker.sock", "RW": True}],
    }


def test_secrets_encrypted_and_preserved(tmp_path):
    store = SettingsStore(tmp_path)
    store.save(PoolSettings(password="secret-sentinel"))
    assert "secret-sentinel" not in store.path.read_text()
    assert "password" not in store.public()
    assert store.public()["password_set"]
    store.save(PoolSettings(desired=2))
    assert SettingsStore(tmp_path).load().password == "secret-sentinel"
    assert store.load().desired == 2


def test_worker_template_has_limits_no_scheduler_no_socket():
    spec = worker_spec(source(), PoolSettings(network="deerflow"), "instance", "token", "/private")
    assert spec["Image"] == "sha256:fixed"
    assert spec["HostConfig"]["Memory"] == 2048 * 1024**2
    assert len(spec["HostConfig"]["Mounts"]) == 1
    assert "SCHEDULED_TASKS_ENABLED=false" in spec["Env"]
    assert not any(e.startswith("WEB_CONCURRENCY=") for e in spec["Env"])
    assert "PortBindings" not in spec["HostConfig"]
    assert spec["Cmd"][-2:] == ["--workers", "1"]


def test_worker_template_uses_snapshot_of_main_runtime():
    spec = worker_spec(
        source(), PoolSettings(network="deerflow"), "instance", "token", "/private",
        runtime_image="sha256:runtime-snapshot",
    )
    assert spec["Image"] == "sha256:runtime-snapshot"


def test_postgres_worker_requires_shared_persistent_chat_archive_directory():
    template = source()
    template["Config"]["Env"].append("DEER_FLOW_HOME=/data")
    spec = worker_spec(
        template, PoolSettings(network="deerflow"), "instance", "token", "/private",
        require_shared_archive=True,
    )
    assert any(mount["Target"] == "/data" and mount["ReadOnly"] is False for mount in spec["HostConfig"]["Mounts"])

    template["Mounts"] = [
        {"Type": "volume", "Name": "other", "Destination": "/other", "RW": True},
    ]
    with pytest.raises(DockerError, match="共享的可写持久卷"):
        worker_spec(
            template, PoolSettings(network="deerflow"), "instance", "token", "/private",
            require_shared_archive=True,
        )


@pytest.mark.asyncio
async def test_live_storage_backend_change_reprobes_and_snapshots_main_runtime(tmp_path, monkeypatch):
    from app.gateway.worker_pool import controller as controller_module

    store = SettingsStore(tmp_path)
    store.save(PoolSettings(
        enabled=True, source_container="main", network="deerflow", desired=1,
    ))
    pool = Controller(store)
    storage_backends = iter(("sqlite", "postgres"))
    discovered = []
    snapshots = []
    template = source()
    template.update({
        "Id": "main-container-id",
        "Image": "sha256:base-image",
        "Config": {"Env": ["AUTH_JWT_SECRET=fixed", "DEER_FLOW_HOME=/data"], "Labels": {}, "User": ""},
        "NetworkSettings": {"Networks": {"deerflow": {}}},
    })

    class Engine:
        def request(self, method, path, body=None):
            assert method == "GET" and path == "/containers/main/json?size=true"
            return template

        def configured_storage_backend(self, source_detail, preferred, working_dir):
            return next(storage_backends)

        def discover_worker_python(self, source_detail, preferred, working_dir):
            discovered.append(working_dir)
            return "/opt/venv/bin/python"

        def snapshot_worker_runtime(self, source_detail, instance):
            image = f"sha256:runtime-{len(snapshots) + 1}"
            snapshots.append((source_detail["Id"], instance))
            return image

    engine = Engine()
    monkeypatch.setattr(controller_module, "DockerEngine", lambda settings: engine)
    try:
        _, sqlite_spec = await pool.inspect_template(store.load())
        _, postgres_spec = await pool.inspect_template(store.load())
        assert sqlite_spec["Image"] == "sha256:runtime-1"
        assert postgres_spec["Image"] == "sha256:runtime-2"
        assert len(discovered) == 2
        assert len(snapshots) == 2
    finally:
        await pool.close()


def test_worker_keeps_service_networks_without_main_aliases_and_uses_private_health():
    template = source()
    template["NetworkSettings"]["Networks"]["knowledge"] = {"Aliases": ["main-only"], "IPAddress": "172.20.0.3"}
    spec = worker_spec(template, PoolSettings(network="deerflow"), "instance", "token", "/private")
    assert spec["NetworkingConfig"]["EndpointsConfig"] == {"deerflow": {}, "knowledge": {}}
    health = spec["Healthcheck"]["Test"]
    assert health[:3] == ["CMD", "/opt/venv/bin/python", "-c"]
    assert "/internal/worker-status" in health[3]
    assert "os.environ['DEER_FLOW_WORKER_TOKEN']" in health[3]


def test_dedicated_private_volume_is_excluded():
    template = source()
    template["Mounts"].append({"Type": "volume", "Name": "secrets", "Destination": "/private", "RW": True})
    spec = worker_spec(template, PoolSettings(network="deerflow"), "instance", "token", "/private")
    assert all(mount["Target"] != "/private" for mount in spec["HostConfig"]["Mounts"])


@pytest.mark.asyncio
async def test_orphaned_run_becomes_terminal_without_retry():
    from app.gateway.worker_pool.manager import WorkerRunManager
    from deerflow.runtime.runs.store.memory import MemoryRunStore

    store = MemoryRunStore()
    await store.put("old", thread_id="thread", status="running")
    manager = WorkerRunManager(store=store)
    assert await manager.get_persisted("old", authorized_thread_id="different-thread") is None
    assert (await store.get("old"))["status"] == "running"
    record = await manager.get_persisted("old", authorized_thread_id="thread")
    assert record.status == "error"
    assert (await store.get("old"))["status"] == "error"
    assert record.task is None


def test_reject_shared_controller_secrets():
    with pytest.raises(DockerError, match="密钥目录"):
        worker_spec(source(), PoolSettings(network="deerflow"), "i", "t", "/data/control")


@pytest.mark.asyncio
async def test_70_simultaneous_runs_are_bounded_and_complete():
    gate = Admission(4, 128, 10)
    maximum = 0
    completed = set()

    async def run(i):
        nonlocal maximum
        await gate.acquire()
        try:
            maximum = max(maximum, gate.active)
            await asyncio.sleep(0.003)
            assert i not in completed
            completed.add(i)
        finally:
            gate.release()

    await asyncio.gather(*(run(i) for i in range(70)))
    assert len(completed) == 70
    assert maximum == 4
    assert gate.active == gate.waiting == 0


@pytest.mark.asyncio
async def test_cancelled_waiter_and_timeout_do_not_leak_capacity():
    gate = Admission(1, 1, 0.03)
    await gate.acquire()
    waiter = asyncio.create_task(gate.acquire())
    await asyncio.sleep(0)
    with pytest.raises(HTTPException) as error:
        await gate.acquire()
    assert error.value.status_code == 429
    waiter.cancel()
    await asyncio.gather(waiter, return_exceptions=True)
    with pytest.raises(HTTPException):
        await gate.acquire()
    gate.release()
    await gate.acquire()
    gate.release()
    assert gate.active == gate.waiting == 0


@pytest.mark.asyncio
async def test_background_run_holds_slot_after_http_returns():
    gate = Admission(1)
    event = asyncio.Event()
    request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(worker_admission=gate)))

    @bounded_worker_run
    async def start(body, thread, request):
        return SimpleNamespace(task=asyncio.create_task(event.wait()))

    record = await start(None, "thread", request)
    assert gate.active == 1
    event.set()
    await record.task
    await asyncio.sleep(0)
    assert gate.active == 0


@pytest.mark.asyncio
async def test_affinity_stable_on_scaleup_fail_closed_and_persisted(tmp_path):
    store = SettingsStore(tmp_path)
    store.save(PoolSettings(enabled=True))
    pool = Controller(store)
    pool.nodes["a"] = {"name": "a", "ready": True, "active": 0, "draining": False}
    try:
        first = await pool.reserve("thread", create=True)
        pool.release("a")
        pool.nodes["b"] = {"name": "b", "ready": True, "active": 0, "draining": False}
        assert (await pool.reserve("thread", create=True))["name"] == first["name"]
        pool.release("a")
        pool.nodes["a"]["ready"] = False
        with pytest.raises(PoolUnavailable):
            await pool.reserve("thread", create=True)
        assert pool._ensure_db().execute("SELECT node FROM affinity WHERE thread='thread'").fetchone()[0] == "a"
    finally:
        await pool.close()
    restarted = Controller(store)
    try:
        with pytest.raises(PoolUnavailable):
            await restarted.reserve("thread", create=True)
    finally:
        await restarted.close()


@pytest.mark.asyncio
async def test_scale_down_waits_for_active_runs_and_requests(tmp_path, monkeypatch):
    store = SettingsStore(tmp_path)
    store.save(PoolSettings(enabled=True, source_container="main", network="deerflow", desired=0))
    pool = Controller(store)
    name = f"dfw-{store.raw()['instance']}-0"
    calls = []

    class Engine:
        def request(self, method, path, body=None):
            calls.append((method, path))
            if path.startswith("/containers/json"):
                return [{"Names": ["/" + name], "Id": "id", "State": "running"}]
            return {}

    async def template(settings):
        return Engine(), {}

    monkeypatch.setattr(pool, "inspect_template", template)
    active = 1

    async def status(request):
        return httpx.Response(200, json={"ready": True, "active": active, "capacity": 4})

    await pool.http.aclose()
    pool.http = httpx.AsyncClient(transport=httpx.MockTransport(status))
    try:
        await pool.reconcile()
        assert not any(method == "POST" for method, _ in calls)
        active = 0
        pool.inflight[name] = 1
        await pool.reconcile()
        assert not any(method == "POST" for method, _ in calls)
        pool.release(name)
        await pool.reconcile()
        assert ("POST", "/containers/id/stop?t=30") in calls
        assert ("DELETE", "/containers/id") in calls
        assert not pool.nodes
    finally:
        await pool.close()


@pytest.mark.asyncio
async def test_changed_runtime_image_drains_and_recreates_old_workers(tmp_path, monkeypatch):
    store = SettingsStore(tmp_path)
    store.save(PoolSettings(
        enabled=True, source_container="main", network="deerflow", desired=1,
    ))
    pool = Controller(store)
    name = f"dfw-{store.raw()['instance']}-0"
    calls = []

    class Engine:
        def request(self, method, path, body=None):
            calls.append((method, path, body))
            if path.startswith("/containers/json"):
                return [{
                    "Names": ["/" + name], "Id": "old-worker", "State": "running",
                    "ImageID": "sha256:old-runtime",
                }]
            if path.startswith("/containers/create?"):
                assert body["Image"] == "sha256:new-runtime"
                return {"Id": "new-worker"}
            return {}

    async def template(_settings):
        return Engine(), {"Image": "sha256:new-runtime", "Entrypoint": ["python"], "WorkingDir": "/app"}

    async def status(_request):
        return httpx.Response(200, json={"ready": True, "active": 0, "capacity": 4})

    monkeypatch.setattr(pool, "inspect_template", template)
    await pool.http.aclose()
    pool.http = httpx.AsyncClient(transport=httpx.MockTransport(status))
    try:
        await pool.reconcile()
        paths = [(method, path) for method, path, _ in calls]
        assert ("POST", "/containers/old-worker/stop?t=30") in paths
        assert ("DELETE", "/containers/old-worker") in paths
        create_index = next(i for i, (_, path) in enumerate(paths) if path.startswith("/containers/create?"))
        assert paths.index(("DELETE", "/containers/old-worker")) < create_index
        assert ("POST", "/containers/new-worker/start") in paths
        assert pool.nodes[name]["id"] == "new-worker"
    finally:
        await pool.close()


@pytest.mark.asyncio
async def test_70_http_clients_distribute_to_four_workers_with_sse(tmp_path, monkeypatch):
    from fastapi import FastAPI

    from app.gateway import deps
    from app.gateway.worker_pool.proxy import WorkerProxy

    store = SettingsStore(tmp_path)
    store.save(PoolSettings(enabled=True, concurrency=32))
    pool = Controller(store)
    for i in range(4):
        name = f"worker-{i}"
        pool.nodes[name] = {"name": name, "ready": True, "active": 0, "capacity": 32, "draining": False, "url": f"http://{name}:8001"}

    async def user(request):
        return SimpleNamespace(id="user")

    async def access(*args, **kwargs):
        return True

    async def metadata(*args, **kwargs):
        return {"metadata": {}}

    monkeypatch.setattr(deps, "get_current_user_from_request", user)
    monkeypatch.setattr(deps, "get_thread_store", lambda request: SimpleNamespace(check_access=access, get=metadata))
    gates = {name: Admission(32) for name in pool.nodes}
    counts = {name: 0 for name in pool.nodes}

    class EventStream(httpx.AsyncByteStream):
        async def __aiter__(self):
            yield b"event: end\ndata: {}\n\n"

    async def worker(request):
        assert request.headers["cookie"] == "session=preserved"
        assert request.headers["host"] == "gateway"
        assert request.headers["x-deerflow-worker-token"] == store.raw()["token"]
        name = request.url.host
        await gates[name].acquire()
        try:
            await asyncio.sleep(0.005)
            counts[name] += 1
            return httpx.Response(200, stream=EventStream(), headers={"content-type": "text/event-stream", "location": "/threads/t/runs/r/stream"})
        finally:
            gates[name].release()

    await pool.http.aclose()
    pool.http = httpx.AsyncClient(transport=httpx.MockTransport(worker))
    app = FastAPI()
    app.state.worker_pool = pool
    app.add_middleware(WorkerProxy)
    try:
        async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://gateway") as client:
            responses = await asyncio.gather(*(client.post(f"/api/threads/t-{i}/runs/stream", headers={"cookie": "session=preserved"}, json={"input": {}}) for i in range(70)))
        assert all(response.status_code == 200 and "event: end" in response.text for response in responses)
        assert all(response.headers["location"] == "/threads/t/runs/r/stream" for response in responses)
        assert sorted(counts.values()) == [17, 17, 18, 18]
        assert all(value == 0 for value in pool.inflight.values())
    finally:
        await pool.close()


@pytest.mark.asyncio
async def test_70_concurrent_runs_fall_back_to_main_without_worker_queue(tmp_path, monkeypatch):
    from fastapi import FastAPI, Request
    from starlette.responses import StreamingResponse

    from app.gateway import deps
    from app.gateway.worker_pool.proxy import WorkerProxy

    store = SettingsStore(tmp_path)
    store.save(PoolSettings(enabled=True, concurrency=4))
    pool = Controller(store)
    pool.nodes["worker-0"] = {
        "name": "worker-0", "ready": True, "active": 0, "capacity": 4,
        "draining": False, "url": "http://worker-0:8001",
    }

    async def user(request):
        return SimpleNamespace(id="user")

    async def access(*args, **kwargs):
        return True

    async def metadata(*args, **kwargs):
        return {"metadata": {}}

    monkeypatch.setattr(deps, "get_current_user_from_request", user)
    monkeypatch.setattr(deps, "get_thread_store", lambda request: SimpleNamespace(check_access=access, get=metadata))
    worker_count = 0
    main_count = 0

    class EventStream(httpx.AsyncByteStream):
        async def __aiter__(self):
            await asyncio.sleep(0)
            yield b"event: end\ndata: {}\n\n"

    async def worker(request):
        nonlocal worker_count
        assert request.headers["x-deerflow-worker-token"] == store.raw()["token"]
        worker_count += 1
        await asyncio.sleep(0.01)
        return httpx.Response(200, stream=EventStream(), headers={"content-type": "text/event-stream"})

    await pool.http.aclose()
    pool.http = httpx.AsyncClient(transport=httpx.MockTransport(worker))
    app = FastAPI()
    app.state.worker_pool = pool
    app.state.checkpointer = object()
    app.state.store = object()
    app.add_middleware(WorkerProxy)

    @app.post("/api/threads/{thread_id}/runs/stream")
    async def main_run(thread_id: str, request: Request):
        nonlocal main_count
        assert request.state.worker_pool_fallback is True
        main_count += 1
        return StreamingResponse(iter([b"event: end\ndata: {}\n\n"]), media_type="text/event-stream")

    try:
        async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://gateway") as client:
            responses = await asyncio.gather(*(
                client.post(f"/api/threads/t-{i}/runs/stream", json={"input": {}})
                for i in range(70)
            ))
        assert all(response.status_code == 200 and "event: end" in response.text for response in responses)
        assert worker_count == 4
        assert main_count == 66
        assert pool._ensure_db().execute("SELECT count(*) FROM affinity").fetchone()[0] == 4
        assert all(value == 0 for value in pool.inflight.values())
    finally:
        await pool.close()


@pytest.mark.asyncio
@pytest.mark.parametrize("worker_thread_active,status_code", [(False, 200), (True, 409)])
async def test_affined_thread_uses_main_fallback_only_when_worker_thread_is_idle(tmp_path, monkeypatch, worker_thread_active, status_code):
    from fastapi import FastAPI, Request

    from app.gateway import deps
    from app.gateway.worker_pool.proxy import WorkerProxy

    store = SettingsStore(tmp_path)
    store.save(PoolSettings(enabled=True, concurrency=1))
    pool = Controller(store)
    pool.nodes["worker-0"] = {
        "name": "worker-0", "ready": True, "active": 1, "capacity": 1,
        "draining": False, "url": "http://worker-0:8001",
    }
    pool._ensure_db().execute("INSERT INTO affinity VALUES (?, ?)", ("thread-1", "worker-0"))
    pool._ensure_db().commit()

    async def user(request):
        return SimpleNamespace(id="user")

    async def access(*args, **kwargs):
        return True

    async def metadata(*args, **kwargs):
        return {"metadata": {}}

    monkeypatch.setattr(deps, "get_current_user_from_request", user)
    monkeypatch.setattr(deps, "get_thread_store", lambda request: SimpleNamespace(check_access=access, get=metadata))

    checked = []

    async def worker(request):
        checked.append(request.url.path)
        return httpx.Response(200, json={"active": worker_thread_active})

    await pool.http.aclose()
    pool.http = httpx.AsyncClient(transport=httpx.MockTransport(worker))
    app = FastAPI()
    app.state.worker_pool = pool
    app.state.checkpointer = object()
    app.state.store = object()
    app.add_middleware(WorkerProxy)

    @app.post("/api/threads/thread-1/runs/stream")
    async def main_run(request: Request):
        assert request.state.worker_pool_fallback is True
        return {"handled_by": "main"}

    try:
        async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://gateway") as client:
            result = await client.post("/api/threads/thread-1/runs/stream", json={"input": {}})
        assert result.status_code == status_code
        assert checked == ["/internal/worker-thread-active/thread-1"]
        if status_code == 200:
            assert result.json() == {"handled_by": "main"}
    finally:
        await pool.close()


def test_unified_pg_and_legacy_override(monkeypatch):
    from app.gateway.worker_pool.preflight import shared_postgres_configured

    cfg = SimpleNamespace(database=SimpleNamespace(backend="postgres"), checkpointer=None)
    assert shared_postgres_configured(cfg)
    cfg.checkpointer = SimpleNamespace(type="sqlite", use_postgres=True)
    assert not shared_postgres_configured(cfg)
    cfg.checkpointer.type = "postgres"
    assert shared_postgres_configured(cfg)
    monkeypatch.setenv("DEER_FLOW_USE_PG_CHECKPOINTER", "false")
    assert not shared_postgres_configured(cfg)


def test_mysql_business_supports_both_chat_backends(monkeypatch, tmp_path):
    from app.gateway.worker_pool.preflight import shared_storage_configured, validate_sqlite_mount

    cfg = SimpleNamespace(database=SimpleNamespace(backend="mysql"), checkpointer=SimpleNamespace(type="postgres", use_postgres=True))
    monkeypatch.delenv("DEER_FLOW_USE_PG_CHECKPOINTER", raising=False)
    assert shared_storage_configured(cfg)
    cfg.checkpointer = SimpleNamespace(type="sqlite", connection_string=str(tmp_path / "chat.db"))
    assert shared_storage_configured(cfg)
    mountinfo = f"31 1 8:1 / {tmp_path.as_posix()} rw - ext4 /dev/sda rw"
    validate_sqlite_mount(cfg, mountinfo)
    with pytest.raises(RuntimeError, match="network filesystem"):
        validate_sqlite_mount(cfg, mountinfo.replace("ext4", "nfs4"))
    with pytest.raises(RuntimeError, match="persistently mounted"):
        validate_sqlite_mount(cfg, "31 1 8:1 / / rw - ext4 /dev/sda rw")
    cfg.checkpointer.connection_string = ":memory:"
    assert not shared_storage_configured(cfg)


@pytest.mark.asyncio
async def test_delayed_health_probe_does_not_hide_existing_run(tmp_path):
    from app.gateway.worker_pool.controller import Controller, PoolUnavailable
    from app.gateway.worker_pool.settings import SettingsStore

    pool = Controller(SettingsStore(tmp_path))
    pool.nodes['worker'] = dict(name='worker', ready=False, ever_ready=True, draining=False, url='http://worker:8001')
    pool._ensure_db().execute("INSERT INTO affinity VALUES (?, ?)", ('thread', 'worker'))
    pool._ensure_db().commit()
    try:
        assert (await pool.reserve('thread', create=False))['name'] == 'worker'
        pool.release('worker')
        with pytest.raises(PoolUnavailable):
            await pool.reserve('thread', create=True)
        pool.nodes['worker']['ever_ready'] = False
        with pytest.raises(PoolUnavailable):
            await pool.reserve('thread', create=False)
    finally:
        await pool.close()


@pytest.mark.asyncio
async def test_reconcile_retains_previous_readiness_during_probe_timeout(tmp_path, monkeypatch):
    import httpx

    store = SettingsStore(tmp_path)
    store.save(PoolSettings(enabled=True, desired=1, source_container='main', network='network'))
    pool = Controller(store)
    name = f"dfw-{store.raw()['instance']}-0"
    pool.nodes[name] = dict(name=name, ready=True, ever_ready=True, active=1, draining=False)
    pool._ensure_db().execute('INSERT INTO affinity VALUES (?, ?)', ('thread', name))
    pool._ensure_db().commit()

    class Engine:
        def request(self, method, path, body=None):
            assert path.startswith('/containers/json')
            return [{'Names': ['/' + name], 'Id': 'id', 'State': 'running'}]

    async def inspect(_settings):
        return Engine(), {}

    async def timeout(*args, **kwargs):
        raise httpx.ReadTimeout('busy')

    monkeypatch.setattr(pool, 'inspect_template', inspect)
    monkeypatch.setattr(pool.http, 'get', timeout)
    try:
        await pool.reconcile()
        await pool.reconcile()
        assert (await pool.reserve('thread', create=False))['name'] == name
        with pytest.raises(PoolUnavailable):
            await pool.reserve('thread', create=True)
    finally:
        await pool.close()


@pytest.mark.asyncio
async def test_read_transport_retries_but_never_replays_model_submission():
    import httpx

    from app.gateway.worker_pool.proxy import send_upstream

    attempts = []

    async def transport(request):
        attempts.append(request.method)
        if request.method == 'POST' or len(attempts) == 1:
            raise httpx.RemoteProtocolError('server closed keep-alive')
        return httpx.Response(200, json={'status': 'success'})

    async with httpx.AsyncClient(transport=httpx.MockTransport(transport)) as client:
        response = await send_upstream(client, 'GET', 'http://worker/run', [], None)
        assert response.status_code == 200
        await response.aclose()
        with pytest.raises(httpx.RemoteProtocolError):
            await send_upstream(client, 'POST', 'http://worker/run', [], b'{}')
    assert attempts == ['GET', 'GET', 'POST']


@pytest.mark.asyncio
async def test_second_gateway_cannot_start_schedulers(tmp_path, monkeypatch):
    from app.gateway.worker_pool.lifecycle import managed_controller
    from app.gateway.worker_pool.settings import PoolSettings, SettingsStore

    monkeypatch.setenv("DEER_FLOW_WORKER_CONTROL_HOME", str(tmp_path))
    async with managed_controller(SimpleNamespace(state=SimpleNamespace())):
        async with managed_controller(SimpleNamespace(state=SimpleNamespace())):
            pass
    assert not (tmp_path / "affinity.sqlite3").exists()
    assert not (tmp_path / "diagnostics.sqlite3").exists()
    assert not (tmp_path / "controller.lock").exists()
    assert not (tmp_path / "settings.json").exists()
    assert not (tmp_path / "secret.key").exists()

    store = SettingsStore(tmp_path)
    store.save(PoolSettings(enabled=True, source_container="main"))
    async with managed_controller(SimpleNamespace(state=SimpleNamespace())) as first:
        assert first is not None
        first.start()
        with pytest.raises(RuntimeError, match="只能由一个主服务进程"):
            async with managed_controller(SimpleNamespace(state=SimpleNamespace())) as second:
                second.start()


@pytest.mark.asyncio
async def test_disabled_pool_ignores_saved_docker_settings(tmp_path, monkeypatch):
    store = SettingsStore(tmp_path)
    store.save(PoolSettings(enabled=False, transport="socket", source_container="main", network="deerflow"))
    pool = Controller(store)

    def unexpected_request(*args, **kwargs):
        pytest.fail("disabled Worker pool must not contact Docker")

    monkeypatch.setattr("app.gateway.worker_pool.controller.DockerEngine.request", unexpected_request)
    pool.start()
    await asyncio.sleep(0)
    assert pool.task is None
    assert pool.db is None
    assert not (tmp_path / "controller.lock").exists()
    assert not (tmp_path / "affinity.sqlite3").exists()
    await pool.close()


@pytest.mark.asyncio
async def test_non_admin_cannot_read_host_credentials(monkeypatch):
    from app.gateway.worker_pool import router

    async def user(request):
        return SimpleNamespace(system_role="user")

    monkeypatch.setattr(router, "get_current_user_from_request", user)
    with pytest.raises(HTTPException) as error:
        await router.controller(SimpleNamespace())
    assert error.value.status_code == 403


@pytest.mark.asyncio
async def test_disabled_proxy_and_reserved_search_are_unchanged(tmp_path):
    from fastapi import FastAPI

    from app.gateway.worker_pool.proxy import WorkerProxy

    pool = Controller(SettingsStore(tmp_path))
    app = FastAPI()
    app.state.worker_pool = pool
    app.add_middleware(WorkerProxy)

    @app.post("/api/threads/search")
    @app.post("/api/threads/t/runs/stream")
    async def endpoint():
        return {"local": True}

    try:
        async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://test") as client:
            assert (await client.post("/api/threads/search")).json() == {"local": True}
            assert (await client.post("/api/threads/t/runs/stream")).json() == {"local": True}
    finally:
        await pool.close()


@pytest.mark.asyncio
@pytest.mark.parametrize("allowed,system,status_code", [(False, False, 404), (True, True, 200)])
async def test_private_and_scheduler_threads_are_not_assigned(tmp_path, monkeypatch, allowed, system, status_code):
    from fastapi import FastAPI

    from app.gateway import deps
    from app.gateway.worker_pool.proxy import WorkerProxy

    store = SettingsStore(tmp_path)
    store.save(PoolSettings(enabled=True))
    pool = Controller(store)

    async def user(request):
        return SimpleNamespace(id="user")

    async def access(*args, **kwargs):
        return allowed

    async def metadata(*args, **kwargs):
        return {"metadata": {"thread_type": "scheduler"} if system else {}}

    monkeypatch.setattr(deps, "get_current_user_from_request", user)
    monkeypatch.setattr(deps, "get_thread_store", lambda request: SimpleNamespace(check_access=access, get=metadata))
    app = FastAPI()
    app.state.worker_pool = pool
    app.add_middleware(WorkerProxy)

    @app.post("/api/threads/t/runs/stream")
    async def endpoint():
        return {"scheduler_stays_local": True}

    try:
        async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://gateway") as client:
            result = await client.post("/api/threads/t/runs/stream")
        assert result.status_code == status_code
        assert pool._ensure_db().execute("SELECT count(*) FROM affinity").fetchone()[0] == 0
    finally:
        await pool.close()


@pytest.mark.asyncio
async def test_worker_boot_disables_background_and_refuses_memory_fallback(monkeypatch):
    from contextlib import asynccontextmanager

    from langgraph.checkpoint.memory import InMemorySaver

    from app.gateway.worker_pool import worker
    from deerflow.config.app_config import AppConfig
    from deerflow.config.database_config import DatabaseConfig

    cfg = AppConfig.model_construct(
        database=DatabaseConfig.model_construct(backend="postgres"),
        checkpointer=None,
        sandbox=SimpleNamespace(use="LocalSandboxProvider"),
    )
    monkeypatch.setattr(worker, "get_app_config", lambda: cfg)

    async def configure(config):
        pass

    monkeypatch.setattr(worker, "configure_runtime_executor", configure)

    @asynccontextmanager
    async def runtime(app, *, background):
        assert background is False
        app.state.checkpointer = InMemorySaver()
        app.state.store = object()
        yield

    monkeypatch.setattr(worker, "langgraph_runtime", runtime)
    app = SimpleNamespace(state=SimpleNamespace())
    with pytest.raises(RuntimeError, match="fallback"):
        async with worker.worker_lifespan(app):
            pytest.fail("memory worker became ready")


@pytest.mark.asyncio
async def test_skill_disable_invalidates_other_worker_cache_immediately(monkeypatch):
    from app.gateway.worker_pool.refresh import SharedSkillRefresh
    from deerflow.agents.lead_agent import prompt
    from deerflow.config import extensions_config

    refresh = SharedSkillRefresh()
    stamp = 1
    calls = []
    monkeypatch.setattr(refresh, "extension_stamp", lambda: stamp)
    monkeypatch.setattr(refresh, "skill_stamp", lambda: ())
    monkeypatch.setattr(extensions_config, "reload_extensions_config", lambda: calls.append("reload"))

    async def invalidate():
        calls.append("refresh")

    monkeypatch.setattr(prompt, "refresh_skills_system_prompt_cache_async", invalidate)
    await refresh.check()
    await refresh.check()
    assert calls == ["reload", "refresh"]
    stamp = 2
    await refresh.check()
    assert calls == ["reload", "refresh", "reload", "refresh"]


@pytest.mark.asyncio
async def test_scaleup_partial_failure_remains_owned(tmp_path, monkeypatch):
    store = SettingsStore(tmp_path)
    store.save(PoolSettings(enabled=True, source_container="main", network="deerflow", desired=1))
    pool = Controller(store)

    class Engine:
        def request(self, method, path, body=None):
            if path.startswith("/containers/json"):
                return []
            if path.startswith("/containers/create?"):
                return {"Id": "created-before-failure"}
            raise DockerError("start failed")

    async def template(settings):
        return Engine(), {}

    monkeypatch.setattr(pool, "inspect_template", template)
    try:
        with pytest.raises(DockerError):
            await pool.reconcile()
        assert len(pool.nodes) == 1
        assert next(iter(pool.nodes.values()))["id"] == "created-before-failure"
        assert not next(iter(pool.nodes.values()))["ready"]
    finally:
        await pool.close()


@pytest.mark.asyncio
async def test_local_preparation_cannot_race_remote_assignment(tmp_path):
    store = SettingsStore(tmp_path)
    store.save(PoolSettings(enabled=True))
    pool = Controller(store)
    request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(worker_pool=pool)))
    entered, finish = asyncio.Event(), asyncio.Event()

    @bounded_worker_run
    async def start(body, thread_id, request):
        entered.set()
        await finish.wait()
        return SimpleNamespace(task=None)

    task = asyncio.create_task(start(None, "thread", request))
    try:
        await entered.wait()
        with pytest.raises(PoolUnavailable, match="主服务执行"):
            await pool.reserve("thread", create=True)
        finish.set()
        await task
        assert not pool.local_threads
    finally:
        task.cancel()
        await asyncio.gather(task, return_exceptions=True)
        await pool.close()