File size: 31,856 Bytes
80d8c84
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
"""Server endpoint tests.

API 02 adds POST /reset endpoint tests.
API 04 adds a smoke test for GET /scenarios.
API 13 adds CORS middleware verification tests.
API 03 adds POST /step endpoint tests.
API 06 adds WebSocket session handler tests.
API 07 adds idle-timeout and graceful disconnect cleanup tests.
"""

from __future__ import annotations

import json
import time
from unittest.mock import patch

import pytest
from fastapi.testclient import TestClient
from starlette.websockets import WebSocketDisconnect

from replicalab.models import ScientistAction
from server.app import app

_EXPECTED_FAMILIES = {"math_reasoning", "ml_benchmark", "finance_trading"}
_EXPECTED_DIFFICULTIES = ["easy", "medium", "hard"]


@pytest.fixture()
def client():
    return TestClient(app)


class TestHealthEndpoint:
    """GET /health β€” API 01."""

    def test_health_returns_200(self, client: TestClient) -> None:
        resp = client.get("/health")
        assert resp.status_code == 200

    def test_health_payload_has_stable_keys(self, client: TestClient) -> None:
        data = client.get("/health").json()
        assert data["status"] == "ok"
        assert data["env"] in ("real", "stub")
        assert "version" in data

    def test_health_version_matches_app(self, client: TestClient) -> None:
        from server.app import app as _app

        data = client.get("/health").json()
        assert data["version"] == _app.version

    def test_health_is_deterministic(self, client: TestClient) -> None:
        r1 = client.get("/health").json()
        r2 = client.get("/health").json()
        assert r1 == r2


class TestRuntimeEndpoint:
    """Local runtime metadata for model-backed Scientist stepping."""

    def test_runtime_defaults_to_baseline_without_api_key(
        self, client: TestClient, monkeypatch
    ) -> None:
        monkeypatch.delenv("REPLICALAB_SCIENTIST_RUNTIME", raising=False)
        monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)

        resp = client.get("/runtime")

        assert resp.status_code == 200
        data = resp.json()
        assert data["scientist_runtime"] == "baseline"
        assert data["scientist_model"] == "baseline-heuristic"
        assert data["agent_step_available"] is True

    def test_runtime_reports_anthropic_when_enabled(
        self, client: TestClient, monkeypatch
    ) -> None:
        monkeypatch.setenv("REPLICALAB_SCIENTIST_RUNTIME", "anthropic")
        monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key")

        resp = client.get("/runtime")

        assert resp.status_code == 200
        data = resp.json()
        assert data["scientist_runtime"] == "anthropic"
        assert data["scientist_ready"] is True
        assert data["agent_step_available"] is True


class TestLogConfig:
    """OBS 02 β€” log level configurability."""

    def test_default_log_level_is_info(self) -> None:
        from replicalab.config import LOG_LEVEL

        # Default when REPLICALAB_LOG_LEVEL is not set
        assert LOG_LEVEL in ("INFO", "DEBUG", "WARNING", "ERROR")

    def test_log_level_env_var_is_respected(self, monkeypatch) -> None:
        """REPLICALAB_LOG_LEVEL env var controls the log level."""
        import importlib

        import replicalab.config as config_mod

        monkeypatch.setenv("REPLICALAB_LOG_LEVEL", "debug")
        importlib.reload(config_mod)

        assert config_mod.LOG_LEVEL == "DEBUG"

        # Restore
        monkeypatch.delenv("REPLICALAB_LOG_LEVEL", raising=False)
        importlib.reload(config_mod)

    def test_log_format_is_readable(self) -> None:
        from replicalab.config import LOG_FORMAT

        assert "%(asctime)s" in LOG_FORMAT
        assert "%(levelname)s" in LOG_FORMAT
        assert "%(name)s" in LOG_FORMAT


class TestRootEndpoint:
    """GET / β€” lightweight landing page for hosted backend deployments."""

    def test_root_returns_200_html(self, client: TestClient) -> None:
        resp = client.get("/")
        assert resp.status_code == 200
        assert "text/html" in resp.headers["content-type"]

    def test_root_mentions_core_api_endpoints(self, client: TestClient) -> None:
        body = client.get("/").text
        # When frontend/dist exists, root serves the SPA; otherwise the API landing
        assert "ReplicaLab" in body
        if "ReplicaLab API" in body:
            assert "GET /health" in body
            assert "POST /reset" in body


class TestWebFallback:
    """GET /web β€” API 19: OpenEnv fallback UI."""

    def test_web_returns_200_html(self, client: TestClient) -> None:
        resp = client.get("/web")
        assert resp.status_code == 200
        assert "text/html" in resp.headers["content-type"]

    def test_web_contains_interactive_controls(self, client: TestClient) -> None:
        body = client.get("/web").text
        assert "ReplicaLab" in body
        assert "btnReset" in body
        assert "btnPropose" in body
        assert "btnAccept" in body
        assert "/reset" in body
        assert "/step" in body

    def test_web_is_self_contained(self, client: TestClient) -> None:
        """Fallback UI must work without external JS/CSS dependencies."""
        body = client.get("/web").text
        assert "<script>" in body
        assert "<style>" in body


class TestScenariosEndpoint:
    """GET /scenarios β€” API 04."""

    def test_returns_200(self, client: TestClient):
        resp = client.get("/scenarios")
        assert resp.status_code == 200

    def test_response_has_scenarios_key(self, client: TestClient):
        data = client.get("/scenarios").json()
        assert "scenarios" in data
        assert isinstance(data["scenarios"], list)

    def test_all_families_present(self, client: TestClient):
        data = client.get("/scenarios").json()
        families = {s["family"] for s in data["scenarios"]}
        assert families == _EXPECTED_FAMILIES

    def test_each_family_has_difficulties(self, client: TestClient):
        data = client.get("/scenarios").json()
        for entry in data["scenarios"]:
            assert entry["difficulties"] == _EXPECTED_DIFFICULTIES

    def test_no_extra_keys(self, client: TestClient):
        data = client.get("/scenarios").json()
        for entry in data["scenarios"]:
            assert set(entry.keys()) == {"family", "difficulties"}


# ---------------------------------------------------------------------------
# POST /reset β€” API 02
# ---------------------------------------------------------------------------


class TestCorsConfiguration:
    """API 13: CORS middleware for local frontend and HF Spaces."""

    def test_preflight_allows_localhost_vite_origin(self, client: TestClient) -> None:
        resp = client.options(
            "/reset",
            headers={
                "Origin": "http://localhost:5173",
                "Access-Control-Request-Method": "POST",
            },
        )

        assert resp.status_code == 200
        assert resp.headers["access-control-allow-origin"] == "http://localhost:5173"
        assert resp.headers["access-control-allow-credentials"] == "true"

    def test_preflight_allows_hf_space_origin(self, client: TestClient) -> None:
        origin = "https://replicalab-demo.hf.space"
        resp = client.options(
            "/health",
            headers={
                "Origin": origin,
                "Access-Control-Request-Method": "GET",
            },
        )

        assert resp.status_code == 200
        assert resp.headers["access-control-allow-origin"] == origin
        assert resp.headers["access-control-allow-credentials"] == "true"

    def test_preflight_rejects_unconfigured_origin(self, client: TestClient) -> None:
        resp = client.options(
            "/reset",
            headers={
                "Origin": "https://evil.example.com",
                "Access-Control-Request-Method": "POST",
            },
        )

        assert resp.status_code == 400
        assert "access-control-allow-origin" not in resp.headers


class TestResetEndpoint:
    """POST /reset β€” API 02."""

    def test_reset_returns_200_with_expected_keys(self, client: TestClient) -> None:
        resp = client.post("/reset", json={"seed": 1})
        assert resp.status_code == 200
        data = resp.json()
        assert "session_id" in data
        assert "episode_id" in data
        assert "observation" in data

    def test_reset_observation_has_both_roles(self, client: TestClient) -> None:
        data = client.post("/reset", json={"seed": 1}).json()
        obs = data["observation"]
        assert "scientist" in obs
        assert "lab_manager" in obs
        assert obs["scientist"]["paper_title"]
        assert obs["lab_manager"]["budget_total"] > 0

    def test_reset_with_explicit_session_id_reuses_slot(
        self, client: TestClient
    ) -> None:
        """Passing session_id reuses the same slot and returns the same id."""
        sid = "my-fixed-session"
        d1 = client.post("/reset", json={"seed": 1, "session_id": sid}).json()
        assert d1["session_id"] == sid

        d2 = client.post("/reset", json={"seed": 2, "session_id": sid}).json()
        assert d2["session_id"] == sid
        # New episode each time
        assert d2["episode_id"] != d1["episode_id"]

    def test_reset_reuse_closes_prior_env(self, client: TestClient) -> None:
        """Resetting with the same session_id produces a fresh episode."""
        sid = "reuse-session"
        d1 = client.post("/reset", json={"seed": 10, "session_id": sid}).json()
        ep1 = d1["episode_id"]

        d2 = client.post("/reset", json={"seed": 20, "session_id": sid}).json()
        ep2 = d2["episode_id"]

        assert ep1 != ep2

    def test_reset_default_params(self, client: TestClient) -> None:
        """Omitting scenario and difficulty uses defaults without error."""
        resp = client.post("/reset", json={"seed": 0})
        assert resp.status_code == 200
        data = resp.json()
        assert data["observation"]["scientist"]["paper_title"]

    def test_reset_custom_scenario_and_difficulty(self, client: TestClient) -> None:
        for family in ("math_reasoning", "ml_benchmark", "finance_trading"):
            for diff in ("easy", "medium", "hard"):
                resp = client.post(
                    "/reset",
                    json={"seed": 42, "scenario": family, "difficulty": diff},
                )
                assert resp.status_code == 200, f"Failed for {family}/{diff}"
                obs = resp.json()["observation"]
                assert obs["scientist"]["paper_title"]
                assert obs["lab_manager"]["budget_total"] > 0

    def test_reset_deterministic_with_same_seed(self, client: TestClient) -> None:
        """Same seed + scenario + difficulty β†’ identical observations."""
        params = {"seed": 99, "scenario": "math_reasoning", "difficulty": "medium"}
        d1 = client.post("/reset", json=params).json()
        d2 = client.post("/reset", json=params).json()

        assert d1["observation"] == d2["observation"]
        # Episode ids differ (new UUID each time)
        assert d1["episode_id"] != d2["episode_id"]


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _reset(client: TestClient, **kwargs) -> dict:
    """Reset and return the response JSON."""
    payload = {"seed": 42, "scenario": "math_reasoning", "difficulty": "easy"}
    payload.update(kwargs)
    resp = client.post("/reset", json=payload)
    assert resp.status_code == 200
    return resp.json()


def _good_action_payload(client: TestClient) -> dict:
    """Build a valid propose_protocol action payload from a fresh scenario."""
    from replicalab.scenarios import generate_scenario

    scenario = generate_scenario(seed=42, template="math_reasoning", difficulty="easy")
    lab = scenario.lab_manager_observation
    spec = scenario.hidden_reference_spec
    return {
        "action_type": "propose_protocol",
        "sample_size": 10,
        "controls": ["baseline", "ablation"],
        "technique": spec.summary[:60] if spec.summary else "replication_plan",
        "duration_days": max(1, min(2, lab.time_limit_days)),
        "required_equipment": (
            list(lab.equipment_available[:1]) if lab.equipment_available else []
        ),
        "required_reagents": (
            list(lab.reagents_in_stock[:1]) if lab.reagents_in_stock else []
        ),
        "questions": [],
        "rationale": (
            f"Plan addresses: {', '.join(spec.required_elements[:2])}. "
            f"Target metric: {spec.target_metric}. "
            f"Target value: {spec.target_value}. "
            "Stay within budget and schedule."
        ),
    }


def _accept_action_payload() -> dict:
    return {
        "action_type": "accept",
        "sample_size": 0,
        "controls": [],
        "technique": "",
        "duration_days": 0,
        "required_equipment": [],
        "required_reagents": [],
        "questions": [],
        "rationale": "",
    }


# ---------------------------------------------------------------------------
# POST /step β€” API 03
# ---------------------------------------------------------------------------


class TestStepEndpoint:
    """POST /step β€” API 03."""

    def test_reset_then_step_happy_path(self, client: TestClient) -> None:
        """Reset, then step with a valid action β†’ 200 with StepResult."""
        reset_data = _reset(client)
        session_id = reset_data["session_id"]

        action = _good_action_payload(client)
        resp = client.post("/step", json={"session_id": session_id, "action": action})

        assert resp.status_code == 200
        data = resp.json()
        assert "observation" in data
        assert "reward" in data
        assert "done" in data
        assert "info" in data
        assert data["done"] is False
        assert data["info"]["error"] is None

    def test_step_invalid_session_returns_404(self, client: TestClient) -> None:
        """Step with a non-existent session_id β†’ 404."""
        action = _good_action_payload(client)
        resp = client.post(
            "/step",
            json={"session_id": "nonexistent-session-id", "action": action},
        )

        assert resp.status_code == 404
        assert "Session not found" in resp.json()["detail"]

    def test_terminal_step_returns_real_reward_breakdown(
        self, client: TestClient
    ) -> None:
        """Propose β†’ accept: terminal step has real reward_breakdown,
        judge_notes, and verdict from the env (not stubs)."""
        reset_data = _reset(client)
        session_id = reset_data["session_id"]

        # Step 1: propose
        action = _good_action_payload(client)
        resp1 = client.post("/step", json={"session_id": session_id, "action": action})
        assert resp1.status_code == 200
        assert resp1.json()["done"] is False

        # Step 2: accept
        resp2 = client.post(
            "/step",
            json={"session_id": session_id, "action": _accept_action_payload()},
        )
        assert resp2.status_code == 200
        data = resp2.json()

        assert data["done"] is True
        assert data["reward"] > 0.0

        info = data["info"]
        assert info["agreement_reached"] is True
        assert info["verdict"] == "accept"
        assert info["judge_notes"] is not None
        assert "rigor" in info["judge_notes"]

        rb = info["reward_breakdown"]
        assert rb is not None
        assert 0.0 <= rb["rigor"] <= 1.0
        assert 0.0 <= rb["feasibility"] <= 1.0
        assert 0.0 <= rb["fidelity"] <= 1.0
        # Verify it's not the old stub 0.8
        assert not (rb["rigor"] == 0.8 and rb["feasibility"] == 0.8 and rb["fidelity"] == 0.8)

    def test_semantic_invalid_action_returns_200_with_error(
        self, client: TestClient
    ) -> None:
        """A semantically invalid action (e.g. duration=999) returns 200
        with info.error set, not a crash or 422."""
        reset_data = _reset(client)
        session_id = reset_data["session_id"]

        bad_action = {
            "action_type": "propose_protocol",
            "sample_size": 5,
            "controls": ["baseline"],
            "technique": "some technique",
            "duration_days": 999,
            "required_equipment": [],
            "required_reagents": [],
            "questions": [],
            "rationale": "Duration is impossibly long for the lab time limit.",
        }
        resp = client.post(
            "/step", json={"session_id": session_id, "action": bad_action}
        )

        assert resp.status_code == 200
        data = resp.json()
        assert data["done"] is False
        assert data["info"]["error"] is not None
        assert "Validation errors" in data["info"]["error"]

    def test_replay_uses_real_judge_data(self, client: TestClient) -> None:
        """After a terminal step, GET /replay/{episode_id} returns
        real judge_notes and verdict, not stub values."""
        reset_data = _reset(client)
        session_id = reset_data["session_id"]
        episode_id = reset_data["episode_id"]

        # Propose then accept
        action = _good_action_payload(client)
        client.post("/step", json={"session_id": session_id, "action": action})
        client.post(
            "/step",
            json={"session_id": session_id, "action": _accept_action_payload()},
        )

        # Fetch replay
        resp = client.get(f"/replay/{episode_id}")
        assert resp.status_code == 200
        replay = resp.json()

        assert replay["agreement_reached"] is True
        assert "rigor" in replay["judge_notes"]
        assert replay["verdict"] == "accept"
        assert replay["reward_breakdown"] is not None
        assert replay["total_reward"] > 0.0
        # Not the old stub string
        assert "Stub audit" not in replay["judge_notes"]

    def test_replay_includes_top_failure_reasons(self, client: TestClient) -> None:
        """Terminal replay records persist the canonical audit reasons."""
        reset_data = _reset(client)
        session_id = reset_data["session_id"]
        episode_id = reset_data["episode_id"]

        # Force a timeout path so the audit builder emits failure reasons.
        for _ in range(6):
            resp = client.post(
                "/step",
                json={"session_id": session_id, "action": _good_action_payload(client)},
            )
            assert resp.status_code == 200
            if resp.json()["done"]:
                break

        replay = client.get(f"/replay/{episode_id}").json()
        assert replay["verdict"] == "timeout"
        assert isinstance(replay["top_failure_reasons"], list)
        assert replay["top_failure_reasons"]
        assert any(
            "round limit" in reason.lower() or "without agreement" in reason.lower()
            for reason in replay["top_failure_reasons"]
        )


class TestAgentStepEndpoint:
    """POST /agent-step uses the configured Scientist runtime."""

    def test_agent_step_runs_runtime_policy(
        self, client: TestClient, monkeypatch
    ) -> None:
        import server.app as server_app

        monkeypatch.setenv("REPLICALAB_SCIENTIST_RUNTIME", "anthropic")
        monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key")
        server_app._SCIENTIST_POLICY_CACHE.clear()

        action = ScientistAction.model_validate(_good_action_payload(client))

        def fake_policy(observation, **kwargs):
            assert observation.paper_title
            assert kwargs["scenario"] == "ml_benchmark"
            assert kwargs["difficulty"] == "medium"
            return action

        monkeypatch.setattr(server_app, "_get_scientist_policy", lambda: fake_policy)

        reset_data = _reset(client, scenario="ml_benchmark", difficulty="medium")
        resp = client.post("/agent-step", json={"session_id": reset_data["session_id"]})

        assert resp.status_code == 200
        data = resp.json()
        assert data["info"]["scientist_runtime"] == "anthropic"
        assert data["info"]["scientist_model"]
        assert data["info"]["scientist_action"]["action_type"] == "propose_protocol"
        assert data["observation"]["scientist"]["round_number"] == 1

    def test_agent_step_invalid_session_returns_404(self, client: TestClient) -> None:
        resp = client.post("/agent-step", json={"session_id": "missing"})

        assert resp.status_code == 404
        assert "Session not found" in resp.json()["detail"]

    def test_agent_step_falls_back_to_baseline_on_runtime_failure(
        self, client: TestClient, monkeypatch
    ) -> None:
        import server.app as server_app

        monkeypatch.setenv("REPLICALAB_SCIENTIST_RUNTIME", "ollama")
        monkeypatch.setattr(
            server_app,
            "_resolve_scientist_action",
            lambda session: (_ for _ in ()).throw(RuntimeError("model timeout")),
        )

        reset_data = _reset(client, scenario="math_reasoning", difficulty="easy")
        resp = client.post("/agent-step", json={"session_id": reset_data["session_id"]})

        assert resp.status_code == 200
        data = resp.json()
        assert data["info"]["scientist_runtime"] == "ollama_fallback"
        assert data["info"]["scientist_error"] == "model timeout"


# ---------------------------------------------------------------------------
# WebSocket handler β€” API 06
# ---------------------------------------------------------------------------


def _ws_send_recv(ws, msg: dict) -> dict:
    """Send a JSON message over the WebSocket and return the parsed response."""
    ws.send_text(json.dumps(msg))
    return json.loads(ws.receive_text())


class TestWebSocket:
    """API 06: WebSocket session handler with isolated env per connection."""

    # -- basic connectivity --------------------------------------------------

    def test_ws_ping_pong(self, client: TestClient) -> None:
        with client.websocket_connect("/ws") as ws:
            resp = _ws_send_recv(ws, {"type": "ping"})
            assert resp["type"] == "pong"

    def test_ws_reset_returns_observation(self, client: TestClient) -> None:
        with client.websocket_connect("/ws") as ws:
            resp = _ws_send_recv(ws, {
                "type": "reset", "seed": 42,
                "scenario": "math_reasoning", "difficulty": "easy",
            })
            assert resp["type"] == "reset_ok"
            assert resp["episode_id"]
            obs = resp["observation"]
            assert obs["scientist"]["paper_title"]
            assert obs["lab_manager"]["budget_total"] > 0

    def test_ws_step_returns_result(self, client: TestClient) -> None:
        action = _good_action_payload(client)
        with client.websocket_connect("/ws") as ws:
            _ws_send_recv(ws, {"type": "reset", "seed": 42})
            resp = _ws_send_recv(ws, {"type": "step", "action": action})

            assert resp["type"] == "step_ok"
            assert resp["done"] is False
            assert resp["reward"] > 0.0
            assert resp["info"]["step_reward_components"]["protocol_delta_bonus"] > 0.0
            assert resp["observation"] is not None

    def test_ws_full_episode_real_reward(self, client: TestClient) -> None:
        """Propose β†’ accept returns real reward breakdown, not stub 0.8."""
        action = _good_action_payload(client)
        with client.websocket_connect("/ws") as ws:
            _ws_send_recv(ws, {"type": "reset", "seed": 42})
            _ws_send_recv(ws, {"type": "step", "action": action})
            resp = _ws_send_recv(ws, {"type": "step", "action": _accept_action_payload()})

            assert resp["type"] == "step_ok"
            assert resp["done"] is True
            assert resp["reward"] > 0.0

            info = resp["info"]
            assert info["agreement_reached"] is True
            assert info["verdict"] == "accept"
            rb = info["reward_breakdown"]
            assert rb is not None
            assert 0.0 <= rb["rigor"] <= 1.0
            assert 0.0 <= rb["feasibility"] <= 1.0
            assert 0.0 <= rb["fidelity"] <= 1.0
            assert not (rb["rigor"] == 0.8 and rb["feasibility"] == 0.8)

    # -- error handling ------------------------------------------------------

    def test_ws_invalid_json(self, client: TestClient) -> None:
        with client.websocket_connect("/ws") as ws:
            ws.send_text("not valid json {{{")
            resp = json.loads(ws.receive_text())
            assert resp["type"] == "error"
            assert "Invalid JSON" in resp["message"]

    def test_ws_missing_action_field(self, client: TestClient) -> None:
        with client.websocket_connect("/ws") as ws:
            _ws_send_recv(ws, {"type": "reset", "seed": 42})
            resp = _ws_send_recv(ws, {"type": "step"})
            assert resp["type"] == "error"
            assert "Missing" in resp["message"]

    def test_ws_invalid_action_payload(self, client: TestClient) -> None:
        """Structurally invalid action (missing required fields) β†’ WS error."""
        with client.websocket_connect("/ws") as ws:
            _ws_send_recv(ws, {"type": "reset", "seed": 42})
            resp = _ws_send_recv(ws, {
                "type": "step",
                "action": {"action_type": "propose_protocol"},
            })
            assert resp["type"] == "error"
            assert "Invalid action" in resp["message"]

    def test_ws_unknown_message_type(self, client: TestClient) -> None:
        with client.websocket_connect("/ws") as ws:
            resp = _ws_send_recv(ws, {"type": "banana"})
            assert resp["type"] == "error"
            assert "Unknown" in resp["message"]

    # -- session isolation ---------------------------------------------------

    def test_ws_session_isolation(self, client: TestClient) -> None:
        """Two WebSocket connections have independent env state."""
        action = _good_action_payload(client)

        with client.websocket_connect("/ws") as ws1:
            r1 = _ws_send_recv(ws1, {"type": "reset", "seed": 1})
            _ws_send_recv(ws1, {"type": "step", "action": action})

            with client.websocket_connect("/ws") as ws2:
                r2 = _ws_send_recv(ws2, {"type": "reset", "seed": 2})

                assert r1["episode_id"] != r2["episode_id"]
                # ws2 is at round 0, ws1 is at round 1
                step2 = _ws_send_recv(ws2, {"type": "step", "action": action})
                assert step2["observation"]["scientist"]["round_number"] == 1

    # -- real-env integration (user-requested) --------------------------------

    def test_ws_semantic_invalid_action_returns_step_ok_with_info_error(
        self, client: TestClient
    ) -> None:
        """A structurally valid but semantically invalid action (e.g.
        duration_days=999) returns step_ok with info.error β€” NOT a
        transport-level WS error frame."""
        with client.websocket_connect("/ws") as ws:
            _ws_send_recv(ws, {"type": "reset", "seed": 42})
            bad_action = {
                "action_type": "propose_protocol",
                "sample_size": 5,
                "controls": ["baseline"],
                "technique": "some technique",
                "duration_days": 999,
                "required_equipment": [],
                "required_reagents": [],
                "questions": [],
                "rationale": "Duration is impossibly long for the lab.",
            }
            resp = _ws_send_recv(ws, {"type": "step", "action": bad_action})

            assert resp["type"] == "step_ok"
            assert resp["done"] is False
            assert resp["info"]["error"] is not None
            assert "Validation errors" in resp["info"]["error"]

    def test_ws_timeout_verdict(self, client: TestClient) -> None:
        """Run to max_rounds without accept β†’ done=True, verdict=timeout,
        reward=0.0. Proves real-env integration."""
        action = _good_action_payload(client)
        with client.websocket_connect("/ws") as ws:
            reset_resp = _ws_send_recv(ws, {"type": "reset", "seed": 42})
            max_rounds = reset_resp["observation"]["scientist"]["max_rounds"]

            resp = None
            for _ in range(max_rounds):
                resp = _ws_send_recv(ws, {"type": "step", "action": action})

            assert resp["done"] is True
            assert resp["info"]["verdict"] == "timeout"
            assert resp["reward"] < 0.0
            assert resp["info"]["reward_breakdown"] is not None
            assert resp["info"]["reward_breakdown"]["penalties"]["timeout"] > 0.0

    def test_ws_terminal_episode_persists_real_replay_log(
        self, client: TestClient
    ) -> None:
        """Complete a WS episode, then verify GET /replay/{episode_id}
        returns real reward_breakdown, judge_notes, and verdict β€”
        not stub strings."""
        action = _good_action_payload(client)
        with client.websocket_connect("/ws") as ws:
            reset_resp = _ws_send_recv(ws, {"type": "reset", "seed": 42})
            episode_id = reset_resp["episode_id"]

            _ws_send_recv(ws, {"type": "step", "action": action})
            _ws_send_recv(ws, {"type": "step", "action": _accept_action_payload()})

        # Fetch replay via REST after WS connection is closed
        replay_resp = client.get(f"/replay/{episode_id}")
        assert replay_resp.status_code == 200
        replay = replay_resp.json()

        assert replay["agreement_reached"] is True
        assert replay["verdict"] == "accept"
        assert replay["total_reward"] > 0.0

        # Real judge_notes, not stub
        assert replay["judge_notes"] != ""
        assert "Stub audit" not in replay["judge_notes"]
        assert "rigor" in replay["judge_notes"]

        # Real reward_breakdown with non-stub scores
        rb = replay["reward_breakdown"]
        assert rb is not None
        assert 0.0 < rb["rigor"] <= 1.0
        assert 0.0 < rb["feasibility"] <= 1.0
        assert 0.0 < rb["fidelity"] <= 1.0
        assert not (rb["rigor"] == 0.8 and rb["feasibility"] == 0.8)

    # -- idle timeout & disconnect cleanup (API 07) -------------------------

    def test_ws_idle_timeout_closes_connection(self, client: TestClient) -> None:
        """API 07: server closes WebSocket after idle timeout (no messages)."""
        with patch("server.app._WS_IDLE_TIMEOUT", 0.5):
            with client.websocket_connect("/ws") as ws:
                # Don't send anything β€” let the server-side timeout fire
                time.sleep(1.0)
                with pytest.raises(WebSocketDisconnect) as exc_info:
                    ws.receive_text()
                assert exc_info.value.code == 1000

    def test_ws_env_closes_on_disconnect(self, client: TestClient) -> None:
        """API 07: env.close() runs in the finally block on disconnect."""
        import server.app as _app

        _original_make_env = _app._make_env
        close_called: list[bool] = []

        def _tracked_make_env():
            env = _original_make_env()
            _original_close = env.close

            def _tracking_close():
                close_called.append(True)
                _original_close()

            env.close = _tracking_close
            return env

        with patch.object(_app, "_make_env", _tracked_make_env):
            with client.websocket_connect("/ws") as ws:
                _ws_send_recv(ws, {"type": "ping"})
            # Context manager exit sends disconnect; server runs finally block
            # TestClient joins the ASGI thread, so close() has already run
            assert len(close_called) == 1