TaskEnvironment protocol + SQLite database env (#11, #9)

README.md

@@ -174,6 +174,7 @@ cat agentfailbench/failures/tool_drift/tool-semantic-drift-001.yaml
 ```text
 parallax/
 ├── agentfailbench/     # benchmark cases + injectors
+│   └── environments/   # TaskEnvironment protocol · customer_api · database (sqlite)
 ├── runtime/            # semantic monitor + detectors
 ├── recovery/           # RecoverAI policies
 ├── baselines/          # comparison methods

agentfailbench/environments/__init__.py

@@ -0,0 +1,7 @@
+"""AgentFailBench environment packages and shared protocol."""
+
+from agentfailbench.environments.base import TaskEnvironment
+from agentfailbench.environments.customer_api import CustomerApiEnv
+from agentfailbench.environments.database import DatabaseEnv
+
+__all__ = ["CustomerApiEnv", "DatabaseEnv", "TaskEnvironment"]

agentfailbench/environments/base.py

@@ -0,0 +1,28 @@
+"""Shared task-environment protocol for AgentFailBench.
+
+Every environment (customer API, database investigation, …) drives an episode
+through the same three hooks so runners and injectors stay env-agnostic.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from runtime.schemas.episode import Action, EnvObservation
+
+
+@runtime_checkable
+class TaskEnvironment(Protocol):
+    """Lifecycle contract every AgentFailBench environment must implement."""
+
+    def reset(self) -> None:
+        """Return the environment to its initial deterministic state."""
+        ...
+
+    def step(self, action: Action) -> EnvObservation:
+        """Apply one agent action and return the transport-level observation."""
+        ...
+
+    def validate_success(self) -> bool:
+        """Return True when the task objective has been achieved."""
+        ...

agentfailbench/environments/customer_api/env.py

@@ -56,7 +56,11 @@ def plan_id_meaning(version: ContractVersion) -> str:
 
 @dataclass
 class CustomerApiEnv:
-    """Deterministic in-memory customer subscription API."""
+    """Deterministic in-memory customer subscription API.
+
+    Implements :class:`~agentfailbench.environments.base.TaskEnvironment`
+    (``reset``, ``step``, ``validate_success``).
+    """
 
     task: TaskSpec
     contract_version: ContractVersion = "v1"

agentfailbench/environments/database/__init__.py

@@ -0,0 +1,5 @@
+"""Deterministic in-memory SQLite investigation environment."""
+
+from agentfailbench.environments.database.env import DatabaseEnv, seed_schema
+
+__all__ = ["DatabaseEnv", "seed_schema"]

agentfailbench/environments/database/env.py

@@ -0,0 +1,150 @@
+"""In-memory SQLite investigation environment for AgentFailBench Milestone 1."""
+
+from __future__ import annotations
+
+import sqlite3
+from dataclasses import dataclass, field
+from typing import Any
+
+from runtime.schemas.episode import Action, EnvObservation, TaskSpec
+
+# Default seed: a tiny CRM suitable for lookup / aggregation tasks.
+SEED_SQL = """
+CREATE TABLE customers (
+    customer_id TEXT PRIMARY KEY,
+    name TEXT NOT NULL,
+    segment TEXT NOT NULL
+);
+CREATE TABLE orders (
+    order_id TEXT PRIMARY KEY,
+    customer_id TEXT NOT NULL REFERENCES customers(customer_id),
+    status TEXT NOT NULL,
+    amount_cents INTEGER NOT NULL
+);
+INSERT INTO customers VALUES
+    ('cust_001', 'Ada Lovelace', 'enterprise'),
+    ('cust_002', 'Grace Hopper', 'smb'),
+    ('cust_003', 'Alan Turing', 'enterprise');
+INSERT INTO orders VALUES
+    ('ord_100', 'cust_001', 'pending', 4500),
+    ('ord_101', 'cust_001', 'shipped', 12000),
+    ('ord_102', 'cust_002', 'pending', 800),
+    ('ord_103', 'cust_002', 'cancelled', 2000),
+    ('ord_104', 'cust_003', 'pending', 9900),
+    ('ord_105', 'cust_003', 'shipped', 1500);
+"""
+
+
+def seed_schema(conn: sqlite3.Connection) -> None:
+    """Create tables and insert the canonical fixture rows."""
+    conn.executescript(SEED_SQL)
+    conn.commit()
+
+
+def _rows_to_dicts(cursor: sqlite3.Cursor) -> list[dict[str, Any]]:
+    columns = [d[0] for d in cursor.description] if cursor.description else []
+    return [dict(zip(columns, row, strict=True)) for row in cursor.fetchall()]
+
+
+@dataclass
+class DatabaseEnv:
+    """Deterministic in-memory SQLite investigation environment.
+
+    Actions
+    -------
+    ``list_tables``
+        Return table names in the schema.
+    ``describe_table``
+        Return column metadata for ``table``.
+    ``execute_sql``
+        Run a single read-only ``SELECT`` (writes are rejected).
+    ``submit_answer``
+        Record the agent's final answer for :meth:`validate_success`.
+    """
+
+    task: TaskSpec
+    conn: sqlite3.Connection = field(default_factory=lambda: sqlite3.connect(":memory:"))
+    call_count: int = 0
+    _submitted: Any = None
+    _seeded: bool = False
+
+    def __post_init__(self) -> None:
+        if not self._seeded:
+            seed_schema(self.conn)
+            self._seeded = True
+
+    def reset(self) -> None:
+        self.conn.close()
+        self.conn = sqlite3.connect(":memory:")
+        seed_schema(self.conn)
+        self._seeded = True
+        self.call_count = 0
+        self._submitted = None
+
+    def step(self, action: Action) -> EnvObservation:
+        self.call_count += 1
+        name = action.name
+        args = action.arguments
+
+        if name == "list_tables":
+            cur = self.conn.execute(
+                "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
+            )
+            tables = [row[0] for row in cur.fetchall()]
+            return EnvObservation(success=True, data={"tables": tables})
+
+        if name == "describe_table":
+            table = str(args.get("table", ""))
+            if not table or not table.replace("_", "").isalnum():
+                return EnvObservation(success=False, status_code=400, error="invalid_table")
+            cur = self.conn.execute(f"PRAGMA table_info({table})")
+            columns = [
+                {"cid": r[0], "name": r[1], "type": r[2], "notnull": r[3], "pk": r[5]}
+                for r in cur.fetchall()
+            ]
+            if not columns:
+                return EnvObservation(success=False, status_code=404, error="table_not_found")
+            return EnvObservation(success=True, data={"table": table, "columns": columns})
+
+        if name == "execute_sql":
+            sql = str(args.get("sql", "")).strip()
+            if not sql:
+                return EnvObservation(success=False, status_code=400, error="empty_sql")
+            # Read-only: reject anything that is not a single SELECT.
+            lowered = sql.lstrip().lower()
+            if not lowered.startswith("select"):
+                return EnvObservation(success=False, status_code=400, error="write_not_allowed")
+            if ";" in sql.rstrip(";"):
+                return EnvObservation(success=False, status_code=400, error="multiple_statements")
+            try:
+                cur = self.conn.execute(sql)
+            except sqlite3.Error as exc:
+                return EnvObservation(
+                    success=False, status_code=400, error=f"sql_error:{exc}", data={}
+                )
+            rows = _rows_to_dicts(cur)
+            return EnvObservation(success=True, data={"rows": rows, "row_count": len(rows)})
+
+        if name == "submit_answer":
+            self._submitted = args.get("value")
+            return EnvObservation(
+                success=True,
+                data={"submitted": self._submitted},
+            )
+
+        return EnvObservation(success=False, status_code=400, error=f"unknown_action:{name}")
+
+    def validate_success(self) -> bool:
+        """True when the agent submitted the task's expected answer."""
+        expected = self.task.expected_answer
+        if expected is None:
+            return False
+        if self._submitted is None:
+            return False
+        # Normalize numeric strings ("15300" vs 15300).
+        if isinstance(expected, (int, float)) and not isinstance(expected, bool):
+            try:
+                return float(self._submitted) == float(expected)
+            except (TypeError, ValueError):
+                return False
+        return str(self._submitted) == str(expected)

agentfailbench/runners/episode.py

@@ -4,9 +4,10 @@
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 from agentfailbench.agents.scripted import ScriptedApiAgent
+from agentfailbench.environments.base import TaskEnvironment
 from agentfailbench.environments.customer_api.env import CustomerApiEnv
 from agentfailbench.failures.base import FailureInjector, FailureInjectorRegistry
 from agentfailbench.failures.memory.injector import StaleMemoryInjector
@@ -99,11 +100,12 @@ def _run_traced(
     case: BenchmarkCaseModel,
     *,
     inject_failure: bool,
-) -> tuple[CustomerApiEnv, ScriptedApiAgent, TraceCollector, FailureInjector]:
+) -> tuple[TaskEnvironment, ScriptedApiAgent, TraceCollector, FailureInjector]:
     task = build_task(case)
-    env = CustomerApiEnv(task=task)
+    api_env = CustomerApiEnv(task=task)
+    env: TaskEnvironment = api_env
     memory = AgentMemory() if _is_memory_case(case) else None
-    agent = ScriptedApiAgent(env=env, memory=memory)
+    agent = ScriptedApiAgent(env=api_env, memory=memory)
     injector, injectors = _build_injector(case, memory=memory)
     collector = TraceCollector(task_id=case.case_id)
 
@@ -191,7 +193,7 @@ def run_episode(
             "injector_triggered": injector.triggered,
             "ground_truth_first_detectable": case.ground_truth.first_detectable_step,
             "ground_truth_root_cause": case.ground_truth.root_cause.value,
-            "contract_final": env.contract_version,
+            "contract_final": cast(CustomerApiEnv, env).contract_version,
             "suite": "memory" if _is_memory_case(case) else "tool_drift",
         },
     )

docs/benchmark-specification.md

@@ -66,12 +66,36 @@ risk:
 - ground-truth root-cause and recovery labels
 - reproducible configurations
 
+## Shared task-environment interface
+
+Every environment implements the ``TaskEnvironment`` protocol
+(``agentfailbench.environments.base``):
+
+| Method | Role |
+| --- | --- |
+| `reset()` | Restore deterministic initial state |
+| `step(action) -> EnvObservation` | Apply one agent action |
+| `validate_success() -> bool` | Task-success validator |
+
+Shipped environments:
+
+| Package | Environment id | Notes |
+| --- | --- | --- |
+| `environments/customer_api` | `customer_service_api` | Versioned subscription tool contracts |
+| `environments/database` | `sqlite_investigation` | In-memory SQLite; `submit_answer` closes the task |
+
+Action / observation schemas live in `runtime.schemas.episode`
+(`Action`, `EnvObservation`, `TaskSpec`).
+
 ## Package layout
 
 ```text
 agentfailbench/
   tasks/
   environments/
+    base.py           # TaskEnvironment protocol
+    customer_api/
+    database/         # in-memory SQLite investigation
   failures/
   labels/
   runners/

runtime/schemas/episode.py

@@ -8,14 +8,20 @@
 
 
 class TaskSpec(BaseModel):
-    """Specification for a benchmark task instance."""
+    """Specification for a benchmark task instance.
+
+    Customer-API tasks use ``customer_id`` / ``target_plan_code``.
+    Database investigation tasks set ``expected_answer`` for
+    :meth:`~agentfailbench.environments.database.env.DatabaseEnv.validate_success`.
+    """
 
     task_id: str
     objective: str
     environment: str
     customer_id: str = "cust_001"
     target_plan_code: str = "GOLD_ANNUAL"
     expected_steps: int = 8
+    expected_answer: Any | None = None
 
 
 class Action(BaseModel):

tests/unit/test_task_environment.py

@@ -0,0 +1,87 @@
+"""Unit tests for the shared TaskEnvironment protocol and DatabaseEnv."""
+
+from __future__ import annotations
+
+from agentfailbench.environments import CustomerApiEnv, DatabaseEnv, TaskEnvironment
+from agentfailbench.environments.customer_api.env import CustomerApiEnv as CustomerApiEnvDirect
+from agentfailbench.environments.database.env import DatabaseEnv as DatabaseEnvDirect
+from runtime.schemas.episode import Action, TaskSpec
+
+
+def _api_task() -> TaskSpec:
+    return TaskSpec(
+        task_id="t-api",
+        objective="update_customer_subscription",
+        environment="customer_service_api",
+    )
+
+
+def _db_task(*, expected_answer: int | str = 15300) -> TaskSpec:
+    return TaskSpec(
+        task_id="t-db",
+        objective="sum_pending_order_amounts",
+        environment="sqlite_investigation",
+        expected_answer=expected_answer,
+        expected_steps=4,
+    )
+
+
+def test_customer_api_satisfies_protocol() -> None:
+    env = CustomerApiEnv(task=_api_task())
+    assert isinstance(env, TaskEnvironment)
+
+
+def test_database_env_satisfies_protocol() -> None:
+    env = DatabaseEnv(task=_db_task())
+    assert isinstance(env, TaskEnvironment)
+
+
+def test_database_clean_episode_succeeds() -> None:
+    """Normal (non-failure) path: explore schema, aggregate, submit answer."""
+    env = DatabaseEnvDirect(task=_db_task(expected_answer=15300))
+    # pending amounts: 4500 + 800 + 9900 = 15300
+    steps = [
+        Action(name="list_tables", arguments={}),
+        Action(name="describe_table", arguments={"table": "orders"}),
+        Action(
+            name="execute_sql",
+            arguments={
+                "sql": "SELECT SUM(amount_cents) AS total FROM orders WHERE status = 'pending'"
+            },
+        ),
+        Action(name="submit_answer", arguments={"value": 15300}),
+    ]
+    for action in steps:
+        obs = env.step(action)
+        assert obs.success, obs.error
+    assert env.validate_success()
+
+
+def test_database_wrong_answer_fails_validator() -> None:
+    env = DatabaseEnv(task=_db_task(expected_answer=15300))
+    env.step(Action(name="submit_answer", arguments={"value": 0}))
+    assert not env.validate_success()
+
+
+def test_database_rejects_writes() -> None:
+    env = DatabaseEnv(task=_db_task())
+    obs = env.step(Action(name="execute_sql", arguments={"sql": "DELETE FROM orders"}))
+    assert not obs.success
+    assert obs.error == "write_not_allowed"
+
+
+def test_database_reset_restores_seed() -> None:
+    env = DatabaseEnv(task=_db_task(expected_answer=3))
+    env.step(Action(name="submit_answer", arguments={"value": 3}))
+    assert env.validate_success()
+    env.reset()
+    assert not env.validate_success()
+    obs = env.step(
+        Action(name="execute_sql", arguments={"sql": "SELECT COUNT(*) AS n FROM customers"})
+    )
+    assert obs.success
+    assert obs.data["rows"][0]["n"] == 3
+
+
+def test_customer_api_still_exports() -> None:
+    assert CustomerApiEnv is CustomerApiEnvDirect