Phase 6: SchemaShiftEnvironment scheduler — full episode lifecycle

server/environment.py

@@ -1,4 +1,363 @@
 """SchemaShiftEnvironment — episode scheduler. reset/step loop with drift ticks + grader.
 
-Will be filled in Phase 6.
+Round 1 bug prevention: step() raises RuntimeError if called before reset(). Never lazy-init.
 """
+from __future__ import annotations
+
+import uuid
+from copy import deepcopy
+from typing import Any
+
+from drift import DriftInjector
+from graders import build_grader, compute_step_shaping
+from models import (
+    Action,
+    DriftReportParams,
+    EpisodeState,
+    HistoryStep,
+    Observation,
+    RewardBreakdown,
+    ToolResponse,
+)
+from scenarios import SCENARIOS
+
+
+def _instantiate_tool(name: str, seed_data: dict) -> Any:
+    """Lazy import so missing stretch tools don't break core scenarios."""
+    if name == "mail":
+        from tools.mail import MailAPI
+        return MailAPI(seed_data)
+    if name == "calendar":
+        from tools.calendar import CalendarAPI
+        return CalendarAPI(seed_data)
+    if name == "crm":
+        from tools.crm import CRMAPI
+        return CRMAPI(seed_data)
+    if name == "chat":
+        from tools.chat import ChatAPI  # type: ignore[attr-defined]
+        return ChatAPI(seed_data)
+    if name == "docs":
+        from tools.docs import DocsAPI  # type: ignore[attr-defined]
+        return DocsAPI(seed_data)
+    raise ValueError(f"Unknown tool: {name}")
+
+
+class SchemaShiftEnvironment:
+    """The SchemaShift RL environment. One instance = one episode at a time."""
+
+    def __init__(self) -> None:
+        self._state: EpisodeState | None = None
+        self._tools: dict[str, Any] = {}
+        self._grader = build_grader()
+
+    # ──────────────────────────────────────────────────────────────
+    # Public API
+    # ──────────────────────────────────────────────────────────────
+
+    def reset(self, task_id: str, seed: int = 0) -> Observation:
+        if task_id not in SCENARIOS:
+            raise ValueError(
+                f"Unknown task_id: {task_id}. Available: {list(SCENARIOS.keys())}"
+            )
+
+        scenario = SCENARIOS[task_id]
+        required_tools = scenario["required_tools"]
+
+        self._tools = {}
+        for tool_name in required_tools:
+            tool_seed = scenario["seed_data"].get(tool_name, {})
+            self._tools[tool_name] = _instantiate_tool(tool_name, tool_seed)
+
+        self._state = EpisodeState(
+            episode_id=str(uuid.uuid4()),
+            task_id=task_id,
+            difficulty=scenario["difficulty"],
+            step=0,
+            max_steps=scenario["max_steps"],
+            token_budget=scenario["token_budget"],
+            token_budget_remaining=scenario["token_budget"],
+            drift_plan=deepcopy(scenario["drift_plan"]),
+            ground_truth_final_state=dict(scenario["ground_truth_final_state"]),
+            agent_state={},
+            history=[],
+            done=False,
+            cumulative_reward=0.0,
+        )
+
+        return self._observation("Episode started.")
+
+    def step(
+        self, action: Action, tokens_used: int = 0
+    ) -> tuple[Observation, RewardBreakdown]:
+        if self._state is None:
+            raise RuntimeError(
+                "Call reset() before step(). "
+                "SchemaShiftEnvironment requires an active episode."
+            )
+        s = self._state
+        if s.done:
+            raise RuntimeError(
+                "Episode already done. Call reset() to start a new episode."
+            )
+
+        s.step += 1
+        s.token_budget_remaining = max(0, s.token_budget_remaining - tokens_used)
+
+        # 1. Apply any scheduled drifts for this step
+        fired_drifts = DriftInjector.tick(s, self._tools)
+
+        # 2. Dispatch action (does NOT mark drift detected — that happens after shaping)
+        response, feedback = self._dispatch_action(action)
+
+        # 3. Compute step shaping BEFORE marking drift detected
+        #    (shaping checks `not d.detected_by_agent` — must run pre-mark)
+        step_shape = compute_step_shaping(s, action, response)
+
+        # 4. Now apply drift-detection mark (so grader sees the detection this step)
+        if action.type == "report_drift" and action.report is not None:
+            self._mark_drift_detected(action.report)
+
+        # 5. Update agent_state from action+response
+        self._update_agent_state(action, response)
+
+        # 6. Log the history step (reward will be filled below)
+        history_step = HistoryStep(
+            step=s.step, action=action, response=response, reward_breakdown=None,
+        )
+        s.history.append(history_step)
+
+        # 7. Check terminal conditions
+        if s.step >= s.max_steps or s.token_budget_remaining <= 0:
+            s.done = True
+        if action.type == "complete_task":
+            s.done = True
+
+        # 8. Run grader (sees marked drifts + updated agent_state + done flag)
+        reward = self._grader(s)
+        reward.step_shaping = step_shape
+        reward.shaped_total += step_shape
+
+        # 9. Log reward into history
+        s.history[-1].reward_breakdown = reward.model_dump()
+        s.cumulative_reward += reward.shaped_total
+
+        # 10. Decorate feedback with drift info and return
+        if fired_drifts:
+            feedback += (
+                f" [DRIFT FIRED: {len(fired_drifts)} drift event(s) on step {s.step}.]"
+            )
+
+        return self._observation(feedback), reward
+
+    # ──────────────────────────────────────────────────────────────
+    # Action dispatch (pure — no state mutation for drift detection)
+    # ──────────────────────────────────────────────────────────────
+
+    def _dispatch_action(self, action: Action) -> tuple[ToolResponse | None, str]:
+        s = self._state
+        assert s is not None
+
+        if action.type == "call_tool":
+            if action.tool_call is None:
+                return None, "Invalid call_tool: missing tool_call params."
+            if action.tool_call.tool not in self._tools:
+                return (
+                    ToolResponse(
+                        ok=False, status=404,
+                        error=f"Tool '{action.tool_call.tool}' not available in this scenario.",
+                    ),
+                    "Tool not available.",
+                )
+            response = self._tools[action.tool_call.tool].call(
+                action.tool_call.endpoint, action.tool_call.params
+            )
+            return (
+                response,
+                f"Called {action.tool_call.tool}.{action.tool_call.endpoint}: status={response.status}",
+            )
+
+        if action.type == "inspect_schema":
+            if action.inspect is None or action.inspect.tool not in self._tools:
+                return (
+                    ToolResponse(
+                        ok=False, status=404,
+                        error="Tool unavailable for inspection.",
+                    ),
+                    "Inspect target missing.",
+                )
+            schema = self._tools[action.inspect.tool].get_schema()
+            return (
+                ToolResponse(ok=True, status=200, body={"schema": schema}),
+                f"Inspected {action.inspect.tool} schema.",
+            )
+
+        if action.type == "retry_with_variant":
+            if action.retry is None or action.retry.tool not in self._tools:
+                return (
+                    ToolResponse(
+                        ok=False, status=404,
+                        error="Retry target unavailable.",
+                    ),
+                    "Retry target missing.",
+                )
+            response = self._tools[action.retry.tool].call(
+                action.retry.endpoint, action.retry.params
+            )
+            return (
+                response,
+                f"Retried {action.retry.tool}.{action.retry.endpoint}: status={response.status}",
+            )
+
+        if action.type == "report_drift":
+            if action.report is None:
+                return None, "Invalid report_drift: missing report params."
+            for d in s.drift_plan:
+                if (d.tool == action.report.tool
+                        and d.kind == action.report.drift_kind
+                        and d.fires_at_step <= s.step
+                        and not d.detected_by_agent):
+                    return (
+                        None,
+                        f"Drift correctly reported: {d.kind} on {d.tool} at step {d.fires_at_step}.",
+                    )
+            return None, "Drift report did not match any undetected fired drift."
+
+        if action.type == "complete_task":
+            summary = action.complete.summary if action.complete else ""
+            s.agent_state["_completion_summary"] = summary
+            self._check_completion_summary(summary)
+            return None, f"Episode marked complete. Summary: {summary[:80]}"
+
+        return None, f"Unknown action type: {action.type}"
+
+    # ──────────────────────────────────────────────────────────────
+    # Post-shaping state mutations
+    # ──────────────────────────────────────────────────────────────
+
+    def _mark_drift_detected(self, report: DriftReportParams) -> None:
+        s = self._state
+        assert s is not None
+        for d in s.drift_plan:
+            if (d.tool == report.tool
+                    and d.kind == report.drift_kind
+                    and d.fires_at_step <= s.step
+                    and not d.detected_by_agent):
+                d.detected_by_agent = True
+                return
+
+    # ──────────────────────────────────────────────────────────────
+    # State tracking — populates agent_state so grader can read it
+    # ──────────────────────────────────────────────────────────────
+
+    def _update_agent_state(
+        self, action: Action, response: ToolResponse | None
+    ) -> None:
+        if response is None or not response.ok:
+            return
+        s = self._state
+        assert s is not None
+        st = s.agent_state
+
+        tool: str | None = None
+        endpoint: str | None = None
+        params: dict = {}
+        if action.type == "call_tool" and action.tool_call is not None:
+            tool = action.tool_call.tool
+            endpoint = action.tool_call.endpoint
+            params = action.tool_call.params
+        elif action.type == "retry_with_variant" and action.retry is not None:
+            tool = action.retry.tool
+            endpoint = action.retry.endpoint
+            params = action.retry.params
+        else:
+            return
+
+        # MAIL ────────────────────────────────────────────────────
+        if tool == "mail":
+            if endpoint in ("send_message", "messages.send"):
+                st["mail.sent_count"] = st.get("mail.sent_count", 0) + 1
+                sent_to = params.get("to", "")
+                st["mail.last_sent_to"] = sent_to
+                subject = str(params.get("subject", "")).lower()
+                if "welcome" in subject:
+                    st["mail.last_subject_contains_welcome"] = True
+                if "all-hands" in subject or "all hands" in subject:
+                    st["mail.last_subject_contains_allhands"] = True
+                recipients: list[str] = st.get("mail.all_recipients", [])
+                if sent_to and sent_to not in recipients:
+                    recipients.append(sent_to)
+                    st["mail.all_recipients"] = recipients
+                e2_required = {"alex@company.com", "jordan@company.com", "sam@company.com"}
+                if e2_required.issubset(set(recipients)):
+                    st["mail.sent_to_all_three_recipients"] = True
+
+        # CALENDAR ────────────────────────────────────────────────
+        if tool == "calendar":
+            if endpoint == "create_event":
+                st["calendar.events_count"] = st.get("calendar.events_count", 0) + 1
+                body = response.body or {}
+                raw = body.get("attendees") or body.get("participants") or []
+                emails: list[str] = []
+                for a in raw:
+                    if isinstance(a, str):
+                        emails.append(a)
+                    elif isinstance(a, dict):
+                        emails.append(a.get("email", ""))
+                st["calendar.last_event_attendees"] = emails
+                if "priya@company.com" in emails and "alex@company.com" in emails:
+                    st["calendar.last_event_has_both_attendees"] = True
+
+        # CRM ─────────────────────────────────────────────────────
+        if tool == "crm":
+            if endpoint in ("update_contact", "contacts.patch"):
+                cid = params.get("contact_id", "")
+                status = params.get("status")
+                if cid and status:
+                    st[f"crm.contact_{cid}_status"] = status
+
+    def _check_completion_summary(self, summary: str) -> None:
+        s = self._state
+        assert s is not None
+        st = s.agent_state
+        gt = s.ground_truth_final_state
+        if "complete_summary_mentions_company" in gt:
+            for c in ("Globex", "Acme", "Initech"):
+                if c.lower() in summary.lower():
+                    st["complete_summary_mentions_company"] = True
+                    break
+
+    # ──────────────────────────────────────────────────────────────
+    # Observation construction
+    # ──────────────────────────────────────────────────────────────
+
+    def _observation(self, feedback: str) -> Observation:
+        s = self._state
+        assert s is not None
+        scenario = SCENARIOS[s.task_id]
+
+        return Observation(
+            episode_id=s.episode_id,
+            task_id=s.task_id,
+            difficulty=s.difficulty,
+            step=s.step,
+            max_steps=s.max_steps,
+            token_budget_remaining=s.token_budget_remaining,
+            task_description=scenario["task_description"],
+            success_criteria=list(scenario["success_criteria"]),
+            tool_schemas={name: t.get_schema() for name, t in self._tools.items()},
+            known_state=dict(s.agent_state),
+            history=list(s.history[-5:]),
+            last_response=s.history[-1].response if s.history else None,
+            drift_events_visible=[
+                {
+                    "tool": d.tool,
+                    "kind": d.kind,
+                    "endpoint": d.endpoint,
+                    "fires_at_step": d.fires_at_step,
+                }
+                for d in s.drift_plan
+                if d.detected_by_agent
+            ],
+            done=s.done,
+            feedback=feedback,
+        )

tests/test_environment.py

@@ -0,0 +1,242 @@
+"""SchemaShiftEnvironment acceptance tests — Phase 6 (end-to-end RL loop)."""
+from __future__ import annotations
+
+import pytest
+
+from models import (
+    Action,
+    CompleteParams,
+    DriftReportParams,
+    InspectParams,
+    RetryParams,
+    ToolCallParams,
+)
+from server.environment import SchemaShiftEnvironment
+
+
+# ──────────────────────────────────────────────────────────────────
+# Test 1 — Round 1 bug prevention
+# ──────────────────────────────────────────────────────────────────
+
+def test_step_before_reset_raises() -> None:
+    env = SchemaShiftEnvironment()
+    action = Action(type="complete_task", complete=CompleteParams(summary="noop"))
+    with pytest.raises(RuntimeError) as excinfo:
+        env.step(action)
+    assert "reset" in str(excinfo.value).lower()
+
+
+# ──────────────────────────────────────────────────────────────────
+# Test 2 — reset returns a valid observation
+# ──────────────────────────────────────────────────────────────────
+
+def test_reset_returns_valid_observation() -> None:
+    env = SchemaShiftEnvironment()
+    obs = env.reset("E1_onboard_new_hire")
+    assert obs.task_id == "E1_onboard_new_hire"
+    assert obs.step == 0
+    assert obs.max_steps == 8
+    assert "mail" in obs.tool_schemas
+    assert "calendar" in obs.tool_schemas
+    assert "crm" not in obs.tool_schemas
+    assert obs.done is False
+    assert obs.difficulty == "easy"
+    assert len(obs.success_criteria) >= 1
+
+
+# ──────────────────────────────────────────────────────────────────
+# Test 3 — reset on unknown task raises
+# ──────────────────────────────────────────────────────────────────
+
+def test_reset_unknown_task_raises() -> None:
+    env = SchemaShiftEnvironment()
+    with pytest.raises(ValueError):
+        env.reset("nonexistent_task")
+
+
+# ──────────────────────────────────────────────────────────────────
+# Test 4 — call_tool success updates agent_state
+# ──────────────────────────────────────────────────────────────────
+
+def test_call_tool_success_updates_state() -> None:
+    env = SchemaShiftEnvironment()
+    env.reset("E1_onboard_new_hire")
+
+    obs, reward = env.step(Action(
+        type="call_tool",
+        tool_call=ToolCallParams(
+            tool="mail",
+            endpoint="send_message",
+            params={
+                "to": "priya@company.com",
+                "subject": "Welcome!",
+                "body": "Welcome to the team.",
+            },
+        ),
+    ))
+    assert obs.last_response is not None
+    assert obs.last_response.ok is True
+    assert obs.known_state["mail.sent_count"] == 1
+    assert obs.known_state["mail.last_sent_to"] == "priya@company.com"
+    assert obs.known_state["mail.last_subject_contains_welcome"] is True
+
+
+# ──────────────────────────────────────────────────────────────────
+# Test 5 — FULL E1 episode with drift → inspect → retry → report → complete
+# ──────────────────────────────────────────────────────────────────
+
+def test_e1_full_episode_with_adaptation() -> None:
+    env = SchemaShiftEnvironment()
+    env.reset("E1_onboard_new_hire")
+
+    # Step 1: send welcome email (pre-drift)
+    obs, r1 = env.step(Action(
+        type="call_tool",
+        tool_call=ToolCallParams(
+            tool="mail", endpoint="send_message",
+            params={"to": "priya@company.com", "subject": "Welcome aboard!",
+                    "body": "Welcome to the team."},
+        ),
+    ))
+    assert obs.last_response is not None and obs.last_response.ok is True
+
+    # Step 2: inspect calendar (pre-drift)
+    obs, r2 = env.step(Action(
+        type="inspect_schema", inspect=InspectParams(tool="calendar"),
+    ))
+    assert obs.last_response is not None and obs.last_response.ok is True
+
+    # Step 3: drift fires at state.step=3; call with stale attendees fails
+    obs, r3 = env.step(Action(
+        type="call_tool",
+        tool_call=ToolCallParams(
+            tool="calendar", endpoint="create_event",
+            params={"title": "New Hire Orientation",
+                    "start": "2026-04-27T10:00:00Z",
+                    "end": "2026-04-27T11:00:00Z",
+                    "attendees": ["priya@company.com", "alex@company.com"]},
+        ),
+    ))
+    assert obs.last_response is not None and obs.last_response.ok is False
+
+    # Step 4: inspect calendar (now shows participants schema)
+    obs, r4 = env.step(Action(
+        type="inspect_schema", inspect=InspectParams(tool="calendar"),
+    ))
+    assert obs.last_response is not None and obs.last_response.ok is True
+    cal_schema = obs.tool_schemas["calendar"]["create_event"]
+    assert "participants" in cal_schema["params"]
+
+    # Step 5: retry with participants format
+    obs, r5 = env.step(Action(
+        type="retry_with_variant",
+        retry=RetryParams(
+            tool="calendar", endpoint="create_event",
+            params={"title": "New Hire Orientation",
+                    "start": "2026-04-27T10:00:00Z",
+                    "end": "2026-04-27T11:00:00Z",
+                    "participants": [
+                        {"email": "priya@company.com", "role": "required"},
+                        {"email": "alex@company.com", "role": "required"},
+                    ]},
+        ),
+    ))
+    assert obs.last_response is not None and obs.last_response.ok is True
+
+    # Step 6: report drift
+    obs, r6 = env.step(Action(
+        type="report_drift",
+        report=DriftReportParams(
+            tool="calendar", drift_kind="field_rename",
+            description="create_event attendees renamed to participants",
+        ),
+    ))
+
+    # Step 7: complete
+    obs, r7 = env.step(Action(
+        type="complete_task",
+        complete=CompleteParams(
+            summary="Onboarded Priya with welcome email and orientation event.",
+        ),
+    ))
+
+    state = env._state
+    assert state is not None
+    assert obs.done is True
+    assert state.agent_state["mail.sent_count"] == 1
+    assert state.agent_state["calendar.events_count"] == 1
+    assert state.agent_state["calendar.last_event_has_both_attendees"] is True
+    assert state.drift_plan[0].detected_by_agent is True
+    assert r7.task_completion == 1.0
+    assert r7.drift_detection == 1.0
+    assert r7.adaptation_quality == 1.0
+    assert r7.shaped_total > 0.5
+    assert r7.binary == 1.0
+
+
+# ──────────────────────────────────────────────────────────────────
+# Test 6 — max_steps terminates episode
+# ──────────────────────────────────────────────────────────────────
+
+def test_max_steps_terminates_episode() -> None:
+    env = SchemaShiftEnvironment()
+    env.reset("E2_meeting_invite_blast")
+    inspect = Action(type="inspect_schema", inspect=InspectParams(tool="mail"))
+    obs = None
+    for _ in range(6):
+        obs, _ = env.step(inspect)
+    assert obs is not None
+    assert obs.done is True
+    assert obs.step == 6
+
+
+# ──────────────────────────────────────────────────────────────────
+# Test 7 — step_shaping +0.10 for inspect after failure
+# ──────────────────────────────────────────────────────────────────
+
+def test_step_shaping_applied_correctly() -> None:
+    env = SchemaShiftEnvironment()
+    env.reset("E1_onboard_new_hire")
+
+    # Step 1: send_message missing required 'body' → 400
+    env.step(Action(
+        type="call_tool",
+        tool_call=ToolCallParams(
+            tool="mail", endpoint="send_message",
+            params={"to": "x@y.com", "subject": "hi"},
+        ),
+    ))
+
+    # Step 2: inspect_schema after failure → +0.10 shaping
+    obs, reward = env.step(Action(
+        type="inspect_schema", inspect=InspectParams(tool="mail"),
+    ))
+    assert reward.step_shaping == pytest.approx(0.10)
+
+
+# ──────────────────────────────────────────────────────────────────
+# Test 8 — dumb retry penalty
+# ──────────────────────────────────────────────────────────────────
+
+def test_dumb_retry_penalty() -> None:
+    env = SchemaShiftEnvironment()
+    env.reset("E1_onboard_new_hire")
+
+    # Step 1: call_tool mail.send_message with only {"to": "x"} → 400
+    env.step(Action(
+        type="call_tool",
+        tool_call=ToolCallParams(
+            tool="mail", endpoint="send_message",
+            params={"to": "x@y.com"},
+        ),
+    ))
+
+    # Step 2: same call again → dumb retry → -0.05 penalty
+    obs, reward = env.step(Action(
+        type="call_tool",
+        tool_call=ToolCallParams(
+            tool="mail", endpoint="send_message",
+            params={"to": "x@y.com"},
+        ),
+    ))
+    assert reward.step_shaping == pytest.approx(-0.05)