Harden env isolation and proxy validation

app.py

@@ -1,336 +1,173 @@
-"""
-FastAPI server exposing the rag-context-optimizer OpenEnv HTTP API.
-"""
-
-from __future__ import annotations
-
+"""
+FastAPI server exposing the rag-context-optimizer OpenEnv HTTP API.
+"""
+
+from __future__ import annotations
+
 from contextlib import asynccontextmanager
 from dataclasses import asdict, is_dataclass
+import os
 from pathlib import Path
 from typing import Any, Literal
-
+from uuid import uuid4
+
 from fastapi import Body, FastAPI, HTTPException, Request
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import HTMLResponse
-from pydantic import BaseModel
-
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse
+from pydantic import BaseModel
+
+from env.corpus import list_corpus_families
 from env.environment import RagContextOptimizerEnv
 from env.models import RagAction
-from env.corpus import list_corpus_families
 from env.prompt_optimizer import CompressionMode, optimize_prompt
 from env.tasks import ALL_TASKS, TASKS_BY_NAME
-
-
+
+
 class ResetRequest(BaseModel):
     task_name: Literal["single_domain_qa", "cross_domain_synthesis", "adversarial_compression"] = "single_domain_qa"
     custom_query: str | None = None
     token_budget: int | None = None
     max_steps: int | None = None
     corpus_family: str | None = None
-
-
+
+
 class OptimizePromptRequest(BaseModel):
     prompt: str
     corpus_family: str | None = None
     compression_mode: CompressionMode = "balanced"
-
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    env = RagContextOptimizerEnv()
-    await env.reset()
-    app.state.env = env
-    yield
-    await app.state.env.close()
-
-
-app = FastAPI(
-    title="rag-context-optimizer",
-    version="1.0.0",
-    description="RAG pipeline optimization environment — minimize tokens, maximize answer quality",
-    lifespan=lifespan,
-)
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-UI_TEMPLATE_PATH = Path(__file__).resolve().parent / "server" / "templates" / "ui.html"
-
-
-
-
-
-@app.middleware("http")
-async def log_requests(request: Request, call_next):
-    print(f"[request] {request.method} {request.url.path}")
-    response = await call_next(request)
-    print(f"[response] {request.method} {request.url.path} -> {response.status_code}")
-    return response
-
-
-@app.get("/", response_class=HTMLResponse)
-async def home_page():
-    return HTMLResponse(
-        UI_TEMPLATE_PATH.read_text(encoding="utf-8"),
-        headers={
-            "Cache-Control": "no-store, max-age=0",
-            "Pragma": "no-cache",
-        },
-    )
-
-
-def _serialize_observation(observation: Any) -> dict[str, Any]:
-    if hasattr(observation, "model_dump"):
-        return observation.model_dump()
-    if is_dataclass(observation):
-        return asdict(observation)
-    return dict(observation)
-
-
-def _serialize_step_result(result: Any, reset: bool = False) -> dict[str, Any]:
-    raw_info = result.info or {}
-    if reset:
-        return {
-            "observation": _serialize_observation(result.observation),
-            "reward": None,
-            "done": False,
-            "info": {},
-        }
-    return {
-        "observation": _serialize_observation(result.observation),
-        "reward": result.reward,
-        "done": result.done,
-        "info": {
-            "grader_breakdown": raw_info.get("grader"),
-            "event": raw_info.get("event"),
-            "passed": raw_info.get("passed"),
-        },
-    }
-
-
-def _is_bad_action_event(event: str | None) -> bool:
-    return event in {
-        "chunk_not_found",
-    }
-
-
-def _tokenize(text: str) -> set[str]:
-    import re
-
-    return set(re.findall(r"[a-z0-9]+", text.lower()))
-
-
-def _content_terms(text: str) -> set[str]:
-    return {term for term in _tokenize(text) if len(term) > 2 and term not in _PROMPT_STOPWORDS}
-
-
-def _clean_output_text(text: str) -> str:
-    import re
-
-    cleaned = text.replace("```", " ").replace("---", " ")
-    cleaned = re.sub(r"\s+", " ", cleaned).strip()
-    cleaned = re.sub(r"[#*_`]+", "", cleaned)
-    cleaned = re.sub(r'\b(title|emoji|colorfrom|colorto|sdk|app_file|pinned)\s*:\s*', "", cleaned, flags=re.IGNORECASE)
-    return cleaned.strip(" -:")
-
-
-def _compact_text(text: str, max_words: int = 28) -> str:
-    words = text.split()
-    if len(words) <= max_words:
-        return text
-    return " ".join(words[:max_words]).rstrip(" ,;:") + " ..."
-
-
-_PROMPT_STOPWORDS = {
-    "a","an","and","are","as","at","be","but","by","can","could","do","does","did",
-    "for","from","had","has","have","how","i","if","in","into","is","it","its","me",
-    "my","of","on","or","our","should","so","than","that","the","their","them","then",
-    "there","these","they","this","to","too","use","using","was","we","were","what",
-    "when","where","which","while","with","without","would","you","your",
-}
-
-
-def _approx_tokens(text: str) -> int:
-    return max(1, len(text.strip()) // 4) if text.strip() else 0
-
-
-def _compress_prompt_text(prompt: str, target_tokens: int) -> str:
-    import re
-
-    raw = " ".join(prompt.strip().split())
-    if not raw:
-        return ""
-
-    tokens = re.findall(r"[A-Za-z0-9][A-Za-z0-9\-_/]*", raw)
-    kept: list[str] = []
-    seen: set[str] = set()
-
-    # Keep “meaningful” tokens: numbers, identifiers, longer words, and acronyms. Drop stopwords.
-    for tok in tokens:
-        low = tok.lower()
-        is_number = low.isdigit()
-        is_identifier = any(ch in tok for ch in ("_", "-", "/")) and len(tok) >= 4
-        is_acronym = tok.isupper() and len(tok) <= 8
-        is_meaningful = is_number or is_identifier or is_acronym or len(low) >= 4
-        if not is_meaningful:
-            continue
-        if low in _PROMPT_STOPWORDS:
-            continue
-        if low in seen:
-            continue
-        seen.add(low)
-        kept.append(tok)
-        if len(kept) >= max(10, target_tokens):
-            break
-
-    if not kept:
-        # Fallback: truncated raw prompt.
-        words = raw.split()
-        return " ".join(words[: max(8, target_tokens)]).rstrip(" ,;:") + (" ..." if len(words) > target_tokens else "")
-
-    # Turn the token list into a copy-paste-ready “goal” sentence.
-    goal = " ".join(kept)
-    goal = re.sub(r"\s+", " ", goal).strip()
-    return goal
-
-
-_INSTRUCTION_PRIORITY_TERMS = {
-    "must","should","only","not","never","always","include","exclude","cite","answer",
-    "return","draft","write","summarize","compare","explain","verify","preserve","focus",
-    "keep","avoid","report","escalate","rollback","refund","incident","customer","security",
-}
-
-
-def _trim_sentence(sentence: str, max_terms: int) -> str:
-    import re
-
-    words = re.findall(r"[A-Za-z0-9][A-Za-z0-9\\-_/]*|[,:;()]", sentence)
-    if not words:
-        return ""
-    kept: list[str] = []
-
-    for index, token in enumerate(words):
-        normalized = re.sub(r"[^A-Za-z0-9]+", "", token).lower()
-        if token in {",", ":", ";", "(", ")"}:
-            if kept and kept[-1] not in {",", ":", ";", "("}:
-                kept.append(token)
-            continue
-        is_priority = normalized in _INSTRUCTION_PRIORITY_TERMS
-        is_meaningful = (
-            normalized.isdigit()
-            or any(ch in token for ch in ("_", "-", "/"))
-            or len(normalized) >= 4
-            or is_priority
-            or index < 3
-        )
-        if not is_meaningful:
-            continue
-        if normalized in _PROMPT_STOPWORDS and not is_priority and index >= 3:
-            continue
-        kept.append(token)
-        if len([word for word in kept if word not in {",", ":", ";", "(", ")"}]) >= max_terms:
-            break
-
-    text = " ".join(kept)
-    text = re.sub(r"\s+([,:;)])", r"\1", text)
-    text = re.sub(r"(\()\s+", r"\1", text)
-    return text.strip(" ,;:")
-
-
-def _rewrite_prompt_text(prompt: str, target_tokens: int) -> str:
-    import re
-
-    raw = " ".join(prompt.strip().split())
-    if not raw:
-        return ""
-
-    sentences = [segment.strip() for segment in re.split(r"(?<=[.!?])\s+|\n+", raw) if segment.strip()]
-    if not sentences:
-        sentences = [raw]
-
-    rewritten: list[str] = []
-    used_terms = 0
-    max_terms = max(8, target_tokens)
-    for index, sentence in enumerate(sentences):
-        remaining = max_terms - used_terms
-        if remaining <= 0:
-            break
-        compact = _trim_sentence(sentence, max(4, remaining if index == 0 else min(remaining, 10)))
-        if not compact:
-            continue
-        rewritten.append(compact)
-        used_terms += len(compact.split())
-        if used_terms >= max_terms:
-            break
-
-    if not rewritten:
-        fallback = _trim_sentence(raw, max_terms)
-        return fallback or raw[: max(16, target_tokens * 4)].strip()
-
-    output = ". ".join(rewritten).strip()
-    if len(rewritten) == 1 and not output.endswith("."):
-        output += "."
-    return output
-
-
-def _fit_citations_into_prompt(base_prompt: str, citation_ids: list[str], input_tokens: int, target_tokens: int, source_prompt: str) -> tuple[str, bool, str | None]:
-    if not citation_ids:
-        return base_prompt, False, "No high-confidence evidence anchors were selected."
-
-    citation_suffix = " Evidence: " + " ".join(f"[{chunk_id}]" for chunk_id in citation_ids[:3])
-    with_all = (base_prompt.rstrip(".") + "." + citation_suffix).strip()
-    if _approx_tokens(with_all) < input_tokens:
-        return with_all, True, None
-
-    one_citation_suffix = " Evidence: " + f"[{citation_ids[0]}]"
-    with_one = (base_prompt.rstrip(".") + "." + one_citation_suffix).strip()
-    if _approx_tokens(with_one) < input_tokens:
-        return with_one, True, None
-
-    tighter_target = max(8, target_tokens - 3)
-    tighter_prompt = _rewrite_prompt_text(source_prompt, tighter_target)
-    tighter_with_one = (tighter_prompt.rstrip(".") + "." + one_citation_suffix).strip()
-    if _approx_tokens(tighter_with_one) < input_tokens:
-        return tighter_with_one, True, None
-
-    return base_prompt, False, "Citations were omitted to keep the optimized prompt shorter than the original. Use the evidence notes below if explicit anchors are required."
-
-
-def _summarize_chunk_for_output(chunk: Any, effective_text: str) -> str:
-    if getattr(chunk, "domain", "").startswith("Project"):
-        keywords = ", ".join(chunk.keywords[:5])
-        domain = chunk.domain.replace("Project ", "").lower()
-        return _compact_text(f"This benchmark's {domain} covers {keywords}.", 24)
-    ranked_sentences = _sentence_rank(" ".join(chunk.keywords), _clean_output_text(effective_text))
-    if ranked_sentences:
-        return _compact_text(_clean_output_text(ranked_sentences[0]))
-    return _compact_text(_clean_output_text(effective_text))
-
-
-def _sentence_rank(query: str, text: str) -> list[str]:
-    import re
-
-    query_terms = _tokenize(query)
-    sentences = [segment.strip() for segment in re.split(r"(?<=[.!?])\s+", text) if segment.strip()]
-    if not sentences:
-        return []
-
-    ranked: list[tuple[float, str]] = []
-    for index, sentence in enumerate(sentences):
-        sentence_terms = _tokenize(sentence)
-        overlap = len(query_terms & sentence_terms)
-        score = (overlap * 2.0) + (0.25 if index == 0 else 0.0)
-        ranked.append((score, sentence))
-
-    ranked.sort(key=lambda item: (-item[0], len(item[1])))
-    return [sentence for _score, sentence in ranked]
-
-
+
+
+class EpisodeStore:
+    def __init__(self, max_episodes: int = 16):
+        self._episodes: dict[str, RagContextOptimizerEnv] = {}
+        self._order: list[str] = []
+        self.latest_episode_id: str | None = None
+        self._max_episodes = max_episodes
+
+    async def close_all(self) -> None:
+        for env in self._episodes.values():
+            await env.close()
+        self._episodes.clear()
+        self._order.clear()
+        self.latest_episode_id = None
+
+    async def create(self, env: RagContextOptimizerEnv) -> str:
+        episode_id = uuid4().hex
+        self._episodes[episode_id] = env
+        self._order.append(episode_id)
+        self.latest_episode_id = episode_id
+
+        while len(self._order) > self._max_episodes:
+            stale_id = self._order.pop(0)
+            stale_env = self._episodes.pop(stale_id, None)
+            if stale_env is not None:
+                await stale_env.close()
+            if self.latest_episode_id == stale_id:
+                self.latest_episode_id = self._order[-1] if self._order else None
+        return episode_id
+
+    def get(self, episode_id: str | None) -> tuple[str, RagContextOptimizerEnv]:
+        resolved_id = episode_id or self.latest_episode_id
+        if resolved_id is None or resolved_id not in self._episodes:
+            raise KeyError("episode_not_found")
+        return resolved_id, self._episodes[resolved_id]
+
+
+def _request_logging_enabled() -> bool:
+    return os.getenv("DEBUG_LOG_REQUESTS", "").strip().lower() in {"1", "true", "yes"}
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app.state.episodes = EpisodeStore()
+    yield
+    await app.state.episodes.close_all()
+
+
+app = FastAPI(
+    title="rag-context-optimizer",
+    version="1.0.0",
+    description="RAG pipeline optimization environment - minimize tokens, maximize answer quality",
+    lifespan=lifespan,
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=False,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+UI_TEMPLATE_PATH = Path(__file__).resolve().parent / "server" / "templates" / "ui.html"
+
+
+@app.middleware("http")
+async def log_requests(request: Request, call_next):
+    should_log = _request_logging_enabled()
+    if should_log:
+        print(f"[request] {request.method} {request.url.path}")
+    response = await call_next(request)
+    if should_log:
+        print(f"[response] {request.method} {request.url.path} -> {response.status_code}")
+    return response
+
+
+@app.get("/", response_class=HTMLResponse)
+async def home_page():
+    return HTMLResponse(
+        UI_TEMPLATE_PATH.read_text(encoding="utf-8"),
+        headers={
+            "Cache-Control": "no-store, max-age=0",
+            "Pragma": "no-cache",
+        },
+    )
+
+
+def _serialize_observation(observation: Any) -> dict[str, Any]:
+    if hasattr(observation, "model_dump"):
+        return observation.model_dump()
+    if is_dataclass(observation):
+        return asdict(observation)
+    return dict(observation)
+
+
+def _serialize_step_result(result: Any, reset: bool = False, episode_id: str | None = None) -> dict[str, Any]:
+    raw_info = result.info or {}
+    payload = {
+        "observation": _serialize_observation(result.observation),
+        "reward": None if reset else result.reward,
+        "done": False if reset else result.done,
+        "info": {} if reset else {
+            "grader_breakdown": raw_info.get("grader"),
+            "event": raw_info.get("event"),
+            "passed": raw_info.get("passed"),
+        },
+    }
+    if episode_id is not None:
+        payload["episode_id"] = episode_id
+    return payload
+
+
+def _is_bad_action_event(event: str | None) -> bool:
+    return event in {"chunk_not_found"}
+
+
+def _episode_store() -> EpisodeStore:
+    episodes = getattr(app.state, "episodes", None)
+    if episodes is None:
+        episodes = EpisodeStore()
+        app.state.episodes = episodes
+    return episodes
+
+
+def _resolve_env(episode_id: str | None) -> tuple[str, RagContextOptimizerEnv]:
+    try:
+        return _episode_store().get(episode_id)
+    except KeyError as exc:
+        raise HTTPException(status_code=404, detail="Episode not found. Call /reset first.") from exc
+
+
 async def _optimize_prompt_backend(
     prompt: str,
     corpus_family: str | None = None,
@@ -346,172 +183,8 @@ async def _optimize_prompt_backend(
         "selected_keywords": result.selected_keywords,
         "optimization_mode": result.optimization_mode,
     }
-    clean_prompt = prompt.strip()
-    env = RagContextOptimizerEnv(
-        task_name="single_domain_qa",
-        query_override=clean_prompt,
-        token_budget_override=800,
-        max_steps_override=6,
-        corpus_family_override=corpus_family,
-    )
-    await env.reset()
-
-    tuning = env._last_tuning or env.context_tuner.tune(clean_prompt, env._available_chunks)
-
-    ranked_candidates = []
-    for chunk in env._available_chunks:
-        tuned = tuning.tuned_scores.get(chunk.chunk_id)
-        score = tuned.final_score if tuned is not None else env.retriever.hybrid_score(clean_prompt, chunk)
-        if score < 0.16:
-            continue
-        ranked_candidates.append((chunk, score, tuned))
-    ranked_candidates.sort(
-        key=lambda item: (
-            -(item[1] / max(item[0].tokens, 1)),
-            -(item[2].citation_prior if item[2] is not None else 0.0),
-            -item[1],
-            item[0].chunk_id,
-        )
-    )
 
-    selected_ids: list[str] = []
-    token_cap = 360
-    running_tokens = 0
-    for chunk, score, tuned in ranked_candidates:
-        if len(selected_ids) >= 4:
-            break
-        if score < 0.22 and selected_ids:
-            break
-        projected = running_tokens + chunk.tokens
-        if projected > token_cap and selected_ids:
-            continue
-        selected_ids.append(chunk.chunk_id)
-        env._selected_chunks.append(chunk.chunk_id)
-        running_tokens += chunk.tokens
-
-    if not selected_ids and ranked_candidates:
-        best_chunk = ranked_candidates[0][0]
-        selected_ids.append(best_chunk.chunk_id)
-        env._selected_chunks.append(best_chunk.chunk_id)
-
-    for chunk_id in list(selected_ids):
-        chunk = env._chunk_map().get(chunk_id)
-        if chunk is None:
-            continue
-        tuned = tuning.tuned_scores.get(chunk_id)
-        score = tuned.final_score if tuned is not None else env.retriever.hybrid_score(clean_prompt, chunk)
-        ratio = tuned.compression_ratio if tuned is not None else 0.5
-        if score >= 0.75:
-            ratio = max(ratio, 0.6)
-        env._compression_ratios[chunk_id] = ratio
-
-    input_tokens = _approx_tokens(clean_prompt)
-    # Target: strictly shorter than input, while preserving more structure for longer prompts.
-    if input_tokens <= 24:
-        target_ratio = 0.85
-    elif input_tokens <= 60:
-        target_ratio = 0.75
-    elif input_tokens <= 120:
-        target_ratio = 0.68
-    else:
-        target_ratio = 0.62
-    target_tokens = max(12, int(input_tokens * target_ratio))
-    target_tokens = min(target_tokens, 80)
-
-    compressed_goal = _rewrite_prompt_text(clean_prompt, target_tokens=target_tokens)
-
-    # Optionally add a tiny amount of distilled context, but only if it still stays shorter overall.
-    distilled_points: list[tuple[str, str]] = []
-    for chunk_id in env._selected_chunks:
-        chunk = env._chunk_map().get(chunk_id)
-        if chunk is None:
-            continue
-        best = _summarize_chunk_for_output(chunk, env._effective_chunk_text(chunk_id))
-        if best and all(existing_point != best for _existing_chunk_id, existing_point in distilled_points):
-            distilled_points.append((chunk_id, best))
-        if len(distilled_points) >= (2 if input_tokens < 80 else 3):
-            break
-
-    lines: list[str] = []
-    lines.append(compressed_goal if compressed_goal else clean_prompt)
-    if distilled_points and input_tokens >= 80:
-        lines.append("")
-        lines.append("Context:")
-        lines.extend([f"- [{chunk_id}] {point}" for chunk_id, point in distilled_points])
-    optimized_prompt = "\n".join(lines).strip()
-
-    # Hard guarantee: never return an “optimized” prompt longer than the input.
-    if input_tokens > 0 and _approx_tokens(optimized_prompt) >= input_tokens:
-        # Enforce by character budget (tokens ~= chars/4).
-        max_chars = max(12, (input_tokens - 1) * 4)
-        optimized_prompt = optimized_prompt[:max_chars].rstrip(" ,;:\n\t")
-        if optimized_prompt and not optimized_prompt.endswith("..."):
-            optimized_prompt = optimized_prompt + " ..."
-        # If still not strictly smaller (very small inputs), trim until it is.
-        while input_tokens > 1 and _approx_tokens(optimized_prompt) >= input_tokens and len(optimized_prompt) > 12:
-            optimized_prompt = optimized_prompt[:-6].rstrip(" ,;:\n\t") + " ..."
-        if input_tokens > 1 and _approx_tokens(optimized_prompt) >= input_tokens:
-            optimized_prompt = _rewrite_prompt_text(clean_prompt, target_tokens=max(5, input_tokens - 1))
-            if optimized_prompt and not optimized_prompt.endswith("...") and _approx_tokens(optimized_prompt) >= input_tokens:
-                optimized_prompt = optimized_prompt[: max(8, (input_tokens - 1) * 4)].strip() + " ..."
-
-    optimized_prompt, citation_ready, citation_guidance = _fit_citations_into_prompt(
-        optimized_prompt,
-        tuning.suggested_citations or list(env._selected_chunks),
-        input_tokens,
-        target_tokens,
-        clean_prompt,
-    )
 
-    original_prompt_tokens = input_tokens
-    optimized_prompt_tokens = _approx_tokens(optimized_prompt)
-    source_tokens = sum(env._chunk_map()[chunk_id].tokens for chunk_id in env._selected_chunks if chunk_id in env._chunk_map())
-    compressed_tokens = sum(env._effective_chunk_tokens(chunk_id) for chunk_id in env._selected_chunks)
-    evidence_terms = _content_terms(" ".join(env._effective_chunk_text(chunk_id) for chunk_id in env._selected_chunks))
-    prompt_terms = _content_terms(optimized_prompt)
-    inline_citations = set(re.findall(r"\[([a-z0-9_]+)\]", optimized_prompt.lower()))
-    grounded_overlap = (len(prompt_terms & evidence_terms) / len(prompt_terms)) if prompt_terms else 0.0
-
-    return {
-        "optimized_prompt": optimized_prompt,
-        "stats": {
-            "selected_chunks": len(env._selected_chunks),
-            "source_tokens": source_tokens,
-            "compressed_context_tokens": compressed_tokens,
-            "original_prompt_tokens": original_prompt_tokens,
-            "optimized_prompt_tokens": optimized_prompt_tokens,
-            "compression_gain": max(0, source_tokens - compressed_tokens),
-        },
-        "grounding": {
-            "citations": tuning.suggested_citations or list(env._selected_chunks),
-            "citation_ready": citation_ready and bool(inline_citations),
-            "citation_guidance": citation_guidance,
-            "grounded_overlap": round(grounded_overlap, 3),
-            "evidence_notes": [
-                {"chunk_id": chunk_id, "note": note}
-                for chunk_id, note in distilled_points
-            ],
-        },
-        "context_tuning": {
-            "mode": tuning.mode,
-            "top_demo_cases": tuning.top_demo_cases,
-            "suggested_citations": tuning.suggested_citations,
-            "token_dropout": tuning.token_dropout,
-            "leave_one_out": tuning.leave_one_out,
-        },
-        "corpus_family": env._corpus_family,
-        "selected_keywords": [
-            keyword
-            for chunk_id in env._selected_chunks
-            for keyword in (env._chunk_map().get(chunk_id).keywords if env._chunk_map().get(chunk_id) else [])
-        ][:10],
-    }
-
-
-
-
-
-
 def _suggest_action(env: RagContextOptimizerEnv) -> dict[str, Any]:
     observation = env._build_observation()
     selected = set(observation.selected_chunks)
@@ -546,8 +219,7 @@ def _suggest_action(env: RagContextOptimizerEnv) -> dict[str, Any]:
             if chunk.keywords:
                 chosen_phrases.append(f"[{chunk.chunk_id}] " + ", ".join(chunk.keywords[:2]))
         answer = (
-            "Grounded answer based on selected evidence: "
-            + "; ".join(chosen_phrases[:3])
+            "Grounded answer based on selected evidence: " + "; ".join(chosen_phrases[:3])
             if chosen_phrases
             else "Grounded answer based on the currently selected evidence."
         )
@@ -559,37 +231,37 @@ def _suggest_action(env: RagContextOptimizerEnv) -> dict[str, Any]:
     for chunk in sorted(
         available,
         key=lambda chunk: (
-            -(score_map.get(chunk.chunk_id).final_score if score_map.get(chunk.chunk_id) else 0.0)
-            / max(chunk.tokens, 1),
+            -(score_map.get(chunk.chunk_id).final_score if score_map.get(chunk.chunk_id) else 0.0) / max(chunk.tokens, 1),
             chunk.tokens,
             chunk.chunk_id,
         ),
     ):
         if chunk.tokens <= remaining_budget:
             return {"action_type": "select_chunk", "chunk_id": chunk.chunk_id}
-
-    if selected_chunks:
-        return {
-            "action_type": "submit_answer",
-            "answer": "Optimized answer based on the currently selected evidence.",
-        }
-    if available:
-        smallest_chunk = min(available, key=lambda chunk: (chunk.tokens, chunk.chunk_id))
-        return {
-            "action_type": "submit_answer",
-            "answer": (
-                "No chunk fits within the current token budget. "
-                f"Increase the budget to at least {smallest_chunk.tokens} tokens or choose a broader budget."
-            ),
-        }
-    return {"action_type": "submit_answer", "answer": "No usable evidence was available."}
-
-
+
+    if selected_chunks:
+        return {
+            "action_type": "submit_answer",
+            "answer": "Optimized answer based on the currently selected evidence.",
+        }
+    if available:
+        smallest_chunk = min(available, key=lambda chunk: (chunk.tokens, chunk.chunk_id))
+        return {
+            "action_type": "submit_answer",
+            "answer": (
+                "No chunk fits within the current token budget. "
+                f"Increase the budget to at least {smallest_chunk.tokens} tokens or choose a broader budget."
+            ),
+        }
+    return {"action_type": "submit_answer", "answer": "No usable evidence was available."}
+
+
 @app.post("/reset")
 async def reset_endpoint(payload: ResetRequest | None = Body(default=None)):
     payload = payload or ResetRequest()
     if payload.task_name not in TASKS_BY_NAME:
         raise HTTPException(status_code=400, detail="Unknown task_name.")
+
     env = RagContextOptimizerEnv(
         task_name=payload.task_name,
         query_override=payload.custom_query,
@@ -597,49 +269,46 @@ async def reset_endpoint(payload: ResetRequest | None = Body(default=None)):
         max_steps_override=payload.max_steps,
         corpus_family_override=payload.corpus_family,
     )
-    app.state.env = env
-    result = await env.reset()
-    return _serialize_step_result(result, reset=True)
-
-
-@app.post("/step")
-async def step_endpoint(action: RagAction):
-    env = getattr(app.state, "env", None)
-    if env is None:
-        raise HTTPException(status_code=400, detail="Environment is not initialized. Call /reset first.")
-
-    result = await env.step(action)
-    event = (result.info or {}).get("event")
-    if _is_bad_action_event(event):
-        raise HTTPException(status_code=400, detail=event)
-    return _serialize_step_result(result, reset=False)
-
-
-@app.get("/state")
-async def state_endpoint():
-    env = getattr(app.state, "env", None)
-    if env is None:
-        raise HTTPException(status_code=400, detail="Environment is not initialized.")
-    return await env.state()
-
-
-@app.get("/health")
-async def health_endpoint():
-    return {"status": "ok", "tasks": [task.name for task in ALL_TASKS]}
-
-
-@app.get("/tasks")
+    result = await env.reset()
+    episode_id = await _episode_store().create(env)
+    return _serialize_step_result(result, reset=True, episode_id=episode_id)
+
+
+@app.post("/step")
+async def step_endpoint(action: RagAction, episode_id: str | None = None):
+    resolved_episode_id, env = _resolve_env(episode_id)
+    result = await env.step(action)
+    event = (result.info or {}).get("event")
+    if _is_bad_action_event(event):
+        raise HTTPException(status_code=400, detail=event)
+    return _serialize_step_result(result, reset=False, episode_id=resolved_episode_id)
+
+
+@app.get("/state")
+async def state_endpoint(episode_id: str | None = None):
+    resolved_episode_id, env = _resolve_env(episode_id)
+    state = await env.state()
+    state["episode_id"] = resolved_episode_id
+    return state
+
+
+@app.get("/health")
+async def health_endpoint():
+    return {"status": "ok", "tasks": [task.name for task in ALL_TASKS]}
+
+
+@app.get("/tasks")
 async def tasks_endpoint():
-    return [
-        {
-            "name": task.name,
-            "description": task.description,
-            "difficulty": task.difficulty,
-            "token_budget": task.token_budget,
-            "query": task.query,
-            "max_steps": task.max_steps,
-        }
-        for task in ALL_TASKS
+    return [
+        {
+            "name": task.name,
+            "description": task.description,
+            "difficulty": task.difficulty,
+            "token_budget": task.token_budget,
+            "query": task.query,
+            "max_steps": task.max_steps,
+        }
+        for task in ALL_TASKS
     ]
 
 
@@ -649,13 +318,11 @@ async def corpus_families_endpoint():
 
 
 @app.post("/optimize-step")
-async def optimize_step_endpoint():
-    env = getattr(app.state, "env", None)
-    if env is None:
-        raise HTTPException(status_code=400, detail="Environment is not initialized. Call /reset first.")
-    return _suggest_action(env)
-
-
+async def optimize_step_endpoint(episode_id: str | None = None):
+    _resolved_episode_id, env = _resolve_env(episode_id)
+    return _suggest_action(env)
+
+
 @app.post("/optimize-prompt")
 async def optimize_prompt_endpoint(payload: OptimizePromptRequest):
     if not payload.prompt.strip():
@@ -665,9 +332,9 @@ async def optimize_prompt_endpoint(payload: OptimizePromptRequest):
         corpus_family=payload.corpus_family,
         compression_mode=payload.compression_mode,
     )
-
-
-if __name__ == "__main__":
-    import uvicorn
-
-    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=False)
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=False)

env/environment.py

@@ -1,532 +1,533 @@
-"""
-Main OpenEnv-style environment for rag-context-optimizer.
-"""
-
-from __future__ import annotations
-
-from dataclasses import asdict, dataclass, is_dataclass, replace
-import os
-from pathlib import Path
-import re
-from typing import Any
-
-from env.corpus import Chunk, load_corpus, resolve_corpus_path
-from env.context_tuner import ContextTunedPlanner
-from env.graders import TaskGrader
-from env.models import ChunkSummary, RagAction, RagObservation
-from env.retriever import HybridRetriever
-from env.tasks import ALL_TASKS, TASKS_BY_NAME, Task
-
-
-@dataclass(slots=True)
-class StepResult:
-    observation: RagObservation
-    reward: float
-    done: bool
-    info: dict[str, Any]
-
-
-class RagContextOptimizerEnv:
-    _PROJECT_STOPWORDS = {
-        "the", "and", "for", "with", "that", "this", "from", "into", "your", "have", "will",
-        "using", "used", "use", "into", "they", "them", "their", "about", "while", "where",
-        "when", "what", "which", "should", "would", "could", "there", "here", "then", "than",
-        "each", "such", "only", "also", "been", "being", "does", "did", "done", "just", "more",
-        "most", "very", "over", "under", "like", "same", "across", "because", "through", "make",
-        "made", "many", "much", "some", "into", "onto", "must", "need", "needs", "task", "tasks",
-        "chunk", "chunks", "query", "prompt", "environment", "agent", "agents", "model", "models",
-    }
-    _PROJECT_QUERY_HINTS = {
-        "openenv", "benchmark", "rag-context-optimizer", "readme", "docker", "fastapi", "api",
-        "endpoint", "inference.py", "app.py", "tasks.py", "graders.py", "environment.py", "repo",
-        "repository", "codebase", "ui", "frontend", "backend", "space", "validator",
-    }
-
-    def __init__(
-        self,
-        task_name: str = "single_domain_qa",
-        query_override: str | None = None,
-        token_budget_override: int | None = None,
-        max_steps_override: int | None = None,
-        corpus_family_override: str | None = None,
-    ):
-        if task_name not in TASKS_BY_NAME:
-            raise ValueError(f"Unknown task_name: {task_name}")
-
+"""
+Main OpenEnv-style environment for rag-context-optimizer.
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, is_dataclass, replace
+import os
+from pathlib import Path
+import re
+from typing import Any
+
+from env.corpus import Chunk, load_corpus, resolve_corpus_path
+from env.context_tuner import ContextTunedPlanner
+from env.graders import TaskGrader
+from env.models import ChunkSummary, RagAction, RagObservation
+from env.retriever import HybridRetriever
+from env.tasks import ALL_TASKS, TASKS_BY_NAME, Task
+
+
+@dataclass(slots=True)
+class StepResult:
+    observation: RagObservation
+    reward: float
+    done: bool
+    info: dict[str, Any]
+
+
+class RagContextOptimizerEnv:
+    _PROJECT_STOPWORDS = {
+        "the", "and", "for", "with", "that", "this", "from", "into", "your", "have", "will",
+        "using", "used", "use", "into", "they", "them", "their", "about", "while", "where",
+        "when", "what", "which", "should", "would", "could", "there", "here", "then", "than",
+        "each", "such", "only", "also", "been", "being", "does", "did", "done", "just", "more",
+        "most", "very", "over", "under", "like", "same", "across", "because", "through", "make",
+        "made", "many", "much", "some", "into", "onto", "must", "need", "needs", "task", "tasks",
+        "chunk", "chunks", "query", "prompt", "environment", "agent", "agents", "model", "models",
+    }
+    _PROJECT_QUERY_HINTS = {
+        "openenv", "benchmark", "rag-context-optimizer", "readme", "docker", "fastapi", "api",
+        "endpoint", "inference.py", "app.py", "tasks.py", "graders.py", "environment.py", "repo",
+        "repository", "codebase", "ui", "frontend", "backend", "space", "validator",
+    }
+
+    def __init__(
+        self,
+        task_name: str = "single_domain_qa",
+        query_override: str | None = None,
+        token_budget_override: int | None = None,
+        max_steps_override: int | None = None,
+        corpus_family_override: str | None = None,
+    ):
+        if task_name not in TASKS_BY_NAME:
+            raise ValueError(f"Unknown task_name: {task_name}")
+
         self._corpus_family = corpus_family_override or os.getenv("RAG_CORPUS_FAMILY") or "enterprise_v1"
         explicit_path = os.getenv("RAG_CORPUS_PATH")
         self._corpus_path = resolve_corpus_path(explicit_path, family=None if explicit_path else self._corpus_family)
         self._all_chunks = load_corpus(self._corpus_path)
         self._query_overridden = bool(query_override and query_override.strip())
-        self._project_chunks = self._load_project_chunks()
+        self._include_project_chunks = os.getenv("ENABLE_PROJECT_CORPUS", "").strip().lower() in {"1", "true", "yes"}
+        self._project_chunks = self._load_project_chunks() if self._include_project_chunks else []
         self.retriever = HybridRetriever(self._all_chunks + self._project_chunks)
-        self.context_tuner = ContextTunedPlanner(
-            self.retriever,
-            self._all_chunks + self._project_chunks,
-            list(ALL_TASKS),
-        )
-        self.grader = TaskGrader()
-        self.task: Task = self._build_task(
-            TASKS_BY_NAME[task_name],
-            query_override=query_override,
-            token_budget_override=token_budget_override,
-            max_steps_override=max_steps_override,
-        )
-
-        self._available_chunks: list[Chunk] = []
-        self._selected_chunks: list[str] = []
-        self._compression_ratios: dict[str, float] = {}
-        self._step_number = 0
-        self._done = False
-        self._last_action_feedback: str | None = None
-        self._last_answer = ""
-        self._last_tuning = None
-
-    @staticmethod
-    def _build_task(
-        base_task: Task,
-        query_override: str | None = None,
-        token_budget_override: int | None = None,
-        max_steps_override: int | None = None,
-    ) -> Task:
-        updated_task = base_task
-        if query_override and query_override.strip():
-            updated_task = replace(updated_task, query=query_override.strip(), domain_filter=None)
-        if token_budget_override is not None and token_budget_override > 0:
-            updated_task = replace(updated_task, token_budget=token_budget_override)
-        if max_steps_override is not None and max_steps_override > 0:
-            updated_task = replace(updated_task, max_steps=max_steps_override)
-        return updated_task
-
-    async def reset(self) -> StepResult:
-        candidate_chunks = self._filter_chunks_for_task(self.task)
-        self._available_chunks = self._rank_chunks_for_query(self.task.query, candidate_chunks)
-        if not self._query_overridden:
-            chunk_by_id = {chunk.chunk_id: chunk for chunk in candidate_chunks}
-            for chunk_id in self.task.required_chunk_ids:
-                chunk = chunk_by_id.get(chunk_id)
-                if chunk and all(existing.chunk_id != chunk_id for existing in self._available_chunks):
-                    self._available_chunks.append(chunk)
-        self._selected_chunks = []
-        self._compression_ratios = {}
-        self._step_number = 0
-        self._done = False
-        self._last_action_feedback = None
-        self._last_answer = ""
-
-        observation = self._build_observation()
-        return StepResult(
-            observation=observation,
-            reward=0.0,
-            done=False,
-            info={"task": self.task.name, "event": "reset"},
-        )
-
-    async def step(self, action: RagAction) -> StepResult:
-        if self._done:
-            return StepResult(
-                observation=self._build_observation(),
-                reward=0.0,
-                done=True,
-                info={"task": self.task.name, "event": "episode_already_done"},
-            )
-
-        reward = 0.0
-        info: dict[str, Any] = {"task": self.task.name, "action_type": action.action_type}
-
-        if action.action_type == "select_chunk":
-            reward, info = self._handle_select(action.chunk_id or "")
-        elif action.action_type == "deselect_chunk":
-            reward, info = self._handle_deselect(action.chunk_id or "")
-        elif action.action_type == "compress_chunk":
-            reward, info = self._handle_compress(action.chunk_id or "", float(action.compression_ratio or 0.0))
-        elif action.action_type == "submit_answer":
-            self._last_answer = action.answer or ""
-            result = self._finalize_submission(reason="submit_answer")
-            self._step_number += 1
-            result.observation.step_number = self._step_number
-            return result
-
-        self._step_number += 1
-
-        if self._step_number >= self.task.max_steps:
-            return self._finalize_submission(reason="max_steps_reached")
-
-        observation = self._build_observation()
-        return StepResult(
-            observation=observation,
-            reward=reward,
-            done=False,
-            info=info,
-        )
-
-    async def state(self) -> dict:
-        selected_chunk_details = []
-        for chunk_id in self._selected_chunks:
-            chunk = self._chunk_map().get(chunk_id)
-            if chunk is None:
-                continue
-            selected_chunk_details.append(
-                {
-                    "chunk_id": chunk.chunk_id,
-                    "domain": chunk.domain,
-                    "original_tokens": chunk.tokens,
-                    "effective_tokens": self._effective_chunk_tokens(chunk_id),
-                    "compression_ratio": round(self._compression_ratios.get(chunk_id, 1.0), 3),
-                    "text": self._effective_chunk_text(chunk_id),
-                    "keywords": chunk.keywords,
-                }
-            )
-        optimized_prompt = self._build_optimized_prompt()
-        return {
-            "task": asdict(self.task) if is_dataclass(self.task) else self.task,
-            "step_number": self._step_number,
-            "done": self._done,
-            "selected_chunks": list(self._selected_chunks),
-            "compression_ratios": dict(self._compression_ratios),
-            "total_tokens_used": self._total_tokens_used(),
-            "token_budget": self.task.token_budget,
-            "last_action_feedback": self._last_action_feedback,
-            "last_answer": self._last_answer,
-            "corpus_family": self._corpus_family,
-            "corpus_path": str(self._corpus_path),
-            "available_chunk_ids": [chunk.chunk_id for chunk in self._available_chunks],
-            "selected_chunk_details": selected_chunk_details,
-            "optimized_prompt_preview": optimized_prompt,
-            "optimized_prompt_tokens": max(1, len(optimized_prompt) // 4) if optimized_prompt else 0,
-            "context_tuning": (
-                {
-                    "mode": self._last_tuning.mode,
-                    "top_demo_cases": self._last_tuning.top_demo_cases,
-                    "suggested_citations": self._last_tuning.suggested_citations,
-                    "token_dropout": self._last_tuning.token_dropout,
-                    "leave_one_out": self._last_tuning.leave_one_out,
-                }
-                if self._last_tuning is not None
-                else None
-            ),
-        }
-
-    async def close(self):
-        self._done = True
-
-    def _filter_chunks_for_task(self, task: Task) -> list[Chunk]:
-        domain_mapping = {
-            "customer_support_operations": "Customer Support Operations",
-            "incident_response_playbooks": "Incident Response Playbooks",
-            "platform_reliability_release_engineering": "Platform Reliability & Release Engineering",
+        self.context_tuner = ContextTunedPlanner(
+            self.retriever,
+            self._all_chunks + self._project_chunks,
+            list(ALL_TASKS),
+        )
+        self.grader = TaskGrader()
+        self.task: Task = self._build_task(
+            TASKS_BY_NAME[task_name],
+            query_override=query_override,
+            token_budget_override=token_budget_override,
+            max_steps_override=max_steps_override,
+        )
+
+        self._available_chunks: list[Chunk] = []
+        self._selected_chunks: list[str] = []
+        self._compression_ratios: dict[str, float] = {}
+        self._step_number = 0
+        self._done = False
+        self._last_action_feedback: str | None = None
+        self._last_answer = ""
+        self._last_tuning = None
+
+    @staticmethod
+    def _build_task(
+        base_task: Task,
+        query_override: str | None = None,
+        token_budget_override: int | None = None,
+        max_steps_override: int | None = None,
+    ) -> Task:
+        updated_task = base_task
+        if query_override and query_override.strip():
+            updated_task = replace(updated_task, query=query_override.strip(), domain_filter=None)
+        if token_budget_override is not None and token_budget_override > 0:
+            updated_task = replace(updated_task, token_budget=token_budget_override)
+        if max_steps_override is not None and max_steps_override > 0:
+            updated_task = replace(updated_task, max_steps=max_steps_override)
+        return updated_task
+
+    async def reset(self) -> StepResult:
+        candidate_chunks = self._filter_chunks_for_task(self.task)
+        self._available_chunks = self._rank_chunks_for_query(self.task.query, candidate_chunks)
+        if not self._query_overridden:
+            chunk_by_id = {chunk.chunk_id: chunk for chunk in candidate_chunks}
+            for chunk_id in self.task.required_chunk_ids:
+                chunk = chunk_by_id.get(chunk_id)
+                if chunk and all(existing.chunk_id != chunk_id for existing in self._available_chunks):
+                    self._available_chunks.append(chunk)
+        self._selected_chunks = []
+        self._compression_ratios = {}
+        self._step_number = 0
+        self._done = False
+        self._last_action_feedback = None
+        self._last_answer = ""
+
+        observation = self._build_observation()
+        return StepResult(
+            observation=observation,
+            reward=0.0,
+            done=False,
+            info={"task": self.task.name, "event": "reset"},
+        )
+
+    async def step(self, action: RagAction) -> StepResult:
+        if self._done:
+            return StepResult(
+                observation=self._build_observation(),
+                reward=0.0,
+                done=True,
+                info={"task": self.task.name, "event": "episode_already_done"},
+            )
+
+        reward = 0.0
+        info: dict[str, Any] = {"task": self.task.name, "action_type": action.action_type}
+
+        if action.action_type == "select_chunk":
+            reward, info = self._handle_select(action.chunk_id or "")
+        elif action.action_type == "deselect_chunk":
+            reward, info = self._handle_deselect(action.chunk_id or "")
+        elif action.action_type == "compress_chunk":
+            reward, info = self._handle_compress(action.chunk_id or "", float(action.compression_ratio or 0.0))
+        elif action.action_type == "submit_answer":
+            self._last_answer = action.answer or ""
+            result = self._finalize_submission(reason="submit_answer")
+            self._step_number += 1
+            result.observation.step_number = self._step_number
+            return result
+
+        self._step_number += 1
+
+        if self._step_number >= self.task.max_steps:
+            return self._finalize_submission(reason="max_steps_reached")
+
+        observation = self._build_observation()
+        return StepResult(
+            observation=observation,
+            reward=reward,
+            done=False,
+            info=info,
+        )
+
+    async def state(self) -> dict:
+        selected_chunk_details = []
+        for chunk_id in self._selected_chunks:
+            chunk = self._chunk_map().get(chunk_id)
+            if chunk is None:
+                continue
+            selected_chunk_details.append(
+                {
+                    "chunk_id": chunk.chunk_id,
+                    "domain": chunk.domain,
+                    "original_tokens": chunk.tokens,
+                    "effective_tokens": self._effective_chunk_tokens(chunk_id),
+                    "compression_ratio": round(self._compression_ratios.get(chunk_id, 1.0), 3),
+                    "text": self._effective_chunk_text(chunk_id),
+                    "keywords": chunk.keywords,
+                }
+            )
+        optimized_prompt = self._build_optimized_prompt()
+        return {
+            "task": asdict(self.task) if is_dataclass(self.task) else self.task,
+            "step_number": self._step_number,
+            "done": self._done,
+            "selected_chunks": list(self._selected_chunks),
+            "compression_ratios": dict(self._compression_ratios),
+            "total_tokens_used": self._total_tokens_used(),
+            "token_budget": self.task.token_budget,
+            "last_action_feedback": self._last_action_feedback,
+            "last_answer": self._last_answer,
+            "corpus_family": self._corpus_family,
+            "corpus_path": str(self._corpus_path),
+            "available_chunk_ids": [chunk.chunk_id for chunk in self._available_chunks],
+            "selected_chunk_details": selected_chunk_details,
+            "optimized_prompt_preview": optimized_prompt,
+            "optimized_prompt_tokens": max(1, len(optimized_prompt) // 4) if optimized_prompt else 0,
+            "context_tuning": (
+                {
+                    "mode": self._last_tuning.mode,
+                    "top_demo_cases": self._last_tuning.top_demo_cases,
+                    "suggested_citations": self._last_tuning.suggested_citations,
+                    "token_dropout": self._last_tuning.token_dropout,
+                    "leave_one_out": self._last_tuning.leave_one_out,
+                }
+                if self._last_tuning is not None
+                else None
+            ),
+        }
+
+    async def close(self):
+        self._done = True
+
+    def _filter_chunks_for_task(self, task: Task) -> list[Chunk]:
+        domain_mapping = {
+            "customer_support_operations": "Customer Support Operations",
+            "incident_response_playbooks": "Incident Response Playbooks",
+            "platform_reliability_release_engineering": "Platform Reliability & Release Engineering",
         }
         if self._query_overridden:
-            if self._is_project_query(task.query):
+            if self._include_project_chunks and self._is_project_query(task.query):
                 return list(self._all_chunks) + list(self._project_chunks)
             return list(self._all_chunks)
-        if task.domain_filter is None:
-            return list(self._all_chunks)
-        normalized = domain_mapping.get(task.domain_filter, task.domain_filter)
-        return [chunk for chunk in self._all_chunks if chunk.domain == normalized]
-
-    def _is_project_query(self, query: str) -> bool:
-        lowered = query.lower()
-        return any(hint in lowered for hint in self._PROJECT_QUERY_HINTS)
-
-    def _rank_chunks_for_query(self, query: str, chunks: list[Chunk], top_k: int = 20) -> list[Chunk]:
-        tuning = self.context_tuner.tune(query, chunks)
-        self._last_tuning = tuning
-        scored = []
+        if task.domain_filter is None:
+            return list(self._all_chunks)
+        normalized = domain_mapping.get(task.domain_filter, task.domain_filter)
+        return [chunk for chunk in self._all_chunks if chunk.domain == normalized]
+
+    def _is_project_query(self, query: str) -> bool:
+        lowered = query.lower()
+        return any(hint in lowered for hint in self._PROJECT_QUERY_HINTS)
+
+    def _rank_chunks_for_query(self, query: str, chunks: list[Chunk], top_k: int = 20) -> list[Chunk]:
+        tuning = self.context_tuner.tune(query, chunks)
+        self._last_tuning = tuning
+        scored = []
         for chunk in chunks:
             tuned = tuning.tuned_scores.get(chunk.chunk_id)
             score = tuned.final_score if tuned is not None else self.retriever.hybrid_score(query, chunk)
-            if self._query_overridden and chunk.domain.startswith("Project"):
+            if self._include_project_chunks and self._query_overridden and chunk.domain.startswith("Project"):
                 score = min(1.0, score + 0.08)
             scored.append((chunk, score))
-        scored.sort(key=lambda item: (-item[1], item[0].tokens, item[0].chunk_id))
-        if not scored:
-            return []
-
-        capped = scored[: max(1, min(top_k * 2, len(scored)))]
-        best_score = capped[0][1]
-        floor = max(0.12, best_score * 0.38)
-        filtered_pairs = [(chunk, score) for chunk, score in capped if score >= floor]
-
-        if self._query_overridden:
+        scored.sort(key=lambda item: (-item[1], item[0].tokens, item[0].chunk_id))
+        if not scored:
+            return []
+
+        capped = scored[: max(1, min(top_k * 2, len(scored)))]
+        best_score = capped[0][1]
+        floor = max(0.12, best_score * 0.38)
+        filtered_pairs = [(chunk, score) for chunk, score in capped if score >= floor]
+
+        if self._include_project_chunks and self._query_overridden:
             project_pairs = [(chunk, score) for chunk, score in filtered_pairs if chunk.domain.startswith("Project")]
             if len(project_pairs) >= 4:
                 filtered_pairs = project_pairs + [
-                    (chunk, score)
-                    for chunk, score in filtered_pairs
-                    if not chunk.domain.startswith("Project")
-                ]
-
-        filtered = [chunk for chunk, _score in filtered_pairs]
-        if not filtered:
-            filtered = [chunk for chunk, _score in capped[: max(1, min(top_k, len(capped)))]]
-
-        return filtered[: max(1, min(top_k, len(filtered)))]
-
-    def _load_project_chunks(self) -> list[Chunk]:
-        root = Path(__file__).resolve().parent.parent
-        chunks: list[Chunk] = []
-        file_specs = [
-            ("Project Documentation", root / "README.md", ["project_docs", "readme"]),
-            ("Project Configuration", root / "openenv.yaml", ["project_docs", "config", "openenv_spec"]),
-            ("Project API", root / "app.py", ["project_docs", "api", "server"]),
-            ("Project Baseline", root / "inference.py", ["project_docs", "baseline", "inference"]),
-            ("Project Environment", root / "env" / "environment.py", ["project_docs", "environment", "state_management"]),
-            ("Project Retrieval", root / "env" / "retriever.py", ["project_docs", "retrieval", "ranking"]),
-            ("Project Grading", root / "env" / "graders.py", ["project_docs", "grading", "reward_design"]),
-            ("Project Tasks", root / "env" / "tasks.py", ["project_docs", "tasks", "difficulty"]),
-            ("Project Validation", root / "validate.py", ["project_docs", "validation", "testing"]),
-        ]
-
-        for domain, path, tags in file_specs:
-            if not path.exists():
-                continue
-            raw_text = path.read_text(encoding="utf-8", errors="ignore")
-            sections = self._chunk_project_text(raw_text)
-            stem = re.sub(r"[^a-z0-9]+", "_", path.stem.lower()).strip("_") or "file"
-            for index, section in enumerate(sections, start=1):
-                keywords = self._extract_project_keywords(section)
-                if not keywords:
-                    keywords = [stem, domain.lower()]
-                chunks.append(
-                    Chunk(
-                        chunk_id=f"project_{stem}_{index:03d}",
-                        domain=domain,
-                        text=section,
-                        tokens=max(30, len(section) // 4),
-                        keywords=keywords[:5],
-                        relevance_tags=tags,
-                    )
-                )
-        return chunks
-
-    def _chunk_project_text(self, raw_text: str, chunk_words: int = 140, stride_words: int = 100) -> list[str]:
-        cleaned = " ".join(raw_text.split())
-        words = cleaned.split()
-        if not words:
-            return []
-        if len(words) <= chunk_words:
-            return [" ".join(words)]
-
-        chunks: list[str] = []
-        start = 0
-        while start < len(words):
-            window = words[start : start + chunk_words]
-            if not window:
-                break
-            chunks.append(" ".join(window))
-            if start + chunk_words >= len(words):
-                break
-            start += stride_words
-        return chunks
-
-    def _extract_project_keywords(self, text: str) -> list[str]:
-        terms = re.findall(r"[a-z0-9_]+", text.lower())
-        counts: dict[str, int] = {}
-        for term in terms:
-            if len(term) < 4 or term in self._PROJECT_STOPWORDS:
-                continue
-            counts[term] = counts.get(term, 0) + 1
-        ranked = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
-        return [term.replace("_", " ") for term, _count in ranked[:8]]
-
-    def _build_observation(self) -> RagObservation:
-        return RagObservation(
-            query=self.task.query,
-            available_chunks=[
-                ChunkSummary(
-                    chunk_id=chunk.chunk_id,
-                    domain=chunk.domain,
-                    tokens=self._effective_chunk_tokens(chunk.chunk_id),
-                    keywords=chunk.keywords,
-                )
-                for chunk in self._available_chunks
-            ],
-            selected_chunks=list(self._selected_chunks),
-            total_tokens_used=self._total_tokens_used(),
-            token_budget=self.task.token_budget,
-            step_number=self._step_number,
-            task_name=self.task.name,
-            last_action_feedback=self._last_action_feedback,
-        )
-
-    def _chunk_map(self) -> dict[str, Chunk]:
-        return {chunk.chunk_id: chunk for chunk in self._available_chunks}
-
-    def _effective_chunk_tokens(self, chunk_id: str) -> int:
-        chunk = self._chunk_map().get(chunk_id)
-        if chunk is None:
-            return 0
-        ratio = self._compression_ratios.get(chunk_id, 1.0)
-        return max(1, int(round(chunk.tokens * ratio)))
-
-    def _total_tokens_used(self) -> int:
-        return sum(self._effective_chunk_tokens(chunk_id) for chunk_id in self._selected_chunks)
-
-    def _effective_chunk_text(self, chunk_id: str) -> str:
-        chunk = self._chunk_map().get(chunk_id)
-        if chunk is None:
-            return ""
-        ratio = self._compression_ratios.get(chunk_id, 1.0)
-        text = " ".join(chunk.text.split())
-        if ratio >= 0.999:
-            return text
-
-        query_terms = self._query_terms(self.task.query)
-        keyword_terms = self._query_terms(" ".join(chunk.keywords))
-        sentences = [segment.strip() for segment in re.split(r"(?<=[.!?])\s+", text) if segment.strip()]
-        if not sentences:
-            return self._truncate_words(text, ratio)
-
-        ranked_sentences: list[tuple[int, float, int, str]] = []
-        for index, sentence in enumerate(sentences):
-            sentence_terms = self._query_terms(sentence)
-            overlap = len(sentence_terms & query_terms)
-            keyword_overlap = len(sentence_terms & keyword_terms)
-            score = (overlap * 2.0) + keyword_overlap + (0.25 if index == 0 else 0.0)
-            ranked_sentences.append((index, score, len(sentence.split()), sentence))
-
-        target_words = max(20, int(len(text.split()) * ratio))
-        chosen: list[tuple[int, str]] = []
-        used_words = 0
-        for index, _score, word_count, sentence in sorted(
-            ranked_sentences,
-            key=lambda item: (-item[1], item[2], item[0]),
-        ):
-            if used_words >= target_words:
-                break
-            chosen.append((index, sentence))
-            used_words += word_count
-
-        if not chosen:
-            return self._truncate_words(text, ratio)
-
-        chosen.sort(key=lambda item: item[0])
-        compressed = " ".join(sentence for _index, sentence in chosen)
-        return self._truncate_words(compressed, ratio)
-
-    @staticmethod
-    def _truncate_words(text: str, ratio: float) -> str:
-        words = text.split()
-        if not words:
-            return ""
-        keep = max(12, int(len(words) * ratio))
-        truncated = " ".join(words[:keep])
-        if keep < len(words):
-            return truncated + " ..."
-        return truncated
-
-    @staticmethod
-    def _query_terms(text: str) -> set[str]:
-        return {token for token in re.findall(r"[a-z0-9]+", text.lower()) if len(token) > 2}
-
-    def _build_optimized_prompt(self) -> str:
-        if not self._selected_chunks:
-            return ""
-        sections = [f"Question: {self.task.query}", "", "Optimized Context:"]
-        for chunk_id in self._selected_chunks:
-            chunk = self._chunk_map().get(chunk_id)
-            if chunk is None:
-                continue
-            sections.append(
-                f"[{chunk.chunk_id} | {self._effective_chunk_tokens(chunk_id)} tokens] {self._effective_chunk_text(chunk_id)}"
-            )
-        return "\n".join(sections).strip()
-
-    def _is_relevant(self, chunk_id: str) -> tuple[bool, float]:
-        chunk = self._chunk_map().get(chunk_id)
-        if chunk is None:
-            return False, 0.0
-        score = self.retriever.hybrid_score(self.task.query, chunk)
-        return score >= 0.3, score
-
-    def _handle_select(self, chunk_id: str) -> tuple[float, dict[str, Any]]:
-        chunk = self._chunk_map().get(chunk_id)
-        if chunk is None:
-            self._last_action_feedback = "chunk_not_found"
-            return -0.1, {"event": "chunk_not_found"}
-        if chunk_id in self._selected_chunks:
-            self._last_action_feedback = "chunk_already_selected"
-            return 0.0, {"event": "chunk_already_selected"}
-
-        projected_tokens = self._total_tokens_used() + self._effective_chunk_tokens(chunk_id)
-        if projected_tokens > self.task.token_budget:
-            self._last_action_feedback = "exceeded_budget"
-            return -0.1, {"event": "exceeded_budget", "chunk_id": chunk_id}
-
-        self._selected_chunks.append(chunk_id)
-        _, score = self._is_relevant(chunk_id)
-        self._last_action_feedback = "chunk_selected"
-        return score * 0.2, {"event": "chunk_selected", "chunk_id": chunk_id, "hybrid_score": score}
-
-    def _handle_deselect(self, chunk_id: str) -> tuple[float, dict[str, Any]]:
-        if chunk_id not in self._selected_chunks:
-            self._last_action_feedback = "chunk_not_selected"
-            return 0.0, {"event": "chunk_not_selected", "chunk_id": chunk_id}
-
-        self._selected_chunks.remove(chunk_id)
-        is_relevant, score = self._is_relevant(chunk_id)
-        self._last_action_feedback = "chunk_deselected"
-        reward = 0.0 if is_relevant else 0.05
-        return reward, {"event": "chunk_deselected", "chunk_id": chunk_id, "hybrid_score": score}
-
-    def _handle_compress(self, chunk_id: str, compression_ratio: float) -> tuple[float, dict[str, Any]]:
-        chunk = self._chunk_map().get(chunk_id)
-        if chunk is None:
-            self._last_action_feedback = "chunk_not_found"
-            return -0.1, {"event": "chunk_not_found", "chunk_id": chunk_id}
-
-        self._compression_ratios[chunk_id] = compression_ratio
-        is_relevant, score = self._is_relevant(chunk_id)
-        reward = 0.03 if is_relevant else 0.0
-        if score >= 0.6 and compression_ratio < 0.4:
-            reward -= 0.05
-            self._last_action_feedback = "overcompressed_relevant_chunk"
-            return reward, {
-                "event": "overcompressed_relevant_chunk",
-                "chunk_id": chunk_id,
-                "hybrid_score": score,
-                "compression_ratio": compression_ratio,
-            }
-
-        self._last_action_feedback = "chunk_compressed"
-        return reward, {
-            "event": "chunk_compressed",
-            "chunk_id": chunk_id,
-            "hybrid_score": score,
-            "compression_ratio": compression_ratio,
-        }
-
-    def _finalize_submission(self, reason: str) -> StepResult:
-        self._done = True
-
-        if not self._selected_chunks:
-            self._last_action_feedback = "no_chunks_selected"
-            observation = self._build_observation()
-            return StepResult(
-                observation=observation,
-                reward=0.0,
-                done=True,
-                info={"event": reason, "grader": None, "passed": False},
-            )
-
-        grader_result = self.grader.grade(
-            selected_chunk_ids=list(self._selected_chunks),
-            answer=self._last_answer,
-            token_budget=self.task.token_budget,
-            total_tokens_used=self._total_tokens_used(),
-            retriever=self.retriever,
-            task=self.task,
-        )
-        self._last_action_feedback = reason
-        observation = self._build_observation()
-        return StepResult(
-            observation=observation,
-            reward=grader_result.score,
-            done=True,
-            info={
-                "event": reason,
-                "grader": grader_result.breakdown,
-                "passed": grader_result.passed,
-            },
-        )
+                    (chunk, score)
+                    for chunk, score in filtered_pairs
+                    if not chunk.domain.startswith("Project")
+                ]
+
+        filtered = [chunk for chunk, _score in filtered_pairs]
+        if not filtered:
+            filtered = [chunk for chunk, _score in capped[: max(1, min(top_k, len(capped)))]]
+
+        return filtered[: max(1, min(top_k, len(filtered)))]
+
+    def _load_project_chunks(self) -> list[Chunk]:
+        root = Path(__file__).resolve().parent.parent
+        chunks: list[Chunk] = []
+        file_specs = [
+            ("Project Documentation", root / "README.md", ["project_docs", "readme"]),
+            ("Project Configuration", root / "openenv.yaml", ["project_docs", "config", "openenv_spec"]),
+            ("Project API", root / "app.py", ["project_docs", "api", "server"]),
+            ("Project Baseline", root / "inference.py", ["project_docs", "baseline", "inference"]),
+            ("Project Environment", root / "env" / "environment.py", ["project_docs", "environment", "state_management"]),
+            ("Project Retrieval", root / "env" / "retriever.py", ["project_docs", "retrieval", "ranking"]),
+            ("Project Grading", root / "env" / "graders.py", ["project_docs", "grading", "reward_design"]),
+            ("Project Tasks", root / "env" / "tasks.py", ["project_docs", "tasks", "difficulty"]),
+            ("Project Validation", root / "validate.py", ["project_docs", "validation", "testing"]),
+        ]
+
+        for domain, path, tags in file_specs:
+            if not path.exists():
+                continue
+            raw_text = path.read_text(encoding="utf-8", errors="ignore")
+            sections = self._chunk_project_text(raw_text)
+            stem = re.sub(r"[^a-z0-9]+", "_", path.stem.lower()).strip("_") or "file"
+            for index, section in enumerate(sections, start=1):
+                keywords = self._extract_project_keywords(section)
+                if not keywords:
+                    keywords = [stem, domain.lower()]
+                chunks.append(
+                    Chunk(
+                        chunk_id=f"project_{stem}_{index:03d}",
+                        domain=domain,
+                        text=section,
+                        tokens=max(30, len(section) // 4),
+                        keywords=keywords[:5],
+                        relevance_tags=tags,
+                    )
+                )
+        return chunks
+
+    def _chunk_project_text(self, raw_text: str, chunk_words: int = 140, stride_words: int = 100) -> list[str]:
+        cleaned = " ".join(raw_text.split())
+        words = cleaned.split()
+        if not words:
+            return []
+        if len(words) <= chunk_words:
+            return [" ".join(words)]
+
+        chunks: list[str] = []
+        start = 0
+        while start < len(words):
+            window = words[start : start + chunk_words]
+            if not window:
+                break
+            chunks.append(" ".join(window))
+            if start + chunk_words >= len(words):
+                break
+            start += stride_words
+        return chunks
+
+    def _extract_project_keywords(self, text: str) -> list[str]:
+        terms = re.findall(r"[a-z0-9_]+", text.lower())
+        counts: dict[str, int] = {}
+        for term in terms:
+            if len(term) < 4 or term in self._PROJECT_STOPWORDS:
+                continue
+            counts[term] = counts.get(term, 0) + 1
+        ranked = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+        return [term.replace("_", " ") for term, _count in ranked[:8]]
+
+    def _build_observation(self) -> RagObservation:
+        return RagObservation(
+            query=self.task.query,
+            available_chunks=[
+                ChunkSummary(
+                    chunk_id=chunk.chunk_id,
+                    domain=chunk.domain,
+                    tokens=self._effective_chunk_tokens(chunk.chunk_id),
+                    keywords=chunk.keywords,
+                )
+                for chunk in self._available_chunks
+            ],
+            selected_chunks=list(self._selected_chunks),
+            total_tokens_used=self._total_tokens_used(),
+            token_budget=self.task.token_budget,
+            step_number=self._step_number,
+            task_name=self.task.name,
+            last_action_feedback=self._last_action_feedback,
+        )
+
+    def _chunk_map(self) -> dict[str, Chunk]:
+        return {chunk.chunk_id: chunk for chunk in self._available_chunks}
+
+    def _effective_chunk_tokens(self, chunk_id: str) -> int:
+        chunk = self._chunk_map().get(chunk_id)
+        if chunk is None:
+            return 0
+        ratio = self._compression_ratios.get(chunk_id, 1.0)
+        return max(1, int(round(chunk.tokens * ratio)))
+
+    def _total_tokens_used(self) -> int:
+        return sum(self._effective_chunk_tokens(chunk_id) for chunk_id in self._selected_chunks)
+
+    def _effective_chunk_text(self, chunk_id: str) -> str:
+        chunk = self._chunk_map().get(chunk_id)
+        if chunk is None:
+            return ""
+        ratio = self._compression_ratios.get(chunk_id, 1.0)
+        text = " ".join(chunk.text.split())
+        if ratio >= 0.999:
+            return text
+
+        query_terms = self._query_terms(self.task.query)
+        keyword_terms = self._query_terms(" ".join(chunk.keywords))
+        sentences = [segment.strip() for segment in re.split(r"(?<=[.!?])\s+", text) if segment.strip()]
+        if not sentences:
+            return self._truncate_words(text, ratio)
+
+        ranked_sentences: list[tuple[int, float, int, str]] = []
+        for index, sentence in enumerate(sentences):
+            sentence_terms = self._query_terms(sentence)
+            overlap = len(sentence_terms & query_terms)
+            keyword_overlap = len(sentence_terms & keyword_terms)
+            score = (overlap * 2.0) + keyword_overlap + (0.25 if index == 0 else 0.0)
+            ranked_sentences.append((index, score, len(sentence.split()), sentence))
+
+        target_words = max(20, int(len(text.split()) * ratio))
+        chosen: list[tuple[int, str]] = []
+        used_words = 0
+        for index, _score, word_count, sentence in sorted(
+            ranked_sentences,
+            key=lambda item: (-item[1], item[2], item[0]),
+        ):
+            if used_words >= target_words:
+                break
+            chosen.append((index, sentence))
+            used_words += word_count
+
+        if not chosen:
+            return self._truncate_words(text, ratio)
+
+        chosen.sort(key=lambda item: item[0])
+        compressed = " ".join(sentence for _index, sentence in chosen)
+        return self._truncate_words(compressed, ratio)
+
+    @staticmethod
+    def _truncate_words(text: str, ratio: float) -> str:
+        words = text.split()
+        if not words:
+            return ""
+        keep = max(12, int(len(words) * ratio))
+        truncated = " ".join(words[:keep])
+        if keep < len(words):
+            return truncated + " ..."
+        return truncated
+
+    @staticmethod
+    def _query_terms(text: str) -> set[str]:
+        return {token for token in re.findall(r"[a-z0-9]+", text.lower()) if len(token) > 2}
+
+    def _build_optimized_prompt(self) -> str:
+        if not self._selected_chunks:
+            return ""
+        sections = [f"Question: {self.task.query}", "", "Optimized Context:"]
+        for chunk_id in self._selected_chunks:
+            chunk = self._chunk_map().get(chunk_id)
+            if chunk is None:
+                continue
+            sections.append(
+                f"[{chunk.chunk_id} | {self._effective_chunk_tokens(chunk_id)} tokens] {self._effective_chunk_text(chunk_id)}"
+            )
+        return "\n".join(sections).strip()
+
+    def _is_relevant(self, chunk_id: str) -> tuple[bool, float]:
+        chunk = self._chunk_map().get(chunk_id)
+        if chunk is None:
+            return False, 0.0
+        score = self.retriever.hybrid_score(self.task.query, chunk)
+        return score >= 0.3, score
+
+    def _handle_select(self, chunk_id: str) -> tuple[float, dict[str, Any]]:
+        chunk = self._chunk_map().get(chunk_id)
+        if chunk is None:
+            self._last_action_feedback = "chunk_not_found"
+            return -0.1, {"event": "chunk_not_found"}
+        if chunk_id in self._selected_chunks:
+            self._last_action_feedback = "chunk_already_selected"
+            return 0.0, {"event": "chunk_already_selected"}
+
+        projected_tokens = self._total_tokens_used() + self._effective_chunk_tokens(chunk_id)
+        if projected_tokens > self.task.token_budget:
+            self._last_action_feedback = "exceeded_budget"
+            return -0.1, {"event": "exceeded_budget", "chunk_id": chunk_id}
+
+        self._selected_chunks.append(chunk_id)
+        _, score = self._is_relevant(chunk_id)
+        self._last_action_feedback = "chunk_selected"
+        return score * 0.2, {"event": "chunk_selected", "chunk_id": chunk_id, "hybrid_score": score}
+
+    def _handle_deselect(self, chunk_id: str) -> tuple[float, dict[str, Any]]:
+        if chunk_id not in self._selected_chunks:
+            self._last_action_feedback = "chunk_not_selected"
+            return 0.0, {"event": "chunk_not_selected", "chunk_id": chunk_id}
+
+        self._selected_chunks.remove(chunk_id)
+        is_relevant, score = self._is_relevant(chunk_id)
+        self._last_action_feedback = "chunk_deselected"
+        reward = 0.0 if is_relevant else 0.05
+        return reward, {"event": "chunk_deselected", "chunk_id": chunk_id, "hybrid_score": score}
+
+    def _handle_compress(self, chunk_id: str, compression_ratio: float) -> tuple[float, dict[str, Any]]:
+        chunk = self._chunk_map().get(chunk_id)
+        if chunk is None:
+            self._last_action_feedback = "chunk_not_found"
+            return -0.1, {"event": "chunk_not_found", "chunk_id": chunk_id}
+
+        self._compression_ratios[chunk_id] = compression_ratio
+        is_relevant, score = self._is_relevant(chunk_id)
+        reward = 0.03 if is_relevant else 0.0
+        if score >= 0.6 and compression_ratio < 0.4:
+            reward -= 0.05
+            self._last_action_feedback = "overcompressed_relevant_chunk"
+            return reward, {
+                "event": "overcompressed_relevant_chunk",
+                "chunk_id": chunk_id,
+                "hybrid_score": score,
+                "compression_ratio": compression_ratio,
+            }
+
+        self._last_action_feedback = "chunk_compressed"
+        return reward, {
+            "event": "chunk_compressed",
+            "chunk_id": chunk_id,
+            "hybrid_score": score,
+            "compression_ratio": compression_ratio,
+        }
+
+    def _finalize_submission(self, reason: str) -> StepResult:
+        self._done = True
+
+        if not self._selected_chunks:
+            self._last_action_feedback = "no_chunks_selected"
+            observation = self._build_observation()
+            return StepResult(
+                observation=observation,
+                reward=0.0,
+                done=True,
+                info={"event": reason, "grader": None, "passed": False},
+            )
+
+        grader_result = self.grader.grade(
+            selected_chunk_ids=list(self._selected_chunks),
+            answer=self._last_answer,
+            token_budget=self.task.token_budget,
+            total_tokens_used=self._total_tokens_used(),
+            retriever=self.retriever,
+            task=self.task,
+        )
+        self._last_action_feedback = reason
+        observation = self._build_observation()
+        return StepResult(
+            observation=observation,
+            reward=grader_result.score,
+            done=True,
+            info={
+                "event": reason,
+                "grader": grader_result.breakdown,
+                "passed": grader_result.passed,
+            },
+        )

env/graders.py

@@ -1,124 +1,125 @@
-"""
-Deterministic graders for rag-context-optimizer tasks.
-"""
-
-from __future__ import annotations
-
-import re
-from dataclasses import dataclass
-
-from env.corpus import Chunk
-from env.retriever import HybridRetriever
-from env.tasks import Task
-
-
-_STOPWORDS = {
-    "a", "an", "and", "are", "as", "at", "be", "because", "by", "for", "from", "how",
-    "if", "in", "into", "is", "it", "its", "of", "on", "or", "that", "the", "their",
-    "them", "there", "these", "this", "to", "was", "were", "what", "when", "where",
-    "which", "while", "with", "within", "without", "you", "your",
-}
-
-
-def _tokenize(text: str) -> set[str]:
-    return set(re.findall(r"[a-z0-9]+", text.lower()))
-
-
-def _content_terms(text: str) -> set[str]:
-    return {term for term in _tokenize(text) if len(term) > 2 and term not in _STOPWORDS}
-
-
-def _extract_citations(text: str) -> list[str]:
-    return re.findall(r"\[([a-z0-9_]+)\]", text.lower())
-
-
-def _normalize_chunk_id(chunk_id: str) -> str:
-    chunk_id = chunk_id.strip()
-    return chunk_id
-
-
-def _normalize_domain_filter(domain_filter: str | None) -> str | None:
-    if domain_filter is None:
-        return None
-    mapping = {
-        "customer_support_operations": "Customer Support Operations",
-        "incident_response_playbooks": "Incident Response Playbooks",
-        "platform_reliability_release_engineering": "Platform Reliability & Release Engineering",
-    }
-    return mapping.get(domain_filter, domain_filter)
-
-
-def _f1_score(selected: set[str], relevant: set[str]) -> float:
-    if not selected and not relevant:
-        return 1.0
-    if not selected or not relevant:
-        return 0.0
-    overlap = len(selected & relevant)
-    if overlap == 0:
-        return 0.0
-    precision = overlap / len(selected)
-    recall = overlap / len(relevant)
-    return 2 * precision * recall / (precision + recall)
-
-
-@dataclass(frozen=True, slots=True)
-class GraderResult:
-    score: float
-    breakdown: dict[str, float]
-    passed: bool
-
-
-class TaskGrader:
-    def _filter_relevant_by_domain(self, relevant_ids: set[str], retriever: HybridRetriever, task: Task) -> set[str]:
-        normalized_domain = _normalize_domain_filter(task.domain_filter)
-        if normalized_domain is None:
-            return relevant_ids
-        allowed_ids = {chunk.chunk_id for chunk in retriever.corpus if chunk.domain == normalized_domain}
-        return relevant_ids & allowed_ids
-
-    def _required_chunks(self, retriever: HybridRetriever, task: Task) -> list[Chunk]:
-        normalized_required = {_normalize_chunk_id(chunk_id) for chunk_id in task.required_chunk_ids}
-        return [chunk for chunk in retriever.corpus if chunk.chunk_id in normalized_required]
-
+"""
+Deterministic graders for rag-context-optimizer tasks.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+from env.corpus import Chunk
+from env.retriever import HybridRetriever
+from env.tasks import Task
+
+
+_STOPWORDS = {
+    "a", "an", "and", "are", "as", "at", "be", "because", "by", "for", "from", "how",
+    "if", "in", "into", "is", "it", "its", "of", "on", "or", "that", "the", "their",
+    "them", "there", "these", "this", "to", "was", "were", "what", "when", "where",
+    "which", "while", "with", "within", "without", "you", "your",
+}
+
+
+def _tokenize(text: str) -> set[str]:
+    return set(re.findall(r"[a-z0-9]+", text.lower()))
+
+
+def _content_terms(text: str) -> set[str]:
+    return {term for term in _tokenize(text) if len(term) > 2 and term not in _STOPWORDS}
+
+
+def _extract_citations(text: str) -> list[str]:
+    return re.findall(r"\[([a-z0-9_]+)\]", text.lower())
+
+
+def _normalize_chunk_id(chunk_id: str) -> str:
+    chunk_id = chunk_id.strip()
+    return chunk_id
+
+
+def _normalize_domain_filter(domain_filter: str | None) -> str | None:
+    if domain_filter is None:
+        return None
+    mapping = {
+        "customer_support_operations": "Customer Support Operations",
+        "incident_response_playbooks": "Incident Response Playbooks",
+        "platform_reliability_release_engineering": "Platform Reliability & Release Engineering",
+    }
+    return mapping.get(domain_filter, domain_filter)
+
+
+def _f1_score(selected: set[str], relevant: set[str]) -> float:
+    if not selected and not relevant:
+        return 1.0
+    if not selected or not relevant:
+        return 0.0
+    overlap = len(selected & relevant)
+    if overlap == 0:
+        return 0.0
+    precision = overlap / len(selected)
+    recall = overlap / len(relevant)
+    return 2 * precision * recall / (precision + recall)
+
+
+@dataclass(frozen=True, slots=True)
+class GraderResult:
+    score: float
+    breakdown: dict[str, float]
+    passed: bool
+
+
+class TaskGrader:
+    def _filter_relevant_by_domain(self, relevant_ids: set[str], retriever: HybridRetriever, task: Task) -> set[str]:
+        normalized_domain = _normalize_domain_filter(task.domain_filter)
+        if normalized_domain is None:
+            return relevant_ids
+        allowed_ids = {chunk.chunk_id for chunk in retriever.corpus if chunk.domain == normalized_domain}
+        return relevant_ids & allowed_ids
+
+    def _required_chunks(self, retriever: HybridRetriever, task: Task) -> list[Chunk]:
+        normalized_required = {_normalize_chunk_id(chunk_id) for chunk_id in task.required_chunk_ids}
+        return [chunk for chunk in retriever.corpus if chunk.chunk_id in normalized_required]
+
     def _answer_quality(self, answer: str, required_chunks: list[Chunk]) -> float:
         answer_terms = _content_terms(answer)
         required_terms = _content_terms(" ".join(chunk.text for chunk in required_chunks))
+        required_terms |= _content_terms(" ".join(" ".join(chunk.keywords) for chunk in required_chunks))
         if not answer_terms or not required_terms:
             return 0.0
         union = answer_terms | required_terms
-        if not union:
-            return 0.0
-        return len(answer_terms & required_terms) / len(union)
-
-    def _citation_accuracy(self, answer: str, selected_chunk_ids: set[str], expected_citation_ids: set[str]) -> float:
-        citations = {_normalize_chunk_id(chunk_id) for chunk_id in _extract_citations(answer)}
-        if not citations:
-            return 0.0
-        valid_citations = citations & selected_chunk_ids
-        precision = len(valid_citations) / len(citations)
-        recall = len(valid_citations & expected_citation_ids) / len(expected_citation_ids) if expected_citation_ids else 1.0
-        return (precision + recall) / 2.0
-
-    def _unsupported_claim_rate(self, answer: str, evidence_chunks: list[Chunk]) -> float:
-        answer_terms = _content_terms(re.sub(r"\[[a-z0-9_]+\]", " ", answer.lower()))
-        evidence_terms = _content_terms(" ".join(chunk.text for chunk in evidence_chunks))
-        if not answer_terms:
-            return 0.0
-        unsupported = answer_terms - evidence_terms
-        return len(unsupported) / len(answer_terms)
-
-    def grade(
-        self,
-        selected_chunk_ids: list[str],
-        answer: str,
-        token_budget: int,
-        total_tokens_used: int,
-        retriever: HybridRetriever,
-        task: Task,
+        if not union:
+            return 0.0
+        return len(answer_terms & required_terms) / len(union)
+
+    def _citation_accuracy(self, answer: str, selected_chunk_ids: set[str], expected_citation_ids: set[str]) -> float:
+        citations = {_normalize_chunk_id(chunk_id) for chunk_id in _extract_citations(answer)}
+        if not citations:
+            return 0.0
+        valid_citations = citations & selected_chunk_ids
+        precision = len(valid_citations) / len(citations)
+        recall = len(valid_citations & expected_citation_ids) / len(expected_citation_ids) if expected_citation_ids else 1.0
+        return (precision + recall) / 2.0
+
+    def _unsupported_claim_rate(self, answer: str, evidence_chunks: list[Chunk]) -> float:
+        answer_terms = _content_terms(re.sub(r"\[[a-z0-9_]+\]", " ", answer.lower()))
+        evidence_terms = _content_terms(" ".join(chunk.text for chunk in evidence_chunks))
+        if not answer_terms:
+            return 0.0
+        unsupported = answer_terms - evidence_terms
+        return len(unsupported) / len(answer_terms)
+
+    def grade(
+        self,
+        selected_chunk_ids: list[str],
+        answer: str,
+        token_budget: int,
+        total_tokens_used: int,
+        retriever: HybridRetriever,
+        task: Task,
     ) -> GraderResult:
         normalized_selected = {_normalize_chunk_id(chunk_id) for chunk_id in selected_chunk_ids}
-        relevant = retriever.get_ground_truth_relevant(task.query, threshold=0.3)
-        relevant = self._filter_relevant_by_domain(relevant, retriever, task)
+        normalized_required = {_normalize_chunk_id(chunk_id) for chunk_id in task.required_chunk_ids}
+        relevant = self._filter_relevant_by_domain(normalized_required, retriever, task)
 
         retrieval_precision = _f1_score(normalized_selected, relevant)
         token_efficiency = 1.0 - (total_tokens_used / token_budget) if total_tokens_used <= token_budget else 0.0
@@ -127,41 +128,41 @@ class TaskGrader:
         required_chunks = self._required_chunks(retriever, task)
         answer_quality = self._answer_quality(answer, required_chunks)
 
-        normalized_required = {_normalize_chunk_id(chunk_id) for chunk_id in task.required_chunk_ids}
         normalized_expected_citations = {
             _normalize_chunk_id(chunk_id)
             for chunk_id in (task.expected_citation_ids or task.required_chunk_ids)
-        }
-        required_chunks_hit = (
-            len(normalized_selected & normalized_required) / len(normalized_required)
-            if normalized_required
-            else 1.0
-        )
-
+        }
+        required_chunks_hit = (
+            len(normalized_selected & normalized_required) / len(normalized_required)
+            if normalized_required
+            else 1.0
+        )
+
         selected_chunks = [
             chunk for chunk in retriever.corpus if chunk.chunk_id in normalized_selected
         ]
+        evidence_chunks = selected_chunks or required_chunks
         citation_accuracy = self._citation_accuracy(answer, normalized_selected, normalized_expected_citations)
-        unsupported_claim_rate = self._unsupported_claim_rate(answer, selected_chunks)
-        hallucination_penalty = min(1.0, unsupported_claim_rate)
-
-        base_score = (
-            0.25 * retrieval_precision
-            + 0.25 * token_efficiency
-            + 0.35 * answer_quality
-            + 0.15 * required_chunks_hit
-        )
-        score = base_score + (0.10 * citation_accuracy) - (0.15 * hallucination_penalty)
-        score = max(0.0, min(1.0, score))
-
-        breakdown = {
-            "retrieval_precision": retrieval_precision,
-            "token_efficiency": token_efficiency,
-            "answer_quality": answer_quality,
-            "required_chunks_hit": required_chunks_hit,
-            "citation_accuracy": citation_accuracy,
-            "unsupported_claim_rate": unsupported_claim_rate,
-            "hallucination_penalty": hallucination_penalty,
-        }
-        passed = score >= 0.7
-        return GraderResult(score=score, breakdown=breakdown, passed=passed)
+        unsupported_claim_rate = self._unsupported_claim_rate(answer, evidence_chunks)
+        hallucination_penalty = min(1.0, unsupported_claim_rate)
+
+        base_score = (
+            0.25 * retrieval_precision
+            + 0.25 * token_efficiency
+            + 0.35 * answer_quality
+            + 0.15 * required_chunks_hit
+        )
+        score = base_score + (0.10 * citation_accuracy) - (0.15 * hallucination_penalty)
+        score = max(0.0, min(1.0, score))
+
+        breakdown = {
+            "retrieval_precision": retrieval_precision,
+            "token_efficiency": token_efficiency,
+            "answer_quality": answer_quality,
+            "required_chunks_hit": required_chunks_hit,
+            "citation_accuracy": citation_accuracy,
+            "unsupported_claim_rate": unsupported_claim_rate,
+            "hallucination_penalty": hallucination_penalty,
+        }
+        passed = score >= 0.7
+        return GraderResult(score=score, breakdown=breakdown, passed=passed)

inference.py

@@ -26,16 +26,16 @@ TASK_SEQUENCE = [
     "adversarial_compression",
 ]
 
-SYSTEM_PROMPT = """You are a baseline RAG context optimizer.
-Read the query and available chunks using chunk_id, keywords, tokens, and domain.
+SYSTEM_PROMPT = """You are a baseline RAG context optimizer.
+Read the query and available chunks using chunk_id, keywords, tokens, and domain.
 Select chunks that maximize keyword overlap with the query.
 Stay under the token budget.
 Compress chunks that are mildly relevant but token-heavy.
 Submit a concise answer once enough useful chunks are selected.
 When you submit an answer, cite selected chunks inline like [support_003] or [incident_002].
 Return only valid JSON matching one of these forms:
-{"action_type":"select_chunk","chunk_id":"support_003"}
-{"action_type":"deselect_chunk","chunk_id":"support_003"}
+{"action_type":"select_chunk","chunk_id":"support_003"}
+{"action_type":"deselect_chunk","chunk_id":"support_003"}
 {"action_type":"compress_chunk","chunk_id":"support_003","compression_ratio":0.5}
 {"action_type":"submit_answer","answer":"Verify outage evidence and the billing ledger before refunding [support_001] [support_003]."}"""
 
@@ -222,12 +222,12 @@ async def _post_json(http_client: httpx.AsyncClient, path: str, payload: dict[st
     return response.json()
 
 
-async def _run_task_http(task_name: str) -> tuple[float, list[float], int]:
+async def _run_task_http(task_name: str) -> tuple[float, list[float], int, bool]:
     rewards: list[float] = []
     steps = 0
     success = False
-    score = 0.0
-    terminal_error: str | None = None
+    score = 0.0
+    terminal_error: str | None = None
     fallback_reason: str | None = None
     model_name = _model_name()
 
@@ -253,7 +253,7 @@ async def _run_task_http(task_name: str) -> tuple[float, list[float], int]:
             flush=True,
         )
         print("[END] success=false steps=0 score=0.000 rewards=")
-        return 0.0, [], 0
+        return 0.0, [], 0, False
 
     try:
         async with httpx.AsyncClient(timeout=30.0) as http_client:
@@ -276,7 +276,7 @@ async def _run_task_http(task_name: str) -> tuple[float, list[float], int]:
                         print(
                             f"[END] success=false steps={steps} score={_clamp_score(score):.3f} rewards={_format_rewards(rewards)}",
                         )
-                        return score, rewards, steps
+                        return score, rewards, steps, False
                     print(
                         f"[warn] Falling back to deterministic policy for {task_name}: {fallback_reason}",
                         file=sys.stderr,
@@ -313,31 +313,34 @@ async def _run_task_http(task_name: str) -> tuple[float, list[float], int]:
                     success = terminal_error is None and fallback_reason is None
                     break
 
-            score = _clamp_score(score)
-            print(
-                f"[END] success={_format_bool(success)} steps={steps} score={score:.3f} rewards={_format_rewards(rewards)}"
-            )
-            return score, rewards, steps
-    except Exception:
-        score = _clamp_score(score)
-        print(
-            f"[END] success=false steps={steps} score={score:.3f} rewards={_format_rewards(rewards)}"
-        )
-        return score, rewards, steps
-
-
-def run_task(task_name: str) -> tuple[float, list[float], int]:
-    return asyncio.run(_run_task_http(task_name))
-
-
-def main() -> None:
-    if RAG_ENV_TASK in TASK_SEQUENCE:
-        tasks = [RAG_ENV_TASK] + [task for task in TASK_SEQUENCE if task != RAG_ENV_TASK]
-    else:
-        tasks = list(TASK_SEQUENCE)
-    for task_name in tasks:
-        run_task(task_name)
-
-
-if __name__ == "__main__":
-    main()
+            score = _clamp_score(score)
+            print(
+                f"[END] success={_format_bool(success)} steps={steps} score={score:.3f} rewards={_format_rewards(rewards)}"
+            )
+            return score, rewards, steps, success
+    except Exception:
+        score = _clamp_score(score)
+        print(
+            f"[END] success=false steps={steps} score={score:.3f} rewards={_format_rewards(rewards)}"
+        )
+        return score, rewards, steps, False
+
+
+def run_task(task_name: str) -> tuple[float, list[float], int, bool]:
+    return asyncio.run(_run_task_http(task_name))
+
+
+def main() -> int:
+    if RAG_ENV_TASK in TASK_SEQUENCE:
+        tasks = [RAG_ENV_TASK] + [task for task in TASK_SEQUENCE if task != RAG_ENV_TASK]
+    else:
+        tasks = list(TASK_SEQUENCE)
+    all_success = True
+    for task_name in tasks:
+        _score, _rewards, _steps, success = run_task(task_name)
+        all_success &= success
+    return 0 if all_success else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/test_api.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+from fastapi.testclient import TestClient
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from app import app
+
+
+client = TestClient(app)
+
+
+def test_reset_accepts_empty_body():
+    response = client.post("/reset")
+    assert response.status_code == 200
+    body = response.json()
+    assert "episode_id" in body
+    assert body["done"] is False
+    assert "observation" in body
+
+
+def test_episode_state_is_isolated():
+    first_reset = client.post("/reset", json={"task_name": "single_domain_qa"})
+    second_reset = client.post("/reset", json={"task_name": "cross_domain_synthesis"})
+    assert first_reset.status_code == 200
+    assert second_reset.status_code == 200
+
+    first_episode = first_reset.json()["episode_id"]
+    second_episode = second_reset.json()["episode_id"]
+    assert first_episode != second_episode
+
+    first_chunk = first_reset.json()["observation"]["available_chunks"][0]["chunk_id"]
+    step = client.post(f"/step?episode_id={first_episode}", json={"action_type": "select_chunk", "chunk_id": first_chunk})
+    assert step.status_code == 200
+    assert step.json()["episode_id"] == first_episode
+
+    first_state = client.get(f"/state?episode_id={first_episode}")
+    second_state = client.get(f"/state?episode_id={second_episode}")
+    assert first_state.status_code == 200
+    assert second_state.status_code == 200
+    assert first_chunk in first_state.json()["selected_chunks"]
+    assert second_state.json()["selected_chunks"] == []

tests/test_inference_proxy.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import json
+import os
+import socket
+import subprocess
+import sys
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+
+import httpx
+
+
+ROOT = Path(__file__).resolve().parents[1]
+PYTHON = ROOT / ".venv" / "Scripts" / "python.exe"
+
+
+def _free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+def test_inference_uses_proxy_api_key():
+    app_port = _free_port()
+    proxy_port = _free_port()
+    requests_seen: list[dict[str, str | None]] = []
+
+    class ProxyHandler(BaseHTTPRequestHandler):
+        def do_POST(self):
+            length = int(self.headers.get("Content-Length", "0"))
+            body = self.rfile.read(length).decode("utf-8")
+            requests_seen.append(
+                {
+                    "path": self.path,
+                    "authorization": self.headers.get("Authorization"),
+                    "body": body,
+                }
+            )
+            payload = {
+                "id": "chatcmpl-test",
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": "proxy-test-model",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": json.dumps(
+                                {
+                                    "action_type": "submit_answer",
+                                    "answer": "Proxy verified [support_003]",
+                                }
+                            ),
+                        },
+                        "finish_reason": "stop",
+                    }
+                ],
+            }
+            encoded = json.dumps(payload).encode("utf-8")
+            self.send_response(200)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(encoded)))
+            self.end_headers()
+            self.wfile.write(encoded)
+
+        def log_message(self, format: str, *args):
+            return
+
+    proxy_server = HTTPServer(("127.0.0.1", proxy_port), ProxyHandler)
+    proxy_thread = threading.Thread(target=proxy_server.serve_forever, daemon=True)
+    proxy_thread.start()
+
+    app_process = subprocess.Popen(
+        [str(PYTHON), "-m", "uvicorn", "app:app", "--host", "127.0.0.1", "--port", str(app_port)],
+        cwd=ROOT,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+
+    try:
+        deadline = time.time() + 20
+        while time.time() < deadline:
+            try:
+                if httpx.get(f"http://127.0.0.1:{app_port}/health", timeout=2).status_code == 200:
+                    break
+            except Exception:
+                time.sleep(0.5)
+
+        env = os.environ.copy()
+        env["RAG_ENV_URL"] = f"http://127.0.0.1:{app_port}"
+        env["RAG_ENV_TASK"] = "single_domain_qa"
+        env["API_BASE_URL"] = f"http://127.0.0.1:{proxy_port}/v1"
+        env["API_KEY"] = "proxy-check-token"
+        env["HF_TOKEN"] = "legacy-should-not-win"
+        result = subprocess.run(
+            [str(PYTHON), "inference.py"],
+            cwd=ROOT,
+            env=env,
+            capture_output=True,
+            text=True,
+            timeout=60,
+        )
+        assert result.returncode == 0
+        assert requests_seen
+        assert requests_seen[0]["path"] == "/v1/chat/completions"
+        assert requests_seen[0]["authorization"] == "Bearer proxy-check-token"
+        assert any(line.startswith("[END]") and "score=" in line for line in result.stdout.splitlines())
+    finally:
+        proxy_server.shutdown()
+        proxy_server.server_close()
+        app_process.terminate()
+        try:
+            app_process.wait(timeout=5)
+        except Exception:
+            app_process.kill()

validate.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
-import json
-import os
-import signal
-import socket
-import subprocess
-import sys
-import time
-from pathlib import Path
+import json
+import os
+import signal
+import socket
+import subprocess
+import sys
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
 
 import httpx
 
@@ -117,24 +119,80 @@ def run_task(client: httpx.Client, base_url: str, task_name: str) -> tuple[bool,
 
 
 def run_inference_script(base_url: str) -> bool:
-    env = os.environ.copy()
-    env["RAG_ENV_URL"] = base_url
-    env["ALLOW_BASELINE_FALLBACK"] = "1"
-    env["API_BASE_URL"] = "http://127.0.0.1:9/v1"
-    env["API_KEY"] = "offline-validation-token"
-    process = subprocess.run(
-        [sys.executable, "inference.py"],
-        cwd=PROJECT_ROOT,
-        capture_output=True,
-        text=True,
-        timeout=120,
-        env=env,
-    )
-    stdout = process.stdout or ""
-    has_start = "[START]" in stdout
-    has_end = "[END]" in stdout
-    end_has_score = " score=" in stdout
-    return process.returncode == 0 and has_start and has_end and end_has_score
+    proxy_port = find_free_port()
+    requests_seen: list[dict[str, str | None]] = []
+
+    class ProxyHandler(BaseHTTPRequestHandler):
+        def do_POST(self):
+            length = int(self.headers.get("Content-Length", "0"))
+            body = self.rfile.read(length).decode("utf-8")
+            requests_seen.append(
+                {
+                    "path": self.path,
+                    "authorization": self.headers.get("Authorization"),
+                    "body": body,
+                }
+            )
+            payload = {
+                "id": "chatcmpl-validate",
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": "validator-proxy",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": json.dumps(
+                                {
+                                    "action_type": "submit_answer",
+                                    "answer": "Validated via proxy [support_003]",
+                                }
+                            ),
+                        },
+                        "finish_reason": "stop",
+                    }
+                ],
+            }
+            encoded = json.dumps(payload).encode("utf-8")
+            self.send_response(200)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(encoded)))
+            self.end_headers()
+            self.wfile.write(encoded)
+
+        def log_message(self, format: str, *args):
+            return
+
+    proxy_server = HTTPServer(("127.0.0.1", proxy_port), ProxyHandler)
+    proxy_thread = threading.Thread(target=proxy_server.serve_forever, daemon=True)
+    proxy_thread.start()
+
+    try:
+        env = os.environ.copy()
+        env["RAG_ENV_URL"] = base_url
+        env.pop("ALLOW_BASELINE_FALLBACK", None)
+        env["API_BASE_URL"] = f"http://127.0.0.1:{proxy_port}/v1"
+        env["API_KEY"] = "offline-validation-token"
+        env["HF_TOKEN"] = "legacy-should-not-win"
+        process = subprocess.run(
+            [sys.executable, "inference.py"],
+            cwd=PROJECT_ROOT,
+            capture_output=True,
+            text=True,
+            timeout=120,
+            env=env,
+        )
+        stdout = process.stdout or ""
+        has_start = "[START]" in stdout
+        has_end = "[END]" in stdout
+        end_has_score = " score=" in stdout
+        proxy_called = any(request["path"] == "/v1/chat/completions" for request in requests_seen)
+        auth_ok = any(request["authorization"] == "Bearer offline-validation-token" for request in requests_seen)
+        return process.returncode == 0 and has_start and has_end and end_has_score and proxy_called and auth_ok
+    finally:
+        proxy_server.shutdown()
+        proxy_server.server_close()
 
 
 def main() -> int: