From ff8ae4152b0e6836afc72dd521a8c3c37e7ac9c3 Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Sun, 4 Oct 2026 19:41:28 -0700 Subject: [PATCH 01/16] Update gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 558ed264..401dfe9f 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,4 @@ sdk/src/asaree_client/_version.py # Private local planning and agent guidance docs/ issues/ +tasks/ From 0ffe1e1c72df83fd1531b63aa006c3510b97557f Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Sun, 4 Oct 2026 20:13:56 -0700 Subject: [PATCH 02/16] Add dataset row input contract --- frontend/src/types/protocols.ts | 10 ++++- src/asaree/services/dataset_row_inputs.py | 48 ++++++++++++++++++++ tests/test_dataset_row_edge_contract.py | 55 +++++++++++++++++++++++ 3 files changed, 112 insertions(+), 1 deletion(-) create mode 100644 src/asaree/services/dataset_row_inputs.py create mode 100644 tests/test_dataset_row_edge_contract.py diff --git a/frontend/src/types/protocols.ts b/frontend/src/types/protocols.ts index 2a992a68..6086f1eb 100644 --- a/frontend/src/types/protocols.ts +++ b/frontend/src/types/protocols.ts @@ -55,13 +55,21 @@ export interface EdgeHandoff { fields?: EdgeHandoffField[] } +export type DatasetInput = + | { mode: 'whole_dataset' } + | { mode: 'per_row'; columns: string[] } + export interface ProtocolEdge { id: string source: string target: string sourceHandle?: string | null targetHandle?: string | null - data?: { handoff?: EdgeHandoff } + data?: { + handoff?: EdgeHandoff + dataset_input?: DatasetInput + [key: string]: unknown + } } export type NodeRunStatus = 'pending' | 'running' | 'completed' | 'failed' | 'skipped' | 'cancelled' diff --git a/src/asaree/services/dataset_row_inputs.py b/src/asaree/services/dataset_row_inputs.py new file mode 100644 index 00000000..7bfb3a9f --- /dev/null +++ b/src/asaree/services/dataset_row_inputs.py @@ -0,0 +1,48 @@ +"""Validation and normalization for Dataset-to-Agent edge input settings.""" + +from __future__ import annotations + +_MISSING = object() + + +class DatasetRowInputError(ValueError): + """A Dataset input configuration that does not match the edge contract.""" + + def __init__(self, code: str, message: str) -> None: + self.code = code + self.message = message + super().__init__(f"dataset_row.{code}: {message}") + + +def normalize_dataset_input(value: object = _MISSING) -> dict[str, str | list[str]]: + """Return a fresh canonical Dataset edge input configuration. + + An omitted configuration keeps the historical whole-dataset behavior. + Explicit JSON null is malformed and must be distinguished from omission. + """ + if value is _MISSING: + return {"mode": "whole_dataset"} + if not isinstance(value, dict): + raise DatasetRowInputError("invalid_configuration", "configuration must be an object") + + unknown_fields = set(value) - {"mode", "columns"} + if unknown_fields: + raise DatasetRowInputError("invalid_configuration", "configuration contains unknown fields") + + mode = value.get("mode") + if mode == "whole_dataset": + if "columns" not in value or value["columns"] == []: + return {"mode": "whole_dataset"} + raise DatasetRowInputError("invalid_columns", "whole_dataset does not accept columns") + + if mode == "per_row": + columns = value.get("columns") + if not isinstance(columns, list) or not columns: + raise DatasetRowInputError("invalid_columns", "per_row requires a non-empty columns array") + if any(not isinstance(column, str) or not column.strip() for column in columns): + raise DatasetRowInputError("invalid_columns", "columns must be non-empty, non-whitespace strings") + if len(set(columns)) != len(columns): + raise DatasetRowInputError("invalid_columns", "columns must be distinct") + return {"mode": "per_row", "columns": list(columns)} + + raise DatasetRowInputError("invalid_mode", "mode must be 'whole_dataset' or 'per_row'") diff --git a/tests/test_dataset_row_edge_contract.py b/tests/test_dataset_row_edge_contract.py new file mode 100644 index 00000000..0b5bf670 --- /dev/null +++ b/tests/test_dataset_row_edge_contract.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import pytest + +from asaree.services.dataset_row_inputs import DatasetRowInputError, normalize_dataset_input + + +def test_missing_configuration_preserves_whole_dataset_behavior() -> None: + assert normalize_dataset_input() == {"mode": "whole_dataset"} + + +def test_per_row_configuration_preserves_column_order_and_spelling() -> None: + configuration = {"mode": "per_row", "columns": ["question", "hint"]} + + assert normalize_dataset_input(configuration) == { + "mode": "per_row", + "columns": ["question", "hint"], + } + assert configuration == {"mode": "per_row", "columns": ["question", "hint"]} + + +@pytest.mark.parametrize( + ("configuration", "code"), + [ + ({"mode": "per_row", "columns": ["question", "question"]}, "invalid_columns"), + ({"mode": "per_row", "columns": []}, "invalid_columns"), + ({"mode": "per_row", "columns": "question"}, "invalid_columns"), + ({"mode": "per_row", "columns": [" "]}, "invalid_columns"), + ({"mode": "whole_dataset", "columns": ["question"]}, "invalid_columns"), + (None, "invalid_configuration"), + ({"mode": "row"}, "invalid_mode"), + ({"mode": "per_row", "columns": ["question"], "extra": True}, "invalid_configuration"), + ], +) +def test_malformed_configuration_is_rejected_with_a_code(configuration: object, code: str) -> None: + with pytest.raises(DatasetRowInputError) as exc_info: + normalize_dataset_input(configuration) + + assert exc_info.value.code == code + assert str(exc_info.value).startswith(f"dataset_row.{code}: ") + + +def test_whole_dataset_empty_columns_are_canonicalized_away() -> None: + assert normalize_dataset_input({"mode": "whole_dataset", "columns": []}) == {"mode": "whole_dataset"} + + +def test_per_row_results_have_fresh_column_lists() -> None: + configuration = {"mode": "per_row", "columns": ["question"]} + + first = normalize_dataset_input(configuration) + second = normalize_dataset_input(configuration) + + assert first is not second + assert first["columns"] is not configuration["columns"] + assert first["columns"] is not second["columns"] From 75bb2f53f8740c7286d69a3ecf055595262d1681 Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Sun, 4 Oct 2026 20:17:48 -0700 Subject: [PATCH 03/16] Resolve dataset row input topology --- src/asaree/services/dataset_row_inputs.py | 148 ++++++++++++++++++++++ tests/test_dataset_row_topology.py | 131 +++++++++++++++++++ 2 files changed, 279 insertions(+) create mode 100644 tests/test_dataset_row_topology.py diff --git a/src/asaree/services/dataset_row_inputs.py b/src/asaree/services/dataset_row_inputs.py index 7bfb3a9f..dc4eeccf 100644 --- a/src/asaree/services/dataset_row_inputs.py +++ b/src/asaree/services/dataset_row_inputs.py @@ -2,6 +2,11 @@ from __future__ import annotations +import uuid +from collections import defaultdict +from contextlib import suppress +from typing import Any + _MISSING = object() @@ -46,3 +51,146 @@ def normalize_dataset_input(value: object = _MISSING) -> dict[str, str | list[st return {"mode": "per_row", "columns": list(columns)} raise DatasetRowInputError("invalid_mode", "mode must be 'whole_dataset' or 'per_row'") + + +_DATASET_NODE_TYPES = frozenset({"dataset"}) +_AGENT_NODE_TYPES = frozenset({"agent"}) +_DATASET_HANDLES = frozenset({"dataset", "resource", "tool"}) + + +def resolve_dataset_row_plan(graph: dict[str, Any], design_spec: dict[str, Any] | None = None) -> dict[str, Any] | None: + """Resolve and validate the Dataset-to-Agent inputs that drive cell rows. + + This is deliberately graph-only: it neither looks up registered datasets nor + changes the supplied graph. Registration identity is checked at runtime. + """ + nodes = {str(node.get("id")): node for node in graph.get("nodes") or [] if node.get("id") is not None} + bindings: list[dict[str, Any]] = [] + drivers: set[str] = set() + whole_by_id: dict[str, list[str]] = defaultdict(list) + whole_by_name: dict[str, list[str]] = defaultdict(list) + + # dataset_input is a reserved edge contract; reject it even where the edge + # would otherwise be ignored by type, handle, or enabled state. + for edge in graph.get("edges") or []: + edge_data = edge.get("data") or {} + if "dataset_input" not in edge_data: + continue + source = nodes.get(str(edge.get("source"))) + target = nodes.get(str(edge.get("target"))) + valid_source = source is not None and source.get("type") in _DATASET_NODE_TYPES + valid_target = target is not None and target.get("type") in _AGENT_NODE_TYPES + if not valid_source or not valid_target: + raise DatasetRowInputError( + "misplaced_configuration", + "dataset_input is only valid on a Dataset-to-Agent edge", + ) + if edge.get("targetHandle") not in _DATASET_HANDLES: + raise DatasetRowInputError("misplaced_configuration", "dataset_input requires a Dataset connector handle") + + for edge in graph.get("edges") or []: + source = nodes.get(str(edge.get("source"))) + target = nodes.get(str(edge.get("target"))) + if ( + source is None + or source.get("type") not in _DATASET_NODE_TYPES + or target is None + or target.get("type") not in _AGENT_NODE_TYPES + ): + continue + if edge.get("targetHandle") not in _DATASET_HANDLES: + continue + + edge_data = edge.get("data") or {} + if edge_data.get("factor_bindings"): + raise DatasetRowInputError("unsupported_edge_factor", "dataset_input edges cannot have factor bindings") + config = (source.get("data") or {}).get("config", _MISSING) + # Validate attached configuration even for disabled source nodes. + normalized = normalize_dataset_input(edge_data.get("dataset_input", _MISSING)) + if config is _MISSING or not isinstance(config, dict): + raise DatasetRowInputError("invalid_dataset_config", "Dataset node configuration must be an object") + enabled = config.get("enabled", True) + if not isinstance(enabled, bool): + raise DatasetRowInputError("invalid_dataset_config", "Dataset enabled must be a boolean") + dataset_name = config.get("dataset_name") + dataset_id = config.get("dataset_id") + if dataset_name is not None and not isinstance(dataset_name, str): + raise DatasetRowInputError("invalid_dataset_config", "dataset_name must be a string") + + if normalized["mode"] == "per_row": + if not isinstance(dataset_id, str): + raise DatasetRowInputError( + "invalid_dataset_id", + "per_row requires a registered dataset UUID in dataset_id", + ) + try: + driver_id = str(uuid.UUID(dataset_id)) + except (ValueError, AttributeError) as exc: + raise DatasetRowInputError("invalid_dataset_id", "per_row dataset_id must be a valid UUID") from exc + if enabled: + # A factor on any part of config can replace/alter the driver. + data = source.get("data") or {} + for field_path in data.get("factor_bindings") or {}: + if field_path == "config" or field_path.startswith("config."): + raise DatasetRowInputError( + "factorized_driver", + "a row-driving Dataset configuration cannot be factor-bound", + ) + drivers.add(driver_id) + bindings.append( + { + "edge_id": str(edge.get("id", "")), + "dataset_node_id": str(source.get("id", "")), + "agent_node_id": str(target.get("id", "")), + "columns": list(normalized["columns"]), + "_dataset_id": driver_id, + } + ) + elif enabled: + if dataset_id is not None: + with suppress(ValueError, TypeError, AttributeError): + whole_by_id[str(uuid.UUID(dataset_id))].append(str(source.get("id", ""))) + if isinstance(dataset_name, str): + whole_by_name[dataset_name].append(str(source.get("id", ""))) + + if len(drivers) > 1: + raise DatasetRowInputError("multiple_drivers", "per_row inputs must use one registered dataset") + if not drivers: + return None + driver_id = next(iter(drivers)) + + driver_nodes = {binding["dataset_node_id"] for binding in bindings} + for node_id in driver_nodes: + data = nodes[node_id].get("data") or {} + for field_path in data.get("factor_bindings") or {}: + if field_path == "config" or field_path.startswith("config."): + raise DatasetRowInputError( + "factorized_driver", + "a row-driving Dataset configuration cannot be factor-bound", + ) + if whole_by_id.get(driver_id): + raise DatasetRowInputError("mixed_driver_modes", "one dataset cannot use both per_row and whole_dataset inputs") + + # Legacy whole-dataset configs can be name-only. In this pure resolver the + # exact matching name is the only identity signal available. + driver_names = { + (nodes[binding["dataset_node_id"]].get("data") or {}).get("config", {}).get("dataset_name") + for binding in bindings + } + if any(name and name in whole_by_name for name in driver_names): + raise DatasetRowInputError("mixed_driver_modes", "one dataset cannot use both per_row and whole_dataset inputs") + + columns_by_agent: dict[str, list[str]] = {} + for binding in bindings: + agent_id = binding["agent_node_id"] + columns = binding["columns"] + if agent_id in columns_by_agent and columns_by_agent[agent_id] != columns: + raise DatasetRowInputError( + "conflicting_agent_columns", + "duplicate row inputs to one Agent must select identical columns", + ) + columns_by_agent[agent_id] = columns + return { + "driver_dataset_id": driver_id, + "bindings": [{key: value for key, value in binding.items() if key != "_dataset_id"} for binding in bindings], + } diff --git a/tests/test_dataset_row_topology.py b/tests/test_dataset_row_topology.py new file mode 100644 index 00000000..196bc192 --- /dev/null +++ b/tests/test_dataset_row_topology.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import copy +import uuid + +import pytest + +from asaree.services.dataset_row_inputs import DatasetRowInputError, resolve_dataset_row_plan + +DATASET_ID = "11111111-1111-4111-8111-111111111111" + + +def graph(*, configs=None, edges=None, extra_nodes=None): + configs = configs or {"d": {"dataset_id": DATASET_ID, "dataset_name": "cohort"}} + nodes = [{"id": node_id, "type": "dataset", "data": {"config": config}} for node_id, config in configs.items()] + nodes.extend({"id": agent, "type": "agent", "data": {}} for agent in {e["target"] for e in edges or []}) + nodes.extend(extra_nodes or []) + return {"nodes": nodes, "edges": edges or []} + + +def edge(source="d", target="a", *, handle="dataset", mode="per_row", columns=None, edge_id="e", **data): + config = {"mode": mode, "columns": columns or ["question"]} if mode == "per_row" else {"mode": mode} + return { + "id": edge_id, + "source": source, + "target": target, + "targetHandle": handle, + "data": {"dataset_input": config, **data}, + } + + +def test_legacy_handles_and_order_are_preserved_without_mutation(): + edges = [edge(edge_id="first", handle="tool"), edge(edge_id="second", handle="resource", target="b")] + value = graph(edges=edges) + before = copy.deepcopy(value) + + plan = resolve_dataset_row_plan(value) + + assert plan == { + "driver_dataset_id": DATASET_ID, + "bindings": [ + {"edge_id": "first", "dataset_node_id": "d", "agent_node_id": "a", "columns": ["question"]}, + {"edge_id": "second", "dataset_node_id": "d", "agent_node_id": "b", "columns": ["question"]}, + ], + } + assert value == before + + +def test_disabled_source_is_ignored_but_its_edge_config_is_validated(): + value = graph(configs={"d": {"dataset_id": DATASET_ID, "enabled": False}}, edges=[edge()]) + assert resolve_dataset_row_plan(value) is None + value["edges"][0]["data"]["dataset_input"] = None + with pytest.raises(DatasetRowInputError, match="invalid_configuration"): + resolve_dataset_row_plan(value) + + +def test_misplaced_dataset_input_is_rejected_even_when_disabled(): + value = graph(edges=[edge()]) + value["edges"][0]["source"] = "ordinary" + value["nodes"].append({"id": "ordinary", "type": "script", "data": {}}) + with pytest.raises(DatasetRowInputError, match="misplaced_configuration"): + resolve_dataset_row_plan(value) + + +def test_alias_nodes_with_same_uuid_can_drive_different_agent_views(): + value = graph( + configs={"d1": {"dataset_id": DATASET_ID}, "d2": {"dataset_id": DATASET_ID}}, + edges=[edge("d1", columns=["question"]), edge("d2", target="b", columns=["grade"], edge_id="e2")], + ) + plan = resolve_dataset_row_plan(value) + assert plan is not None + assert plan["driver_dataset_id"] == DATASET_ID + assert [b["columns"] for b in plan["bindings"]] == [["question"], ["grade"]] + + +def test_two_drivers_and_mixed_access_are_rejected(): + other = str(uuid.uuid4()) + value = graph( + configs={"d": {"dataset_id": DATASET_ID}, "d2": {"dataset_id": other}}, + edges=[edge(), edge("d2", target="b", edge_id="e2")], + ) + with pytest.raises(DatasetRowInputError, match="multiple_drivers"): + resolve_dataset_row_plan(value) + value = graph(edges=[edge(), edge(mode="whole_dataset", edge_id="whole")]) + with pytest.raises(DatasetRowInputError, match="mixed_driver_modes"): + resolve_dataset_row_plan(value) + + +def test_duplicate_edges_collapse_only_when_columns_match(): + value = graph(edges=[edge(), edge(edge_id="e2")]) + assert len(resolve_dataset_row_plan(value)["bindings"]) == 2 + value["edges"][1]["data"]["dataset_input"]["columns"] = ["answer"] + with pytest.raises(DatasetRowInputError, match="conflicting_agent_columns"): + resolve_dataset_row_plan(value) + + +def test_row_mode_requires_dataset_uuid_not_only_name(): + value = graph(configs={"d": {"dataset_name": "cohort"}}, edges=[edge()]) + with pytest.raises(DatasetRowInputError, match="invalid_dataset_id"): + resolve_dataset_row_plan(value) + + +def test_prompt_factor_is_allowed_but_driver_factor_is_forbidden(): + value = graph(edges=[edge()]) + value["nodes"].append( + { + "id": "a", + "type": "agent", + "data": { + "config": {"system_prompt": "base"}, + "factor_bindings": {"config.system_prompt": "Prompt"}, + }, + } + ) + plan = resolve_dataset_row_plan(value, {"factors": [{"name": "Prompt", "levels": ["base", "other"]}]}) + assert plan is not None + value["nodes"][0]["data"]["factor_bindings"] = {"config": "Dataset"} + with pytest.raises(DatasetRowInputError, match="factorized_driver"): + resolve_dataset_row_plan(value) + + +def test_factor_applied_graph_must_keep_driver_and_ordered_binding_columns(): + base = graph(edges=[edge()]) + plan = resolve_dataset_row_plan(base) + changed = copy.deepcopy(base) + changed["edges"][0]["data"]["dataset_input"]["columns"] = ["answer"] + applied_plan = resolve_dataset_row_plan(changed) + assert applied_plan is not None + assert applied_plan["driver_dataset_id"] == plan["driver_dataset_id"] + assert applied_plan["bindings"][0]["columns"] != plan["bindings"][0]["columns"] + assert plan["bindings"][0]["columns"] == ["question"] From a3222e896db2307516bf2176a840f45fe8d8080d Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Sun, 4 Oct 2026 21:19:30 -0700 Subject: [PATCH 04/16] feat: add dataset row publication and storage foundations --- CONTEXT.md | 7 +- scripts/implement-asaree-row.sh | 64 +++++ src/asaree/api/datasets.py | 58 ++++- src/asaree/api/experiments.py | 2 +- src/asaree/api/protocols.py | 27 ++- src/asaree/migrations/env.py | 1 + ...b2c3d4e5f_persist_factorial_row_results.py | 86 +++++++ src/asaree/models/__init__.py | 3 + src/asaree/models/factorial_row_result.py | 53 ++++ src/asaree/models/protocol_run.py | 9 + src/asaree/services/dataset_row_csv.py | 129 ++++++++++ src/asaree/services/protocol_execution.py | 5 + src/asaree/services/protocol_graph_schema.py | 21 +- src/asaree/services/protocol_revisions.py | 94 +++++++- tests/test_dataset_row_csv.py | 105 ++++++++ tests/test_dataset_row_publication.py | 165 +++++++++++++ tests/test_dataset_row_schema.py | 135 +++++++++++ tests/test_dataset_row_storage.py | 227 ++++++++++++++++++ tests/test_protocol_graph_schema.py | 33 +++ 19 files changed, 1204 insertions(+), 20 deletions(-) create mode 100755 scripts/implement-asaree-row.sh create mode 100644 src/asaree/migrations/versions/0a1b2c3d4e5f_persist_factorial_row_results.py create mode 100644 src/asaree/models/factorial_row_result.py create mode 100644 src/asaree/services/dataset_row_csv.py create mode 100644 tests/test_dataset_row_csv.py create mode 100644 tests/test_dataset_row_publication.py create mode 100644 tests/test_dataset_row_schema.py create mode 100644 tests/test_dataset_row_storage.py diff --git a/CONTEXT.md b/CONTEXT.md index eb9c5554..f3eeae8e 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -21,7 +21,10 @@ One unique factor combination together with all its planned replicates. _Avoid_: Condition, treatment (accepted in prose, never in code) **Replicate**: -One run of one cell. Counts of pending/run/scored are replicate counts, never cell counts. +One execution unit planned under a cell. A row-mode parent Replicate owns one Row execution per original row; counts of pending/run/scored are replicate counts, never cell or row counts. Rows are neither factors nor extra stochastic replicates. + +**Row execution**: +One complete protocol execution for a (cell, original dataset row, replicate) tuple. **Primary metric**: The optional declared metric Results uses to rank cells. An editable experiment may have none; when declared, it must be unique. @@ -113,4 +116,4 @@ An agent node with its critic gate. `approved` and `revisions_used` are recorded _Avoid_: Reviewer, supervisor **Attempt**: -One ProtocolRun for a replicate; later attempts supersede earlier ones but every attempt's own facts stay immutable. +One ProtocolRun under either a whole-dataset Replicate or a stable row slot; later attempts supersede earlier ones but every attempt's own facts stay immutable. A run may also retain a nullable dataset-row snapshot when no row slot exists, as for previews. diff --git a/scripts/implement-asaree-row.sh b/scripts/implement-asaree-row.sh new file mode 100755 index 00000000..f34982cd --- /dev/null +++ b/scripts/implement-asaree-row.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash + +set -uo pipefail + +repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" +task_dir="$repo_root/tasks/asaree-row" +start_number="${1:-005}" + +if [[ ! -f "$repo_root/.env" ]]; then + printf 'Missing environment file: %s\n' "$repo_root/.env" >&2 + exit 1 +fi + +set -a +# shellcheck disable=SC1091 +source "$repo_root/.env" +set +a + +# Motoro and ASAREE use different names for the same core database URL. +if [[ -n "${ASAREE_DATABASE_URL:-}" ]]; then + export DATABASE_URL="$ASAREE_DATABASE_URL" +fi + +if [[ ! "$start_number" =~ ^[0-9]+$ ]]; then + printf 'Usage: %s [starting task number]\n' "${0##*/}" >&2 + printf 'Example: %s 005\n' "${0##*/}" >&2 + exit 2 +fi + +start_number=$((10#$start_number)) +found=0 + +shopt -s nullglob +tasks=("$task_dir"/ASAREE-ROW-L*.yaml) + +cd -- "$repo_root" + +for task in "${tasks[@]}"; do + task_name="${task##*/}" + task_number="${task_name#ASAREE-ROW-L}" + task_number="${task_number%.yaml}" + + [[ "$task_number" =~ ^[0-9]+$ ]] || continue + ((10#$task_number >= start_number)) || continue + + found=1 + printf '\n==> Implementing %s\n' "$task_name" + + if swe implement "$task"; then + printf '==> Completed %s\n' "$task_name" + else + status=$? + printf '==> Stopping: %s failed (exit %d)\n' "$task_name" "$status" >&2 + exit "$status" + fi +done + +if ((found == 0)); then + printf 'No ASAREE row tasks found at or after L%03d in %s\n' \ + "$start_number" "$task_dir" >&2 + exit 1 +fi + +printf '\nAll selected ASAREE row tasks passed.\n' diff --git a/src/asaree/api/datasets.py b/src/asaree/api/datasets.py index 52026c5e..290b16a5 100644 --- a/src/asaree/api/datasets.py +++ b/src/asaree/api/datasets.py @@ -15,9 +15,12 @@ from fastapi import APIRouter, File, Form, HTTPException, UploadFile from pydantic import BaseModel +from sqlalchemy.exc import IntegrityError from asaree.deps import CurrentUser, DbSession from asaree.models.dataset_workspace_event import WorkspaceEventType +from asaree.services.dataset_row_csv import DatasetRowCsvError, read_row_source +from asaree.services.dataset_row_inputs import DatasetRowInputError from asaree.services.dataset_workspace_events import list_events, record_event from asaree.services.datasets import ( DatasetNameConflictError, @@ -34,6 +37,19 @@ router = APIRouter(prefix="/datasets", tags=["datasets"]) +def _integrity_constraint_name(exc: IntegrityError) -> str | None: + """Read a constraint name through SQLAlchemy's asyncpg exception wrapper.""" + for candidate in (exc.orig, getattr(exc.orig, "__cause__", None)): + if candidate is None: + continue + constraint = getattr(candidate, "constraint_name", None) or getattr( + getattr(candidate, "diag", None), "constraint_name", None + ) + if constraint is not None: + return constraint + return None + + async def _get_owned_dataset(db: DbSession, dataset_id: uuid.UUID, user: CurrentUser) -> Any: dataset = await get_dataset(db, dataset_id) if dataset is None or dataset.owner_id != user.id: @@ -79,6 +95,13 @@ class DatasetResponse(BaseModel): created_at: datetime | None = None +class DatasetRowSchemaResponse(BaseModel): + dataset_id: uuid.UUID + raw_sha256: str + columns: list[str] + row_count: int + + def _dataset_response(d: Any) -> DatasetResponse: return DatasetResponse( id=d.id, @@ -221,10 +244,43 @@ async def get_dataset_endpoint(dataset_id: uuid.UUID, db: DbSession, user: Curre return _dataset_response(dataset) +@router.get("/{dataset_id}/row-schema", response_model=DatasetRowSchemaResponse) +async def get_dataset_row_schema_endpoint( + dataset_id: uuid.UUID, db: DbSession, user: CurrentUser +) -> DatasetRowSchemaResponse: + dataset = await _get_owned_dataset(db, dataset_id, user) + try: + source = read_row_source( + dataset_id=str(dataset.id), + raw_path=dataset.raw_path, + raw_sha256=dataset.raw_sha256, + ) + except DatasetRowCsvError as exc: + raise HTTPException( + status_code=422, + detail=str(DatasetRowInputError(exc.code, exc.message)), + ) from exc + return DatasetRowSchemaResponse( + dataset_id=dataset.id, + raw_sha256=source.raw_sha256, + columns=list(source.columns), + row_count=len(source.rows), + ) + + @router.delete("/{dataset_id}", status_code=204) async def delete_dataset_endpoint(dataset_id: uuid.UUID, db: DbSession, user: CurrentUser) -> None: await _get_owned_dataset(db, dataset_id, user) - await delete_dataset(db, dataset_id) + try: + await delete_dataset(db, dataset_id) + except IntegrityError as exc: + if _integrity_constraint_name(exc) != "fk_factorial_row_results_dataset": + raise + await db.rollback() + raise HTTPException( + status_code=409, + detail="Dataset is referenced by stored row results and cannot be deleted", + ) from exc @router.post("/{dataset_id}/workspace-events", response_model=WorkspaceEventResponse, status_code=201) diff --git a/src/asaree/api/experiments.py b/src/asaree/api/experiments.py index d3fd4a5e..6615f049 100644 --- a/src/asaree/api/experiments.py +++ b/src/asaree/api/experiments.py @@ -569,7 +569,7 @@ async def import_experiment_definition_endpoint( allow_preserved_bindings=False, ) if published_graph is not None: - await publish_protocol(db, protocol) + await publish_protocol(db, protocol, owner_id=user.id) if graph != published_graph: protocol.graph = graph await db.flush() diff --git a/src/asaree/api/protocols.py b/src/asaree/api/protocols.py index d2eb7236..851e141f 100644 --- a/src/asaree/api/protocols.py +++ b/src/asaree/api/protocols.py @@ -368,15 +368,6 @@ async def publish_protocol_endpoint(protocol_id: uuid.UUID, user: CurrentUser, d """Make the current autosaved canvas the immutable version future runs use.""" protocol = await _get_owned_protocol(db, protocol_id, user) published = await get_published_revision(db, protocol) - if published is not None and is_draft_published(protocol, published): - return await _protocol_response(db, protocol) - if protocol.experiment_id: - experiment = await get_experiment(db, protocol.experiment_id) - if experiment is not None and experiment.locked_at is not None: - raise HTTPException( - status_code=409, - detail="Experiment is locked. Unlock it before publishing a changed canvas.", - ) try: # Strategy before shape: the acyclic requirement is a pipeline # requirement, and a peer_collaboration canvas isn't run as one -- see @@ -401,7 +392,23 @@ async def publish_protocol_endpoint(protocol_id: uuid.UUID, user: CurrentUser, d raise ProtocolValidationError("; ".join(issue.message for issue in blocking_issues)) except (ProtocolValidationError, ValueError) as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc - await publish_protocol(db, protocol) + if published is not None and is_draft_published(protocol, published): + try: + await publish_protocol(db, protocol, owner_id=user.id) + except ValueError as exc: + raise HTTPException(status_code=422, detail=str(exc)) from exc + return await _protocol_response(db, protocol) + if protocol.experiment_id: + experiment = await get_experiment(db, protocol.experiment_id) + if experiment is not None and experiment.locked_at is not None: + raise HTTPException( + status_code=409, + detail="Experiment is locked. Unlock it before publishing a changed canvas.", + ) + try: + await publish_protocol(db, protocol, owner_id=user.id) + except ValueError as exc: + raise HTTPException(status_code=422, detail=str(exc)) from exc return await _protocol_response(db, protocol) diff --git a/src/asaree/migrations/env.py b/src/asaree/migrations/env.py index 609d035c..7be61c2a 100644 --- a/src/asaree/migrations/env.py +++ b/src/asaree/migrations/env.py @@ -26,6 +26,7 @@ import asaree.models.experiment_design_revision # noqa: F401 import asaree.models.factorial_cell # noqa: F401 import asaree.models.factorial_replicate_result # noqa: F401 +import asaree.models.factorial_row_result # noqa: F401 import asaree.models.password_reset_token # noqa: F401 import asaree.models.protocol # noqa: F401 import asaree.models.protocol_revision # noqa: F401 diff --git a/src/asaree/migrations/versions/0a1b2c3d4e5f_persist_factorial_row_results.py b/src/asaree/migrations/versions/0a1b2c3d4e5f_persist_factorial_row_results.py new file mode 100644 index 00000000..0ac504ae --- /dev/null +++ b/src/asaree/migrations/versions/0a1b2c3d4e5f_persist_factorial_row_results.py @@ -0,0 +1,86 @@ +"""persist stable row slots and optional attempt row snapshots + +Revision ID: 0a1b2c3d4e5f +Revises: b2c7e4d91a60 +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +revision: str = "0a1b2c3d4e5f" +down_revision: str | None = "b2c7e4d91a60" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.create_table( + "factorial_row_results", + sa.Column("id", postgresql.UUID(as_uuid=True), nullable=False), + sa.Column("replicate_result_id", postgresql.UUID(as_uuid=True), nullable=False), + sa.Column("protocol_revision_id", postgresql.UUID(as_uuid=True), nullable=False), + sa.Column("dataset_id", postgresql.UUID(as_uuid=True), nullable=False), + sa.Column("raw_sha256", sa.String(length=64), nullable=False), + sa.Column("row_index", sa.Integer(), nullable=False), + sa.Column("run_id", postgresql.UUID(as_uuid=True), nullable=True), + sa.Column("workspace_id", sa.String(length=255), nullable=True), + sa.Column("metric_values", postgresql.JSONB(astext_type=sa.Text()), nullable=True), + sa.Column("artifacts", postgresql.JSONB(astext_type=sa.Text()), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.CheckConstraint("row_index >= 0", name="ck_factorial_row_results_row_index_nonnegative"), + sa.ForeignKeyConstraint( + ["dataset_id"], ["registered_datasets.id"], name="fk_factorial_row_results_dataset", ondelete="RESTRICT" + ), + sa.ForeignKeyConstraint( + ["protocol_revision_id"], ["protocol_revisions.id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["replicate_result_id"], ["factorial_replicate_results.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "replicate_result_id", + "protocol_revision_id", + "dataset_id", + "raw_sha256", + "row_index", + name="uq_factorial_row_results_row_identity", + ), + ) + op.create_index("ix_factorial_row_results_dataset_id", "factorial_row_results", ["dataset_id"]) + op.create_index( + "ix_factorial_row_results_protocol_revision_id", "factorial_row_results", ["protocol_revision_id"] + ) + op.create_index( + "ix_factorial_row_results_replicate_result_id", "factorial_row_results", ["replicate_result_id"] + ) + op.create_index("ix_factorial_row_results_run_id", "factorial_row_results", ["run_id"]) + op.add_column("protocol_runs", sa.Column("row_result_id", postgresql.UUID(as_uuid=True), nullable=True)) + op.add_column("protocol_runs", sa.Column("dataset_row", postgresql.JSONB(astext_type=sa.Text()), nullable=True)) + op.create_foreign_key( + "fk_protocol_runs_row_result_id_factorial_row_results", + "protocol_runs", + "factorial_row_results", + ["row_result_id"], + ["id"], + ondelete="SET NULL", + ) + op.create_index("ix_protocol_runs_row_result_id", "protocol_runs", ["row_result_id"]) + + +def downgrade() -> None: + op.drop_index("ix_protocol_runs_row_result_id", table_name="protocol_runs") + op.drop_constraint("fk_protocol_runs_row_result_id_factorial_row_results", "protocol_runs", type_="foreignkey") + op.drop_column("protocol_runs", "dataset_row") + op.drop_column("protocol_runs", "row_result_id") + op.drop_index("ix_factorial_row_results_run_id", table_name="factorial_row_results") + op.drop_index("ix_factorial_row_results_replicate_result_id", table_name="factorial_row_results") + op.drop_index("ix_factorial_row_results_protocol_revision_id", table_name="factorial_row_results") + op.drop_index("ix_factorial_row_results_dataset_id", table_name="factorial_row_results") + op.drop_table("factorial_row_results") diff --git a/src/asaree/models/__init__.py b/src/asaree/models/__init__.py index e69de29b..120d912e 100644 --- a/src/asaree/models/__init__.py +++ b/src/asaree/models/__init__.py @@ -0,0 +1,3 @@ +"""Model package registration for ASAREE metadata consumers.""" + +from asaree.models.factorial_row_result import FactorialRowResult # noqa: F401 diff --git a/src/asaree/models/factorial_row_result.py b/src/asaree/models/factorial_row_result.py new file mode 100644 index 00000000..976d7b9e --- /dev/null +++ b/src/asaree/models/factorial_row_result.py @@ -0,0 +1,53 @@ +"""Immutable identity and latest projection for one original dataset row.""" + +from __future__ import annotations + +import uuid +from typing import Any + +from sqlalchemy import CheckConstraint, ForeignKey, Integer, String, UniqueConstraint +from sqlalchemy.dialects.postgresql import JSONB, UUID +from sqlalchemy.orm import Mapped, mapped_column + +from asaree.models.base import Base, TimestampMixin, generate_uuid + + +class FactorialRowResult(Base, TimestampMixin): + __tablename__ = "factorial_row_results" + __table_args__ = ( + CheckConstraint("row_index >= 0", name="ck_factorial_row_results_row_index_nonnegative"), + UniqueConstraint( + "replicate_result_id", + "protocol_revision_id", + "dataset_id", + "raw_sha256", + "row_index", + name="uq_factorial_row_results_row_identity", + ), + ) + + id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=generate_uuid) + replicate_result_id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), + ForeignKey("factorial_replicate_results.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + protocol_revision_id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), ForeignKey("protocol_revisions.id", ondelete="CASCADE"), nullable=False, index=True + ) + dataset_id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), + ForeignKey("registered_datasets.id", ondelete="RESTRICT", name="fk_factorial_row_results_dataset"), + nullable=False, + index=True, + ) + raw_sha256: Mapped[str] = mapped_column(String(64), nullable=False) + row_index: Mapped[int] = mapped_column(Integer, nullable=False) + run_id: Mapped[uuid.UUID | None] = mapped_column(UUID(as_uuid=True), nullable=True, index=True) + workspace_id: Mapped[str | None] = mapped_column(String(255), nullable=True) + metric_values: Mapped[dict[str, Any] | None] = mapped_column(JSONB, nullable=True) + artifacts: Mapped[dict[str, Any] | None] = mapped_column(JSONB, nullable=True) + + +__all__ = ["FactorialRowResult"] diff --git a/src/asaree/models/protocol_run.py b/src/asaree/models/protocol_run.py index 6771f7b9..e48c008b 100644 --- a/src/asaree/models/protocol_run.py +++ b/src/asaree/models/protocol_run.py @@ -93,6 +93,15 @@ class ProtocolRun(Base, TimestampMixin): nullable=True, index=True, ) + # Stable row slot for row-mode attempts. Preview runs may carry dataset_row + # without a slot; legacy and whole-dataset attempts leave both null. + row_result_id: Mapped[uuid.UUID | None] = mapped_column( + UUID(as_uuid=True), + ForeignKey("factorial_row_results.id", ondelete="SET NULL"), + nullable=True, + index=True, + ) + dataset_row: Mapped[dict[str, Any] | None] = mapped_column(JSONB, nullable=True) # Which design revision this run's cell belonged to when it was planned # (null for a plain graph run, same as replicate_label). Without it, a result # arriving after the user regenerated the design would be written against diff --git a/src/asaree/services/dataset_row_csv.py b/src/asaree/services/dataset_row_csv.py new file mode 100644 index 00000000..5b37377b --- /dev/null +++ b/src/asaree/services/dataset_row_csv.py @@ -0,0 +1,129 @@ +"""Read verified original Dataset CSV rows and write selected-row snapshots.""" + +from __future__ import annotations + +import csv +import hashlib +import io +import re +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path + + +class DatasetRowCsvError(ValueError): + """A source CSV or requested row snapshot is invalid.""" + + def __init__(self, code: str, message: str) -> None: + self.code = code + self.message = message + super().__init__(f"dataset_row.{code}: {message}") + + +@dataclass(frozen=True) +class RowSource: + dataset_id: str + raw_sha256: str + columns: tuple[str, ...] + rows: tuple[tuple[str, ...], ...] + + +_SHA256_RE = re.compile(r"[0-9a-f]{64}\Z") + + +def read_row_source( + *, + dataset_id: str, + raw_path: str | None, + raw_sha256: str | None, + expected_sha256: str | None = None, +) -> RowSource: + """Read and parse only the registered original file after verifying its bytes.""" + if raw_path is None or not isinstance(raw_sha256, str) or not _SHA256_RE.fullmatch(raw_sha256): + raise DatasetRowCsvError("source_unavailable", "registered original CSV or SHA-256 is unavailable") + if expected_sha256 is not None and ( + not isinstance(expected_sha256, str) or not _SHA256_RE.fullmatch(expected_sha256) + ): + raise DatasetRowCsvError("source_hash_mismatch", "expected SHA-256 is invalid") + + try: + original_bytes = Path(raw_path).read_bytes() + except OSError as exc: + raise DatasetRowCsvError("source_unavailable", "registered original CSV cannot be read") from exc + + actual_sha256 = hashlib.sha256(original_bytes).hexdigest() + if actual_sha256 != raw_sha256 or (expected_sha256 is not None and actual_sha256 != expected_sha256): + raise DatasetRowCsvError("source_hash_mismatch", "original CSV does not match its registered SHA-256") + + try: + text = original_bytes.decode("utf-8-sig") + reader = csv.reader(io.StringIO(text, newline=""), delimiter=",", strict=True) + records = (record for record in reader if record) + header = next(records, None) + if header is None: + raise DatasetRowCsvError("invalid_csv", "CSV header is missing") + if any(not column.strip() for column in header) or len(set(header)) != len(header): + raise DatasetRowCsvError("invalid_csv", "CSV headers must be distinct and nonempty") + rows: list[tuple[str, ...]] = [] + for record in records: + if len(record) != len(header): + raise DatasetRowCsvError("invalid_csv", "CSV record has a different number of fields than its header") + rows.append(tuple(record)) + except UnicodeDecodeError as exc: + raise DatasetRowCsvError("invalid_csv", "CSV is not valid UTF-8") from exc + except csv.Error as exc: + raise DatasetRowCsvError("invalid_csv", "CSV syntax is malformed") from exc + + return RowSource(dataset_id, actual_sha256, tuple(header), tuple(rows)) + + +def project_row(source: RowSource, *, row_index: int, columns: Sequence[str]) -> dict[str, object]: + """Return an ordered, exact-string projection of one zero-based data record.""" + if isinstance(row_index, bool) or not isinstance(row_index, int): + raise DatasetRowCsvError("row_out_of_range", "row index must be an integer") + if not isinstance(columns, Sequence) or isinstance(columns, (str, bytes)): + raise DatasetRowCsvError("invalid_columns", "columns must be a nonempty sequence") + selected = list(columns) + if ( + not selected + or any(not isinstance(column, str) or not column.strip() for column in selected) + or len(set(selected)) != len(selected) + or any(column not in source.columns for column in selected) + ): + raise DatasetRowCsvError("invalid_columns", "columns must be distinct existing nonempty headers") + if row_index < 0 or row_index >= len(source.rows): + code = "empty_source" if not source.rows else "row_out_of_range" + raise DatasetRowCsvError(code, "row index is outside the source") + row = source.rows[row_index] + indices = [source.columns.index(column) for column in selected] + return { + "dataset_id": source.dataset_id, + "raw_sha256": source.raw_sha256, + "row_index": row_index, + "columns": selected, + "values": {column: row[index] for column, index in zip(selected, indices, strict=True)}, + } + + +def materialize_row_csv(view: dict, *, destination: Path) -> Path: + """Write a one-row selected-column CSV snapshot to an explicit run path.""" + columns = view.get("columns") + values = view.get("values") + if ( + not isinstance(columns, list) + or not columns + or any(not isinstance(column, str) for column in columns) + or len(set(columns)) != len(columns) + or not isinstance(values, dict) + or any(column not in values or not isinstance(values[column], str) for column in columns) + ): + raise DatasetRowCsvError("invalid_columns", "row view must contain distinct columns and string values") + try: + destination.parent.mkdir(parents=True, exist_ok=True) + with destination.open("w", encoding="utf-8", newline="") as stream: + writer = csv.writer(stream, delimiter=",", lineterminator="\n") + writer.writerow(columns) + writer.writerow([values[column] for column in columns]) + except OSError as exc: + raise DatasetRowCsvError("source_unavailable", "row snapshot destination cannot be written") from exc + return destination diff --git a/src/asaree/services/protocol_execution.py b/src/asaree/services/protocol_execution.py index d712a174..e3651e70 100644 --- a/src/asaree/services/protocol_execution.py +++ b/src/asaree/services/protocol_execution.py @@ -58,6 +58,7 @@ from asaree.services import prompt_references from asaree.services.agent_cards import AgentCard, build_agent_card from asaree.services.coordination import coordination_strategy_slug +from asaree.services.dataset_row_inputs import DatasetRowInputError, resolve_dataset_row_plan from asaree.services.dataset_workspaces import ( WorkspaceSeedError, fetch_owned_registration, @@ -1047,6 +1048,10 @@ def topological_order(graph: dict[str, Any], *, require_acyclic: bool = True) -> unreachable are appended in declaration order, since with the sort's premise gone there is no order left to claim. """ + try: + resolve_dataset_row_plan(graph) + except DatasetRowInputError as exc: + raise ProtocolValidationError(str(exc)) from exc _, downstream, _ = _adjacency(graph) nodes, ordered, complete = _kahn_order(graph) if not nodes: diff --git a/src/asaree/services/protocol_graph_schema.py b/src/asaree/services/protocol_graph_schema.py index f036ed20..445b1883 100644 --- a/src/asaree/services/protocol_graph_schema.py +++ b/src/asaree/services/protocol_graph_schema.py @@ -7,6 +7,8 @@ from copy import deepcopy from typing import Any +from asaree.services.dataset_row_inputs import normalize_dataset_input + _LEGACY_MODEL_NODE_TYPES = { "llm_anthropic": "model_anthropic", "llm_openai": "model_openai", @@ -67,16 +69,25 @@ def functional_protocol_graph(graph: dict[str, Any]) -> dict[str, Any]: for node in nodes or [] if isinstance(node, dict) ] - functional_edges = [ - { + functional_edges = [] + for edge in edges or []: + if not isinstance(edge, dict): + continue + functional_edge = { "source": edge.get("source"), "target": edge.get("target"), "sourceHandle": edge.get("sourceHandle"), "targetHandle": edge.get("targetHandle"), } - for edge in edges or [] - if isinstance(edge, dict) - ] + data = edge.get("data") + if isinstance(data, dict) and "dataset_input" in data: + try: + configured = normalize_dataset_input(data["dataset_input"]) + except ValueError: + configured = None + if configured is not None and configured.get("mode") != "whole_dataset": + functional_edge["data"] = {"dataset_input": configured} + functional_edges.append(functional_edge) def canonical_json(value: dict[str, Any]) -> str: return json.dumps(value, sort_keys=True, separators=(",", ":")) diff --git a/src/asaree/services/protocol_revisions.py b/src/asaree/services/protocol_revisions.py index 636382a7..22a062ab 100644 --- a/src/asaree/services/protocol_revisions.py +++ b/src/asaree/services/protocol_revisions.py @@ -9,10 +9,14 @@ from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession +from asaree.models.dataset import RegisteredDataset from asaree.models.factorial_replicate_result import FactorialReplicateResult from asaree.models.protocol import Protocol from asaree.models.protocol_revision import ProtocolRevision from asaree.models.protocol_run import ProtocolRun +from asaree.services.dataset_row_csv import DatasetRowCsvError, read_row_source +from asaree.services.dataset_row_inputs import DatasetRowInputError, resolve_dataset_row_plan +from asaree.services.datasets import get_dataset_by_name from asaree.services.protocol_graph_schema import functional_protocol_graph_hash @@ -26,8 +30,96 @@ async def get_published_revision(db: AsyncSession, protocol: Protocol) -> Protoc return await get_revision(db, protocol.published_revision_id) -async def publish_protocol(db: AsyncSession, protocol: Protocol) -> ProtocolRevision: +async def publish_protocol( + db: AsyncSession, protocol: Protocol, *, owner_id: uuid.UUID | None = None +) -> ProtocolRevision: """Freeze the protocol's current draft as its next production revision.""" + plan = resolve_dataset_row_plan(protocol.graph) + if plan is not None: + if owner_id is None: + raise DatasetRowInputError("source_unavailable", "registered row source is unavailable") + driver_id = uuid.UUID(plan["driver_dataset_id"]) + registration = ( + await db.execute( + select(RegisteredDataset).where( + RegisteredDataset.id == driver_id, + RegisteredDataset.owner_id == owner_id, + ) + ) + ).scalar_one_or_none() + if registration is None: + raise DatasetRowInputError("source_unavailable", "registered row source is unavailable") + nodes = {str(node.get("id")): node for node in protocol.graph.get("nodes", []) if isinstance(node, dict)} + for binding in plan["bindings"]: + driver_node = nodes[binding["dataset_node_id"]] + driver_config = (driver_node.get("data") or {}).get("config") or {} + if driver_config.get("dataset_name") != registration.name: + raise DatasetRowInputError("invalid_dataset_config", "dataset_id and dataset_name do not agree") + # Resolve whole-context aliases as well as row drivers. Name-only legacy + # nodes are owner-scoped and still identify the registered row source. + for node in nodes.values(): + if node.get("type") != "dataset": + continue + config = (node.get("data") or {}).get("config") or {} + if config.get("enabled", True) is False: + continue + node_id = str(node.get("id")) + connected = any( + edge.get("source") == node_id + and edge.get("targetHandle") in {"dataset", "resource", "tool"} + and (edge.get("data") or {}).get("dataset_input", {}).get("mode", "whole_dataset") == "whole_dataset" + for edge in protocol.graph.get("edges", []) + if isinstance(edge, dict) and isinstance(edge.get("data") or {}, dict) + ) + if not connected: + continue + configured_id = config.get("dataset_id") + name = config.get("dataset_name") + resolved_by_id = None + if configured_id is not None: + try: + candidate = ( + await db.execute( + select(RegisteredDataset).where( + RegisteredDataset.id == uuid.UUID(configured_id), + RegisteredDataset.owner_id == owner_id, + ) + ) + ).scalar_one_or_none() + except (ValueError, TypeError, AttributeError): + candidate = None + if candidate is not None: + resolved_by_id = candidate + resolved_by_name = ( + await get_dataset_by_name(db, name, owner_id=owner_id) + if isinstance(name, str) and name + else None + ) + if configured_id is not None and resolved_by_id is None: + # In a row publication, a configured ID is an identity claim. + # Do not silently fall back to a different name match (and do + # not reveal whether the ID belongs to another user). + raise DatasetRowInputError("source_unavailable", "registered dataset source is unavailable") + if resolved_by_id is not None and name is not None and name != resolved_by_id.name: + raise DatasetRowInputError("invalid_dataset_config", "dataset_id and dataset_name do not agree") + if ( + resolved_by_id is not None + and resolved_by_name is not None + and resolved_by_id.id != resolved_by_name.id + ): + raise DatasetRowInputError("invalid_dataset_config", "dataset_id and dataset_name do not agree") + resolved = resolved_by_id or resolved_by_name + if resolved is not None and resolved.id == registration.id: + raise DatasetRowInputError( + "mixed_driver_modes", "one dataset cannot use both per_row and whole_dataset inputs" + ) + source = read_row_source( + dataset_id=str(registration.id), raw_path=registration.raw_path, raw_sha256=registration.raw_sha256 + ) + for binding in plan["bindings"]: + missing = [column for column in binding["columns"] if column not in source.columns] + if missing: + raise DatasetRowCsvError("invalid_columns", "selected columns are not present in the registered CSV") published = await get_published_revision(db, protocol) if published is not None and is_draft_published(protocol, published): return published diff --git a/tests/test_dataset_row_csv.py b/tests/test_dataset_row_csv.py new file mode 100644 index 00000000..97ffc779 --- /dev/null +++ b/tests/test_dataset_row_csv.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +import hashlib +from pathlib import Path + +import pytest + +from asaree.services.dataset_row_csv import ( + DatasetRowCsvError, + materialize_row_csv, + project_row, + read_row_source, +) + + +def _register(path: Path, content: bytes) -> tuple[str, str]: + path.write_bytes(content) + return str(path), hashlib.sha256(content).hexdigest() + + +def test_reads_original_rows_and_projects_only_authorized_columns(tmp_path: Path) -> None: + original = ( + 'question,reference,note\n' + '"What is\nUnicode?",,café\n' + 'duplicate,,雪\n' + 'duplicate,,雪\n' + ).encode() + raw_path, digest = _register(tmp_path / "source.csv", original) + source = read_row_source(dataset_id="dataset-1", raw_path=raw_path, raw_sha256=digest, expected_sha256=digest) + + assert source.columns == ("question", "reference", "note") + assert source.rows == (("What is\nUnicode?", "", "café"), ("duplicate", "", "雪"), ("duplicate", "", "雪")) + first = project_row(source, row_index=1, columns=["question", "reference"]) + second = project_row(source, row_index=1, columns=["question", "note"]) + assert first["values"] == {"question": "duplicate", "reference": ""} + assert second["values"] == {"question": "duplicate", "note": "雪"} + assert first["row_index"] == second["row_index"] == 1 + + +def test_hash_mutation_is_rejected(tmp_path: Path) -> None: + raw_path, digest = _register(tmp_path / "source.csv", b"question\na\n") + Path(raw_path).write_bytes(b"question\nb\n") + + with pytest.raises(DatasetRowCsvError) as exc_info: + read_row_source(dataset_id="dataset-1", raw_path=raw_path, raw_sha256=digest) + + assert exc_info.value.code == "source_hash_mismatch" + + +def test_accepts_bom_and_ignores_blank_physical_records(tmp_path: Path) -> None: + content = b'\xef\xbb\xbfquestion,reference\n\n"",\n\nhello,world\n' + raw_path, digest = _register(tmp_path / "source.csv", content) + source = read_row_source(dataset_id="dataset-1", raw_path=raw_path, raw_sha256=digest) + + assert source.columns == ("question", "reference") + assert source.rows == (("", ""), ("hello", "world")) + + +@pytest.mark.parametrize( + "content", + [ + b"", + b"\n\n", + b"question,question\na,b\n", + b"question, \na,b\n", + b"question\n\xff\n", + b"question,reference\na\n", + b'question\n"unterminated\n', + ], +) +def test_invalid_csv_sources_are_rejected(tmp_path: Path, content: bytes) -> None: + raw_path, digest = _register(tmp_path / "source.csv", content) + + with pytest.raises(DatasetRowCsvError) as exc_info: + read_row_source(dataset_id="dataset-1", raw_path=raw_path, raw_sha256=digest) + + assert exc_info.value.code == "invalid_csv" + + +def test_empty_data_source_is_valid_but_cannot_project_a_row(tmp_path: Path) -> None: + raw_path, digest = _register(tmp_path / "source.csv", b"question\n\n") + source = read_row_source(dataset_id="dataset-1", raw_path=raw_path, raw_sha256=digest) + + assert source.rows == () + with pytest.raises(DatasetRowCsvError) as exc_info: + project_row(source, row_index=0, columns=["question"]) + assert exc_info.value.code == "empty_source" + + +def test_materialization_preserves_original_and_writes_exactly_one_selected_row(tmp_path: Path) -> None: + original = b'question,reference,note\n"multi\nline",,caf\xc3\xa9\nother,x,y\n' + raw_path, digest = _register(tmp_path / "source.csv", original) + source = read_row_source(dataset_id="dataset-1", raw_path=raw_path, raw_sha256=digest) + view = project_row(source, row_index=0, columns=["note", "question"]) + destination = tmp_path / "run" / "snapshot.csv" + + assert materialize_row_csv(view, destination=destination) == destination + assert Path(raw_path).read_bytes() == original + snapshot = read_row_source( + dataset_id="dataset-1", + raw_path=str(destination), + raw_sha256=hashlib.sha256(destination.read_bytes()).hexdigest(), + ) + assert snapshot.columns == ("note", "question") + assert snapshot.rows == (("café", "multi\nline"),) diff --git a/tests/test_dataset_row_publication.py b/tests/test_dataset_row_publication.py new file mode 100644 index 00000000..1de1ae4d --- /dev/null +++ b/tests/test_dataset_row_publication.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import hashlib +import uuid +from collections.abc import AsyncIterator +from copy import deepcopy +from pathlib import Path + +import pytest +import pytest_asyncio +from sqlalchemy import func, select + +import asaree.models.dataset # noqa: F401 +import asaree.models.experiment # noqa: F401 +from asaree.models.database import dispose_engine, get_session +from asaree.models.dataset import RegisteredDataset +from asaree.models.protocol_revision import ProtocolRevision +from asaree.models.user import User +from asaree.security.passwords import hash_password +from asaree.services.dataset_row_csv import DatasetRowCsvError +from asaree.services.dataset_row_inputs import DatasetRowInputError +from asaree.services.protocol_revisions import publish_protocol +from asaree.services.protocols import create_protocol, delete_protocol, update_protocol + + +@pytest_asyncio.fixture(autouse=True) +async def _dispose_engine() -> AsyncIterator[None]: + yield + await dispose_engine() + + +@pytest.mark.asyncio +async def test_row_publication_revisions_and_source_validation(tmp_path: Path) -> None: + owner_id = uuid.uuid4() + other_owner_id = uuid.uuid4() + dataset_id = uuid.uuid4() + csv_path = tmp_path / "rows.csv" + contents = b"question,answer\nfirst,one\nsecond,two\n" + csv_path.write_bytes(contents) + digest = hashlib.sha256(contents).hexdigest() + async with get_session() as db: + owner = User( + id=owner_id, + email="test@test.com", + hashed_password=hash_password("test"), + display_name="Row Publisher", + ) + other = User( + id=other_owner_id, + email="other@test.com", + hashed_password=hash_password("test"), + display_name="Other Publisher", + ) + db.add_all([owner, other]) + await db.flush() + dataset = RegisteredDataset( + id=dataset_id, + name="cohort", + raw_path=str(csv_path), + raw_sha256=digest, + owner_id=owner_id, + ) + other_dataset = RegisteredDataset( + id=uuid.uuid4(), + name="other-cohort", + raw_path=str(csv_path), + raw_sha256=digest, + owner_id=owner_id, + ) + db.add_all([dataset, other_dataset]) + graph = { + "nodes": [ + { + "id": "d", + "type": "dataset", + "data": {"config": {"dataset_id": str(dataset_id), "dataset_name": "cohort"}}, + }, + {"id": "a", "type": "agent", "data": {}}, + {"id": "b", "type": "agent", "data": {}}, + ], + "edges": [ + { + "id": "e1", + "source": "d", + "target": "a", + "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["question"]}}, + }, + { + "id": "e2", + "source": "d", + "target": "b", + "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["answer"]}}, + }, + ], + } + protocol = await create_protocol(db, name="row publication", owner_id=owner_id, graph=graph) + first = await publish_protocol(db, protocol, owner_id=owner_id) + original_snapshot = deepcopy(first.graph) + + changed = deepcopy(graph) + changed["edges"][0]["data"]["dataset_input"]["columns"] = ["answer", "question"] + await update_protocol(db, protocol.id, fields={"graph": changed}) + second = await publish_protocol(db, protocol, owner_id=owner_id) + assert second.revision == first.revision + 1 + assert first.graph == original_snapshot + + whole = deepcopy(changed) + whole["edges"] = [{"id": "e1", "source": "d", "target": "a", "targetHandle": "dataset"}] + await update_protocol(db, protocol.id, fields={"graph": whole}) + third = await publish_protocol(db, protocol, owner_id=owner_id) + explicit_whole = deepcopy(whole) + explicit_whole["edges"][0]["data"] = {"dataset_input": {"mode": "whole_dataset"}} + await update_protocol(db, protocol.id, fields={"graph": explicit_whole}) + assert (await publish_protocol(db, protocol, owner_id=owner_id)).id == third.id + + cosmetic = deepcopy(explicit_whole) + cosmetic["nodes"][0]["position"] = {"x": 200, "y": 300} + cosmetic["edges"][0]["id"] = "new-edge-id" + await update_protocol(db, protocol.id, fields={"graph": cosmetic}) + assert (await publish_protocol(db, protocol, owner_id=owner_id)).id == third.id + + invalid = deepcopy(graph) + invalid["edges"][0]["data"]["dataset_input"]["columns"] = ["missing"] + await update_protocol(db, protocol.id, fields={"graph": invalid}) + with pytest.raises(DatasetRowCsvError, match="invalid_columns"): + await publish_protocol(db, protocol, owner_id=owner_id) + unknown = deepcopy(graph) + unknown["nodes"][0]["data"]["config"]["dataset_id"] = str(uuid.uuid4()) + await update_protocol(db, protocol.id, fields={"graph": unknown}) + with pytest.raises(DatasetRowInputError, match="source_unavailable"): + await publish_protocol(db, protocol, owner_id=owner_id) + unowned = deepcopy(graph) + unowned["nodes"][0]["data"]["config"]["dataset_id"] = str(dataset_id) + await update_protocol(db, protocol.id, fields={"graph": unowned}) + with pytest.raises(DatasetRowInputError, match="source_unavailable"): + await publish_protocol(db, protocol, owner_id=other_owner_id) + + mismatched_alias = deepcopy(graph) + mismatched_alias["nodes"].append( + { + "id": "whole-alias", + "type": "dataset", + "data": {"config": {"dataset_id": str(uuid.uuid4()), "dataset_name": "other-cohort"}}, + } + ) + mismatched_alias["edges"].append( + {"id": "whole-edge", "source": "whole-alias", "target": "a", "targetHandle": "dataset"} + ) + await update_protocol(db, protocol.id, fields={"graph": mismatched_alias}) + with pytest.raises(DatasetRowInputError, match="source_unavailable"): + await publish_protocol(db, protocol, owner_id=owner_id) + + revision_count = ( + await db.execute( + select(func.count()).select_from(ProtocolRevision).where(ProtocolRevision.protocol_id == protocol.id) + ) + ).scalar_one() + assert revision_count == 3 + await delete_protocol(db, protocol.id) + await db.delete(dataset) + await db.delete(other_dataset) + await db.delete(owner) + await db.delete(other) diff --git a/tests/test_dataset_row_schema.py b/tests/test_dataset_row_schema.py new file mode 100644 index 00000000..1f435bd6 --- /dev/null +++ b/tests/test_dataset_row_schema.py @@ -0,0 +1,135 @@ +"""Authenticated Dataset row metadata API tests using real CSV files.""" + +from __future__ import annotations + +import hashlib +import uuid +from pathlib import Path +from types import SimpleNamespace + +import httpx +import pytest +from fastapi import FastAPI + +from asaree.api import datasets +from asaree.deps import get_current_user +from asaree.models.database import get_db + + +@pytest.fixture +def row_schema_app(monkeypatch: pytest.MonkeyPatch) -> tuple[FastAPI, uuid.UUID, dict[str, object]]: + user_id = uuid.uuid4() + dataset_id = uuid.uuid4() + state: dict[str, object] = {} + + async def _auth() -> SimpleNamespace: + return SimpleNamespace(id=user_id) + + async def _db() -> None: + return None + + async def _get_dataset(_db_session: object, requested_id: uuid.UUID) -> object: + assert requested_id == dataset_id + return state["dataset"] + + app = FastAPI() + app.include_router(datasets.router, prefix="/api") + app.dependency_overrides[get_current_user] = _auth + app.dependency_overrides[get_db] = _db + monkeypatch.setattr(datasets, "get_dataset", _get_dataset) + state["dataset_id"] = dataset_id + state["user_id"] = user_id + return app, dataset_id, state + + +def _dataset(dataset_id: uuid.UUID, owner_id: uuid.UUID, path: Path, content: bytes) -> SimpleNamespace: + return SimpleNamespace( + id=dataset_id, + owner_id=owner_id, + name="private-dataset-name", + raw_path=str(path), + raw_sha256=hashlib.sha256(content).hexdigest(), + ) + + +async def _get_schema(app: FastAPI, dataset_id: uuid.UUID) -> httpx.Response: + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + return await client.get(f"/api/datasets/{dataset_id}/row-schema") + + +async def test_row_schema_returns_ordered_metadata_without_source_disclosure( + tmp_path: Path, row_schema_app: tuple[FastAPI, uuid.UUID, dict[str, object]] +) -> None: + app, dataset_id, state = row_schema_app + content = b"subject,value,label\n001,alpha,yes\n002,beta,no\n003,gamma,yes\n" + raw_path = tmp_path / "private-source-file.csv" + raw_path.write_bytes(content) + state["dataset"] = _dataset(dataset_id, state["user_id"], raw_path, content) + + response = await _get_schema(app, dataset_id) + + assert response.status_code == 200 + assert response.json() == { + "dataset_id": str(dataset_id), + "raw_sha256": hashlib.sha256(content).hexdigest(), + "columns": ["subject", "value", "label"], + "row_count": 3, + } + assert "alpha" not in response.text + assert "private-dataset-name" not in response.text + assert str(raw_path) not in response.text + + +async def test_row_schema_keeps_unowned_dataset_concealed( + tmp_path: Path, row_schema_app: tuple[FastAPI, uuid.UUID, dict[str, object]] +) -> None: + app, dataset_id, state = row_schema_app + content = b"id,value\n1,secret\n" + raw_path = tmp_path / "readable-unowned.csv" + raw_path.write_bytes(content) + state["dataset"] = _dataset(dataset_id, uuid.uuid4(), raw_path, content) + + response = await _get_schema(app, dataset_id) + + assert response.status_code == 404 + assert "secret" not in response.text + assert str(raw_path) not in response.text + + +async def test_row_schema_rejects_corrupted_registered_hash( + tmp_path: Path, row_schema_app: tuple[FastAPI, uuid.UUID, dict[str, object]] +) -> None: + app, dataset_id, state = row_schema_app + content = b"id,value\n1,secret\n" + raw_path = tmp_path / "corrupted.csv" + raw_path.write_bytes(content + b"2,changed\n") + state["dataset"] = _dataset(dataset_id, state["user_id"], raw_path, content) + + response = await _get_schema(app, dataset_id) + + assert response.status_code == 422 + assert response.json() == { + "detail": "dataset_row.source_hash_mismatch: original CSV does not match its registered SHA-256" + } + assert str(raw_path) not in response.text + assert "secret" not in response.text + + +async def test_row_schema_accepts_header_only_csv( + tmp_path: Path, row_schema_app: tuple[FastAPI, uuid.UUID, dict[str, object]] +) -> None: + app, dataset_id, state = row_schema_app + content = b"first,second\n" + raw_path = tmp_path / "header-only.csv" + raw_path.write_bytes(content) + state["dataset"] = _dataset(dataset_id, state["user_id"], raw_path, content) + + response = await _get_schema(app, dataset_id) + + assert response.status_code == 200 + assert response.json() == { + "dataset_id": str(dataset_id), + "raw_sha256": hashlib.sha256(content).hexdigest(), + "columns": ["first", "second"], + "row_count": 0, + } diff --git a/tests/test_dataset_row_storage.py b/tests/test_dataset_row_storage.py new file mode 100644 index 00000000..9f0f0267 --- /dev/null +++ b/tests/test_dataset_row_storage.py @@ -0,0 +1,227 @@ +"""PostgreSQL persistence tests for stable dataset-row slots.""" + +from __future__ import annotations + +import os +import uuid +from collections.abc import AsyncIterator +from datetime import UTC, datetime + +import httpx +import pytest +import pytest_asyncio +from fastapi import FastAPI +from sqlalchemy import delete, select +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from asaree.api import datasets as datasets_api +from asaree.deps import get_current_user +from asaree.models.database import get_db +from asaree.models.dataset import RegisteredDataset +from asaree.models.experiment import ResearchExperiment +from asaree.models.experiment_design_revision import ExperimentDesignRevision +from asaree.models.factorial_cell import FactorialCell +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol import Protocol +from asaree.models.protocol_revision import ProtocolRevision +from asaree.models.protocol_run import ProtocolRun +from asaree.models.user import User +from asaree.security.passwords import hash_password + + +@pytest_asyncio.fixture +async def storage_graph() -> AsyncIterator[tuple[AsyncSession, dict[str, uuid.UUID]]]: + url = os.environ["ASAREE_PRODUCT_DATABASE_URL"] + engine = create_async_engine(url) + sessions = async_sessionmaker(engine, expire_on_commit=False) + ids = { + name: uuid.uuid4() + for name in ("user", "dataset", "experiment", "design", "cell", "replicate", "protocol", "protocol_revision") + } + async with sessions() as db: + db.add( + User( + id=ids["user"], + email="test@test.com", + hashed_password=hash_password("test"), + display_name="Row test", + ) + ) + await db.flush() + db.add_all( + [ + RegisteredDataset(id=ids["dataset"], name=f"row-{ids['dataset']}", owner_id=ids["user"]), + ResearchExperiment(id=ids["experiment"], name=f"row-{ids['experiment']}", owner_id=ids["user"]), + Protocol(id=ids["protocol"], name=f"row-{ids['protocol']}", owner_id=ids["user"], graph={}), + ] + ) + await db.flush() + db.add_all( + [ + ExperimentDesignRevision(id=ids["design"], experiment_id=ids["experiment"], revision=1), + ProtocolRevision( + id=ids["protocol_revision"], + protocol_id=ids["protocol"], + revision=1, + graph={}, + published_at=datetime.now(UTC), + ), + ] + ) + await db.flush() + db.add( + FactorialCell( + id=ids["cell"], + experiment_id=ids["experiment"], + design_revision_id=ids["design"], + cell_label="cell-1", + factor_values={}, + ) + ) + await db.flush() + replicate = FactorialReplicateResult( + id=ids["replicate"], cell_id=ids["cell"], replicate_number=1, replicate_label="replicate-1" + ) + db.add(replicate) + await db.flush() + yield db, ids + await db.rollback() + await engine.dispose() + + +def _slot(ids: dict[str, uuid.UUID], row_index: int = 0) -> FactorialRowResult: + return FactorialRowResult( + replicate_result_id=ids["replicate"], + protocol_revision_id=ids["protocol_revision"], + dataset_id=ids["dataset"], + raw_sha256="a" * 64, + row_index=row_index, + ) + + +async def test_composite_identity_and_duplicate_valued_positions(storage_graph) -> None: + db, ids = storage_graph + first = _slot(ids, 0) + second = _slot(ids, 1) + db.add_all([first, second]) + await db.flush() + assert first.id != second.id + + db.add(_slot(ids, 0)) + with pytest.raises(IntegrityError): + await db.flush() + await db.rollback() + + +async def test_negative_row_index_is_rejected(storage_graph) -> None: + db, ids = storage_graph + db.add(_slot(ids, -1)) + with pytest.raises(IntegrityError): + await db.flush() + await db.rollback() + + +async def test_legacy_run_nullable_fields_and_row_slot_cascade(storage_graph) -> None: + db, ids = storage_graph + legacy = ProtocolRun(protocol_id=ids["protocol"], owner_id=ids["user"], node_runs={}) + preview = ProtocolRun( + protocol_id=ids["protocol"], + owner_id=ids["user"], + node_runs={}, + dataset_row={"dataset_id": str(ids["dataset"]), "raw_sha256": "a" * 64, "row_index": 0}, + ) + slot = _slot(ids) + db.add_all([legacy, preview, slot]) + await db.flush() + assert legacy.row_result_id is None and legacy.dataset_row is None + assert preview.row_result_id is None and preview.dataset_row is not None + linked = ProtocolRun( + protocol_id=ids["protocol"], + owner_id=ids["user"], + node_runs={}, + row_result_id=slot.id, + dataset_row={"dataset_id": str(ids["dataset"]), "raw_sha256": "a" * 64, "row_index": 0}, + ) + db.add(linked) + await db.flush() + await db.delete(slot) + await db.flush() + await db.refresh(linked) + assert linked.row_result_id is None + assert linked.id is not None + + +async def test_replicate_deletion_cascades_slots_and_nulls_run_link(storage_graph) -> None: + db, ids = storage_graph + slot = _slot(ids) + run = ProtocolRun( + protocol_id=ids["protocol"], owner_id=ids["user"], node_runs={}, row_result_id=slot.id + ) + db.add_all([slot, run]) + await db.flush() + await db.execute(delete(FactorialReplicateResult).where(FactorialReplicateResult.id == ids["replicate"])) + await db.flush() + await db.refresh(run) + assert run.row_result_id is None + assert await db.scalar(select(FactorialRowResult.id).where(FactorialRowResult.id == slot.id)) is None + + + +async def test_design_deletion_cascades_slots_and_nulls_run_link(storage_graph) -> None: + db, ids = storage_graph + slot = _slot(ids) + run = ProtocolRun( + protocol_id=ids["protocol"], owner_id=ids["user"], node_runs={}, row_result_id=slot.id + ) + db.add_all([slot, run]) + await db.flush() + await db.execute(delete(ExperimentDesignRevision).where(ExperimentDesignRevision.id == ids["design"])) + await db.flush() + await db.refresh(run) + assert run.row_result_id is None + assert run.id is not None + + +async def test_protocol_revision_cascade_and_dataset_restrict(storage_graph) -> None: + db, ids = storage_graph + slot = _slot(ids) + db.add(slot) + await db.flush() + await db.execute(delete(ProtocolRevision).where(ProtocolRevision.id == ids["protocol_revision"])) + await db.flush() + assert await db.scalar(select(FactorialRowResult.id).where(FactorialRowResult.id == slot.id)) is None + + # Recreate a revision and slot to verify dataset RESTRICT independently. + revision = ProtocolRevision( + id=uuid.uuid4(), protocol_id=ids["protocol"], revision=2, graph={}, published_at=datetime.now(UTC) + ) + db.add(revision) + await db.flush() + retained = FactorialRowResult( + replicate_result_id=ids["replicate"], + protocol_revision_id=revision.id, + dataset_id=ids["dataset"], + raw_sha256="b" * 64, + row_index=0, + ) + db.add(retained) + await db.flush() + app = FastAPI() + app.include_router(datasets_api.router, prefix="/api") + + async def current_user() -> User: + return await db.get(User, ids["user"]) + + async def db_session(): + yield db + + app.dependency_overrides[get_current_user] = current_user + app.dependency_overrides[get_db] = db_session + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://test" + ) as client: + response = await client.delete(f"/api/datasets/{ids['dataset']}") + assert response.status_code == 409 + assert response.json()["detail"] == "Dataset is referenced by stored row results and cannot be deleted" diff --git a/tests/test_protocol_graph_schema.py b/tests/test_protocol_graph_schema.py index 0c08c3b9..34427095 100644 --- a/tests/test_protocol_graph_schema.py +++ b/tests/test_protocol_graph_schema.py @@ -1,5 +1,7 @@ from __future__ import annotations +import hashlib +import json from copy import deepcopy from asaree.services.protocol_graph_schema import functional_protocol_graph_hash, normalize_protocol_graph @@ -117,3 +119,34 @@ def test_functional_hash_changes_for_node_configuration_and_edge_reconnection() assert functional_protocol_graph_hash(configured) != functional_protocol_graph_hash(graph) assert functional_protocol_graph_hash(reconnected) != functional_protocol_graph_hash(graph) + + +def test_dataset_input_hash_preserves_legacy_whole_hash_and_ordered_row_columns() -> None: + graph = { + "nodes": [{"id": "d", "type": "dataset", "data": {"config": {}}}], + "edges": [{"id": "e", "source": "d", "target": "a", "targetHandle": "dataset"}], + } + omitted = functional_protocol_graph_hash(graph) + legacy_functional_graph = { + "nodes": [{"id": "d", "type": "dataset", "data": {"config": {}}}], + "edges": [{"source": "d", "target": "a", "sourceHandle": None, "targetHandle": "dataset"}], + } + legacy_bytes = json.dumps(legacy_functional_graph, sort_keys=True, separators=(",", ":")).encode() + assert omitted == hashlib.sha256(legacy_bytes).hexdigest() + explicit = deepcopy(graph) + explicit["edges"][0]["data"] = {"dataset_input": {"mode": "whole_dataset"}} + assert functional_protocol_graph_hash(explicit) == omitted + + row = deepcopy(graph) + row["edges"][0]["data"] = {"dataset_input": {"mode": "per_row", "columns": ["question", "answer"]}} + reversed_row = deepcopy(row) + reversed_row["edges"][0]["data"]["dataset_input"]["columns"].reverse() + assert functional_protocol_graph_hash(row) != omitted + assert functional_protocol_graph_hash(row) != functional_protocol_graph_hash(reversed_row) + + cosmetic = deepcopy(row) + cosmetic["edges"][0]["id"] = "replacement" + cosmetic["edges"][0]["data"]["transient"] = True + cosmetic["edges"][0]["data"]["unrelated"] = {"metadata": "ignored"} + cosmetic["nodes"][0]["position"] = {"x": 900, "y": 300} + assert functional_protocol_graph_hash(cosmetic) == functional_protocol_graph_hash(row) From 66c036e4dc9b43b50c2cbc7ce7bc829a57393f4a Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Sun, 4 Oct 2026 21:37:18 -0700 Subject: [PATCH 05/16] feat: add dataset row planning and test users --- scripts/seed_row_test_users.py | 34 ++ src/asaree/api/protocols.py | 13 + src/asaree/services/dataset_row_planning.py | 99 ++++++ src/asaree/services/factorial_row_results.py | 228 ++++++++++++ src/asaree/services/protocol_execution.py | 130 +++++++ src/asaree/services/protocol_runs.py | 136 ++++++++ tests/test_dataset_row_candidates.py | 151 ++++++++ tests/test_dataset_row_publication.py | 41 ++- tests/test_dataset_row_queries.py | 346 +++++++++++++++++++ tests/test_dataset_row_run_identity.py | 254 ++++++++++++++ tests/test_dataset_row_storage.py | 19 +- 11 files changed, 1424 insertions(+), 27 deletions(-) create mode 100644 scripts/seed_row_test_users.py create mode 100644 src/asaree/services/dataset_row_planning.py create mode 100644 src/asaree/services/factorial_row_results.py create mode 100644 tests/test_dataset_row_candidates.py create mode 100644 tests/test_dataset_row_queries.py create mode 100644 tests/test_dataset_row_run_identity.py diff --git a/scripts/seed_row_test_users.py b/scripts/seed_row_test_users.py new file mode 100644 index 00000000..3fc90d77 --- /dev/null +++ b/scripts/seed_row_test_users.py @@ -0,0 +1,34 @@ +"""Create or refresh the persistent development users used by row tests.""" + +from __future__ import annotations + +import asyncio + +from asaree.models.database import dispose_engine, get_session +from asaree.services.users import create_user, get_user_by_email, set_password + +PASSWORD = "Test1234" +USERS = ( + ("test@test.com", "Row Test User"), + ("other@test.com", "Row Test Other User"), +) + + +async def seed() -> None: + try: + async with get_session() as db: + for email, display_name in USERS: + user = await get_user_by_email(db, email) + if user is None: + user = await create_user(db, email=email, password=PASSWORD, display_name=display_name) + else: + await set_password(db, user, new_password=PASSWORD) + user.display_name = display_name + user.is_active = True + print("Seeded row-test users: test@test.com, other@test.com") + finally: + await dispose_engine() + + +if __name__ == "__main__": + asyncio.run(seed()) diff --git a/src/asaree/api/protocols.py b/src/asaree/api/protocols.py index 851e141f..4a1d5ecd 100644 --- a/src/asaree/api/protocols.py +++ b/src/asaree/api/protocols.py @@ -22,6 +22,7 @@ validate_experiment_measurement_plan, ) from asaree.services.experiments import get_experiment +from asaree.services.dataset_row_inputs import resolve_dataset_row_plan from asaree.services.factor_bindings import validate_factor_bindings from asaree.services.protocol_execution import ( ProtocolValidationError, @@ -105,6 +106,8 @@ class ProtocolRunResponse(BaseModel): error: str | None replicate_label: str | None replicate_result_id: uuid.UUID | None + row_result_id: uuid.UUID | None = None + dataset_row: dict[str, Any] | None = None factor_values: dict[str, Any] | None design_revision_id: uuid.UUID | None protocol_revision_id: uuid.UUID | None @@ -118,11 +121,14 @@ class ProtocolRunResponse(BaseModel): observations: list[dict[str, Any]] = Field(default_factory=list) artifacts: list[dict[str, Any]] = Field(default_factory=list) + def _protocol_run_response(run: Any) -> ProtocolRunResponse: response = ProtocolRunResponse.model_validate(run) measurement = (run.attempt_result or {}).get("measurement") if isinstance(run.attempt_result, dict) else None return response.model_copy( update={ + "row_result_id": run.row_result_id, + "dataset_row": run.dataset_row, "observations": list(measurement.get("observations") or []) if isinstance(measurement, dict) else [], "artifacts": list(measurement.get("artifacts") or []) if isinstance(measurement, dict) else [], } @@ -226,6 +232,7 @@ class ProtocolRevisionResponse(BaseModel): graph: dict[str, Any] published_at: datetime + class CreateProtocolRunRequest(BaseModel): # Omitted/null -- today's ad-hoc, un-substituted whole-graph run. Set -- # runs that one already-generated replicate for real, its cell's factor_values @@ -258,6 +265,8 @@ class CellRunBatchResponse(BaseModel): skipped: int protocol_revision_id: uuid.UUID protocol_revision: int + consumption_mode: str = "whole_dataset" + row_result_ids: list[uuid.UUID] = Field(default_factory=list) class PromptPreviewRequest(BaseModel): @@ -587,6 +596,10 @@ async def create_cell_runs_endpoint( skipped=skipped, protocol_revision_id=revision.id, protocol_revision=revision.revision, + consumption_mode=( + "per_row" if resolve_dataset_row_plan(revision.graph) is not None else "whole_dataset" + ), + row_result_ids=[r.row_result_id for r in runs if r.row_result_id is not None], ) diff --git a/src/asaree/services/dataset_row_planning.py b/src/asaree/services/dataset_row_planning.py new file mode 100644 index 00000000..28c3fcfc --- /dev/null +++ b/src/asaree/services/dataset_row_planning.py @@ -0,0 +1,99 @@ +"""Pure planning of row executions for factorial replicate parents. + +Callers remain responsible for the published graph's existing graph and +measurement guards. This module only verifies that factor substitution leaves +the resolved row driver and its ordered Agent/column bindings unchanged. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from copy import deepcopy +from typing import Any + +from asaree.services.dataset_row_csv import RowSource +from asaree.services.dataset_row_inputs import DatasetRowInputError, resolve_dataset_row_plan +from asaree.services.protocol_execution import apply_factor_bindings + + +def enumerate_row_candidates( + *, + parents: Sequence[dict], + source: RowSource, + graph: dict, + design_spec: dict | None, + protocol_revision_id: str, +) -> list[dict]: + """Expand sorted replicate parents across original source rows. + + Earlier graph and measurement validation remains the caller's + responsibility. No database access, scheduling, or input mutation occurs. + """ + if not source.rows: + raise DatasetRowInputError("empty_source", "row source must contain at least one data row") + + original_plan = resolve_dataset_row_plan(graph, design_spec) + if original_plan is None: + raise DatasetRowInputError("driver_mismatch", "published graph has no per-row Dataset driver") + if original_plan["driver_dataset_id"] != source.dataset_id: + raise DatasetRowInputError("driver_mismatch", "row source does not match the published graph driver") + + if not isinstance(protocol_revision_id, str) or not protocol_revision_id: + raise DatasetRowInputError("missing_revision_id", "protocol revision id is required") + + seen: set[str] = set() + ordered_parents = [] + required = ( + "replicate_result_id", + "cell_id", + "cell_label", + "replicate_label", + "replicate_number", + "design_revision_id", + "factor_values", + ) + for parent in parents: + if any(key not in parent for key in required): + raise DatasetRowInputError("invalid_parent", "row parent is missing required identity or factor fields") + identity = parent["replicate_result_id"] + if not isinstance(identity, str) or not identity: + raise DatasetRowInputError("invalid_parent", "replicate result id is required") + if identity in seen: + raise DatasetRowInputError("duplicate_parent", "replicate result ids must be unique") + seen.add(identity) + if not isinstance(parent["design_revision_id"], str) or not parent["design_revision_id"]: + raise DatasetRowInputError("missing_revision_id", "design revision id is required for every parent") + if isinstance(parent["replicate_number"], bool) or not isinstance(parent["replicate_number"], int): + raise DatasetRowInputError("invalid_parent", "replicate number must be an integer") + if not isinstance(parent["factor_values"], dict): + raise DatasetRowInputError("invalid_parent", "factor values must be an object") + ordered_parents.append(parent) + + ordered_parents.sort(key=lambda parent: ( + parent["cell_label"], + parent["replicate_number"], + parent["replicate_result_id"], + )) + + candidates: list[dict[str, Any]] = [] + for parent in ordered_parents: + applied_graph = apply_factor_bindings(graph, parent["factor_values"]) + applied_plan = resolve_dataset_row_plan(applied_graph, design_spec) + if applied_plan is None or ( + applied_plan["driver_dataset_id"], applied_plan["bindings"] + ) != (original_plan["driver_dataset_id"], original_plan["bindings"]): + raise DatasetRowInputError( + "varying_topology", + "factor substitution must preserve the row driver and ordered Agent/column bindings", + ) + + for row_index in range(len(source.rows)): + candidate = deepcopy(parent) + candidate["protocol_revision_id"] = protocol_revision_id + candidate["dataset_row"] = { + "dataset_id": source.dataset_id, + "raw_sha256": source.raw_sha256, + "row_index": row_index, + } + candidates.append(candidate) + return candidates diff --git a/src/asaree/services/factorial_row_results.py b/src/asaree/services/factorial_row_results.py new file mode 100644 index 00000000..c73fe12f --- /dev/null +++ b/src/asaree/services/factorial_row_results.py @@ -0,0 +1,228 @@ +"""Scoped persistence queries for stable dataset row slots and their attempts.""" + +from __future__ import annotations + +import re +import uuid + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from asaree.models.dataset import RegisteredDataset +from asaree.models.experiment import ResearchExperiment +from asaree.models.experiment_design_revision import ExperimentDesignRevision +from asaree.models.factorial_cell import FactorialCell +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol import Protocol +from asaree.models.protocol_revision import ProtocolRevision +from asaree.models.protocol_run import ProtocolRun +from asaree.services.design_revisions import get_current_revision + +_SHA256 = re.compile(r"^[0-9a-f]{64}$") + + +async def _revision_id( + db: AsyncSession, experiment_id: uuid.UUID, design_revision_id: uuid.UUID | None +) -> uuid.UUID | None: + if design_revision_id is None: + current = await get_current_revision(db, experiment_id) + return current.id if current is not None else None + valid = await db.scalar( + select(ExperimentDesignRevision.id).where( + ExperimentDesignRevision.id == design_revision_id, + ExperimentDesignRevision.experiment_id == experiment_id, + ) + ) + return valid + + +def _slot_scope(experiment_id: uuid.UUID, revision_id: uuid.UUID, protocol_revision_id: uuid.UUID): + return ( + FactorialRowResult.replicate_result_id == FactorialReplicateResult.id, + FactorialReplicateResult.cell_id == FactorialCell.id, + ProtocolRevision.id == FactorialRowResult.protocol_revision_id, + Protocol.experiment_id == experiment_id, + FactorialCell.experiment_id == experiment_id, + FactorialCell.design_revision_id == revision_id, + FactorialRowResult.protocol_revision_id == protocol_revision_id, + ) + + +async def get_row_result( + db: AsyncSession, + *, + experiment_id: uuid.UUID, + row_result_id: uuid.UUID, + design_revision_id: uuid.UUID | None = None, + protocol_revision_id: uuid.UUID, +) -> FactorialRowResult | None: + revision_id = await _revision_id(db, experiment_id, design_revision_id) + if revision_id is None: + return None + return ( + await db.execute( + select(FactorialRowResult) + .join(FactorialReplicateResult, FactorialRowResult.replicate_result_id == FactorialReplicateResult.id) + .join(FactorialCell, FactorialReplicateResult.cell_id == FactorialCell.id) + .join(ProtocolRevision, ProtocolRevision.id == FactorialRowResult.protocol_revision_id) + .join(Protocol, Protocol.id == ProtocolRevision.protocol_id) + .where( + *_slot_scope(experiment_id, revision_id, protocol_revision_id), + FactorialRowResult.id == row_result_id, + ) + ) + ).scalar_one_or_none() + + +async def list_row_results( + db: AsyncSession, + *, + experiment_id: uuid.UUID, + design_revision_id: uuid.UUID | None = None, + protocol_revision_id: uuid.UUID, +) -> list[FactorialRowResult]: + revision_id = await _revision_id(db, experiment_id, design_revision_id) + if revision_id is None: + return [] + return list( + ( + await db.execute( + select(FactorialRowResult) + .join(FactorialReplicateResult, FactorialRowResult.replicate_result_id == FactorialReplicateResult.id) + .join(FactorialCell, FactorialReplicateResult.cell_id == FactorialCell.id) + .join(ProtocolRevision, ProtocolRevision.id == FactorialRowResult.protocol_revision_id) + .join(Protocol, Protocol.id == ProtocolRevision.protocol_id) + .where(*_slot_scope(experiment_id, revision_id, protocol_revision_id)) + .order_by( + FactorialCell.cell_label, + FactorialReplicateResult.replicate_number, + FactorialRowResult.row_index, + FactorialRowResult.id, + ) + ) + ) + .scalars() + .all() + ) + + +async def ensure_row_result( + db: AsyncSession, + *, + experiment_id: uuid.UUID, + design_revision_id: uuid.UUID, + protocol_revision_id: uuid.UUID, + replicate_result_id: uuid.UUID, + dataset_id: uuid.UUID, + raw_sha256: str, + row_index: int, +) -> FactorialRowResult: + def invalid() -> ValueError: + return ValueError("invalid_row_scope") + + if ( + not all( + isinstance(value, uuid.UUID) + for value in (experiment_id, design_revision_id, protocol_revision_id, replicate_result_id, dataset_id) + ) + or not isinstance(raw_sha256, str) + or _SHA256.fullmatch(raw_sha256) is None + or isinstance(row_index, bool) + or not isinstance(row_index, int) + or row_index < 0 + ): + raise invalid() + + current = await get_current_revision(db, experiment_id) + parent = await db.scalar( + select(FactorialReplicateResult.id) + .join(FactorialCell, FactorialReplicateResult.cell_id == FactorialCell.id) + .where( + FactorialReplicateResult.id == replicate_result_id, + FactorialCell.experiment_id == experiment_id, + FactorialCell.design_revision_id == design_revision_id, + ) + ) + published = await db.scalar( + select(ProtocolRevision.id) + .join(Protocol, ProtocolRevision.protocol_id == Protocol.id) + .where( + ProtocolRevision.id == protocol_revision_id, + Protocol.experiment_id == experiment_id, + Protocol.published_revision_id == protocol_revision_id, + ) + ) + owned_dataset = await db.scalar( + select(RegisteredDataset.id) + .join(ResearchExperiment, ResearchExperiment.id == experiment_id) + .where(RegisteredDataset.id == dataset_id, RegisteredDataset.owner_id == ResearchExperiment.owner_id) + ) + if ( + current is None + or current.id != design_revision_id + or parent is None + or published is None + or owned_dataset is None + ): + raise invalid() + + existing = ( + await db.execute( + select(FactorialRowResult).where( + FactorialRowResult.replicate_result_id == replicate_result_id, + FactorialRowResult.protocol_revision_id == protocol_revision_id, + FactorialRowResult.dataset_id == dataset_id, + FactorialRowResult.raw_sha256 == raw_sha256, + FactorialRowResult.row_index == row_index, + ) + ) + ).scalar_one_or_none() + if existing is not None: + return existing + result = FactorialRowResult( + replicate_result_id=replicate_result_id, + protocol_revision_id=protocol_revision_id, + dataset_id=dataset_id, + raw_sha256=raw_sha256, + row_index=row_index, + ) + db.add(result) + await db.flush() + return result + + +async def list_row_attempts( + db: AsyncSession, + *, + experiment_id: uuid.UUID, + row_result_id: uuid.UUID, + design_revision_id: uuid.UUID | None = None, + protocol_revision_id: uuid.UUID, +) -> list[ProtocolRun]: + revision_id = await _revision_id(db, experiment_id, design_revision_id) + if revision_id is None: + return [] + return list( + ( + await db.execute( + select(ProtocolRun) + .join(FactorialRowResult, ProtocolRun.row_result_id == FactorialRowResult.id) + .join(FactorialReplicateResult, FactorialRowResult.replicate_result_id == FactorialReplicateResult.id) + .join(FactorialCell, FactorialReplicateResult.cell_id == FactorialCell.id) + .join(ProtocolRevision, ProtocolRevision.id == FactorialRowResult.protocol_revision_id) + .join(Protocol, Protocol.id == ProtocolRevision.protocol_id) + .where( + *_slot_scope(experiment_id, revision_id, protocol_revision_id), + ProtocolRun.row_result_id == row_result_id, + ProtocolRun.protocol_revision_id == protocol_revision_id, + ) + .order_by(ProtocolRun.created_at, ProtocolRun.id) + ) + ) + .scalars() + .all() + ) + + +__all__ = ["ensure_row_result", "get_row_result", "list_row_attempts", "list_row_results"] diff --git a/src/asaree/services/protocol_execution.py b/src/asaree/services/protocol_execution.py index e3651e70..286886e8 100644 --- a/src/asaree/services/protocol_execution.py +++ b/src/asaree/services/protocol_execution.py @@ -54,11 +54,17 @@ from asaree.config import get_settings from asaree.models.database import get_session +from asaree.models.dataset import RegisteredDataset +from asaree.models.experiment import ResearchExperiment +from asaree.models.protocol import Protocol +from asaree.models.protocol_revision import ProtocolRevision from asaree.models.protocol_run import ProtocolRun from asaree.services import prompt_references from asaree.services.agent_cards import AgentCard, build_agent_card from asaree.services.coordination import coordination_strategy_slug +from asaree.services.dataset_row_csv import DatasetRowCsvError, project_row, read_row_source from asaree.services.dataset_row_inputs import DatasetRowInputError, resolve_dataset_row_plan +from asaree.services.dataset_row_planning import enumerate_row_candidates from asaree.services.dataset_workspaces import ( WorkspaceSeedError, fetch_owned_registration, @@ -76,6 +82,7 @@ from asaree.services.experiments import get_experiment from asaree.services.factor_bindings import validate_factor_bindings from asaree.services.factorial_cells import get_replicate, list_replicates, upsert_replicate +from asaree.services.factorial_row_results import ensure_row_result, list_row_attempts from asaree.services.protocol_revisions import get_revision from asaree.services.protocol_runs import ( TERMINAL_PROTOCOL_RUN_STATUSES, @@ -4521,6 +4528,129 @@ async def plan_cell_runs( f"proposed {impact.proposed_cell_count} cells/{impact.proposed_replicate_count} replicates." ) + # A row batch is pinned to an immutable published revision. Validate the + # complete ownership/scope relationship and the caller-supplied snapshot + # before the shared current-design parents can create any row slots. + revision = await db.get(ProtocolRevision, protocol_revision_id) if protocol_revision_id else None + protocol = await db.get(Protocol, protocol_id) + if ( + revision is not None + and protocol is not None + and revision.protocol_id == protocol_id + and protocol.owner_id == owner_id + and protocol.experiment_id == experiment_id + and protocol.published_revision_id == revision.id + and revision.graph == graph + ): + try: + original_row_plan = resolve_dataset_row_plan(revision.graph, design_spec) + except DatasetRowInputError as exc: + raise ProtocolValidationError(str(exc)) from exc + if original_row_plan is not None: + if rerun_replicate_labels: + raise ProtocolValidationError("Row executions cannot be rerun through the cell batch; use row retry.") + registration = ( + await db.execute( + select(RegisteredDataset).where( + RegisteredDataset.id == uuid.UUID(original_row_plan["driver_dataset_id"]), + RegisteredDataset.owner_id == owner_id, + ) + ) + ).scalar_one_or_none() + experiment_owner = await db.scalar( + select(ResearchExperiment.owner_id).where(ResearchExperiment.id == experiment_id) + ) + if registration is None or experiment_owner != owner_id: + raise ProtocolValidationError("Registered row source is unavailable.") + try: + source = read_row_source( + dataset_id=str(registration.id), + raw_path=registration.raw_path, + raw_sha256=registration.raw_sha256, + ) + except DatasetRowCsvError as exc: + raise ProtocolValidationError(str(exc)) from exc + parents = await list_replicates(db, experiment_id=experiment_id) + labels = {parent.replicate_label for parent in parents} + requested_labels = labels if replicate_labels is None else replicate_labels + unknown = requested_labels - labels + if unknown: + raise ProtocolValidationError(f"Unknown replicate label(s): {', '.join(sorted(unknown))}.") + selected_parents = [parent for parent in parents if parent.replicate_label in requested_labels] + parent_values = [ + { + "replicate_result_id": str(parent.id), + "cell_id": str(parent.cell_id), + "cell_label": parent.cell_label, + "replicate_label": parent.replicate_label, + "replicate_number": parent.replicate_number, + "design_revision_id": str(parent.design_revision_id), + "factor_values": parent.factor_values or {}, + } + for parent in selected_parents + ] + try: + candidates = enumerate_row_candidates( + parents=parent_values, + source=source, + graph=revision.graph, + design_spec=design_spec, + protocol_revision_id=str(revision.id), + ) + planned = [] + for candidate in candidates: + patched_graph = apply_factor_bindings(revision.graph, candidate["factor_values"]) + patched_plan = resolve_dataset_row_plan(patched_graph, design_spec) + if patched_plan is None or patched_plan["driver_dataset_id"] != str(registration.id): + raise DatasetRowInputError("driver_mismatch", "factor substitution changed the row driver") + columns = list(dict.fromkeys( + column for binding in patched_plan["bindings"] for column in binding["columns"] + )) + view = project_row(source, row_index=candidate["dataset_row"]["row_index"], columns=columns) + planned.append((candidate, view)) + except (DatasetRowInputError, DatasetRowCsvError) as exc: + raise ProtocolValidationError(str(exc)) from exc + + created: list[ProtocolRun] = [] + skipped = 0 + for candidate, view in planned: + row_result = await ensure_row_result( + db, + experiment_id=experiment_id, + design_revision_id=uuid.UUID(candidate["design_revision_id"]), + protocol_revision_id=revision.id, + replicate_result_id=uuid.UUID(candidate["replicate_result_id"]), + dataset_id=registration.id, + raw_sha256=source.raw_sha256, + row_index=candidate["dataset_row"]["row_index"], + ) + attempts = await list_row_attempts( + db, + experiment_id=experiment_id, + row_result_id=row_result.id, + design_revision_id=uuid.UUID(candidate["design_revision_id"]), + protocol_revision_id=revision.id, + ) + if attempts: + skipped += 1 + continue + run = await create_protocol_run( + db, + protocol_id=protocol_id, + owner_id=owner_id, + replicate_label=candidate["replicate_label"], + factor_values=candidate["factor_values"], + replicate_result_id=uuid.UUID(candidate["replicate_result_id"]), + design_revision_id=uuid.UUID(candidate["design_revision_id"]), + protocol_revision_id=revision.id, + row_result_id=row_result.id, + dataset_row=view, + ) + created.append(run) + return created, skipped + elif resolve_dataset_row_plan(revision.graph if revision is not None else graph, design_spec) is not None: + raise ProtocolValidationError("Published protocol revision does not match the requested protocol and graph.") + # Current design only -- list_replicates scopes to the experiment's current # revision, so a superseded design's replicates are neither counted nor run. replicates = await list_replicates(db, experiment_id=experiment_id) diff --git a/src/asaree/services/protocol_runs.py b/src/asaree/services/protocol_runs.py index 14e807a1..803358e2 100644 --- a/src/asaree/services/protocol_runs.py +++ b/src/asaree/services/protocol_runs.py @@ -19,10 +19,13 @@ from sqlalchemy.ext.asyncio import AsyncSession from asaree.models.experiment import ResearchExperiment +from asaree.models.factorial_cell import FactorialCell from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult from asaree.models.protocol import Protocol from asaree.models.protocol_revision import ProtocolRevision from asaree.models.protocol_run import ProtocolRun +from asaree.services.design_revisions import get_current_revision from asaree.services.factorial_cells import get_replicate, list_replicates from asaree.services.measurement_engine import MeasurementEvaluation, normalize_measurement_plan from asaree.services.measurement_migration import normalize_experiment_measurement_plan @@ -86,6 +89,8 @@ async def create_protocol_run( target_node_id: str | None = None, design_revision_id: uuid.UUID | None = None, protocol_revision_id: uuid.UUID | None = None, + row_result_id: uuid.UUID | None = None, + dataset_row: dict[str, Any] | None = None, is_test_run: bool = False, ) -> ProtocolRun: """``replicate_label``/``factor_values``/``design_revision_id`` are set together @@ -97,6 +102,29 @@ async def create_protocol_run( ``target_node_id`` is set only for a single-node "Play" run (see ``ProtocolRun`` model's own comment) -- mutually exclusive with replicate_label/factor_values in practice, though nothing enforces that here.""" + if dataset_row is not None: + _validate_dataset_row_snapshot(dataset_row) + row_slot: FactorialRowResult | None = None + if row_result_id is not None: + if ( + dataset_row is None + or replicate_result_id is None + or design_revision_id is None + or protocol_revision_id is None + ): + raise ValueError("invalid_row_scope") + row_slot, parent, cell = await _validate_row_attempt( + db, + protocol_id=protocol_id, + owner_id=owner_id, + row_result_id=row_result_id, + replicate_result_id=replicate_result_id, + replicate_label=replicate_label, + factor_values=factor_values, + design_revision_id=design_revision_id, + protocol_revision_id=protocol_revision_id, + dataset_row=dataset_row, + ) measurement_plan_snapshot: dict[str, Any] | None = None reference_values: dict[str, Any] = {} protocol = await db.get(Protocol, protocol_id) @@ -129,6 +157,8 @@ async def create_protocol_run( replicate_label=replicate_label, factor_values=factor_values, replicate_result_id=replicate_result_id, + row_result_id=row_result_id, + dataset_row=dataset_row, target_node_id=target_node_id, design_revision_id=design_revision_id, protocol_revision_id=protocol_revision_id, @@ -140,12 +170,35 @@ async def create_protocol_run( else {} ), **({"reference_values": reference_values} if reference_values else {}), + **( + { + "row_provenance": { + "row_result_id": str(row_result_id) if row_result_id is not None else None, + "replicate_result_id": ( + str(replicate_result_id) if replicate_result_id is not None else None + ), + "design_revision_id": str(design_revision_id) if design_revision_id is not None else None, + "protocol_revision_id": str(protocol_revision_id), + "dataset_row": dict(dataset_row), + } + } + if dataset_row is not None + else {} + ), } or None ), ) db.add(run) await db.flush() + if row_slot is not None: + row_slot.run_id = run.id + row_slot.workspace_id = None + row_slot.metric_values = None + row_slot.artifacts = None + await db.flush() + await db.refresh(run) + return run if replicate_result_id is not None: # A planned run is a new attempt for this stable replicate slot. Its # result projection must immediately become "latest attempt only" -- @@ -176,6 +229,89 @@ async def create_protocol_run( return run +def _validate_dataset_row_snapshot(snapshot: dict[str, Any]) -> None: + import re + + if not isinstance(snapshot, dict) or set(snapshot) != { + "dataset_id", + "raw_sha256", + "row_index", + "columns", + "values", + }: + raise ValueError("invalid_dataset_row_snapshot") + try: + uuid.UUID(snapshot["dataset_id"]) + except (ValueError, TypeError, AttributeError) as exc: + raise ValueError("invalid_dataset_row_snapshot") from exc + if ( + not isinstance(snapshot["dataset_id"], str) + or not isinstance(snapshot["raw_sha256"], str) + or re.fullmatch(r"[0-9a-f]{64}", snapshot["raw_sha256"]) is None + or isinstance(snapshot["row_index"], bool) + or not isinstance(snapshot["row_index"], int) + or snapshot["row_index"] < 0 + or not isinstance(snapshot["columns"], list) + or not snapshot["columns"] + or any(not isinstance(column, str) or not column for column in snapshot["columns"]) + or len(set(snapshot["columns"])) != len(snapshot["columns"]) + or not isinstance(snapshot["values"], dict) + or set(snapshot["values"]) != set(snapshot["columns"]) + or any(not isinstance(value, str) for value in snapshot["values"].values()) + ): + raise ValueError("invalid_dataset_row_snapshot") + + +async def _validate_row_attempt( + db: AsyncSession, + *, + protocol_id: uuid.UUID, + owner_id: uuid.UUID, + row_result_id: uuid.UUID, + replicate_result_id: uuid.UUID, + replicate_label: str | None, + factor_values: dict[str, Any] | None, + design_revision_id: uuid.UUID, + protocol_revision_id: uuid.UUID, + dataset_row: dict[str, Any], +) -> tuple[FactorialRowResult, FactorialReplicateResult, FactorialCell]: + statement = ( + select(FactorialRowResult, FactorialReplicateResult, FactorialCell) + .join(FactorialReplicateResult, FactorialReplicateResult.id == FactorialRowResult.replicate_result_id) + .join(FactorialCell, FactorialCell.id == FactorialReplicateResult.cell_id) + .where(FactorialRowResult.id == row_result_id) + ) + result = (await db.execute(statement)).one_or_none() + if result is None: + raise ValueError("invalid_row_scope") + slot, parent, cell = result + current = await get_current_revision(db, cell.experiment_id) + revision = await db.get(ProtocolRevision, protocol_revision_id) + protocol = await db.get(Protocol, protocol_id) + experiment = await db.get(ResearchExperiment, cell.experiment_id) + if ( + cell.design_revision_id != design_revision_id + or current is None + or current.id != design_revision_id + or parent.id != replicate_result_id + or parent.replicate_label != replicate_label + or dict(cell.factor_values or {}) != dict(factor_values or {}) + or slot.protocol_revision_id != protocol_revision_id + or slot.dataset_id != uuid.UUID(dataset_row["dataset_id"]) + or slot.raw_sha256 != dataset_row["raw_sha256"] + or slot.row_index != dataset_row["row_index"] + or revision is None + or revision.protocol_id != protocol_id + or protocol is None + or protocol.experiment_id != cell.experiment_id + or protocol.owner_id != owner_id + or experiment is None + or experiment.owner_id != owner_id + ): + raise ValueError("invalid_row_scope") + return slot, parent, cell + + async def create_test_run( db: AsyncSession, *, protocol_id: uuid.UUID, owner_id: uuid.UUID, protocol_revision_id: uuid.UUID ) -> ProtocolRun: diff --git a/tests/test_dataset_row_candidates.py b/tests/test_dataset_row_candidates.py new file mode 100644 index 00000000..608f416e --- /dev/null +++ b/tests/test_dataset_row_candidates.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import copy +import hashlib + +import pytest + +from asaree.services.dataset_row_csv import read_row_source +from asaree.services.dataset_row_inputs import DatasetRowInputError +from asaree.services.dataset_row_planning import enumerate_row_candidates + +DATASET_ID = "11111111-1111-4111-8111-111111111111" + + +def make_graph() -> dict: + return { + "nodes": [ + { + "id": "dataset", + "type": "dataset", + "data": {"config": {"dataset_id": DATASET_ID}}, + }, + { + "id": "agent", + "type": "agent", + "data": { + "config": {"system_prompt": "Base prompt"}, + "factor_bindings": {"config.system_prompt": "Prompt"}, + }, + }, + ], + "edges": [ + { + "id": "dataset-agent", + "source": "dataset", + "target": "agent", + "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["question"]}}, + } + ], + } + + +def make_source(tmp_path, content: str = "question\nrepeat\nrepeat\nlast\n"): + path = tmp_path / "original.csv" + raw = content.encode() + path.write_bytes(raw) + return read_row_source( + dataset_id=DATASET_ID, + raw_path=str(path), + raw_sha256=hashlib.sha256(raw).hexdigest(), + ) + + +def make_parents() -> list[dict]: + return [ + { + "replicate_result_id": f"result-{level}-{replicate}", + "cell_id": f"cell-{level}", + "cell_label": f"Prompt={level}", + "replicate_label": f"Prompt={level} / replicate {replicate}", + "replicate_number": replicate, + "design_revision_id": "design-revision", + "factor_values": {"Prompt": f"Prompt for level {level}"}, + } + for level in range(5) + for replicate in range(1, 3) + ] + + +def test_expands_five_factor_cells_and_two_parents_across_three_original_rows(tmp_path): + candidates = enumerate_row_candidates( + parents=make_parents(), + source=make_source(tmp_path), + graph=make_graph(), + design_spec={"factors": [{"name": "Prompt", "levels": [str(i) for i in range(5)]}]}, + protocol_revision_id="published-revision", + ) + + assert len(candidates) == 30 + assert len({candidate["cell_id"] for candidate in candidates}) == 5 + assert len({candidate["replicate_result_id"] for candidate in candidates}) == 10 + assert [candidate["dataset_row"]["row_index"] for candidate in candidates[:3]] == [0, 1, 2] + assert all(candidate["protocol_revision_id"] == "published-revision" for candidate in candidates) + + +def test_duplicate_row_values_keep_distinct_original_positions(tmp_path): + source = make_source(tmp_path) + candidates = enumerate_row_candidates( + parents=make_parents()[:1], + source=source, + graph=make_graph(), + design_spec=None, + protocol_revision_id="published-revision", + ) + + assert source.rows[0] == source.rows[1] + assert [candidate["dataset_row"]["row_index"] for candidate in candidates[:2]] == [0, 1] + + +def test_candidate_order_does_not_depend_on_parent_input_order(tmp_path): + arguments = { + "source": make_source(tmp_path), + "graph": make_graph(), + "design_spec": None, + "protocol_revision_id": "published-revision", + } + + forward = enumerate_row_candidates(parents=make_parents(), **arguments) + backward = enumerate_row_candidates(parents=list(reversed(make_parents())), **arguments) + + assert forward == backward + assert [candidate["replicate_result_id"] for candidate in forward[:6]] == [ + "result-0-1", "result-0-1", "result-0-1", + "result-0-2", "result-0-2", "result-0-2", + ] + + +def test_prompt_factor_substitution_is_allowed_and_inputs_are_unchanged(tmp_path): + parents = make_parents()[:1] + graph = make_graph() + parents_before, graph_before = copy.deepcopy(parents), copy.deepcopy(graph) + + candidates = enumerate_row_candidates( + parents=parents, + source=make_source(tmp_path), + graph=graph, + design_spec=None, + protocol_revision_id="published-revision", + ) + + assert len(candidates) == 3 + assert parents == parents_before + assert graph == graph_before + + +def test_rejects_empty_source_and_driver_mutation(tmp_path): + empty = make_source(tmp_path, "question\n") + with pytest.raises(DatasetRowInputError, match="empty_source"): + enumerate_row_candidates( + parents=make_parents(), source=empty, graph=make_graph(), design_spec=None, + protocol_revision_id="published-revision", + ) + + graph = make_graph() + graph["nodes"][0]["data"]["factor_bindings"] = {"config.enabled": "Driver enabled"} + with pytest.raises(DatasetRowInputError, match="factorized_driver"): + enumerate_row_candidates( + parents=make_parents(), source=make_source(tmp_path), graph=graph, design_spec=None, + protocol_revision_id="published-revision", + ) diff --git a/tests/test_dataset_row_publication.py b/tests/test_dataset_row_publication.py index 1de1ae4d..45f503da 100644 --- a/tests/test_dataset_row_publication.py +++ b/tests/test_dataset_row_publication.py @@ -9,6 +9,7 @@ import pytest import pytest_asyncio from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession import asaree.models.dataset # noqa: F401 import asaree.models.experiment # noqa: F401 @@ -16,11 +17,23 @@ from asaree.models.dataset import RegisteredDataset from asaree.models.protocol_revision import ProtocolRevision from asaree.models.user import User -from asaree.security.passwords import hash_password from asaree.services.dataset_row_csv import DatasetRowCsvError from asaree.services.dataset_row_inputs import DatasetRowInputError from asaree.services.protocol_revisions import publish_protocol from asaree.services.protocols import create_protocol, delete_protocol, update_protocol +from asaree.services.users import create_user, get_user_by_email, set_password + + +async def _row_test_user( + db: AsyncSession, *, email: str, display_name: str +) -> tuple[User, bool]: + user = await get_user_by_email(db, email) + if user is None: + return await create_user(db, email=email, password="Test1234", display_name=display_name), True + await set_password(db, user, new_password="Test1234") + user.display_name = display_name + user.is_active = True + return user, False @pytest_asyncio.fixture(autouse=True) @@ -31,28 +44,20 @@ async def _dispose_engine() -> AsyncIterator[None]: @pytest.mark.asyncio async def test_row_publication_revisions_and_source_validation(tmp_path: Path) -> None: - owner_id = uuid.uuid4() - other_owner_id = uuid.uuid4() dataset_id = uuid.uuid4() csv_path = tmp_path / "rows.csv" contents = b"question,answer\nfirst,one\nsecond,two\n" csv_path.write_bytes(contents) digest = hashlib.sha256(contents).hexdigest() async with get_session() as db: - owner = User( - id=owner_id, - email="test@test.com", - hashed_password=hash_password("test"), - display_name="Row Publisher", + owner, owner_created = await _row_test_user( + db, email="test@test.com", display_name="Row Publisher" ) - other = User( - id=other_owner_id, - email="other@test.com", - hashed_password=hash_password("test"), - display_name="Other Publisher", + other, other_created = await _row_test_user( + db, email="other@test.com", display_name="Other Publisher" ) - db.add_all([owner, other]) - await db.flush() + owner_id = owner.id + other_owner_id = other.id dataset = RegisteredDataset( id=dataset_id, name="cohort", @@ -161,5 +166,7 @@ async def test_row_publication_revisions_and_source_validation(tmp_path: Path) - await delete_protocol(db, protocol.id) await db.delete(dataset) await db.delete(other_dataset) - await db.delete(owner) - await db.delete(other) + if owner_created: + await db.delete(owner) + if other_created: + await db.delete(other) diff --git a/tests/test_dataset_row_queries.py b/tests/test_dataset_row_queries.py new file mode 100644 index 00000000..97ed8803 --- /dev/null +++ b/tests/test_dataset_row_queries.py @@ -0,0 +1,346 @@ +"""PostgreSQL tests for experiment and revision scoped row result queries.""" + +from __future__ import annotations + +import os +import uuid +from collections.abc import AsyncIterator +from datetime import UTC, datetime + +import pytest +import pytest_asyncio +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from asaree.models.dataset import RegisteredDataset +from asaree.models.experiment import ResearchExperiment +from asaree.models.experiment_design_revision import ExperimentDesignRevision +from asaree.models.factorial_cell import FactorialCell +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol import Protocol +from asaree.models.protocol_revision import ProtocolRevision +from asaree.models.protocol_run import ProtocolRun +from asaree.services.factorial_row_results import ( + ensure_row_result, + get_row_result, + list_row_attempts, + list_row_results, +) +from asaree.services.users import create_user, get_user_by_email, set_password + + +@pytest_asyncio.fixture +async def row_graph() -> AsyncIterator[tuple[AsyncSession, dict[str, uuid.UUID]]]: + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + sessions = async_sessionmaker(engine, expire_on_commit=False) + ids = { + key: uuid.uuid4() + for key in ( + "user", + "dataset", + "experiment", + "other_experiment", + "design", + "old_design", + "other_design", + "protocol", + "other_protocol", + "publication", + "other_publication", + "replicate_a", + "replicate_b", + "other_replicate", + "slot_a", + "slot_b", + "other_slot", + "historical_replicate", + "historical_slot", + "other_user", + "foreign_dataset", + ) + } + async with sessions() as db: + user = await get_user_by_email(db, "test@test.com") + if user is None: + user = await create_user(db, email="test@test.com", password="Test1234", display_name="Row test") + else: + await set_password(db, user, new_password="Test1234") + user.display_name = "Row test" + user.is_active = True + ids["user"] = user.id + db.add_all( + [ + RegisteredDataset(id=ids["dataset"], name=f"rows-{ids['dataset']}", owner_id=ids["user"]), + ResearchExperiment(id=ids["experiment"], name=f"exp-{ids['experiment']}", owner_id=ids["user"]), + ResearchExperiment( + id=ids["other_experiment"], name=f"exp-{ids['other_experiment']}", owner_id=ids["user"] + ), + ] + ) + other_user = await get_user_by_email(db, "other@test.com") + if other_user is None: + other_user = await create_user( + db, email="other@test.com", password="Test1234", display_name="Other row test" + ) + else: + await set_password(db, other_user, new_password="Test1234") + other_user.display_name = "Other row test" + other_user.is_active = True + ids["other_user"] = other_user.id + db.add( + RegisteredDataset( + id=ids["foreign_dataset"], name=f"foreign-{ids['foreign_dataset']}", owner_id=ids["other_user"] + ) + ) + await db.flush() + db.add_all( + [ + ExperimentDesignRevision( + id=ids["old_design"], experiment_id=ids["experiment"], revision=1, superseded_at=datetime.now(UTC) + ), + ExperimentDesignRevision(id=ids["design"], experiment_id=ids["experiment"], revision=2), + ExperimentDesignRevision(id=ids["other_design"], experiment_id=ids["other_experiment"], revision=1), + Protocol( + id=ids["protocol"], + name=f"protocol-{ids['protocol']}", + owner_id=ids["user"], + experiment_id=ids["experiment"], + graph={}, + ), + Protocol( + id=ids["other_protocol"], + name=f"protocol-{ids['other_protocol']}", + owner_id=ids["user"], + experiment_id=ids["other_experiment"], + graph={}, + ), + ] + ) + await db.flush() + db.add_all( + [ + ProtocolRevision( + id=ids["publication"], + protocol_id=ids["protocol"], + revision=1, + graph={}, + published_at=datetime.now(UTC), + ), + ProtocolRevision( + id=ids["other_publication"], + protocol_id=ids["other_protocol"], + revision=1, + graph={}, + published_at=datetime.now(UTC), + ), + ] + ) + await db.flush() + protocol = await db.get(Protocol, ids["protocol"]) + other_protocol = await db.get(Protocol, ids["other_protocol"]) + protocol.published_revision_id = ids["publication"] + other_protocol.published_revision_id = ids["other_publication"] + await db.flush() + cells = [ + FactorialCell( + id=uuid.uuid4(), + experiment_id=ids["experiment"], + design_revision_id=ids["design"], + cell_label=label, + factor_values={}, + ) + for label in ("b", "a") + ] + old_cell = FactorialCell( + id=uuid.uuid4(), + experiment_id=ids["experiment"], + design_revision_id=ids["old_design"], + cell_label="historic", + factor_values={}, + ) + other_cell = FactorialCell( + id=uuid.uuid4(), + experiment_id=ids["other_experiment"], + design_revision_id=ids["other_design"], + cell_label="a", + factor_values={}, + ) + db.add_all([*cells, old_cell, other_cell]) + await db.flush() + reps = [ + FactorialReplicateResult( + id=ids["replicate_a"], cell_id=cells[0].id, replicate_number=1, replicate_label="same-label" + ), + FactorialReplicateResult( + id=ids["replicate_b"], cell_id=cells[1].id, replicate_number=1, replicate_label="same-label" + ), + FactorialReplicateResult( + id=ids["other_replicate"], cell_id=other_cell.id, replicate_number=1, replicate_label="same-label" + ), + FactorialReplicateResult( + id=ids["historical_replicate"], cell_id=old_cell.id, replicate_number=1, replicate_label="same-label" + ), + ] + db.add_all(reps) + await db.flush() + slots = [ + FactorialRowResult( + id=ids["slot_a"], + replicate_result_id=ids["replicate_a"], + protocol_revision_id=ids["publication"], + dataset_id=ids["dataset"], + raw_sha256="a" * 64, + row_index=2, + ), + FactorialRowResult( + id=ids["slot_b"], + replicate_result_id=ids["replicate_b"], + protocol_revision_id=ids["publication"], + dataset_id=ids["dataset"], + raw_sha256="a" * 64, + row_index=0, + ), + FactorialRowResult( + id=ids["other_slot"], + replicate_result_id=ids["other_replicate"], + protocol_revision_id=ids["other_publication"], + dataset_id=ids["dataset"], + raw_sha256="a" * 64, + row_index=0, + ), + FactorialRowResult( + id=ids["historical_slot"], + replicate_result_id=ids["historical_replicate"], + protocol_revision_id=ids["publication"], + dataset_id=ids["dataset"], + raw_sha256="a" * 64, + row_index=0, + ), + ] + db.add_all(slots) + await db.flush() + yield db, ids + await db.rollback() + await engine.dispose() + + +async def test_scoped_lookup_order_and_forged_mixed_scopes(row_graph) -> None: + db, ids = row_graph + rows = await list_row_results(db, experiment_id=ids["experiment"], protocol_revision_id=ids["publication"]) + assert [row.id for row in rows] == [ids["slot_b"], ids["slot_a"]] + assert ( + await get_row_result( + db, + experiment_id=ids["experiment"], + row_result_id=ids["other_slot"], + protocol_revision_id=ids["other_publication"], + ) + is None + ) + assert ( + await list_row_results( + db, + experiment_id=ids["experiment"], + design_revision_id=ids["other_design"], + protocol_revision_id=ids["publication"], + ) + == [] + ) + assert ( + await list_row_results(db, experiment_id=ids["experiment"], protocol_revision_id=ids["other_publication"]) == [] + ) + assert ( + await get_row_result( + db, + experiment_id=ids["experiment"], + row_result_id=ids["slot_a"], + design_revision_id=ids["old_design"], + protocol_revision_id=ids["publication"], + ) + is None + ) + historical = await list_row_results( + db, + experiment_id=ids["experiment"], + design_revision_id=ids["old_design"], + protocol_revision_id=ids["publication"], + ) + assert [row.id for row in historical] == [ids["historical_slot"]] + + +async def test_history_is_grouped_by_slot_identity_not_replicate_label(row_graph) -> None: + db, ids = row_graph + earlier = datetime(2025, 1, 1, tzinfo=UTC) + later = datetime(2025, 1, 2, tzinfo=UTC) + first_attempt = ProtocolRun( + id=uuid.UUID(int=10), + protocol_id=ids["protocol"], + owner_id=ids["user"], + node_runs={}, + row_result_id=ids["slot_a"], + protocol_revision_id=ids["publication"], + created_at=later, + ) + second_attempt = ProtocolRun( + id=uuid.UUID(int=11), + protocol_id=ids["protocol"], + owner_id=ids["user"], + node_runs={}, + row_result_id=ids["slot_a"], + protocol_revision_id=ids["publication"], + created_at=earlier, + ) + other_slot_attempt = ProtocolRun( + protocol_id=ids["protocol"], + owner_id=ids["user"], + node_runs={}, + row_result_id=ids["slot_b"], + protocol_revision_id=ids["publication"], + created_at=earlier, + ) + db.add_all([first_attempt, second_attempt, other_slot_attempt]) + await db.flush() + attempts = await list_row_attempts( + db, experiment_id=ids["experiment"], row_result_id=ids["slot_a"], protocol_revision_id=ids["publication"] + ) + assert [attempt.id for attempt in attempts] == [second_attempt.id, first_attempt.id] + assert all(attempt.row_result_id == ids["slot_a"] for attempt in attempts) + other_attempts = await list_row_attempts( + db, + experiment_id=ids["experiment"], + row_result_id=ids["slot_b"], + protocol_revision_id=ids["publication"], + ) + assert [attempt.id for attempt in other_attempts] == [other_slot_attempt.id] + + +async def test_ensure_is_idempotent_and_checks_parent_dataset_and_publication(row_graph) -> None: + db, ids = row_graph + kwargs = dict( + experiment_id=ids["experiment"], + design_revision_id=ids["design"], + protocol_revision_id=ids["publication"], + replicate_result_id=ids["replicate_a"], + dataset_id=ids["dataset"], + raw_sha256="b" * 64, + row_index=4, + ) + first = await ensure_row_result(db, **kwargs) + second = await ensure_row_result(db, **kwargs) + assert first.id == second.id + invalid_scopes = [ + {**kwargs, "design_revision_id": ids["old_design"]}, + {**kwargs, "replicate_result_id": ids["other_replicate"]}, + {**kwargs, "protocol_revision_id": ids["other_publication"]}, + {**kwargs, "protocol_revision_id": uuid.uuid4()}, + {**kwargs, "dataset_id": ids["foreign_dataset"]}, + ] + for bad in invalid_scopes: + with pytest.raises(ValueError, match="invalid_row_scope"): + await ensure_row_result(db, **bad) + for bad_value in ("A" * 64, "bad", "a" * 63): + with pytest.raises(ValueError, match="invalid_row_scope"): + await ensure_row_result(db, **{**kwargs, "raw_sha256": bad_value}) + for bad_value in (-1, True, 1.5): + with pytest.raises(ValueError, match="invalid_row_scope"): + await ensure_row_result(db, **{**kwargs, "row_index": bad_value}) diff --git a/tests/test_dataset_row_run_identity.py b/tests/test_dataset_row_run_identity.py new file mode 100644 index 00000000..0ea2b6df --- /dev/null +++ b/tests/test_dataset_row_run_identity.py @@ -0,0 +1,254 @@ +"""Database-backed identity tests for row-slot ProtocolRun attempts.""" + +from __future__ import annotations + +import hashlib +import os +import uuid +from collections.abc import AsyncIterator +from datetime import UTC, datetime + +import pytest +import pytest_asyncio +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +import asaree.api.protocols as protocol_api +from asaree.models.dataset import RegisteredDataset +from asaree.models.experiment import ResearchExperiment +from asaree.models.experiment_design_revision import ExperimentDesignRevision +from asaree.models.factorial_cell import FactorialCell +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol import Protocol +from asaree.models.protocol_revision import ProtocolRevision +from asaree.models.protocol_run import ProtocolRun +from asaree.services.dataset_row_csv import project_row, read_row_source +from asaree.services.protocol_runs import create_protocol_run +from asaree.services.users import create_user, get_user_by_email, set_password + + +@pytest_asyncio.fixture +async def row_graph(tmp_path) -> AsyncIterator[tuple[AsyncSession, dict]]: + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + async with engine.begin() as connection: + sessions = async_sessionmaker(bind=connection, expire_on_commit=False) + async with sessions() as db: + user_id, experiment_id, dataset_id, protocol_id = (uuid.uuid4() for _ in range(4)) + design_id, protocol_revision_id = uuid.uuid4(), uuid.uuid4() + cell_id, replicate_id = uuid.uuid4(), uuid.uuid4() + raw = b"subject,value\na,1\nb,2\n" + raw_path = tmp_path / "source.csv" + raw_path.write_bytes(raw) + digest = hashlib.sha256(raw).hexdigest() + source = read_row_source(dataset_id=str(dataset_id), raw_path=str(raw_path), raw_sha256=digest) + workspace = tmp_path / "workspace" + workspace.mkdir() + (workspace / "trace.json").write_text('{"run":"fixture"}', encoding="utf-8") + user = await get_user_by_email(db, "test@test.com") + if user is None: + user = await create_user(db, email="test@test.com", password="Test1234", display_name="Row test") + else: + await set_password(db, user, new_password="Test1234") + user.display_name = "Row test" + user.is_active = True + user_id = user.id + experiment = ResearchExperiment(id=experiment_id, name=f"row-{experiment_id}", owner_id=user_id) + dataset = RegisteredDataset( + id=dataset_id, name=f"rows-{dataset_id}", owner_id=user_id, raw_path=str(raw_path), raw_sha256=digest + ) + protocol = Protocol( + id=protocol_id, name=f"protocol-{protocol_id}", owner_id=user_id, experiment_id=experiment_id, graph={} + ) + db.add_all([experiment, dataset, protocol]) + await db.flush() + design = ExperimentDesignRevision(id=design_id, experiment_id=experiment_id, revision=1) + revision = ProtocolRevision( + id=protocol_revision_id, protocol_id=protocol_id, revision=1, graph={}, published_at=datetime.now(UTC) + ) + db.add_all([design, revision]) + await db.flush() + protocol.published_revision_id = protocol_revision_id + cell = FactorialCell( + id=cell_id, + experiment_id=experiment_id, + design_revision_id=design_id, + cell_label="cell-1", + factor_values={"condition": "same"}, + ) + db.add(cell) + await db.flush() + replicate = FactorialReplicateResult( + id=replicate_id, + cell_id=cell_id, + replicate_number=1, + replicate_label="replicate-1", + run_id=uuid.uuid4(), + workspace_id="parent-workspace", + metric_values={"old": 1}, + artifacts={"old": True}, + ) + db.add(replicate) + await db.flush() + rows = [ + FactorialRowResult( + replicate_result_id=replicate_id, + protocol_revision_id=protocol_revision_id, + dataset_id=dataset_id, + raw_sha256=digest, + row_index=index, + run_id=uuid.uuid4(), + workspace_id=f"row-workspace-{index}", + metric_values={"old": index}, + artifacts={"old": index}, + ) + for index in range(2) + ] + db.add_all(rows) + await db.flush() + context = { + "owner_id": user_id, + "experiment_id": experiment_id, + "dataset_id": dataset_id, + "protocol_id": protocol_id, + "design_id": design_id, + "protocol_revision_id": protocol_revision_id, + "replicate_id": replicate_id, + "replicate": replicate, + "rows": rows, + "digest": digest, + "workspace": workspace, + "source": source, + } + try: + yield db, context + finally: + await db.rollback() + await engine.dispose() + + +def _snapshot(ctx: dict, row_index: int) -> dict: + return project_row(ctx["source"], row_index=row_index, columns=["subject", "value"]) + + +async def _create(db: AsyncSession, ctx: dict, index: int) -> ProtocolRun: + return await create_protocol_run( + db, + protocol_id=ctx["protocol_id"], + owner_id=ctx["owner_id"], + replicate_label="replicate-1", + factor_values={"condition": "same"}, + replicate_result_id=ctx["replicate_id"], + row_result_id=ctx["rows"][index].id, + dataset_row=_snapshot(ctx, index), + design_revision_id=ctx["design_id"], + protocol_revision_id=ctx["protocol_revision_id"], + ) + + +async def test_row_slots_have_independent_attempt_identity_and_projection(row_graph) -> None: + db, ctx = row_graph + assert (ctx["workspace"] / "trace.json").is_file() + parent_before = ( + ctx["replicate"].run_id, + ctx["replicate"].workspace_id, + ctx["replicate"].metric_values, + ctx["replicate"].artifacts, + ) + old_row_run_id = ctx["rows"][0].run_id + first = await _create(db, ctx, 0) + second = await _create(db, ctx, 1) + assert first.id != second.id + assert first.replicate_label == second.replicate_label == "replicate-1" + assert first.row_result_id != second.row_result_id + assert first.attempt_result["row_provenance"]["row_result_id"] == str(ctx["rows"][0].id) + assert first.attempt_result["row_provenance"]["dataset_row"] == _snapshot(ctx, 0) + assert ctx["rows"][0].run_id == first.id and ctx["rows"][1].run_id == second.id + assert ( + ctx["rows"][0].workspace_id is None + and ctx["rows"][0].metric_values is None + and ctx["rows"][0].artifacts is None + ) + assert parent_before == ( + ctx["replicate"].run_id, + ctx["replicate"].workspace_id, + ctx["replicate"].metric_values, + ctx["replicate"].artifacts, + ) + assert old_row_run_id != first.id + serialized = protocol_api._protocol_run_response(first) + assert serialized.row_result_id == first.row_result_id + assert serialized.dataset_row == _snapshot(ctx, 0) + + old_attempt_snapshot = dict(first.attempt_result) + replacement = await _create(db, ctx, 0) + assert ctx["rows"][0].run_id == replacement.id + assert first.attempt_result == old_attempt_snapshot + assert await db.scalar(select(ProtocolRun.id).where(ProtocolRun.id == first.id)) == first.id + assert parent_before == ( + ctx["replicate"].run_id, + ctx["replicate"].workspace_id, + ctx["replicate"].metric_values, + ctx["replicate"].artifacts, + ) + + +async def test_mismatched_row_scope_snapshot_and_factors_rejected_before_run_insert(row_graph) -> None: + db, ctx = row_graph + for kwargs in ( + {"dataset_row": {**_snapshot(ctx, 0), "row_index": 1}}, + {"factor_values": {"condition": "other"}}, + {"design_revision_id": uuid.uuid4()}, + {"protocol_revision_id": uuid.uuid4()}, + {"replicate_result_id": uuid.uuid4()}, + ): + args = { + "protocol_id": ctx["protocol_id"], + "owner_id": ctx["owner_id"], + "replicate_label": "replicate-1", + "factor_values": {"condition": "same"}, + "replicate_result_id": ctx["replicate_id"], + "row_result_id": ctx["rows"][0].id, + "dataset_row": _snapshot(ctx, 0), + "design_revision_id": ctx["design_id"], + "protocol_revision_id": ctx["protocol_revision_id"], + } + args.update(kwargs) + with pytest.raises(ValueError): + await create_protocol_run(db, **args) + assert await db.scalar(select(ProtocolRun.id).where(ProtocolRun.row_result_id == ctx["rows"][0].id)) is None + + with pytest.raises(ValueError, match="invalid_dataset_row_snapshot"): + await create_protocol_run( + db, + protocol_id=ctx["protocol_id"], + owner_id=ctx["owner_id"], + dataset_row={**_snapshot(ctx, 0), "extra": True}, + ) + + +async def test_snapshot_only_and_legacy_runs_have_no_factorial_projection(row_graph) -> None: + db, ctx = row_graph + parent_before = ( + ctx["replicate"].run_id, + ctx["replicate"].workspace_id, + ctx["replicate"].metric_values, + ctx["replicate"].artifacts, + ) + preview = await create_protocol_run( + db, + protocol_id=ctx["protocol_id"], + owner_id=ctx["owner_id"], + target_node_id="node-preview", + dataset_row=_snapshot(ctx, 0), + ) + legacy = await create_protocol_run(db, protocol_id=ctx["protocol_id"], owner_id=ctx["owner_id"]) + assert preview.row_result_id is None and preview.replicate_result_id is None + assert preview.attempt_result["row_provenance"]["row_result_id"] is None + assert legacy.row_result_id is None and legacy.dataset_row is None and legacy.replicate_result_id is None + assert parent_before == ( + ctx["replicate"].run_id, + ctx["replicate"].workspace_id, + ctx["replicate"].metric_values, + ctx["replicate"].artifacts, + ) diff --git a/tests/test_dataset_row_storage.py b/tests/test_dataset_row_storage.py index 9f0f0267..8395da9f 100644 --- a/tests/test_dataset_row_storage.py +++ b/tests/test_dataset_row_storage.py @@ -28,7 +28,7 @@ from asaree.models.protocol_revision import ProtocolRevision from asaree.models.protocol_run import ProtocolRun from asaree.models.user import User -from asaree.security.passwords import hash_password +from asaree.services.users import create_user, get_user_by_email, set_password @pytest_asyncio.fixture @@ -41,15 +41,14 @@ async def storage_graph() -> AsyncIterator[tuple[AsyncSession, dict[str, uuid.UU for name in ("user", "dataset", "experiment", "design", "cell", "replicate", "protocol", "protocol_revision") } async with sessions() as db: - db.add( - User( - id=ids["user"], - email="test@test.com", - hashed_password=hash_password("test"), - display_name="Row test", - ) - ) - await db.flush() + user = await get_user_by_email(db, "test@test.com") + if user is None: + user = await create_user(db, email="test@test.com", password="Test1234", display_name="Row test") + else: + await set_password(db, user, new_password="Test1234") + user.display_name = "Row test" + user.is_active = True + ids["user"] = user.id db.add_all( [ RegisteredDataset(id=ids["dataset"], name=f"row-{ids['dataset']}", owner_id=ids["user"]), From a06831b6f5427b1a9711106f32c30621420d762c Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Sun, 4 Oct 2026 21:43:46 -0700 Subject: [PATCH 06/16] fix: complete dataset row batch planning --- src/asaree/api/protocols.py | 2 +- src/asaree/services/dataset_row_planning.py | 5 +- tests/test_dataset_row_batch_planning.py | 172 ++++++++++++++++++++ 3 files changed, 177 insertions(+), 2 deletions(-) create mode 100644 tests/test_dataset_row_batch_planning.py diff --git a/src/asaree/api/protocols.py b/src/asaree/api/protocols.py index 4a1d5ecd..1ebf0aa4 100644 --- a/src/asaree/api/protocols.py +++ b/src/asaree/api/protocols.py @@ -17,12 +17,12 @@ from pydantic import BaseModel, ConfigDict, Field from asaree.deps import CurrentUser, DbSession +from asaree.services.dataset_row_inputs import resolve_dataset_row_plan from asaree.services.experiment_measurements import ( blocking_measurement_plan_issues, validate_experiment_measurement_plan, ) from asaree.services.experiments import get_experiment -from asaree.services.dataset_row_inputs import resolve_dataset_row_plan from asaree.services.factor_bindings import validate_factor_bindings from asaree.services.protocol_execution import ( ProtocolValidationError, diff --git a/src/asaree/services/dataset_row_planning.py b/src/asaree/services/dataset_row_planning.py index 28c3fcfc..e8e24430 100644 --- a/src/asaree/services/dataset_row_planning.py +++ b/src/asaree/services/dataset_row_planning.py @@ -13,7 +13,6 @@ from asaree.services.dataset_row_csv import RowSource from asaree.services.dataset_row_inputs import DatasetRowInputError, resolve_dataset_row_plan -from asaree.services.protocol_execution import apply_factor_bindings def enumerate_row_candidates( @@ -29,6 +28,10 @@ def enumerate_row_candidates( Earlier graph and measurement validation remains the caller's responsibility. No database access, scheduling, or input mutation occurs. """ + # Imported lazily to keep this pure helper independent at module import + # time: protocol_execution calls us while it is itself being imported. + from asaree.services.protocol_execution import apply_factor_bindings + if not source.rows: raise DatasetRowInputError("empty_source", "row source must contain at least one data row") diff --git a/tests/test_dataset_row_batch_planning.py b/tests/test_dataset_row_batch_planning.py new file mode 100644 index 00000000..ca2e9c5e --- /dev/null +++ b/tests/test_dataset_row_batch_planning.py @@ -0,0 +1,172 @@ +"""PostgreSQL integration coverage for row-mode production batches.""" + +from __future__ import annotations + +import hashlib +import os +import uuid +from collections.abc import AsyncIterator +from pathlib import Path + +import pytest +import pytest_asyncio +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +import asaree.api.protocols as protocol_api +from asaree.api.protocols import CellRunBatchRequest, create_cell_runs_endpoint +from asaree.models.dataset import RegisteredDataset +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.protocol_run import ProtocolRun +from asaree.services.experiments import create_experiment +from asaree.services.factorial_cells import upsert_replicate +from asaree.services.protocol_revisions import publish_protocol +from asaree.services.protocols import create_protocol +from asaree.services.users import get_user_by_email + + +@pytest_asyncio.fixture +async def row_batch(tmp_path: Path) -> AsyncIterator[tuple[AsyncSession, dict]]: + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + async with engine.begin() as connection: + sessions = async_sessionmaker(bind=connection, expire_on_commit=False) + async with sessions() as db: + user = await get_user_by_email(db, "test@test.com") + assert user is not None, "run scripts/seed_row_test_users.py before row database tests" + raw = b"question,answer\n" + b"prompt,response\n" * 30 + raw_path = tmp_path / "batch.csv" + raw_path.write_bytes(raw) + dataset_id = uuid.uuid4() + dataset = RegisteredDataset( + id=dataset_id, + name=f"batch-{dataset_id}", + owner_id=user.id, + raw_path=str(raw_path), + raw_sha256=hashlib.sha256(raw).hexdigest(), + ) + db.add(dataset) + experiment = await create_experiment( + db, + name=f"row-batch-{uuid.uuid4()}", + owner_id=user.id, + design_spec={"factors": [], "metrics": []}, + ) + await upsert_replicate( + db, + experiment_id=experiment.id, + replicate_label="batch-parent", + fields={"factor_values": {}}, + ) + graph = { + "nodes": [ + { + "id": "dataset", + "type": "dataset", + "data": { + "config": {"dataset_id": str(dataset_id), "dataset_name": dataset.name} + }, + }, + {"id": "agent", "type": "agent", "data": {}}, + {"id": "model", "type": "model_openai", "data": {"config": {}}}, + ], + "edges": [ + { + "id": "dataset-agent", + "source": "dataset", + "target": "agent", + "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["question"]}}, + }, + { + "id": "model-agent", + "source": "model", + "target": "agent", + "targetHandle": "model", + }, + ], + } + protocol = await create_protocol( + db, + name=f"row-protocol-{uuid.uuid4()}", + owner_id=user.id, + experiment_id=experiment.id, + graph=graph, + ) + revision = await publish_protocol(db, protocol, owner_id=user.id) + protocol.published_revision_id = revision.id + await db.flush() + yield db, {"user": user, "protocol": protocol, "revision": revision} + await engine.dispose() + + +@pytest.mark.asyncio +async def test_batch_persists_each_original_row_once(row_batch, monkeypatch: pytest.MonkeyPatch) -> None: + db, ctx = row_batch + queued: list[uuid.UUID] = [] + + async def enqueue(run_id: uuid.UUID) -> None: + queued.append(run_id) + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + first = await create_cell_runs_endpoint( + ctx["protocol"].id, ctx["user"], db, CellRunBatchRequest() + ) + assert len(first.protocol_run_ids) == 30 + assert len(first.row_result_ids) == 30 + assert first.replicate_labels == ["batch-parent"] * 30 + assert first.consumption_mode == "per_row" + assert first.skipped == 0 + failed_run = await db.get(ProtocolRun, first.protocol_run_ids[0]) + completed_run = await db.get(ProtocolRun, first.protocol_run_ids[1]) + assert failed_run is not None and completed_run is not None + failed_run.status = "failed" + completed_run.status = "completed" + parent = await db.scalar( + select(FactorialReplicateResult).where(FactorialReplicateResult.replicate_label == "batch-parent") + ) + assert parent is not None and parent.run_id is None and parent.metric_values is None + + second = await create_cell_runs_endpoint( + ctx["protocol"].id, ctx["user"], db, CellRunBatchRequest() + ) + assert second.protocol_run_ids == [] + assert second.row_result_ids == [] + assert second.skipped == 30 + assert len(queued) == 30 + assert await db.scalar( + select(func.count()).select_from(ProtocolRun).where( + ProtocolRun.protocol_id == ctx["protocol"].id, + ProtocolRun.is_test_run.is_(False), + ) + ) == 30 + + +@pytest.mark.asyncio +async def test_whole_dataset_batch_response_keeps_legacy_mode(row_batch, monkeypatch: pytest.MonkeyPatch) -> None: + db, ctx = row_batch + async def enqueue(_run_id: uuid.UUID) -> None: + return None + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + protocol = ctx["protocol"] + protocol.graph = { + "nodes": [ + {"id": "agent", "type": "agent", "data": {}}, + {"id": "model", "type": "model_openai", "data": {"config": {}}}, + ], + "edges": [ + { + "id": "model-agent", + "source": "model", + "target": "agent", + "targetHandle": "model", + } + ], + } + revision = await publish_protocol(db, protocol, owner_id=ctx["user"].id) + protocol.published_revision_id = revision.id + # This regression checks response mode only; the row parent is still valid + # for the existing whole-dataset planner path. + result = await create_cell_runs_endpoint(protocol.id, ctx["user"], db) + assert result.consumption_mode == "whole_dataset" + assert result.row_result_ids == [] From b69866297349cd440755f1bd1633079e5cb31e2c Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Mon, 5 Oct 2026 11:42:30 -0700 Subject: [PATCH 07/16] Fix duplicate row result ID in serialized attempt claims --- src/asaree/services/factorial_row_results.py | 110 ++++++++++- tests/test_dataset_row_claims.py | 196 +++++++++++++++++++ 2 files changed, 301 insertions(+), 5 deletions(-) create mode 100644 tests/test_dataset_row_claims.py diff --git a/src/asaree/services/factorial_row_results.py b/src/asaree/services/factorial_row_results.py index c73fe12f..bb82ade8 100644 --- a/src/asaree/services/factorial_row_results.py +++ b/src/asaree/services/factorial_row_results.py @@ -5,7 +5,8 @@ import re import uuid -from sqlalchemy import select +from sqlalchemy import exists, select, update +from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from asaree.models.dataset import RegisteredDataset @@ -18,8 +19,54 @@ from asaree.models.protocol_revision import ProtocolRevision from asaree.models.protocol_run import ProtocolRun from asaree.services.design_revisions import get_current_revision +from asaree.services.protocol_runs import TERMINAL_PROTOCOL_RUN_STATUSES _SHA256 = re.compile(r"^[0-9a-f]{64}$") +_ROW_PROJECTION_FIELDS = frozenset({"workspace_id", "metric_values", "artifacts"}) + + +async def project_row_attempt( + db: AsyncSession, *, row_result_id: uuid.UUID, run_id: uuid.UUID, fields: dict +) -> bool: + """Atomically update a row slot only while ``run_id`` still owns it. + + The update is scoped through the owning cell and both pinned revisions. + The caller cannot use an arbitrary row UUID to project into another + experiment, and an earlier current-attempt read is never authorization. + """ + if not isinstance(fields, dict) or not fields or fields.keys() - _ROW_PROJECTION_FIELDS: + raise ValueError("invalid_row_projection_fields") + if not isinstance(row_result_id, uuid.UUID) or not isinstance(run_id, uuid.UUID): + raise ValueError("invalid_row_scope") + + valid_scope = exists( + select(FactorialRowResult.id) + .join( + FactorialReplicateResult, + FactorialReplicateResult.id == FactorialRowResult.replicate_result_id, + ) + .join(FactorialCell, FactorialCell.id == FactorialReplicateResult.cell_id) + .join(ProtocolRevision, ProtocolRevision.id == FactorialRowResult.protocol_revision_id) + .join(Protocol, Protocol.id == ProtocolRevision.protocol_id) + .join(ProtocolRun, ProtocolRun.id == run_id) + .where( + FactorialRowResult.id == row_result_id, + FactorialRowResult.run_id == run_id, + ProtocolRun.row_result_id == FactorialRowResult.id, + ProtocolRun.replicate_result_id == FactorialReplicateResult.id, + ProtocolRun.design_revision_id == FactorialCell.design_revision_id, + ProtocolRun.protocol_revision_id == FactorialRowResult.protocol_revision_id, + ProtocolRun.protocol_id == Protocol.id, + Protocol.experiment_id == FactorialCell.experiment_id, + ) + ) + result = await db.execute( + update(FactorialRowResult) + .where(FactorialRowResult.id == row_result_id, valid_scope) + .values(**fields) + .execution_options(synchronize_session=False) + ) + return result.rowcount == 1 async def _revision_id( @@ -187,9 +234,55 @@ def invalid() -> ValueError: raw_sha256=raw_sha256, row_index=row_index, ) - db.add(result) - await db.flush() - return result + # A concurrent planner can allocate the same identity after our lookup. + # Isolate the insert in a savepoint so a uniqueness race does not poison + # unrelated work in the caller's transaction. + try: + async with db.begin_nested(): + db.add(result) + await db.flush() + return result + except IntegrityError: + existing = ( + await db.execute( + select(FactorialRowResult).where( + FactorialRowResult.replicate_result_id == replicate_result_id, + FactorialRowResult.protocol_revision_id == protocol_revision_id, + FactorialRowResult.dataset_id == dataset_id, + FactorialRowResult.raw_sha256 == raw_sha256, + FactorialRowResult.row_index == row_index, + ) + ) + ).scalar_one() + return existing + + +async def claim_row_attempt( + db: AsyncSession, + *, + row_result_id: uuid.UUID, + expected_run_id: uuid.UUID | None, + create_kwargs: dict, +) -> ProtocolRun | None: + """Atomically claim the latest-attempt slot for a row execution.""" + slot = await db.scalar( + select(FactorialRowResult) + .where(FactorialRowResult.id == row_result_id) + .with_for_update() + .execution_options(populate_existing=True) + ) + if slot is None or slot.run_id != expected_run_id: + return None + if expected_run_id is not None: + previous = await db.get(ProtocolRun, expected_run_id, populate_existing=True) + if previous is None or previous.status not in TERMINAL_PROTOCOL_RUN_STATUSES: + return None + if previous.status not in {"failed", "cancelled"}: + return None + from asaree.services.protocol_runs import create_protocol_run + + # The locked slot defines the identity, even when callers include it in kwargs. + return await create_protocol_run(db, **{**create_kwargs, "row_result_id": row_result_id}) async def list_row_attempts( @@ -225,4 +318,11 @@ async def list_row_attempts( ) -__all__ = ["ensure_row_result", "get_row_result", "list_row_attempts", "list_row_results"] +__all__ = [ + "claim_row_attempt", + "ensure_row_result", + "get_row_result", + "list_row_attempts", + "list_row_results", + "project_row_attempt", +] diff --git a/tests/test_dataset_row_claims.py b/tests/test_dataset_row_claims.py new file mode 100644 index 00000000..e7712a92 --- /dev/null +++ b/tests/test_dataset_row_claims.py @@ -0,0 +1,196 @@ +"""PostgreSQL integration checks for serialized dataset row attempt claims.""" + +from __future__ import annotations + +import asyncio +import hashlib +import os +import uuid +from collections.abc import AsyncIterator +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest +import pytest_asyncio +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +import asaree.api.protocols as protocol_api +from asaree.api.protocols import CellRunBatchRequest, create_cell_runs_endpoint +from asaree.models.dataset import RegisteredDataset +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol_run import ProtocolRun +from asaree.services.experiments import create_experiment +from asaree.services.factorial_cells import upsert_replicate +from asaree.services.factorial_row_results import claim_row_attempt +from asaree.services.protocol_revisions import publish_protocol +from asaree.services.protocols import create_protocol +from asaree.services.users import get_user_by_email + + +@pytest.mark.parametrize("include_row_result_id", [False, True]) +async def test_claim_retry_forwards_row_result_id_once(monkeypatch, include_row_result_id) -> None: + slot_id, previous_run_id = uuid.uuid4(), uuid.uuid4() + slot = FactorialRowResult(id=slot_id, run_id=previous_run_id) + previous = ProtocolRun(id=previous_run_id, status="failed") + db = AsyncMock(spec=AsyncSession) + db.scalar.return_value = slot + db.get.return_value = previous + create_run = AsyncMock(return_value=ProtocolRun(id=uuid.uuid4(), row_result_id=slot_id)) + monkeypatch.setattr("asaree.services.protocol_runs.create_protocol_run", create_run) + create_kwargs = {"protocol_id": uuid.uuid4(), "owner_id": uuid.uuid4()} + if include_row_result_id: + create_kwargs["row_result_id"] = slot_id + original_kwargs = create_kwargs.copy() + + run = await claim_row_attempt( + db, + row_result_id=slot_id, + expected_run_id=previous_run_id, + create_kwargs=create_kwargs, + ) + + assert run is create_run.return_value + create_run.assert_awaited_once_with(db, **{**create_kwargs, "row_result_id": slot_id}) + assert create_kwargs == original_kwargs + + +@pytest_asyncio.fixture +async def row_batch(tmp_path: Path) -> AsyncIterator[tuple[AsyncSession, dict]]: + """Create row-batch setup committed for independently locking sessions.""" + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + sessions = async_sessionmaker(engine, expire_on_commit=False) + try: + async with sessions.begin() as setup: + user = await get_user_by_email(setup, "test@test.com") + assert user is not None, "run scripts/seed_row_test_users.py before row database tests" + raw = b"question,answer\n" + b"prompt,response\n" * 30 + raw_path = tmp_path / "batch.csv" + raw_path.write_bytes(raw) + dataset_id = uuid.uuid4() + dataset = RegisteredDataset( + id=dataset_id, + name=f"batch-{dataset_id}", + owner_id=user.id, + raw_path=str(raw_path), + raw_sha256=hashlib.sha256(raw).hexdigest(), + ) + setup.add(dataset) + experiment = await create_experiment( + setup, + name=f"row-batch-{uuid.uuid4()}", + owner_id=user.id, + design_spec={"factors": [], "metrics": []}, + ) + await upsert_replicate( + setup, + experiment_id=experiment.id, + replicate_label="batch-parent", + fields={"factor_values": {}}, + ) + graph = { + "nodes": [ + { + "id": "dataset", + "type": "dataset", + "data": {"config": {"dataset_id": str(dataset_id), "dataset_name": dataset.name}}, + }, + {"id": "agent", "type": "agent", "data": {}}, + {"id": "model", "type": "model_openai", "data": {"config": {}}}, + ], + "edges": [ + { + "id": "dataset-agent", + "source": "dataset", + "target": "agent", + "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["question"]}}, + }, + { + "id": "model-agent", + "source": "model", + "target": "agent", + "targetHandle": "model", + }, + ], + } + protocol = await create_protocol( + setup, + name=f"row-protocol-{uuid.uuid4()}", + owner_id=user.id, + experiment_id=experiment.id, + graph=graph, + ) + revision = await publish_protocol(setup, protocol, owner_id=user.id) + protocol_id, user_id, revision_id = protocol.id, user.id, revision.id + + async with sessions() as db: + user = await db.get(type(user), user_id) + protocol = await db.get(type(protocol), protocol_id) + revision = await db.get(type(revision), revision_id) + assert user is not None and protocol is not None and revision is not None + yield db, {"user": user, "protocol": protocol, "revision": revision} + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_claim_waits_for_committed_claim_and_never_duplicates(row_batch) -> None: + db, ctx = row_batch + async def no_enqueue(_run_id: uuid.UUID) -> None: + return None + + original_enqueue = protocol_api.enqueue_protocol_run + protocol_api.enqueue_protocol_run = no_enqueue + try: + result = await create_cell_runs_endpoint(ctx["protocol"].id, ctx["user"], db, CellRunBatchRequest()) + finally: + protocol_api.enqueue_protocol_run = original_enqueue + slot_id = result.row_result_ids[0] + protocol = ctx["protocol"] + run = await db.get(ProtocolRun, result.protocol_run_ids[0]) + assert run is not None + + # The slot row lock is the production claim boundary. The second session + # waits at that lock and observes the first transaction's latest run id. + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + sessions = async_sessionmaker(engine, expire_on_commit=False) + locked = asyncio.Event() + second_started = asyncio.Event() + async with sessions() as first, sessions() as second: + async with first.begin(): + slot = await first.scalar( + select(FactorialRowResult).where(FactorialRowResult.id == slot_id).with_for_update() + ) + assert slot is not None + locked.set() + + async def competing_claim(): + async with second.begin(): + second_started.set() + return await claim_row_attempt( + second, + row_result_id=slot_id, + expected_run_id=None, + create_kwargs={ + "protocol_id": protocol.id, + "owner_id": ctx["user"].id, + "replicate_label": run.replicate_label, + "factor_values": run.factor_values, + "replicate_result_id": run.replicate_result_id, + "design_revision_id": run.design_revision_id, + "protocol_revision_id": run.protocol_revision_id, + "dataset_row": run.dataset_row, + }, + ) + + task = asyncio.create_task(competing_claim()) + await locked.wait() + await second_started.wait() + assert await task is None + async with sessions() as check: + count = await check.scalar( + select(func.count()).select_from(ProtocolRun).where(ProtocolRun.row_result_id == slot_id) + ) + assert count == 1 + await engine.dispose() From d9f242659e59b49ebb2675991ab34235adf664a0 Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Mon, 5 Oct 2026 16:33:37 -0700 Subject: [PATCH 08/16] feat: complete per-row dataset execution and results --- frontend/src/api/client.ts | 48 +- frontend/src/api/datasetRows.test.ts | 56 ++ .../protocol/AgentNodeInspector.tsx | 10 +- .../protocol/DatasetRowActivation.test.tsx | 36 + .../protocol/DatasetRowHistory.test.tsx | 34 + .../protocol/DatasetRowPlay.test.tsx | 46 ++ .../protocol/DatasetRowPromptPreview.test.tsx | 22 + .../protocol/DatasetRowResults.test.tsx | 44 ++ .../components/protocol/DatasetRowResults.tsx | 67 ++ .../protocol/DatasetRowRetry.test.tsx | 20 + .../protocol/DatasetRowRuns.test.tsx | 38 ++ .../components/protocol/DatasetRowRuns.tsx | 48 ++ .../protocol/DatasetRowSelector.tsx | 6 + .../protocol/DatasetRowTestRun.test.tsx | 76 +++ .../protocol/ExperimentSidePanel.tsx | 2 + .../protocol/PromptPreviewPanel.tsx | 19 +- .../components/protocol/ProtocolCanvas.tsx | 49 +- .../protocol/ProtocolCanvasContext.tsx | 1 + .../src/components/protocol/ResultsTab.tsx | 34 +- frontend/src/components/protocol/RunsTab.tsx | 14 +- .../components/protocol/TestRunResults.tsx | 3 +- .../components/protocol/cells/CellsTab.tsx | 10 +- .../protocol/edges/DatasetInputPanel.test.tsx | 43 ++ .../protocol/edges/DatasetInputPanel.tsx | 43 ++ .../protocol/edges/InteractEdge.tsx | 14 +- .../protocol/useDatasetRowSelection.ts | 16 + frontend/src/lib/datasetRows.ts | 39 ++ frontend/src/lib/protocolGraph.test.ts | 7 + frontend/src/lib/protocolGraph.ts | 10 +- frontend/src/lib/reportedMetrics.ts | 7 + frontend/src/pages/ProtocolCanvasPage.tsx | 9 +- frontend/src/types/experiments.ts | 56 ++ frontend/src/types/protocols.ts | 24 + src/asaree/api/experiments.py | 172 ++++- src/asaree/api/protocols.py | 198 +++++- src/asaree/mcp_servers/script_server.py | 57 +- src/asaree/mcp_servers/workspace_server.py | 72 ++ src/asaree/script_context.py | 10 + src/asaree/services/agent_messenger.py | 27 + src/asaree/services/csv_export.py | 190 +++++- src/asaree/services/dataset_row_workspaces.py | 179 +++++ src/asaree/services/design_generation.py | 25 +- src/asaree/services/experiment_run_results.py | 405 ++++++++++- src/asaree/services/factorial_analysis.py | 27 +- src/asaree/services/factorial_row_results.py | 12 +- src/asaree/services/protocol_execution.py | 641 ++++++++++++++++-- src/asaree/services/protocol_revisions.py | 8 + src/asaree/services/protocol_runs.py | 203 +++++- src/asaree/services/reported_metrics.py | 3 +- src/asaree/services/runtime_metrics.py | 28 +- src/asaree/services/test_run_results.py | 5 + tests/test_dataset_row_acceptance.py | 433 ++++++++++++ tests/test_dataset_row_analysis_exclusion.py | 121 ++++ tests/test_dataset_row_batch_planning.py | 63 ++ tests/test_dataset_row_csv_export.py | 139 ++++ tests/test_dataset_row_executor.py | 194 ++++++ tests/test_dataset_row_finalization.py | 415 ++++++++++++ tests/test_dataset_row_node_play.py | 386 +++++++++++ tests/test_dataset_row_prompt_preview.py | 133 ++++ tests/test_dataset_row_prompts.py | 142 ++++ tests/test_dataset_row_results_projection.py | 393 +++++++++++ tests/test_dataset_row_retry.py | 405 +++++++++++ tests/test_dataset_row_script_inputs.py | 224 ++++++ tests/test_dataset_row_single_run.py | 204 ++++++ tests/test_dataset_row_stale_writes.py | 301 ++++++++ tests/test_dataset_row_test_run.py | 308 +++++++++ tests/test_dataset_row_tool_bindings.py | 109 +++ tests/test_dataset_row_workspace_lifecycle.py | 129 ++++ tests/test_protocol_runs.py | 1 + tests/test_reported_metrics.py | 12 + .../src/asaree_workspace_core/__init__.py | 2 + .../src/asaree_workspace_core/context.py | 7 + 72 files changed, 7140 insertions(+), 194 deletions(-) create mode 100644 frontend/src/api/datasetRows.test.ts create mode 100644 frontend/src/components/protocol/DatasetRowActivation.test.tsx create mode 100644 frontend/src/components/protocol/DatasetRowHistory.test.tsx create mode 100644 frontend/src/components/protocol/DatasetRowPlay.test.tsx create mode 100644 frontend/src/components/protocol/DatasetRowPromptPreview.test.tsx create mode 100644 frontend/src/components/protocol/DatasetRowResults.test.tsx create mode 100644 frontend/src/components/protocol/DatasetRowResults.tsx create mode 100644 frontend/src/components/protocol/DatasetRowRetry.test.tsx create mode 100644 frontend/src/components/protocol/DatasetRowRuns.test.tsx create mode 100644 frontend/src/components/protocol/DatasetRowRuns.tsx create mode 100644 frontend/src/components/protocol/DatasetRowSelector.tsx create mode 100644 frontend/src/components/protocol/DatasetRowTestRun.test.tsx create mode 100644 frontend/src/components/protocol/edges/DatasetInputPanel.test.tsx create mode 100644 frontend/src/components/protocol/edges/DatasetInputPanel.tsx create mode 100644 frontend/src/components/protocol/useDatasetRowSelection.ts create mode 100644 frontend/src/lib/datasetRows.ts create mode 100644 frontend/src/lib/reportedMetrics.ts create mode 100644 src/asaree/services/dataset_row_workspaces.py create mode 100644 tests/test_dataset_row_acceptance.py create mode 100644 tests/test_dataset_row_analysis_exclusion.py create mode 100644 tests/test_dataset_row_csv_export.py create mode 100644 tests/test_dataset_row_executor.py create mode 100644 tests/test_dataset_row_finalization.py create mode 100644 tests/test_dataset_row_node_play.py create mode 100644 tests/test_dataset_row_prompt_preview.py create mode 100644 tests/test_dataset_row_prompts.py create mode 100644 tests/test_dataset_row_results_projection.py create mode 100644 tests/test_dataset_row_retry.py create mode 100644 tests/test_dataset_row_script_inputs.py create mode 100644 tests/test_dataset_row_single_run.py create mode 100644 tests/test_dataset_row_stale_writes.py create mode 100644 tests/test_dataset_row_test_run.py create mode 100644 tests/test_dataset_row_tool_bindings.py create mode 100644 tests/test_dataset_row_workspace_lifecycle.py diff --git a/frontend/src/api/client.ts b/frontend/src/api/client.ts index cd8ee2c5..0503eb29 100644 --- a/frontend/src/api/client.ts +++ b/frontend/src/api/client.ts @@ -29,7 +29,7 @@ import type { import type { LLMConnectionCheck, LLMProvider, LLMSetting, LLMSettingModelsResponse } from '@/types/llmSettings' import type { McpOAuthAuthorization, McpServer } from '@/types/mcpServers' import type { OkfBundle, OkfDocument } from '@/types/okf' -import type { CellRunBatch, PromptPreview, Protocol, ProtocolGraph, ProtocolRevision, ProtocolRun, TestRun } from '@/types/protocols' +import type { CellRunBatch, DatasetRowSchema, PromptPreview, Protocol, ProtocolGraph, ProtocolRevision, ProtocolRun, TestRun } from '@/types/protocols' import type { Run, RunStep } from '@/types/runs' import type { Skill, SkillListResponse, SkillUrlPreview } from '@/types/skills' @@ -181,6 +181,28 @@ export const tokenApi = { revoke: (id: string) => request(`/auth/me/tokens/${id}`, { method: 'DELETE' }), } +export interface ResultsScope { + protocol_id?: string + design_revision_id?: string + protocol_revision_id?: string +} + +export interface RowSelection { row_index?: number | null } + +function scopeQuery(scope?: ResultsScope): string { + const params = new URLSearchParams() + for (const [key, value] of Object.entries(scope ?? {})) if (value !== undefined) params.set(key, value) + return params.size ? `?${params}` : '' +} + +function rowSelection(options?: RowSelection): RowSelection | undefined { + const index = options?.row_index + if (index !== undefined && index !== null && (!Number.isInteger(index) || index < 0)) { + throw new RangeError('Source row must be a nonnegative integer') + } + return options +} + export const experimentsApi = { list: (opts?: { includeArchived?: boolean }) => request(opts?.includeArchived ? '/experiments?include_archived=true' : '/experiments'), @@ -255,7 +277,8 @@ export const experimentsApi = { // anywhere in the experiment (see services.csv_export.replicates_to_csv) -- // a Blob, not JSON, so callers hand it straight to URL.createObjectURL. downloadReplicatesCsv: (id: string) => requestBlob(`/experiments/${id}/replicates.csv`), - downloadRunResultsCsv: (id: string) => requestBlob(`/experiments/${id}/run-results.csv`), + downloadRunResultsCsv: (id: string, scope?: ResultsScope) => requestBlob(`/experiments/${id}/run-results.csv${scopeQuery(scope)}`), + getRunResultsSchema: (id: string, scope?: ResultsScope) => request>(`/experiments/${id}/run-results.schema.json${scopeQuery(scope)}`), // Materializes one cell per combination and its replicate-result children. // declared factors, returning the current design's replicates. If the new design // isn't the same set of cells as the current one, the current design @@ -275,7 +298,7 @@ export const experimentsApi = { getResults: (id: string) => request(`/experiments/${id}/results`), // A current-design scorecard and the per-cell/per-replicate evidence behind // it. Unlike getResults(), no balanced-factorial assumptions are required. - getRunResults: (id: string) => request(`/experiments/${id}/run-results`), + getRunResults: (id: string, scope?: ResultsScope) => request(`/experiments/${id}/run-results${scopeQuery(scope)}`), } export const protocolsApi = { @@ -292,20 +315,21 @@ export const protocolsApi = { // status "pending"; poll getRun for progress. cellLabel runs that one // already-generated cell for real (its own factor_values substituted in) // instead of today's ad-hoc, un-substituted whole-graph run. - run: (id: string, cellLabel?: string | null) => - request(`/protocols/${id}/runs`, { method: 'POST', body: { replicate_label: cellLabel ?? null } }), - testRun: (id: string) => request(`/protocols/${id}/test-runs`, { method: 'POST' }), + run: (id: string, cellLabel?: string | null, options?: RowSelection) => + request(`/protocols/${id}/runs`, { method: 'POST', body: { replicate_label: cellLabel ?? null, ...rowSelection(options) } }), + testRun: (id: string, options?: RowSelection) => request(`/protocols/${id}/test-runs`, { method: 'POST', body: rowSelection(options) }), getLatestTestRun: (id: string) => request(`/protocols/${id}/test-runs/latest`), // The per-node Play icon -- 422 if the node has upstream input or isn't a // runnable Agent (see validate_single_node_runnable). Same polling shape // as a plain run (getRun), just with node_runs carrying only this one key. - runNode: (id: string, nodeId: string) => request(`/protocols/${id}/nodes/${nodeId}/run`, { method: 'POST' }), + runNode: (id: string, nodeId: string, options?: RowSelection) => request(`/protocols/${id}/nodes/${nodeId}/run`, { method: 'POST', body: rowSelection(options) }), // Read-only despite the POST: the graph goes in the body because the canvas // being previewed is the one on screen, including edits autosave hasn't // flushed yet. Creates no run of any kind. 422 for a node that isn't an // agent, since only an agent is ever given a prompt. - promptPreview: (id: string, nodeId: string, graph: ProtocolGraph) => - request(`/protocols/${id}/nodes/${nodeId}/prompt-preview`, { method: 'POST', body: { graph } }), + promptPreview: (id: string, nodeId: string, graph: ProtocolGraph, options?: RowSelection) => + request(`/protocols/${id}/nodes/${nodeId}/prompt-preview`, { method: 'POST', body: { graph, ...rowSelection(options) } }), + listRevisions: (id: string) => request(`/protocols/${id}/revisions`), getRevision: (id: string, revisionId: string) => request(`/protocols/${id}/revisions/${revisionId}`), getRun: (id: string, runId: string) => request(`/protocols/${id}/runs/${runId}`), // Queued work is cancelled immediately. Active work raises @@ -316,19 +340,21 @@ export const protocolsApi = { // "Run all cells" -- 422 if there's no linked experiment or the graph // doesn't have exactly one final node; fans out one ProtocolRun per pending // replicate. The optional list explicitly re-runs selected completed rows. - runCells: (id: string, options?: { replicateLabels?: string[]; rerunReplicateLabels?: string[] }) => + runCells: (id: string, options?: { replicateLabels?: string[]; rerunReplicateLabels?: string[]; retry_row_result_ids?: string[] }) => request(`/protocols/${id}/cell-runs`, { method: 'POST', body: options ? { replicate_labels: options.replicateLabels, - rerun_replicate_labels: options.rerunReplicateLabels ?? [], + rerun_replicate_labels: options.retry_row_result_ids ? undefined : options.rerunReplicateLabels ?? [], + retry_row_result_ids: options.retry_row_result_ids, } : undefined, }), } export const datasetsApi = { + getRowSchema: (id: string) => request(`/datasets/${id}/row-schema`), // Owner-scoped, same convention as mcpServersApi.list -- backs the canvas's // dataset browser (DatasetBrowserPanel) and the Dataset node inspector's // read-out of whichever dataset the node is bound to. diff --git a/frontend/src/api/datasetRows.test.ts b/frontend/src/api/datasetRows.test.ts new file mode 100644 index 00000000..d2271cda --- /dev/null +++ b/frontend/src/api/datasetRows.test.ts @@ -0,0 +1,56 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { datasetsApi, experimentsApi, protocolsApi } from './client' + +const fetchMock = vi.fn() +beforeEach(() => { + localStorage.clear() + vi.stubGlobal('fetch', fetchMock) + fetchMock.mockImplementation(async () => new Response(JSON.stringify({ dataset_row: null, row_result_id: null }), { status: 200 })) +}) +afterEach(() => { vi.unstubAllGlobals(); vi.clearAllMocks() }) +const lastBody = () => { + const body = fetchMock.mock.calls.at(-1)?.[1]?.body + return body === undefined ? undefined : JSON.parse(body) +} +describe('row HTTP contracts', () => { + it('retains omitted legacy bodies and explicit null', async () => { + await protocolsApi.testRun('p') + expect(lastBody()).toBeUndefined() + await protocolsApi.runNode('p', 'a') + expect(lastBody()).toBeUndefined() + await protocolsApi.testRun('p', { row_index: null }) + expect(lastBody()).toEqual({ row_index: null }) + }) + it('sends zero-based selections on each execution interface', async () => { + await protocolsApi.testRun('p', { row_index: 2 }) + expect(lastBody()).toEqual({ row_index: 2 }) + await protocolsApi.runNode('p', 'a', { row_index: 2 }) + expect(lastBody()).toEqual({ row_index: 2 }) + await protocolsApi.run('p', null, { row_index: 2 }) + expect(lastBody()).toEqual({ replicate_label: null, row_index: 2 }) + const graph = { nodes: [], edges: [] } + await protocolsApi.promptPreview('p', 'a', graph, { row_index: 2 }) + expect(lastBody()).toEqual({ graph, row_index: 2 }) + expect(() => protocolsApi.testRun('p', { row_index: -1 })).toThrow(RangeError) + }) + it('retries exactly one row without a replicate rerun payload', async () => { + await protocolsApi.runCells('p', { retry_row_result_ids: ['slot'] }) + expect(lastBody()).toEqual({ retry_row_result_ids: ['slot'] }) + }) + it('uses identical encoded scope on projection and exports', async () => { + const scope = { protocol_id: 'p /', design_revision_id: 'd', protocol_revision_id: 'r' } + await experimentsApi.getRunResults('e', scope) + const query = fetchMock.mock.calls.at(-1)?.[0].split('?')[1] + await experimentsApi.downloadRunResultsCsv('e', scope) + expect(fetchMock.mock.calls.at(-1)?.[0].split('?')[1]).toBe(query) + await experimentsApi.getRunResultsSchema('e', scope) + expect(fetchMock.mock.calls.at(-1)?.[0].split('?')[1]).toBe(query) + }) + it('fetches row metadata and preserves returned identities', async () => { + const schema = { dataset_id: 'id', raw_sha256: 'hash', columns: ['question'], row_count: 3 } + fetchMock.mockResolvedValueOnce(new Response(JSON.stringify(schema))) + expect(await datasetsApi.getRowSchema('id')).toEqual(schema) + expect(fetchMock.mock.calls.at(-1)?.[0]).toContain('/datasets/id/row-schema') + expect(await protocolsApi.testRun('p')).toMatchObject({ dataset_row: null, row_result_id: null }) + }) +}) diff --git a/frontend/src/components/protocol/AgentNodeInspector.tsx b/frontend/src/components/protocol/AgentNodeInspector.tsx index 7778acd1..f1d5babe 100644 --- a/frontend/src/components/protocol/AgentNodeInspector.tsx +++ b/frontend/src/components/protocol/AgentNodeInspector.tsx @@ -67,6 +67,7 @@ export function AgentNodeInspector({ handoffPeers, wiredOutputParserLabel, fetchPromptPreview, + previewGraph, nodeRun, onChange, onDelete, @@ -90,7 +91,8 @@ export function AgentNodeInspector({ // Assembles the real prompt server-side against the live canvas. A callback // rather than an id pair because the graph it posts is the unsaved one on // screen, which only ProtocolCanvas holds. - fetchPromptPreview: (nodeId: string) => Promise + previewGraph?: import('@/types/protocols').ProtocolGraph + fetchPromptPreview: (nodeId: string, rowIndex?: number) => Promise nodeRun?: NodeRunState onChange: (nodeId: string, data: AgentNodeData) => void onDelete: (nodeId: string) => void @@ -228,8 +230,10 @@ export function AgentNodeInspector({ // The panel holds the last text it assembled; on a node switch // that text describes the previous node, so it starts over. key={node.id} - signature={JSON.stringify(data)} - fetchPreview={() => fetchPromptPreview(node.id)} + signature={JSON.stringify({ data, graph: previewGraph })} + graph={previewGraph} + nodeId={node.id} + fetchPreview={rowIndex => fetchPromptPreview(node.id, rowIndex)} /> {/* Below the design-time preview, because it supersedes it: a placeholder proves nothing about a run that actually happened. */} diff --git a/frontend/src/components/protocol/DatasetRowActivation.test.tsx b/frontend/src/components/protocol/DatasetRowActivation.test.tsx new file mode 100644 index 00000000..e78b5b23 --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowActivation.test.tsx @@ -0,0 +1,36 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { fireEvent, render, screen } from '@testing-library/react' +import { beforeEach, expect, it, vi } from 'vitest' +import { datasetsApi } from '@/api/client' +import { toPersistedGraph } from '@/lib/protocolGraph' +import { Position } from '@xyflow/react' +import { InteractEdge } from './edges/InteractEdge' + +const state = vi.hoisted(() => ({ edges: [] as any[], nodes: [] as any[], setEdges: vi.fn() })) +vi.mock('@xyflow/react', async importOriginal => ({ ...await importOriginal(), BaseEdge: () => null, EdgeLabelRenderer: ({ children }: any) => children, EdgeToolbar: ({ children }: any) => children, useNodes: () => state.nodes, useEdges: () => state.edges, useReactFlow: () => ({ getNode: (id: string) => state.nodes.find(node => node.id === id), setEdges: state.setEdges }) })) +vi.mock('./ProtocolCanvasContext', () => ({ useProtocolCanvasActions: () => ({ requestEdgeInsert: vi.fn(), experimentLocked: false }) })) +beforeEach(() => { + state.nodes = [{ id: 'd', type: 'dataset', position: { x: 0, y: 0 }, data: { config: { dataset_id: '11111111-1111-4111-8111-111111111111' } } }, { id: 'a', type: 'agent', position: { x: 0, y: 0 }, data: { config: {} } }] + state.edges = [{ id: 'edge', source: 'd', target: 'a', targetHandle: 'dataset', data: { custom: 'keep' } }] + state.setEdges.mockImplementation(fn => { state.edges = fn(state.edges) }) + vi.spyOn(datasetsApi, 'getRowSchema').mockResolvedValue({ dataset_id: 'id', raw_sha256: 'hash', row_count: 3, columns: ['question','reference'] }) +}) +function mount(handle = 'dataset') { + render() +} +it('Dataset edge settings apply through production serialization and reload', async () => { + mount() + fireEvent.click(screen.getByRole('button', { name: 'Dataset input settings' })) + await screen.findByText(/3 original rows/) + expect(screen.getByLabelText('Whole dataset')).toBeChecked() + fireEvent.click(screen.getByLabelText('Per row')) + expect(screen.getByRole('button', { name: 'Apply' })).toBeDisabled() + fireEvent.click(screen.getByRole('checkbox', { name: 'question' })) + fireEvent.click(screen.getByRole('button', { name: 'Apply' })) + const graph = toPersistedGraph(state.nodes, state.edges) + expect(JSON.parse(JSON.stringify(graph)).edges[0].data).toEqual({ custom: 'keep', dataset_input: { mode: 'per_row', columns: ['question'] } }) +}) +it('unrelated typed edge has no Dataset settings', () => { + mount('model') + expect(screen.queryByRole('button', { name: 'Dataset input settings' })).not.toBeInTheDocument() +}) diff --git a/frontend/src/components/protocol/DatasetRowHistory.test.tsx b/frontend/src/components/protocol/DatasetRowHistory.test.tsx new file mode 100644 index 00000000..11d492ac --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowHistory.test.tsx @@ -0,0 +1,34 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { fireEvent, render, screen, waitFor } from '@testing-library/react' +import { afterEach, expect, it, vi } from 'vitest' +import { experimentsApi, protocolsApi } from '@/api/client' +import type { Experiment, ExperimentRunResults } from '@/types/experiments' +import { ResultsInspectorPanel, ResultsTab } from './ResultsTab' +afterEach(() => vi.restoreAllMocks()) +vi.mock('./NodeRunOutputPanel', () => ({ RunStepTrace: ({ runId }: { runId: string }) =>

Trace {runId}

, ReceivedPromptPanel: () => null, UnresolvedReferencesNote: () => null })) +it('preserves failed attempt output while inspecting the latest row', async () => { + const row = { row_result_id: 'slot', replicate_number: 1, dataset_row: { dataset_id: 'dataset-uuid', raw_sha256: 'original-hash', row_index: 2 }, design_revision_id: 'design', protocol_revision_id: 'publication', latest_attempt: { run_id: 'B' }, attempts: [{ run_id: 'A', status: 'failed', current: false, error: 'old failure', node_runs: { agent: { run_id: 'trace-A', output_text: 'old output' } }, attempt_result: {} }, { run_id: 'B', status: 'completed', current: true, node_runs: { agent: { run_id: 'trace-B', output_text: 'new output' } }, attempt_result: {} }] } + vi.spyOn(experimentsApi, 'getRunResults').mockResolvedValue({ row_results: [row] } as unknown as ExperimentRunResults) + const scope = { protocol_id: 'p', design_revision_id: 'design', protocol_revision_id: 'publication' } + render() + await screen.findByText('new output') + expect(screen.getByText('original-hash')).toBeInTheDocument() + fireEvent.change(screen.getByRole('combobox', { name: 'Inspect attempt' }), { target: { value: 'A' } }) + expect(screen.getByText('old output')).toBeInTheDocument() + expect(screen.getByRole('alert')).toHaveTextContent('old failure') + expect(screen.getByText('Trace trace-A')).toBeInTheDocument() + expect(experimentsApi.getRunResults).toHaveBeenCalledWith('e', scope) +}) + +it('history selectors request their exact design and publication scope', async () => { + vi.spyOn(experimentsApi, 'listDesignRevisions').mockResolvedValue([{ id: 'design-old', revision: 1 }] as never) + vi.spyOn(protocolsApi, 'listRevisions').mockResolvedValue([{ id: 'publication-old', revision: 1 }] as never) + const results = vi.spyOn(experimentsApi, 'getRunResults').mockResolvedValue({ consumption_mode: 'per_row', row_results: [], row_summary: null } as unknown as ExperimentRunResults) + render() + await screen.findByRole('combobox', { name: 'Results design revision' }) + await waitFor(() => expect(screen.getAllByText('Revision 1')).toHaveLength(2)) + fireEvent.change(screen.getByRole('combobox', { name: 'Results design revision' }), { target: { value: 'design-old' } }) + fireEvent.change(await screen.findByRole('combobox', { name: 'Results published revision' }), { target: { value: 'publication-old' } }) + await waitFor(() => expect(results).toHaveBeenLastCalledWith('e', { protocol_id: 'p', design_revision_id: 'design-old', protocol_revision_id: 'publication-old' })) + expect(await screen.findByText('Read-only historical scope')).toBeInTheDocument() +}) diff --git a/frontend/src/components/protocol/DatasetRowPlay.test.tsx b/frontend/src/components/protocol/DatasetRowPlay.test.tsx new file mode 100644 index 00000000..50e28992 --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowPlay.test.tsx @@ -0,0 +1,46 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { ReactFlowProvider } from '@xyflow/react' +import { fireEvent, render, screen, waitFor } from '@testing-library/react' +import { MemoryRouter } from 'react-router-dom' +import { afterEach, beforeEach, expect, it, vi } from 'vitest' +import { datasetsApi, protocolsApi } from '@/api/client' +import { defaultAgentNodeData, defaultAnthropicModelNodeData, defaultDatasetNodeData, defaultReasonActPatternNodeData, type ProtocolGraph } from '@/types/protocols' +import { ProtocolCanvas } from './ProtocolCanvas' + +vi.mock('./PythonCodeEditor', () => ({ PythonCodeEditor: () => null })) +const graph: ProtocolGraph = { nodes: [ + { id: 'a', type: 'agent', position: { x: 100, y: 100 }, data: defaultAgentNodeData('Prediction') }, + { id: 'm', type: 'model_anthropic', position: { x: 0, y: 0 }, data: defaultAnthropicModelNodeData() }, + { id: 'd', type: 'dataset', position: { x: 0, y: 200 }, data: { ...defaultDatasetNodeData(), config: { ...defaultDatasetNodeData().config, dataset_id: '11111111-1111-4111-8111-111111111111', dataset_name: 'original' } } }, + { id: 'pattern', type: 'reason_act', position: { x: 0, y: 300 }, data: defaultReasonActPatternNodeData() }, + ], edges: [ { id: 'model', source: 'm', target: 'a', targetHandle: 'model' }, { id: 'pattern', source: 'pattern', target: 'a', targetHandle: 'architectural_pattern' }, { id: 'dataset', source: 'd', target: 'a', targetHandle: 'dataset', data: { dataset_input: { mode: 'per_row', columns: ['question'] } } } ] } +function setup() { + vi.stubGlobal('DOMMatrixReadOnly', class { m22 = 1 }) + vi.spyOn(HTMLElement.prototype, 'getBoundingClientRect').mockReturnValue({ x: 0, y: 0, top: 0, left: 0, right: 1000, bottom: 800, width: 1000, height: 800, toJSON: () => ({}) }) + vi.stubGlobal('ResizeObserver', class { private callback: ResizeObserverCallback; constructor(callback: ResizeObserverCallback) { this.callback = callback } observe(target: Element) { this.callback([{ target, contentRect: target.getBoundingClientRect() } as ResizeObserverEntry], this as unknown as ResizeObserver) } unobserve() {} disconnect() {} }) + vi.spyOn(protocolsApi, 'get').mockResolvedValue({ id: 'p', published_revision_id: 'r' } as never) + vi.spyOn(protocolsApi, 'getRevision').mockResolvedValue({ id: 'r', graph } as never) + vi.spyOn(datasetsApi, 'getRowSchema').mockResolvedValue({ dataset_id: 'd', raw_sha256: 'hash', row_count: 3, columns: ['question','reference'] }) + vi.spyOn(protocolsApi, 'listRuns').mockResolvedValue([]) + vi.spyOn(protocolsApi, 'getLatestTestRun').mockRejectedValue(new Error('No test')) + vi.spyOn(protocolsApi, 'update').mockResolvedValue({} as never) +} +beforeEach(setup) +afterEach(() => { vi.restoreAllMocks(); vi.unstubAllGlobals() }) +function mount() { + render(
) +} +it('node hover Play pins selected row and submits exactly one node request', async () => { + const play = vi.spyOn(protocolsApi, 'runNode').mockRejectedValue(new Error('mock queue')) + const testRun = vi.spyOn(protocolsApi, 'testRun') + mount() + const button = document.querySelector('button[aria-label^="Run this node"]') + expect(button).not.toBeNull() + fireEvent.click(button!) + const selector = await screen.findByRole('spinbutton', { name: 'Source row' }) + await waitFor(() => expect(selector).not.toBeDisabled()) + fireEvent.change(selector, { target: { value: '3' } }) + fireEvent.click(screen.getByRole('button', { name: /Run anyway|Confirm & run/ })) + await waitFor(() => expect(play).toHaveBeenCalledWith('p', 'a', { row_index: 2 })) + expect(testRun).not.toHaveBeenCalled() +}) diff --git a/frontend/src/components/protocol/DatasetRowPromptPreview.test.tsx b/frontend/src/components/protocol/DatasetRowPromptPreview.test.tsx new file mode 100644 index 00000000..e4f03334 --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowPromptPreview.test.tsx @@ -0,0 +1,22 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { fireEvent, render, screen, waitFor } from '@testing-library/react' +import { afterEach, expect, it, vi } from 'vitest' +import { datasetsApi, protocolsApi } from '@/api/client' +import { defaultAgentNodeData, defaultDatasetNodeData, type ProtocolGraph } from '@/types/protocols' +import { PromptPreviewPanel } from './PromptPreviewPanel' + +afterEach(() => vi.restoreAllMocks()) +it('draft preview selection requests only text with graph and index', async () => { + const graph: ProtocolGraph = { nodes: [{ id: 'agent', type: 'agent', position: { x: 0, y: 0 }, data: defaultAgentNodeData() }, { id: 'dataset', type: 'dataset', position: { x: 0, y: 0 }, data: { ...defaultDatasetNodeData(), config: { ...defaultDatasetNodeData().config, dataset_id: '11111111-1111-4111-8111-111111111111' } } }], edges: [{ id: 'e', source: 'dataset', target: 'agent', targetHandle: 'dataset', data: { dataset_input: { mode: 'per_row', columns: ['question'] } } }] } + vi.spyOn(datasetsApi, 'getRowSchema').mockResolvedValue({ dataset_id: 'draft', raw_sha256: 'draft-hash', columns: ['question'], row_count: 3 }) + const preview = vi.spyOn(protocolsApi, 'promptPreview').mockImplementation(async (_id, _node, _graph, selection) => ({ text: 'Authorized question', dataset_row: { dataset_id: 'draft', raw_sha256: 'draft-hash', row_index: selection?.row_index ?? 0, columns: ['question'], values: { question: 'Authorized question' } } })) + const run = vi.spyOn(protocolsApi, 'testRun') + render( protocolsApi.promptPreview('p', 'agent', graph, { row_index: rowIndex })} />) + fireEvent.click(screen.getByRole('button', { name: 'Prompt this agent will receive' })) + await screen.findByText('Authorized question') + expect(preview).toHaveBeenCalledWith('p', 'agent', graph, { row_index: 0 }) + fireEvent.change(screen.getByRole('spinbutton', { name: 'Source row' }), { target: { value: '3' } }) + await waitFor(() => expect(preview).toHaveBeenLastCalledWith('p', 'agent', graph, { row_index: 2 })) + await screen.findByText(/Source row 3 · draft · draft-hash/) + expect(run).not.toHaveBeenCalled() +}) diff --git a/frontend/src/components/protocol/DatasetRowResults.test.tsx b/frontend/src/components/protocol/DatasetRowResults.test.tsx new file mode 100644 index 00000000..ec195a9b --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowResults.test.tsx @@ -0,0 +1,44 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { fireEvent, render, screen } from '@testing-library/react' +import { afterEach, expect, it, vi } from 'vitest' +import { experimentsApi } from '@/api/client' +import type { Experiment, ExperimentRunResults } from '@/types/experiments' +import { ResultsTab } from './ResultsTab' +import { DatasetRowResults } from './DatasetRowResults' + +afterEach(() => vi.restoreAllMocks()) +it('renders opaque values and distinct row output selections without ranking', async () => { + const results = { consumption_mode: 'per_row', row_summary: { planned: 4, completed: 3, failed: 1, missing_reported: 1, metric_coverage: {} }, row_results: ['12', { answer: 2 }, null, undefined].map((value,index) => ({ row_result_id: `row-${index}`, factor_values: { Prompt: 'level' }, dataset_row: { row_index: index }, replicate_number: 2, status: index === 3 ? 'failed' : 'completed', latest_attempt: { run_id: `run-${index}` }, measurement: { observations: [{ metric_id: 'answer', status: index === 3 ? 'unavailable' : 'measured', value }] } })) } as unknown as ExperimentRunResults + vi.spyOn(experimentsApi, 'getRunResults').mockResolvedValue(results) + const experiment = { id: 'e', name: 'Rows', design_spec: { metrics: [{ id: 'answer', name: 'Answer', kind: 'custom' }] } } as unknown as Experiment + const inspect = vi.fn() + render() + await screen.findByText('null') + expect(screen.getByText('12')).toBeInTheDocument() + expect(screen.getByText('{"answer":2}')).toBeInTheDocument() + expect(screen.getByText('unavailable')).toBeInTheDocument() + expect(screen.queryByText('Best current result')).not.toBeInTheDocument() + fireEvent.click(screen.getAllByRole('button', { name: 'Inspect row' })[2]) + expect(inspect).toHaveBeenCalledWith({ type: 'row', rowResultId: 'row-2', scope: { protocol_id: 'p' } }) +}) + +it('paginates dense rows, maximizes with Escape, and exports the selected scope', async () => { + const scope = { protocol_id: 'p', protocol_revision_id: 'old', design_revision_id: 'design' } + const results = { row_results: Array.from({ length: 21 }, (_, index) => ({ row_result_id: `slot-${index}`, factor_values: {}, dataset_row: { row_index: index }, replicate_number: 1, status: 'completed', latest_attempt: { run_id: `run-${index}` }, measurement: { observations: [] } })) } as unknown as ExperimentRunResults + const csv = vi.spyOn(experimentsApi, 'downloadRunResultsCsv').mockResolvedValue(new Blob(['rows'])) + const create = vi.fn(() => 'blob:csv') + Object.defineProperty(URL, 'createObjectURL', { configurable: true, value: create }) + Object.defineProperty(URL, 'revokeObjectURL', { configurable: true, value: vi.fn() }) + vi.spyOn(HTMLAnchorElement.prototype, 'click').mockImplementation(() => {}) + render() + expect(screen.getAllByRole('button', { name: 'Inspect row' })).toHaveLength(20) + fireEvent.click(screen.getByRole('button', { name: 'Next' })) + expect(screen.getAllByRole('button', { name: 'Inspect row' })).toHaveLength(1) + fireEvent.click(screen.getByRole('button', { name: 'Maximize Results' })) + expect(document.body.style.overflow).toBe('hidden') + fireEvent.keyDown(document, { key: 'Escape' }) + expect(document.body.style.overflow).not.toBe('hidden') + fireEvent.click(screen.getByRole('button', { name: 'Download CSV' })) + expect(csv).toHaveBeenCalledWith('e', scope) + await screen.findByRole('button', { name: 'Download CSV' }) +}) diff --git a/frontend/src/components/protocol/DatasetRowResults.tsx b/frontend/src/components/protocol/DatasetRowResults.tsx new file mode 100644 index 00000000..a16ea0ee --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowResults.tsx @@ -0,0 +1,67 @@ +import { useEffect, useState } from 'react' +import { experimentsApi, type ResultsScope } from '@/api/client' +import { Button } from '@/components/ui/button' +import { reportedMetricText } from '@/lib/reportedMetrics' +import type { Experiment, ExperimentRunResults, RowResult } from '@/types/experiments' +import { RunStepTrace } from './NodeRunOutputPanel' + +export function DatasetRowResults({ results, experiment, scope, onInspect }: { results: ExperimentRunResults; experiment: Experiment; scope: ResultsScope; onInspect: (rowId: string) => void }) { + const [page, setPage] = useState(0) + const [maximized, setMaximized] = useState(false) + useEffect(() => { + if (!maximized) return + const previous = document.body.style.overflow + document.body.style.overflow = 'hidden' + const close = (event: KeyboardEvent) => { if (event.key === 'Escape') setMaximized(false) } + document.addEventListener('keydown', close) + return () => { document.body.style.overflow = previous; document.removeEventListener('keydown', close) } + }, [maximized]) + const [error, setError] = useState(null) + const [downloading, setDownloading] = useState(false) + const rows = results.row_results ?? [] + const currentPage = Math.min(page, Math.max(0, Math.ceil(rows.length / 20) - 1)) + const visibleRows = rows.slice(currentPage * 20, (currentPage + 1) * 20) + const factors = [...new Set(rows.flatMap(row => Object.keys(row.factor_values)))] + const declared = experiment.design_spec?.metrics?.filter(metric => metric.kind === 'custom') ?? [] + const observed = rows.flatMap(row => row.measurement?.observations ?? []).filter(item => item.producer?.kind !== 'runtime') + const metrics = [...new Map([...observed.map(item => ({ id: item.metric_id, name: item.metric_name })), ...declared].map(metric => [metric.id, metric])).values()] + async function download() { + setDownloading(true) + try { + const blob = await experimentsApi.downloadRunResultsCsv(experiment.id, scope) + const url = URL.createObjectURL(blob) + const link = document.createElement('a') + link.href = url + link.download = 'row-results.csv' + link.click() + URL.revokeObjectURL(url) + } catch (error) { setError(error instanceof Error ? error.message : 'CSV unavailable') } + finally { setDownloading(false) } + } + return
+

Per-row outcomes are reported individually. Aggregation, ranking and factorial analysis are unavailable.

+ {results.row_summary && <>

{results.row_summary.planned} planned · {results.row_summary.completed} completed · {results.row_summary.failed} failed · {results.row_summary.missing_reported} missing reported

{JSON.stringify(results.row_summary.metric_coverage, null, 2)}
} + + + {error &&

{error}

} +
{[...factors, 'Source row', 'Replicate', 'Execution status', 'Reported coverage', ...metrics.map(metric => metric.name), 'Output'].map((heading, index) => )}{visibleRows.map(row => { + const observations = row.measurement?.observations ?? [] + const measured = metrics.filter(metric => observations.some(item => item.metric_id === metric.id && item.status === 'measured')).length + return {factors.map(factor => )}{metrics.map(metric => )} + })}
{heading}
{String(row.factor_values[factor] ?? '')}{row.dataset_row.row_index + 1}{row.replicate_number}{row.status}{measured}/{metrics.length}{reportedMetricText(observations.find(item => item.metric_id === metric.id))}
+ {rows.length > 20 &&
Page {currentPage + 1} of {Math.ceil(rows.length / 20)}
} + {!rows.length &&

No row executions planned in this scope.

} +
+} + +export function DatasetRowDetail({ row, onClose }: { row: RowResult; onClose: () => void }) { + const [attemptId, setAttemptId] = useState(row.latest_attempt?.run_id) + const attempt = row.attempts.find(item => item.run_id === attemptId) ?? row.latest_attempt + const provenance = attempt?.attempt_result?.row_provenance + return +} diff --git a/frontend/src/components/protocol/DatasetRowRetry.test.tsx b/frontend/src/components/protocol/DatasetRowRetry.test.tsx new file mode 100644 index 00000000..b47e6715 --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowRetry.test.tsx @@ -0,0 +1,20 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { fireEvent, render, screen, waitFor } from '@testing-library/react' +import { afterEach, expect, it, vi } from 'vitest' +import { protocolsApi } from '@/api/client' +import type { ExperimentRunResults } from '@/types/experiments' +import type { Protocol } from '@/types/protocols' +import { DatasetRowRuns } from './DatasetRowRuns' + +afterEach(() => vi.restoreAllMocks()) +it('retries only the failed row and keeps completed unscored disabled', async () => { + const retry = vi.spyOn(protocolsApi, 'runCells').mockRejectedValue(new Error('invalid_retry_target')) + const results = { row_summary: { cell_count: 1, parent_replicate_count: 1, row_count: 2, expected: 2 }, row_results: ['failed','completed'].map((status,index) => ({ row_result_id: `slot-${index}`, cell_label: 'same', dataset_row: { row_index: index }, replicate_number: 1, status, protocol_revision_id: 'r', latest_attempt: { run_id: `run-${index}` } })) } as unknown as ExperimentRunResults + const protocol = { id: 'p', graph: { nodes: [], edges: [] }, published_revision_id: 'r' } as unknown as Protocol + render() + const buttons = screen.getAllByRole('button', { name: 'Retry' }) + expect(buttons[1]).toBeDisabled() + fireEvent.click(buttons[0]) + await waitFor(() => expect(retry).toHaveBeenCalledWith('p', { retry_row_result_ids: ['slot-0'] })) + expect(await screen.findByRole('alert')).toHaveTextContent('invalid_retry_target') +}) diff --git a/frontend/src/components/protocol/DatasetRowRuns.test.tsx b/frontend/src/components/protocol/DatasetRowRuns.test.tsx new file mode 100644 index 00000000..29459249 --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowRuns.test.tsx @@ -0,0 +1,38 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { fireEvent, render, screen, waitFor } from '@testing-library/react' +import { afterEach, expect, it, vi } from 'vitest' +import { experimentsApi, protocolsApi } from '@/api/client' +import type { ExperimentRunResults } from '@/types/experiments' +import type { Protocol } from '@/types/protocols' +import { RunsTab } from './RunsTab' + +export const rowResults = { + consumption_mode: 'per_row', cells: [], replicates: [], overview: {}, + row_summary: { cell_count: 5, parent_replicate_count: 10, row_count: 3, expected: 30, planned: 30, pending: 1, running: 2, completed: 26, failed: 1, cancelled: 0, scored: 25, missing_reported: 1, metric_coverage: {} }, + row_results: ['failed', 'completed', 'finalizing'].map((status, index) => ({ row_result_id: `slot-${index}`, cell_label: 'duplicate', replicate_label: 'same', replicate_number: 1, dataset_row: { row_index: index, dataset_id: 'd', raw_sha256: 'hash' }, status, protocol_revision_id: 'r', latest_attempt: { run_id: `run-${index}` }, attempts: [] })), +} as unknown as ExperimentRunResults +export const protocol = { name: 'Rows', description: null, created_at: '', updated_at: '', id: 'p', experiment_id: 'e', published_revision_id: 'r', published_revision: 1, has_unpublished_changes: false, graph: { nodes: [], edges: [] } } as Protocol +afterEach(() => vi.restoreAllMocks()) +export function mountRows(inspect = vi.fn(), results = rowResults) { + vi.spyOn(experimentsApi, 'listReplicates').mockResolvedValue([]) + vi.spyOn(experimentsApi, 'listTrials').mockResolvedValue([]) + vi.spyOn(experimentsApi, 'getRunResults').mockResolvedValue(results) + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }) + render() + return client +} +it('renders authoritative forecast and distinct row identities', async () => { + const inspect = vi.fn() + mountRows(inspect) + expect(await screen.findByText('5 cells / 10 replicates / 3 rows / 30 executions')).toBeInTheDocument() + expect(screen.getByText(/26 completed.*1 failed.*1 missing reported/)).toBeInTheDocument() + fireEvent.click(screen.getAllByRole('button', { name: 'Output' })[1]) + expect(inspect).toHaveBeenCalledWith('slot-1') + expect(experimentsApi.getRunResults).toHaveBeenCalledWith('e', { protocol_id: 'p' }) +}) +it('cancels exactly the active row run', async () => { + const cancel = vi.spyOn(protocolsApi, 'cancelRun').mockResolvedValue({} as never) + mountRows() + fireEvent.click(await screen.findByRole('button', { name: 'Cancel' })) + await waitFor(() => expect(cancel).toHaveBeenCalledWith('p', 'run-2')) +}) diff --git a/frontend/src/components/protocol/DatasetRowRuns.tsx b/frontend/src/components/protocol/DatasetRowRuns.tsx new file mode 100644 index 00000000..f9826de8 --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowRuns.tsx @@ -0,0 +1,48 @@ +import { useState } from 'react' +import { useMutation, useQueryClient } from '@tanstack/react-query' +import { protocolsApi } from '@/api/client' +import { Button } from '@/components/ui/button' +import type { ExperimentRunResults, RowSummary } from '@/types/experiments' +import type { Protocol } from '@/types/protocols' +import { RunConfirmDialog } from './RunConfirmDialog' + +export function rowExecutionForecast(summary: RowSummary): string { + return `${summary.cell_count} cells / ${summary.parent_replicate_count} replicates / ${summary.row_count ?? 'unknown'} rows / ${summary.expected ?? 'unknown'} executions` +} + +export function DatasetRowRuns({ results, protocol, blocked, onInspect }: { + results: ExperimentRunResults; protocol: Protocol; blocked: boolean; onInspect?: (rowId: string) => void +}) { + const queryClient = useQueryClient() + const [confirm, setConfirm] = useState(false) + const [retrying, setRetrying] = useState(null) + const [page, setPage] = useState(0) + const run = useMutation({ + mutationFn: () => protocolsApi.runCells(protocol.id), + onSuccess: () => { setConfirm(false); queryClient.invalidateQueries({ queryKey: ['experiments', protocol.experiment_id] }) }, + }) + const summary = results.row_summary + const rows = results.row_results ?? [] + const currentPage = Math.min(page, Math.max(0, Math.ceil(rows.length / 20) - 1)) + const visibleRows = rows.slice(currentPage * 20, (currentPage + 1) * 20) + const cancel = useMutation({ mutationFn: (runId: string) => protocolsApi.cancelRun(protocol.id, runId), onSuccess: () => queryClient.invalidateQueries({ queryKey: ['experiments', protocol.experiment_id] }) }) + const retry = useMutation({ + mutationFn: (rowId: string) => protocolsApi.runCells(protocol.id, { retry_row_result_ids: [rowId] }), + onMutate: rowId => setRetrying(rowId), + onSuccess: () => queryClient.invalidateQueries({ queryKey: ['experiments', protocol.experiment_id] }), + onSettled: () => setRetrying(null), + }) + return
+

{summary ? rowExecutionForecast(summary) : 'Execution forecast unavailable: original source metadata unavailable.'}

+ {summary &&

{summary.planned} planned · {summary.pending} pending · {summary.running} running · {summary.completed} completed · {summary.failed} failed · {summary.cancelled} cancelled · {summary.missing_reported} missing reported

} + + {[run.error, cancel.error, retry.error].filter(Boolean).map((error, index) =>

{error?.message}

)} +
{['Cell', 'Source row', 'Replicate', 'Status', 'Publication', 'Actions'].map(label => )}{visibleRows.map(row => )}
{label}
{row.cell_label}{row.dataset_row.row_index + 1}{row.replicate_number}{row.status}{row.protocol_revision_id} + + {row.latest_attempt && ['pending', 'running', 'finalizing'].includes(row.status) && } + +
+ {rows.length > 20 &&
Page {currentPage + 1} of {Math.ceil(rows.length / 20)}
} + {confirm && setConfirm(false)} onConfirm={() => run.mutate()} hasUnpublishedChanges={protocol.has_unpublished_changes} publishedRevision={protocol.published_revision} confirmDisabled={blocked || run.isPending} additionalContent={

{summary && rowExecutionForecast(summary)}

} />} +
+} diff --git a/frontend/src/components/protocol/DatasetRowSelector.tsx b/frontend/src/components/protocol/DatasetRowSelector.tsx new file mode 100644 index 00000000..314cb839 --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowSelector.tsx @@ -0,0 +1,6 @@ +import { Input } from '@/components/ui/input' + +export function DatasetRowSelector({ rowIndex, onChange, rowCount, disabled = false }: { rowIndex: number; onChange: (zeroBased: number) => void; rowCount: number; disabled?: boolean }) { + const valid = Number.isInteger(rowIndex) && rowIndex >= 0 && rowIndex < rowCount + return +} diff --git a/frontend/src/components/protocol/DatasetRowTestRun.test.tsx b/frontend/src/components/protocol/DatasetRowTestRun.test.tsx new file mode 100644 index 00000000..882562f0 --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowTestRun.test.tsx @@ -0,0 +1,76 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { ReactFlowProvider } from '@xyflow/react' +import { fireEvent, render, renderHook, screen, waitFor, act } from '@testing-library/react' +import { MemoryRouter } from 'react-router-dom' +import { afterEach, beforeEach, expect, it, vi } from 'vitest' +import { datasetsApi, protocolsApi } from '@/api/client' +import { defaultAgentNodeData, defaultAnthropicModelNodeData, defaultDatasetNodeData, defaultReasonActPatternNodeData, type ProtocolGraph } from '@/types/protocols' +import { ProtocolCanvas } from './ProtocolCanvas' +import { useDatasetRowSelection } from './useDatasetRowSelection' + +vi.mock('./PythonCodeEditor', () => ({ PythonCodeEditor: () => null })) +const graph: ProtocolGraph = { nodes: [ + { id: 'a', type: 'agent', position: { x: 100, y: 100 }, data: defaultAgentNodeData('Prediction') }, + { id: 'm', type: 'model_anthropic', position: { x: 0, y: 0 }, data: defaultAnthropicModelNodeData() }, + { id: 'd', type: 'dataset', position: { x: 0, y: 200 }, data: { ...defaultDatasetNodeData(), config: { ...defaultDatasetNodeData().config, dataset_id: '11111111-1111-4111-8111-111111111111', dataset_name: 'original' } } }, + { id: 'pattern', type: 'reason_act', position: { x: 0, y: 300 }, data: defaultReasonActPatternNodeData() }, + ], edges: [ { id: 'model', source: 'm', target: 'a', targetHandle: 'model' }, { id: 'pattern', source: 'pattern', target: 'a', targetHandle: 'architectural_pattern' }, { id: 'dataset', source: 'd', target: 'a', targetHandle: 'dataset', data: { dataset_input: { mode: 'per_row', columns: ['question'] } } } ] } +function setup() { + vi.stubGlobal('DOMMatrixReadOnly', class { m22 = 1 }) + vi.spyOn(HTMLElement.prototype, 'getBoundingClientRect').mockReturnValue({ x: 0, y: 0, top: 0, left: 0, right: 1000, bottom: 800, width: 1000, height: 800, toJSON: () => ({}) }) + vi.stubGlobal('ResizeObserver', class { private callback: ResizeObserverCallback; constructor(callback: ResizeObserverCallback) { this.callback = callback } observe(target: Element) { this.callback([{ target, contentRect: target.getBoundingClientRect() } as ResizeObserverEntry], this as unknown as ResizeObserver) } unobserve() {} disconnect() {} }) + vi.spyOn(protocolsApi, 'get').mockResolvedValue({ id: 'p', published_revision_id: 'r' } as never) + vi.spyOn(protocolsApi, 'getRevision').mockResolvedValue({ id: 'r', graph } as never) + vi.spyOn(datasetsApi, 'getRowSchema').mockResolvedValue({ dataset_id: 'd', raw_sha256: 'hash', row_count: 3, columns: ['question','reference'] }) + vi.spyOn(protocolsApi, 'listRuns').mockResolvedValue([]) + vi.spyOn(protocolsApi, 'getLatestTestRun').mockRejectedValue(new Error('No test')) + vi.spyOn(protocolsApi, 'update').mockResolvedValue({} as never) +} +beforeEach(setup) +afterEach(() => { vi.restoreAllMocks(); vi.unstubAllGlobals() }) +function mount() { + render(
) +} +it('canvas Test Run selects published source row 3 and submits only index 2', async () => { + const start = vi.spyOn(protocolsApi, 'testRun').mockRejectedValue(new Error('mock queue')) + mount() + fireEvent.click(await screen.findByRole('button', { name: /Test Run/i })) + const selector = await screen.findByRole('spinbutton', { name: 'Source row' }) + await waitFor(() => expect(selector).not.toBeDisabled()) + expect(selector).toHaveValue(1) + fireEvent.change(selector, { target: { value: '3' } }) + expect(selector).toHaveValue(3) + expect(start).not.toHaveBeenCalled() + const buttons = screen.getAllByRole('button', { name: 'Test Run' }) + fireEvent.click(buttons.at(-1)!) + await waitFor(() => expect(start).toHaveBeenCalledWith('p', { row_index: 2 })) +}) +it('source row control validates bounds and converts one-based display', async () => { + const { DatasetRowSelector } = await import('./DatasetRowSelector') + const change = vi.fn() + render() + fireEvent.change(screen.getByRole('spinbutton'), { target: { value: '3' } }) + expect(change).toHaveBeenCalledWith(2) +}) + +it('published selection resets after publication changes and rejects invalid bounds and columns', async () => { + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }) + const wrapper = ({ children }: { children: React.ReactNode }) => {children} + const { result, rerender } = renderHook(({ scope, source }) => useDatasetRowSelection(source, scope), { wrapper, initialProps: { scope: 'publication-1', source: graph } }) + await waitFor(() => expect(result.current.error).toBeNull()) + act(() => result.current.setRowIndex(2)) + expect(result.current.rowIndex).toBe(2) + rerender({ scope: 'publication-2', source: graph }) + await waitFor(() => expect(result.current.error).toBeNull()) + expect(result.current.rowIndex).toBe(0) + act(() => result.current.setRowIndex(3)) + expect(result.current.error).toMatch(/outside/) + const invalid = structuredClone(graph) + invalid.edges[2].data = { dataset_input: { mode: 'per_row', columns: ['missing'] } } + act(() => result.current.setRowIndex(0)) + rerender({ scope: 'publication-2', source: invalid }) + expect(result.current.error).toMatch(/columns are absent/) + rerender({ scope: 'whole', source: { ...graph, edges: graph.edges.filter(edge => edge.id !== 'dataset') } }) + expect(result.current.binding).toBeUndefined() + expect(result.current.error).toBeNull() +}) diff --git a/frontend/src/components/protocol/ExperimentSidePanel.tsx b/frontend/src/components/protocol/ExperimentSidePanel.tsx index cfc5e79c..84ba47b0 100644 --- a/frontend/src/components/protocol/ExperimentSidePanel.tsx +++ b/frontend/src/components/protocol/ExperimentSidePanel.tsx @@ -256,10 +256,12 @@ export function ExperimentSidePanel({ protocol={protocol} regenerationRequired={regenerationRequired} unboundFactors={unboundFactors} + onViewRowResult={rowResultId => { onResultSelection({ type: 'row', rowResultId, scope: { protocol_id: protocol?.id } }); openPanel('results') }} onViewResult={viewReplicateResult} /> ) : ( Promise + fetchPreview: (rowIndex?: number) => Promise }) { + const row = useDatasetRowSelection(graph, 'draft', nodeId) + const hasRowBinding = !!row.binding + const rowSignature = `${row.schema?.dataset_id}:${row.schema?.raw_sha256}:${row.rowIndex}:${row.error}` const [open, setOpen] = useState(false) const [preview, setPreview] = useState(null) const [error, setError] = useState(null) @@ -42,13 +52,14 @@ export function PromptPreviewPanel({ useEffect(() => { if (!open) return + if (hasRowBinding && row.error) { setError(row.error); setPreview(null); setLoading(false); return } let cancelled = false setLoading(true) // Debounced, because `signature` changes on every keystroke in the prompt // field above and each change is a request. const timer = setTimeout(() => { fetchRef - .current() + .current(hasRowBinding ? row.rowIndex : undefined) .then((result) => { if (cancelled) return setPreview(result) @@ -67,7 +78,7 @@ export function PromptPreviewPanel({ cancelled = true clearTimeout(timer) } - }, [open, signature, nonce]) + }, [open, signature, nonce, rowSignature, hasRowBinding, row.rowIndex, row.error]) const placeholders = preview ? (preview.text.match(PLACEHOLDER_RE) ?? []).length : 0 @@ -100,6 +111,8 @@ export function PromptPreviewPanel({ )} + {open && row.binding && } + {open && preview?.dataset_row &&

Source row {preview.dataset_row.row_index + 1} · {preview.dataset_row.dataset_id} · {preview.dataset_row.raw_sha256}

} {!open ? (

The exact user message, assembled the way a run assembles it — your prompt with every reference resolved, diff --git a/frontend/src/components/protocol/ProtocolCanvas.tsx b/frontend/src/components/protocol/ProtocolCanvas.tsx index 3e65e3bd..0a7f2cf5 100644 --- a/frontend/src/components/protocol/ProtocolCanvas.tsx +++ b/frontend/src/components/protocol/ProtocolCanvas.tsx @@ -1,3 +1,6 @@ +import { DatasetRowSelector } from './DatasetRowSelector' +import { useDatasetRowSelection } from './useDatasetRowSelection' +import { rowBindingForNode } from '@/lib/datasetRows' import { forwardRef, useCallback, useEffect, useImperativeHandle, useMemo, useRef, useState } from 'react' import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query' import { @@ -609,12 +612,12 @@ export const ProtocolCanvas = forwardRef protocolsApi.publish(protocolId), - onSuccess: (published) => { + onSuccess: async (published) => { if (published.experiment_id) { queryClient.setQueryData(protocolForExperimentQueryKey(published.experiment_id), published) } - confirmPendingRun() + const revision = await protocolsApi.getRevision(protocolId, published.published_revision_id!) + const binding = rowBindingForNode(revision.graph, pendingRunConfirm?.type === 'node' ? pendingRunConfirm.nodeId : undefined) + queryClient.setQueryData(['protocols', protocolId, 'row-publication', published.published_revision], published) + queryClient.setQueryData(['protocols', protocolId, 'row-published-graph', published.published_revision_id], revision) + // The new publication can change the driver. A new source begins at its + // first row, then later selections use that publication's verified schema. + if (pendingRunConfirm?.type === 'node') { + await protocolsApi.runNode(protocolId, pendingRunConfirm.nodeId, binding ? { row_index: 0 } : undefined).then(run => { setRunId(run.id); setRunResultsOpen(true); setTestResultsOpen(false) }) + } else { + await protocolsApi.testRun(protocolId, binding ? { row_index: 0 } : undefined).then(run => { setRunId(run.id); setTestResultsOpen(true); setRunResultsOpen(false); queryClient.setQueryData(['protocols', protocolId, 'test-run'], run) }) + } + setPendingRunConfirm(null) }, }) @@ -640,8 +654,17 @@ export const ProtocolCanvas = forwardRef protocolsApi.get(protocolId) }) + const publishedGraphQuery = useQuery({ queryKey: ['protocols', protocolId, 'row-published-graph', publishedProtocolQuery.data?.published_revision_id], queryFn: () => protocolsApi.getRevision(protocolId, publishedProtocolQuery.data!.published_revision_id!), enabled: !!publishedProtocolQuery.data?.published_revision_id }) + const publishedRow = useDatasetRowSelection(publishedGraphQuery.data?.graph, publishedProtocolQuery.data?.published_revision_id ?? '') + const pendingHasRow = !!publishedGraphQuery.data && !!rowBindingForNode(publishedGraphQuery.data.graph, pendingRunConfirm?.type === 'node' ? pendingRunConfirm.nodeId : undefined) + const publishedSourceError = publishedGraphQuery.isError || publishedProtocolQuery.isError ? 'Published row source unavailable.' : publishedProtocolQuery.isLoading || (!!publishedProtocolQuery.data?.published_revision_id && publishedGraphQuery.isLoading) ? 'Loading published inputs…' : null + const testRunMutation = useMutation({ - mutationFn: () => protocolsApi.testRun(protocolId), + mutationFn: (rowIndex: number) => { + if (publishedSourceError || (publishedRow.binding && publishedRow.error)) throw new Error(publishedSourceError ?? publishedRow.error!) + return protocolsApi.testRun(protocolId, publishedRow.binding ? { row_index: rowIndex } : undefined) + }, onSuccess: (run) => { setRunId(run.id) setRunResultsOpen(false) @@ -657,7 +680,11 @@ export const ProtocolCanvas = forwardRef protocolsApi.runNode(protocolId, nodeId), + mutationFn: ({ nodeId, rowIndex }: { nodeId: string; rowIndex: number }) => { + const binding = publishedGraphQuery.data && rowBindingForNode(publishedGraphQuery.data.graph, nodeId) + if (publishedSourceError || (binding && publishedRow.error)) throw new Error(publishedSourceError ?? publishedRow.error!) + return protocolsApi.runNode(protocolId, nodeId, binding ? { row_index: rowIndex } : undefined) + }, onSuccess: (run) => { setRunId(run.id) setTestResultsOpen(false) @@ -696,6 +723,7 @@ export const ProtocolCanvas = forwardRef + (nodeId: string, rowIndex?: number) => protocolsApi.promptPreview(protocolId, nodeId, { nodes: nodes as unknown as ProtocolNode[], edges: edges as unknown as ProtocolEdge[], - }), + }, rowIndex === undefined ? undefined : { row_index: rowIndex }), [protocolId, nodes, edges], ) @@ -1390,6 +1418,7 @@ export const ProtocolCanvas = forwardRef ({ + experimentLocked, requestConnectorAdd, requestMainEdgeAdd, requestEdgeInsert, @@ -1400,6 +1429,7 @@ export const ProtocolCanvas = forwardRef publishAndRunMutation.mutate()} confirmLabel={pendingRunConfirm.type === 'graph' ? 'Test Run' : undefined} + confirmDisabled={!!publishedSourceError || (pendingHasRow && !!publishedRow.error)} + additionalContent={pendingHasRow ?

Published source{hasUnpublishedChanges ? ' (canvas has unpublished edits)' : ''}

{publishedRow.error &&

{publishedRow.error}

}
: publishedSourceError ?

{publishedSourceError}

: undefined} /> )} {factorPickerNodeId && ( diff --git a/frontend/src/components/protocol/ProtocolCanvasContext.tsx b/frontend/src/components/protocol/ProtocolCanvasContext.tsx index 9939ac0a..5bd3e109 100644 --- a/frontend/src/components/protocol/ProtocolCanvasContext.tsx +++ b/frontend/src/components/protocol/ProtocolCanvasContext.tsx @@ -54,6 +54,7 @@ export interface EdgeInsertRequest { } interface ProtocolCanvasActions { + experimentLocked?: boolean // Opens the "+" side panel filtered to whatever node type(s) fill this // slot (today always exactly one per slot, but the panel-based flow // is deliberately a picker rather than instant-creating, so this diff --git a/frontend/src/components/protocol/ResultsTab.tsx b/frontend/src/components/protocol/ResultsTab.tsx index 73d2c393..7d672a34 100644 --- a/frontend/src/components/protocol/ResultsTab.tsx +++ b/frontend/src/components/protocol/ResultsTab.tsx @@ -1,7 +1,7 @@ import { useState } from 'react' import { useQuery } from '@tanstack/react-query' import { AlertTriangle, ChevronDown, ChevronRight, CircleDollarSign, Clock3, Coins, Cpu, Download, ExternalLink, Trophy, X } from 'lucide-react' -import { experimentsApi } from '@/api/client' +import { experimentsApi, protocolsApi, type ResultsScope } from '@/api/client' import { Badge } from '@/components/ui/badge' import { Button } from '@/components/ui/button' import { Card } from '@/components/ui/card' @@ -12,6 +12,7 @@ import { displayFactorLevel, formatMetricLabel, formatMetricValue } from '@/lib/ import { OBSERVATION_LABELS } from '@/lib/measurementPlan' import { sanitizeFilename } from '@/lib/utils' import type { DesignMetric, EvaluationArtifact, Experiment, HistoricalRun, MetricObservation, ObsoleteRun, ResultCell, ResultNodeRun, ResultReplicate, SupersededRun } from '@/types/experiments' +import { DatasetRowDetail, DatasetRowResults } from './DatasetRowResults' import { InfoTooltip } from './InfoTooltip' import { ReceivedPromptPanel, RunStepTrace, UnresolvedReferencesNote } from './NodeRunOutputPanel' @@ -417,6 +418,7 @@ function latestObsoleteRun(replicate: ResultReplicate | null): ObsoleteRun | nul export type ResultsSelection = | { type: 'cell'; cellLabel: string } | { type: 'replicate'; replicateLabel: string } + | { type: 'row'; rowResultId: string; scope: ResultsScope } function resultsMetricKey(metric: DesignMetric): string { return metric.kind === 'runtime' && metric.catalogKey ? metric.catalogKey : metric.name @@ -450,12 +452,17 @@ export function ResultsInspectorPanel({ }) { const [expandedObsoleteRuns, setExpandedObsoleteRuns] = useState>(() => new Set()) const resultsQuery = useQuery({ - queryKey: ['experiments', experimentId, 'run-results'], - queryFn: () => experimentsApi.getRunResults(experimentId), + queryKey: ['experiments', experimentId, 'run-results', selection?.type === 'row' ? selection.scope : null], + queryFn: () => experimentsApi.getRunResults(experimentId, selection?.type === 'row' ? selection.scope : undefined), enabled: selection !== null, refetchInterval: 5000, }) if (!selection) return null + if (selection.type === 'row') { + const row = resultsQuery.data?.row_results?.find(item => item.row_result_id === selection.rowResultId) + if (!row) return + return + } const replicate = selection.type === 'replicate' ? resultsQuery.data?.replicates.find((candidate) => candidate.replicate_label === selection.replicateLabel) ?? null @@ -597,11 +604,13 @@ export function ResultsInspectorPanel({ } export function ResultsTab({ + protocolId, experimentId, experimentName, experiment, onSelectResult, }: { + protocolId?: string experimentId: string experimentName: string experiment: Experiment @@ -610,9 +619,21 @@ export function ResultsTab({ const [metricPreference, setMetricPreference] = useState(null) const [expandedResultCells, setExpandedResultCells] = useState>(() => new Set()) const [downloading, setDownloading] = useState(false) - const resultsQuery = useQuery({ queryKey: ['experiments', experimentId, 'run-results'], queryFn: () => experimentsApi.getRunResults(experimentId), refetchInterval: 5000 }) + const [designRevisionId, setDesignRevisionId] = useState('') + const [protocolRevisionId, setProtocolRevisionId] = useState('') + const designs = useQuery({ queryKey: ['experiments', experimentId, 'design-revisions'], queryFn: () => experimentsApi.listDesignRevisions(experimentId), enabled: !!protocolId }) + const publications = useQuery({ queryKey: ['protocols', protocolId, 'revisions'], queryFn: () => protocolsApi.listRevisions(protocolId!), enabled: !!protocolId }) + const scope: ResultsScope = { protocol_id: protocolId, design_revision_id: designRevisionId || undefined, protocol_revision_id: protocolRevisionId || undefined } + const scopeControls = protocolId ?
+ + + {(designRevisionId || protocolRevisionId) &&

Read-only historical scope

} +
: null + const resultsQuery = useQuery({ queryKey: ['experiments', experimentId, 'run-results', scope], queryFn: () => experimentsApi.getRunResults(experimentId, scope), refetchInterval: 5000 }) if (resultsQuery.isLoading) return
- if (resultsQuery.isError || !resultsQuery.data) return

Could not load this experiment’s results.

+ if (resultsQuery.isError || !resultsQuery.data) return <>{scopeControls}

{resultsQuery.error?.message ?? 'Could not load this experiment’s results.'}

+ + if (resultsQuery.data.consumption_mode === 'per_row') return <>{scopeControls} onSelectResult({ type: 'row', rowResultId, scope })} /> const { overview, cells, replicates, metric_types: metricTypes, metric_aggregations: metricAggregations, metric_directions: metricDirections, primary_metric: primaryMetric, primary_metric_direction: primaryMetricDirection } = resultsQuery.data const metricKeys = orderedResultsMetricKeys(resultsQuery.data.metric_keys, experiment) @@ -652,7 +673,7 @@ export function ResultsTab({ async function downloadResults() { setDownloading(true) try { - const blob = await experimentsApi.downloadRunResultsCsv(experimentId) + const blob = await experimentsApi.downloadRunResultsCsv(experimentId, scope) const url = URL.createObjectURL(blob) const link = document.createElement('a') link.href = url @@ -666,6 +687,7 @@ export function ResultsTab({ return (
+ {scopeControls}
{metricKeys.length > 0 && (
diff --git a/frontend/src/components/protocol/RunsTab.tsx b/frontend/src/components/protocol/RunsTab.tsx index 912bf051..f494e539 100644 --- a/frontend/src/components/protocol/RunsTab.tsx +++ b/frontend/src/components/protocol/RunsTab.tsx @@ -12,6 +12,7 @@ import { factorBindingDiscrepancies, type FactorBindingDiscrepancy } from '@/lib import { protocolForExperimentQueryKey } from '@/lib/protocolGraph' import type { Experiment, ResultCell, ResultReplicate, Trial } from '@/types/experiments' import type { Protocol } from '@/types/protocols' +import { DatasetRowRuns } from './DatasetRowRuns' import { RunConfirmDialog } from './RunConfirmDialog' import { WarningBadge } from './nodes/WarningBadge' @@ -729,12 +730,14 @@ export function RunsTab({ regenerationRequired, unboundFactors, onViewResult, + onViewRowResult, }: { experimentId: string designSpec: Experiment['design_spec'] protocol: Protocol | undefined regenerationRequired: boolean unboundFactors: string[] + onViewRowResult?: (rowId: string) => void onViewResult: (replicateLabel: string) => void }) { const [expandedCells, setExpandedCells] = useState>(() => new Set()) @@ -751,11 +754,18 @@ export function RunsTab({ refetchInterval: 3000, }) const resultsQuery = useQuery({ - queryKey: ['experiments', experimentId, 'run-results'], - queryFn: () => experimentsApi.getRunResults(experimentId), + queryKey: ['experiments', experimentId, 'run-results', protocol?.id, protocol?.published_revision_id], + queryFn: () => experimentsApi.getRunResults(experimentId, { protocol_id: protocol?.id }), refetchInterval: 5000, }) + if (protocol && resultsQuery.data?.consumption_mode === 'per_row') { + return 0 || factorBindingDiscrepancies(designSpec, protocol.graph).length > 0} onInspect={onViewRowResult} /> + } + + if (protocol && resultsQuery.isLoading) return + if (protocol && resultsQuery.isError) return

Could not load the published execution scope.

+ if (replicatesQuery.isLoading) { return (
diff --git a/frontend/src/components/protocol/TestRunResults.tsx b/frontend/src/components/protocol/TestRunResults.tsx index 6be9242b..0d899c76 100644 --- a/frontend/src/components/protocol/TestRunResults.tsx +++ b/frontend/src/components/protocol/TestRunResults.tsx @@ -34,7 +34,7 @@ function formatObservationValue(observation: MetricObservation): string { if (typeof observation.value === 'boolean') return observation.value ? 'Yes' : 'No' if (typeof observation.value === 'number') return new Intl.NumberFormat(undefined, { maximumFractionDigits: 4 }).format(observation.value) if (typeof observation.value === 'string') return observation.value - return observation.value == null ? '—' : JSON.stringify(observation.value) + return JSON.stringify(observation.value) } function staleReason(reasons: TestRun['freshness']['reasons']): string { @@ -159,6 +159,7 @@ export function TestRunResults({ {!collapsed && + {run.dataset_row &&

Source row {run.dataset_row.row_index + 1} · {run.dataset_row.dataset_id} · {run.dataset_row.raw_sha256}

} {run.freshness.out_of_date && (

Out of date

{staleReason(run.freshness.reasons)}

diff --git a/frontend/src/components/protocol/cells/CellsTab.tsx b/frontend/src/components/protocol/cells/CellsTab.tsx index 8f8da801..ccd01050 100644 --- a/frontend/src/components/protocol/cells/CellsTab.tsx +++ b/frontend/src/components/protocol/cells/CellsTab.tsx @@ -52,7 +52,7 @@ function CellsBody({ experiment, cells }: { experiment: Experiment; cells: Repli * API. The table is mounted in exactly one of the two places at a time, so * there's never a second copy carrying its own divergent sort/page state. */ -export function CellsTab({ experiment }: { experiment: Experiment }) { +export function CellsTab({ experiment, perRow = false }: { experiment: Experiment; perRow?: boolean }) { const [fullscreen, setFullscreen] = useState(false) // null = the current design. Set from DesignHistory to inspect a superseded // revision's cells; owned here so the heatmap, the table, the scored tally @@ -113,7 +113,7 @@ export function CellsTab({ experiment }: { experiment: Experiment }) { ) } - const cells = replicatesQuery.data + const cells = perRow ? replicatesQuery.data?.map(replicate => ({ ...replicate, metric_values: null })) : replicatesQuery.data const cellCount = cells ? groupReplicatesIntoCells(cells).length : 0 const scoredReplicates = cells?.filter((replicate) => replicate.metric_values).length ?? 0 const viewingHistory = revisionId !== null @@ -138,7 +138,7 @@ export function CellsTab({ experiment }: { experiment: Experiment }) {

) } else { - body = + body = perRow ?

Per-row observations are available in Results. Parent scores and heatmaps do not apply to this publication.

: } // A superseded revision is a record of what was, not something to keep @@ -165,7 +165,7 @@ export function CellsTab({ experiment }: { experiment: Experiment }) { Cells — {cellCount} {cellCount === 1 ? 'cell' : 'cells'} · {scoredReplicates}/{cells?.length ?? 0} replicates scored

- +
+} diff --git a/frontend/src/components/protocol/edges/InteractEdge.tsx b/frontend/src/components/protocol/edges/InteractEdge.tsx index b0af7931..73332f12 100644 --- a/frontend/src/components/protocol/edges/InteractEdge.tsx +++ b/frontend/src/components/protocol/edges/InteractEdge.tsx @@ -4,11 +4,13 @@ import { BaseEdge, EdgeLabelRenderer, EdgeToolbar, getBezierPath, useReactFlow, // uses, so "remove this thing" looks identical whether the thing is a node or // an edge. An X here also collided with the two other X's on the canvas // (dismissing a panel, unbinding a factor), neither of which deletes anything. -import { Filter, Plus, Trash2 } from 'lucide-react' +import { Filter, Plus, Settings2, Trash2 } from 'lucide-react' import { Popover, PopoverContent, PopoverTrigger } from '@/components/ui/popover' import { describeHandoff } from '@/lib/promptReferences' import type { EdgeHandoff } from '@/types/protocols' import { useProtocolCanvasActions } from '../ProtocolCanvasContext' +import { Button } from '@/components/ui/button' +import { DatasetInputPanel } from './DatasetInputPanel' import { EdgeHandoffPanel } from './EdgeHandoffPanel' // Every edge's look is decided here and nowhere else -- no edge in a persisted @@ -87,15 +89,17 @@ export function InteractEdge({ markerEnd, }: EdgeProps) { const [hovered, setHovered] = useState(false) + const [datasetOpen, setDatasetOpen] = useState(false) const [handoffOpen, setHandoffOpen] = useState(false) const { setEdges, getNode } = useReactFlow() - const { requestEdgeInsert } = useProtocolCanvasActions() + const { requestEdgeInsert, experimentLocked } = useProtocolCanvasActions() const [edgePath, labelX, labelY] = getBezierPath({ sourceX, sourceY, sourcePosition, targetX, targetY, targetPosition }) const isMainEdge = !sourceHandleId && !targetHandleId const isDirectedFlow = data?.directedFlow === true const isPatternEdge = targetHandleId === 'architectural_pattern' // "What passes" is only a choice on a main edge into an Agent: a Critic Gate // always reviews the whole answer (validate_edge_handoffs). + const canConfigureDataset = getNode(source)?.type === 'dataset' && getNode(target)?.type === 'agent' && ['dataset', 'resource', 'tool'].includes(targetHandleId ?? '') const canNarrow = isMainEdge && getNode(target)?.type === 'agent' const handoffLabel = describeHandoff((data as { handoff?: EdgeHandoff } | undefined)?.handoff) @@ -137,7 +141,7 @@ export function InteractEdge({
)} - +
setHovered(true)} @@ -154,6 +158,10 @@ export function InteractEdge({ )} + {canConfigureDataset && + }> + + } {canNarrow && ( datasetsApi.getRowSchema(id!), enabled: !!id, retry: false }) + const key = `${scope}:${id}:${schema.data?.raw_sha256}:${nodeId ?? ''}` + const [selection, setSelection] = useState({ key, index: 0 }) + const rowIndex = selection.key === key ? selection.index : 0 + const error = !binding ? null : !id ? 'Registered row source unavailable.' : schema.isLoading ? 'Loading original rows…' : schema.isError ? schema.error.message : !schema.data ? 'Original row metadata unavailable.' : !schema.data.row_count ? 'Original Dataset has no rows.' : !Number.isInteger(rowIndex) || rowIndex < 0 || rowIndex >= schema.data.row_count ? 'Source row is outside the original Dataset.' : binding.input.mode === 'per_row' && binding.input.columns.some(column => !schema.data.columns.includes(column)) ? 'Selected columns are absent from the original CSV.' : graph ? datasetRowTopologyError(graph) : null + return { binding, schema: schema.data, rowIndex, setRowIndex: (index: number) => setSelection({ key, index }), error } +} diff --git a/frontend/src/lib/datasetRows.ts b/frontend/src/lib/datasetRows.ts new file mode 100644 index 00000000..57aff1ef --- /dev/null +++ b/frontend/src/lib/datasetRows.ts @@ -0,0 +1,39 @@ +import type { ProtocolGraph, ProtocolNode } from '@/types/protocols' + +const handles = new Set(['dataset', 'resource', 'tool']) +export function datasetRowBindings(graph: ProtocolGraph) { + const nodes = new Map(graph.nodes.map(node => [node.id, node])) + return graph.edges.flatMap(edge => { + const source = nodes.get(edge.source) + const target = nodes.get(edge.target) + const config = source?.data.config as { enabled?: boolean; dataset_id?: string; dataset_name?: string } | undefined + if (source?.type !== 'dataset' || target?.type !== 'agent' || !handles.has(edge.targetHandle ?? '') || config?.enabled === false) return [] + return [{ edge, source, config, input: edge.data?.dataset_input ?? { mode: 'whole_dataset' as const } }] + }) +} + +export function rowBindingForNode(graph: ProtocolGraph, nodeId?: string) { + return datasetRowBindings(graph).find(binding => binding.input.mode === 'per_row' && (!nodeId || binding.edge.target === nodeId)) +} + +export function datasetRowTopologyError(graph: ProtocolGraph): string | null { + const bindings = datasetRowBindings(graph) + const rows = bindings.filter(binding => binding.input.mode === 'per_row') + if (!rows.length) return null + const id = rows[0].config?.dataset_id + if (!id || !/^[\da-f]{8}-[\da-f]{4}-[\da-f]{4}-[\da-f]{4}-[\da-f]{12}$/i.test(id)) return 'Per row requires a registered dataset UUID.' + if (rows.some(binding => binding.config?.dataset_id?.toLowerCase() !== id.toLowerCase())) return 'Per row inputs must share one Dataset driver.' + const names = new Set(rows.map(binding => binding.config?.dataset_name).filter(Boolean)) + if (bindings.some(binding => binding.input.mode === 'whole_dataset' && (binding.config?.dataset_id?.toLowerCase() === id.toLowerCase() || names.has(binding.config?.dataset_name)))) return 'One Dataset cannot use both Whole dataset and Per row.' + const agents = new Map() + for (const binding of rows) { + if (binding.input.mode !== 'per_row') continue + if (!binding.input.columns.length || new Set(binding.input.columns).size !== binding.input.columns.length || binding.input.columns.some(column => !column.trim())) return 'Select distinct, nonempty columns.' + const key = JSON.stringify(binding.input.columns) + if (agents.has(binding.edge.target) && agents.get(binding.edge.target) !== key) return 'Duplicate Dataset connections must select identical columns.' + agents.set(binding.edge.target, key) + const factors = (binding.source.data as ProtocolNode['data']).factor_bindings + if (Object.keys(factors ?? {}).some(path => path === 'config' || path.startsWith('config.'))) return 'The row driver cannot be factor-bound.' + } + return null +} diff --git a/frontend/src/lib/protocolGraph.test.ts b/frontend/src/lib/protocolGraph.test.ts index d30e2d5b..9f24cb4f 100644 --- a/frontend/src/lib/protocolGraph.test.ts +++ b/frontend/src/lib/protocolGraph.test.ts @@ -29,3 +29,10 @@ describe('mergeProtocolSaveIntoCache', () => { }) }) }) + +it('preserves dataset input and unrelated metadata through canvas serialization', async () => { + const { toPersistedGraph } = await import('./protocolGraph') + const graph = toPersistedGraph([], [{ id: 'edge', source: 'd', target: 'a', targetHandle: 'dataset', data: { directedFlow: false, custom: 'keep', dataset_input: { mode: 'per_row', columns: ['question'] } } }]) + expect(graph.edges[0].data).toEqual({ custom: 'keep', dataset_input: { mode: 'per_row', columns: ['question'] } }) + expect(JSON.parse(JSON.stringify(graph))).toEqual(graph) +}) diff --git a/frontend/src/lib/protocolGraph.ts b/frontend/src/lib/protocolGraph.ts index f9c43feb..310d8748 100644 --- a/frontend/src/lib/protocolGraph.ts +++ b/frontend/src/lib/protocolGraph.ts @@ -96,16 +96,18 @@ export function toPersistedGraph(nodes: Node[], edges: Edge[]): ProtocolGraph { data: n.data as AgentNodeData | McpToolNodeData | CriticGateNodeData, })), edges: edges.map((e) => { - // Only the handoff is persisted from edge data -- the rest of it - // (`directedFlow`) is display state the canvas derives on render. - const handoff = (e.data as ProtocolEdge['data'])?.handoff + // Preserve saved input settings and other authored metadata; display + // bookkeeping is reconstructed when the canvas loads. + const data = { ...e.data } as NonNullable + delete data.directedFlow + if (data.handoff?.mode === 'full') delete data.handoff return { id: e.id, source: e.source, target: e.target, sourceHandle: e.sourceHandle, targetHandle: e.targetHandle, - ...(handoff && handoff.mode !== 'full' ? { data: { handoff } } : {}), + ...(Object.keys(data).length ? { data } : {}), } }), } diff --git a/frontend/src/lib/reportedMetrics.ts b/frontend/src/lib/reportedMetrics.ts new file mode 100644 index 00000000..26473be8 --- /dev/null +++ b/frontend/src/lib/reportedMetrics.ts @@ -0,0 +1,7 @@ +import type { MetricObservation } from '@/types/experiments' + +export function reportedMetricText(observation?: MetricObservation): string { + if (!observation) return 'Unavailable' + if (observation.status !== 'measured') return observation.status.replaceAll('_', ' ') + return typeof observation.value === 'string' ? observation.value : JSON.stringify(observation.value) +} diff --git a/frontend/src/pages/ProtocolCanvasPage.tsx b/frontend/src/pages/ProtocolCanvasPage.tsx index bf60ac59..25be6f52 100644 --- a/frontend/src/pages/ProtocolCanvasPage.tsx +++ b/frontend/src/pages/ProtocolCanvasPage.tsx @@ -108,7 +108,8 @@ function TopBarStat({ icon: Icon, value, title, accent }: { icon: LucideIcon; va // Same design_type gate the Cells tab uses -- cells/metric_values are a // factorial concept, so a future non-factorial experiment gets no chips // rather than a nonsensical "0/0 scored". -function TopBarStats({ experiment, cells, obsoleteRunCount = 0 }: { experiment: Experiment; cells: Replicate[] | undefined; obsoleteRunCount?: number }) { +function TopBarStats({ experiment, cells, obsoleteRunCount = 0, rowSummary }: { experiment: Experiment; cells: Replicate[] | undefined; obsoleteRunCount?: number; rowSummary?: import('@/types/experiments').RowSummary | null }) { + if (rowSummary) return if (experiment.design_type !== 'factorial' || !cells) return null const scored = cells.filter((c) => c.metric_values).length const cellCount = groupReplicatesIntoCells(cells).length @@ -259,8 +260,8 @@ export function ProtocolCanvasPage() { enabled: !!experimentId, }) const runResultsQuery = useQuery({ - queryKey: ['experiments', experimentId, 'run-results'], - queryFn: () => experimentsApi.getRunResults(experimentId!), + queryKey: ['experiments', experimentId, 'run-results', protocolQuery.data?.id, protocolQuery.data?.published_revision_id], + queryFn: () => experimentsApi.getRunResults(experimentId!, { protocol_id: protocolQuery.data?.id }), enabled: !!experimentId, refetchInterval: 5000, }) @@ -286,7 +287,7 @@ export function ProtocolCanvasPage() { {experimentQuery.data && } {experimentQuery.data?.locked_at && Locked} - {experimentQuery.data && } + {experimentQuery.data && }
{protocolQuery.data && experimentId && ( <> diff --git a/frontend/src/types/experiments.ts b/frontend/src/types/experiments.ts index 92fe91e5..d5a564b9 100644 --- a/frontend/src/types/experiments.ts +++ b/frontend/src/types/experiments.ts @@ -1,3 +1,4 @@ +import type { DatasetRowIdentity, DatasetRowSnapshot, NodeRunState, Conversation, ProtocolRun } from './protocols' export interface DesignFactor { name: string levels: unknown[] @@ -458,7 +459,62 @@ export interface RunResultsOverview { reported_cost_count: number } +export interface RowAttempt { + run_id: string + status: ProtocolRun['status'] + error: string | null + started_at: string | null + completed_at: string | null + workspace_id: string | null + attempt_result: Record | null + node_runs: Record + conversation: Conversation | null + dataset_row: DatasetRowSnapshot | null + protocol_revision_id: string | null + design_revision_id: string | null + current: boolean +} + +export interface RowResult { + row_result_id: string + cell_id: string + cell_label: string + factor_values: Record + replicate_result_id: string + replicate_label: string + replicate_number: number + dataset_row: DatasetRowIdentity + design_revision_id: string + protocol_revision_id: string + status: ProtocolRun['status'] + workspace_id: string | null + metric_values: Record | null + measurement: { observations: MetricObservation[]; artifacts: EvaluationArtifact[] } | null + artifacts: Record | null + latest_attempt: RowAttempt | null + attempts: RowAttempt[] +} + +export interface RowSummary { + cell_count: number + parent_replicate_count: number + row_count: number | null + expected: number | null + planned: number + pending: number + running: number + completed: number + failed: number + cancelled: number + scored: number + missing_reported: number + metric_coverage: Record> +} + export interface ExperimentRunResults { + consumption_mode?: 'whole_dataset' | 'per_row' + row_results?: RowResult[] + row_summary?: RowSummary | null overview: RunResultsOverview metric_keys: string[] metric_types: Record diff --git a/frontend/src/types/protocols.ts b/frontend/src/types/protocols.ts index 6086f1eb..ef8c96d1 100644 --- a/frontend/src/types/protocols.ts +++ b/frontend/src/types/protocols.ts @@ -138,7 +138,27 @@ export interface Conversation { messages: ConversationMessage[] } +export interface DatasetRowIdentity { + dataset_id: string + raw_sha256: string + row_index: number +} + +export interface DatasetRowSnapshot extends DatasetRowIdentity { + columns: string[] + values: Record +} + +export interface DatasetRowSchema { + dataset_id: string + raw_sha256: string + row_count: number + columns: string[] +} + export interface ProtocolRun { + dataset_row?: DatasetRowSnapshot | null + row_result_id?: string | null id: string protocol_id: string // `limit_reached` is conversation-mode only: the agents were still talking @@ -182,6 +202,7 @@ export interface TestRunResourceUsage { } export interface TestRun { + dataset_row?: DatasetRowSnapshot | null id: string protocol_id: string status: ProtocolRun['status'] @@ -212,6 +233,7 @@ export interface TestRun { // `` placeholder wherever upstream output would go. Nothing // is created; this is a rendering, not a resource. export interface PromptPreview { + dataset_row?: DatasetRowSnapshot | null text: string } @@ -227,6 +249,8 @@ export interface ProtocolRevision { // not-yet-completed replicate. skipped is how many replicates already had // metrics or a completed run and were left alone (resume semantics). export interface CellRunBatch { + consumption_mode?: 'whole_dataset' | 'per_row' + row_result_ids?: string[] protocol_run_ids: string[] replicate_labels: string[] skipped: number diff --git a/src/asaree/api/experiments.py b/src/asaree/api/experiments.py index 6615f049..8b0a7842 100644 --- a/src/asaree/api/experiments.py +++ b/src/asaree/api/experiments.py @@ -35,7 +35,7 @@ measurement_plan_issue_is_blocking, validate_experiment_measurement_plan, ) -from asaree.services.experiment_run_results import summarize_experiment_run_results +from asaree.services.experiment_run_results import RunResultsProjectionError, summarize_experiment_run_results from asaree.services.experiments import ( create_experiment, create_untitled_experiment, @@ -963,14 +963,24 @@ class AnalyzeFactorialRequest(BaseModel): @router.post("/{experiment_id}/analyze") async def analyze_factorial_endpoint( - experiment_id: uuid.UUID, body: AnalyzeFactorialRequest, user: CurrentUser, db: DbSession + experiment_id: uuid.UUID, body: AnalyzeFactorialRequest, user: CurrentUser, db: DbSession, + protocol_id: uuid.UUID | None = None, + design_revision_id: uuid.UUID | None = None, + protocol_revision_id: uuid.UUID | None = None, ) -> dict[str, Any]: """Failure homogeneity, factorial effects (Freedman-Lane + max-stat FWER), estimated marginal means, non-inferiority vs. the reference condition (BCa bootstrap + Holm), and heteroscedasticity diagnostics — computed fresh from this experiment's current replicate results, not persisted.""" experiment = await _get_owned_experiment(db, experiment_id, user) - declared_primary = declared_primary_metric((experiment.design_spec or {}).get("metrics")) + selection = await _factorial_analysis_selection( + db, experiment, protocol_id=protocol_id, design_revision_id=design_revision_id, + protocol_revision_id=protocol_revision_id, + ) + if selection["consumption_mode"] == "per_row": + return _row_analysis_unavailable() + selected_design = selection["design_spec"] + declared_primary = declared_primary_metric((selected_design or {}).get("metrics")) if declared_primary is None: raise HTTPException(status_code=422, detail="This experiment has no declared primary metric.") if declared_primary["name"] != body.primary_metric: @@ -978,7 +988,7 @@ async def analyze_factorial_endpoint( status_code=422, detail=f"{body.primary_metric!r} is not the experiment's declared primary metric.", ) - replicates = await list_replicates(db, experiment_id=experiment_id) + replicates = await list_replicates(db, experiment_id=experiment_id, revision_id=design_revision_id) try: analysis = ( analyze_binary_factorial @@ -1016,6 +1026,42 @@ class ResultsResponse(BaseModel): best_condition: dict[str, Any] | None +def _row_analysis_unavailable() -> dict[str, Any]: + return { + "available": False, + "reason": "Per-row executions are available for inspection and export; factorial analysis is not supported.", + "analysis": None, + "best_condition": None, + } + + +async def _factorial_analysis_selection( + db: DbSession, + experiment: Any, + *, + protocol_id: uuid.UUID | None, + design_revision_id: uuid.UUID | None, + protocol_revision_id: uuid.UUID | None, +) -> dict[str, Any]: + """Resolve mode and design from the same publication/revision selectors as Results.""" + try: + projection = await summarize_experiment_run_results( + db, experiment_id=experiment.id, design_spec=experiment.design_spec, + protocol_id=protocol_id, design_revision_id=design_revision_id, + protocol_revision_id=protocol_revision_id, + ) + except RunResultsProjectionError as exc: + status = 422 if str(exc) == "ambiguous_protocol" else 404 + raise HTTPException(status_code=status, detail=str(exc)) from exc + design_spec = experiment.design_spec + if design_revision_id is not None: + revision = await get_revision(db, design_revision_id) + if revision is None or revision.experiment_id != experiment.id: + raise HTTPException(status_code=404, detail="design_revision_not_found") + design_spec = revision.design_spec + return {"consumption_mode": projection["consumption_mode"], "design_spec": design_spec} + + class RunResultsResponse(BaseModel): """The general-purpose results scorecard, unlike the optional factorial analysis.""" @@ -1028,44 +1074,85 @@ class RunResultsResponse(BaseModel): primary_metric_direction: str | None cells: list[dict[str, Any]] replicates: list[dict[str, Any]] + consumption_mode: str + row_results: list[dict[str, Any]] + row_summary: dict[str, Any] | None @router.get("/{experiment_id}/results", response_model=ResultsResponse) async def get_experiment_results_endpoint( - experiment_id: uuid.UUID, user: CurrentUser, db: DbSession + experiment_id: uuid.UUID, user: CurrentUser, db: DbSession, + protocol_id: uuid.UUID | None = None, + design_revision_id: uuid.UUID | None = None, + protocol_revision_id: uuid.UUID | None = None, ) -> ResultsResponse: """See services.factorial_analysis.analyze_experiment_design -- this endpoint is a thin pass-through, all the real derivation/wrapping logic lives there so it's unit-testable without a request/response cycle.""" experiment = await _get_owned_experiment(db, experiment_id, user) - replicates = await list_replicates(db, experiment_id=experiment_id) - result = analyze_experiment_design(experiment.design_spec, replicates) + selection = await _factorial_analysis_selection( + db, experiment, protocol_id=protocol_id, design_revision_id=design_revision_id, + protocol_revision_id=protocol_revision_id, + ) + if selection["consumption_mode"] == "per_row": + return ResultsResponse(**_row_analysis_unavailable()) + replicates = await list_replicates(db, experiment_id=experiment_id, revision_id=design_revision_id) + result = analyze_experiment_design(selection["design_spec"], replicates, consumption_mode="whole_dataset") return ResultsResponse(**result) @router.get("/{experiment_id}/run-results", response_model=RunResultsResponse) async def get_experiment_run_results_endpoint( - experiment_id: uuid.UUID, user: CurrentUser, db: DbSession + experiment_id: uuid.UUID, + user: CurrentUser, + db: DbSession, + protocol_id: uuid.UUID | None = None, + design_revision_id: uuid.UUID | None = None, + protocol_revision_id: uuid.UUID | None = None, ) -> RunResultsResponse: - """Operational and outcome summary for the experiment's current cells. - - This remains available for every experiment, including one without the - balanced 2-level design needed by the statistical ``/results`` endpoint. - """ + """Operational Results for the selected published protocol and design.""" experiment = await _get_owned_experiment(db, experiment_id, user) - return RunResultsResponse( - **(await summarize_experiment_run_results(db, experiment_id=experiment_id, design_spec=experiment.design_spec)) - ) + try: + results = await summarize_experiment_run_results( + db, + experiment_id=experiment_id, + design_spec=experiment.design_spec, + protocol_id=protocol_id, + design_revision_id=design_revision_id, + protocol_revision_id=protocol_revision_id, + ) + except RunResultsProjectionError as exc: + if str(exc) == "ambiguous_protocol": + raise HTTPException(status_code=422, detail="ambiguous_protocol") from exc + raise HTTPException(status_code=404, detail=str(exc)) from exc + return RunResultsResponse(**results) @router.get("/{experiment_id}/run-results.csv") -async def export_run_results_csv_endpoint(experiment_id: uuid.UUID, user: CurrentUser, db: DbSession) -> Response: +async def export_run_results_csv_endpoint( + experiment_id: uuid.UUID, + user: CurrentUser, + db: DbSession, + protocol_id: uuid.UUID | None = None, + design_revision_id: uuid.UUID | None = None, + protocol_revision_id: uuid.UUID | None = None, +) -> Response: """Download the same enriched data shown on the Results tab.""" experiment = await _get_owned_experiment(db, experiment_id, user) - results = await summarize_experiment_run_results( - db, experiment_id=experiment_id, design_spec=experiment.design_spec + try: + results = await summarize_experiment_run_results( + db, experiment_id=experiment_id, design_spec=experiment.design_spec, + protocol_id=protocol_id, design_revision_id=design_revision_id, protocol_revision_id=protocol_revision_id, + ) + except RunResultsProjectionError as exc: + if str(exc) == "ambiguous_protocol": + raise HTTPException(status_code=422, detail="ambiguous_protocol") from exc + raise HTTPException(status_code=404, detail=str(exc)) from exc + row_mode = results["consumption_mode"] == "per_row" + csv_text = result_rows_to_csv( + results["row_results"] if row_mode else results["replicates"], experiment.design_spec, + consumption_mode=results["consumption_mode"], ) - csv_text = result_rows_to_csv(results["replicates"], experiment.design_spec) filename = _UNSAFE_FILENAME_CHAR.sub("_", experiment.name.strip()) or "experiment" return Response( content=csv_text, @@ -1075,14 +1162,29 @@ async def export_run_results_csv_endpoint(experiment_id: uuid.UUID, user: Curren @router.get("/{experiment_id}/run-results.schema.json") -async def get_run_results_schema_endpoint(experiment_id: uuid.UUID, user: CurrentUser, db: DbSession) -> dict[str, Any]: +async def get_run_results_schema_endpoint( + experiment_id: uuid.UUID, + user: CurrentUser, + db: DbSession, + protocol_id: uuid.UUID | None = None, + design_revision_id: uuid.UUID | None = None, + protocol_revision_id: uuid.UUID | None = None, +) -> dict[str, Any]: """Machine-readable factor encoding and outcome types for Results CSV consumers.""" experiment = await _get_owned_experiment(db, experiment_id, user) - results = await summarize_experiment_run_results( - db, experiment_id=experiment_id, design_spec=experiment.design_spec - ) + try: + results = await summarize_experiment_run_results( + db, experiment_id=experiment_id, design_spec=experiment.design_spec, + protocol_id=protocol_id, design_revision_id=design_revision_id, protocol_revision_id=protocol_revision_id, + ) + except RunResultsProjectionError as exc: + if str(exc) == "ambiguous_protocol": + raise HTTPException(status_code=422, detail="ambiguous_protocol") from exc + raise HTTPException(status_code=404, detail=str(exc)) from exc return result_rows_schema( - results["replicates"], results["metric_types"], results["metric_aggregations"], experiment.design_spec + results["row_results"] if results["consumption_mode"] == "per_row" else results["replicates"], + results["metric_types"], results["metric_aggregations"], experiment.design_spec, + consumption_mode=results["consumption_mode"], ) @@ -1139,12 +1241,28 @@ async def list_replicates_endpoint( @router.get("/{experiment_id}/replicates.csv") -async def export_replicates_csv_endpoint(experiment_id: uuid.UUID, user: CurrentUser, db: DbSession) -> Response: +async def export_replicates_csv_endpoint( + experiment_id: uuid.UUID, user: CurrentUser, db: DbSession, + protocol_id: uuid.UUID | None = None, + design_revision_id: uuid.UUID | None = None, + protocol_revision_id: uuid.UUID | None = None, +) -> Response: """One row per replicate that's actually run, one column per factor_values/ metric_values key seen across them -- see services.csv_export (replicates_that_ran / replicates_to_csv).""" experiment = await _get_owned_experiment(db, experiment_id, user) - replicates = await list_replicates(db, experiment_id=experiment_id) + selection = await _factorial_analysis_selection( + db, experiment, protocol_id=protocol_id, design_revision_id=design_revision_id, + protocol_revision_id=protocol_revision_id, + ) + if selection["consumption_mode"] == "per_row": + raise HTTPException(status_code=422, detail={ + "code": "use_row_results_export", + "message": "Per-row executions are available in the run-results.csv export.", + "href": f"/experiments/{experiment_id}/run-results.csv", + "link_text": "run-results.csv", + }) + replicates = await list_replicates(db, experiment_id=experiment_id, revision_id=design_revision_id) csv_text = replicates_to_csv(replicates_that_ran(replicates), design_spec=experiment.design_spec) filename = _UNSAFE_FILENAME_CHAR.sub("_", experiment.name.strip()) or "experiment" return Response( diff --git a/src/asaree/api/protocols.py b/src/asaree/api/protocols.py index 1ebf0aa4..1447f719 100644 --- a/src/asaree/api/protocols.py +++ b/src/asaree/api/protocols.py @@ -14,10 +14,14 @@ from typing import Any from fastapi import APIRouter, HTTPException -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, StrictInt +from sqlalchemy import select from asaree.deps import CurrentUser, DbSession -from asaree.services.dataset_row_inputs import resolve_dataset_row_plan +from asaree.models.dataset import RegisteredDataset +from asaree.models.protocol_revision import ProtocolRevision +from asaree.services.dataset_row_csv import DatasetRowCsvError, project_row, read_row_source +from asaree.services.dataset_row_inputs import DatasetRowInputError, resolve_dataset_row_plan from asaree.services.experiment_measurements import ( blocking_measurement_plan_issues, validate_experiment_measurement_plan, @@ -46,6 +50,7 @@ from asaree.services.protocol_runs import ( create_protocol_run, create_test_run, + fail_protocol_run, get_protocol_run, list_protocol_runs, request_protocol_run_cancellation, @@ -174,6 +179,7 @@ class TestRunResponse(BaseModel): status: str error: str | None protocol_revision_id: uuid.UUID | None + dataset_row: dict[str, Any] | None created_at: datetime updated_at: datetime observations: list[dict[str, Any]] @@ -193,7 +199,8 @@ async def _test_run_response(db: DbSession, run: Any, protocol: Any, experiment: protocol_id=run.protocol_id, status=run.status, error=run.error, - protocol_revision_id=run.protocol_revision_id, + protocol_revision_id=result.protocol_revision_id, + dataset_row=result.dataset_row, created_at=run.created_at, updated_at=run.updated_at, observations=result.observations, @@ -240,10 +247,23 @@ class CreateProtocolRunRequest(BaseModel): # the same as one entry of "Run all cells" but picked by name instead of # running every not-yet-completed replicate at once. replicate_label: str | None = None + row_index: StrictInt | None = Field(default=None, ge=0) + + +class TestRunRequest(BaseModel): + """Optional original Dataset row selection for a row-mode Test Run.""" + + row_index: StrictInt | None = Field(default=None, ge=0) + + +class NodePlayRequest(BaseModel): + """Optional original Dataset row selection for an eligible node Play.""" + + row_index: StrictInt | None = Field(default=None, ge=0) class CellRunBatchRequest(BaseModel): - """Previously completed replicates the user explicitly chose to run again. + """Options for a cell batch or an explicit row-slot retry. Omit this body for the normal resume behavior: every never-completed replicate runs, while completed ones remain skipped. @@ -253,6 +273,7 @@ class CellRunBatchRequest(BaseModel): # a run from one cell to just that cell's replicates. replicate_labels: list[str] | None = None rerun_replicate_labels: list[str] = [] + retry_row_result_ids: list[uuid.UUID] | None = None class CellRunBatchResponse(BaseModel): @@ -274,10 +295,12 @@ class PromptPreviewRequest(BaseModel): # is on screen, including edits autosave has not flushed yet; omitted, the # stored draft stands in. Nothing is written either way. graph: dict[str, Any] | None = None + row_index: StrictInt | None = None class PromptPreviewResponse(BaseModel): text: str + dataset_row: dict[str, Any] | None = None async def _get_owned_protocol(db: DbSession, protocol_id: uuid.UUID, user: CurrentUser) -> Any: @@ -421,6 +444,18 @@ async def publish_protocol_endpoint(protocol_id: uuid.UUID, user: CurrentUser, d return await _protocol_response(db, protocol) +@router.get("/{protocol_id}/revisions", response_model=list[ProtocolRevisionResponse]) +async def list_protocol_revisions_endpoint( + protocol_id: uuid.UUID, user: CurrentUser, db: DbSession, +) -> list[ProtocolRevisionResponse]: + await _get_owned_protocol(db, protocol_id, user) + revisions = await db.scalars( + select(ProtocolRevision).where(ProtocolRevision.protocol_id == protocol_id) + .order_by(ProtocolRevision.revision.desc()) + ) + return [ProtocolRevisionResponse.model_validate(revision) for revision in revisions] + + @router.get("/{protocol_id}/revisions/{revision_id}", response_model=ProtocolRevisionResponse) async def get_protocol_revision_endpoint( protocol_id: uuid.UUID, revision_id: uuid.UUID, user: CurrentUser, db: DbSession @@ -455,8 +490,15 @@ async def create_protocol_run_endpoint( protocol = await _get_owned_protocol(db, protocol_id, user) revision = await _require_published_revision(db, protocol) replicate_label = body.replicate_label if body else None + row_index = body.row_index if body else None try: - if replicate_label: + row_mode = resolve_dataset_row_plan(revision.graph) is not None + except DatasetRowInputError as exc: + raise HTTPException(status_code=422, detail=str(exc)) from exc + if row_index is not None and not row_mode: + raise HTTPException(status_code=422, detail="no_row_driver") + try: + if replicate_label or row_mode: run = await plan_single_replicate_run( db, protocol_id=protocol_id, @@ -464,6 +506,8 @@ async def create_protocol_run_endpoint( owner_id=user.id, graph=revision.graph, replicate_label=replicate_label, + row_index=row_index, + snapshot_only=replicate_label is None, protocol_revision_id=revision.id, ) else: @@ -478,12 +522,20 @@ async def create_protocol_run_endpoint( ) except ProtocolValidationError as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc + if row_mode: + await db.commit() + await db.refresh(run) await enqueue_protocol_run(run.id) return _protocol_run_response(run) @router.post("/{protocol_id}/test-runs", response_model=TestRunResponse, status_code=201) -async def create_test_run_endpoint(protocol_id: uuid.UUID, user: CurrentUser, db: DbSession) -> TestRunResponse: +async def create_test_run_endpoint( + protocol_id: uuid.UUID, + user: CurrentUser, + db: DbSession, + body: TestRunRequest | None = None, +) -> TestRunResponse: """Start the experiment's single current canvas Test Run.""" protocol = await _get_owned_protocol(db, protocol_id, user) revision = await _require_published_revision(db, protocol) @@ -496,9 +548,49 @@ async def create_test_run_endpoint(protocol_id: uuid.UUID, user: CurrentUser, db revision.graph, require_acyclic=not is_conversation_strategy(experiment.design_spec if experiment else None), ) - run = await create_test_run(db, protocol_id=protocol.id, owner_id=user.id, protocol_revision_id=revision.id) + dataset_row = None + try: + row_plan = resolve_dataset_row_plan(revision.graph, experiment.design_spec if experiment else None) + except DatasetRowInputError as exc: + raise HTTPException(status_code=422, detail=str(exc)) from exc + row_index = body.row_index if body is not None else None + if row_plan is None and row_index is not None: + raise HTTPException(status_code=422, detail="no_row_driver") + if row_plan is not None: + registration = ( + await db.execute( + select(RegisteredDataset).where( + RegisteredDataset.id == uuid.UUID(row_plan["driver_dataset_id"]), + RegisteredDataset.owner_id == user.id, + ) + ) + ).scalar_one_or_none() + if registration is None: + raise HTTPException(status_code=422, detail="Registered row source is unavailable.") + try: + source = read_row_source( + dataset_id=str(registration.id), + raw_path=registration.raw_path, + raw_sha256=registration.raw_sha256, + ) + columns = list( + dict.fromkeys(column for binding in row_plan["bindings"] for column in binding["columns"]) + ) + selected_index = 0 if row_index is None else row_index + dataset_row = project_row(source, row_index=selected_index, columns=columns) + except DatasetRowCsvError as exc: + raise HTTPException(status_code=422, detail=str(exc)) from exc + run = await create_test_run( + db, + protocol_id=protocol.id, + owner_id=user.id, + protocol_revision_id=revision.id, + dataset_row=dataset_row, + ) except (ProtocolValidationError, ValueError) as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc + if run.dataset_row is not None: + await db.commit() await enqueue_protocol_run(run.id) return await _test_run_response(db, run, protocol, experiment) @@ -523,7 +615,11 @@ async def get_latest_test_run_endpoint(protocol_id: uuid.UUID, user: CurrentUser @router.post("/{protocol_id}/nodes/{node_id}/run", response_model=ProtocolRunResponse, status_code=201) async def run_single_node_endpoint( - protocol_id: uuid.UUID, node_id: str, user: CurrentUser, db: DbSession + protocol_id: uuid.UUID, + node_id: str, + user: CurrentUser, + db: DbSession, + body: NodePlayRequest | None = None, ) -> ProtocolRunResponse: """The canvas's per-node Play icon -- runs one Agent node in isolation. Scoped to a node with no upstream input (validated up front, same @@ -534,11 +630,53 @@ async def run_single_node_endpoint( revision = await _require_published_revision(db, protocol) try: validate_single_node_runnable(revision.graph, node_id) + row_plan = resolve_dataset_row_plan(revision.graph) except ProtocolValidationError as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc + except DatasetRowInputError as exc: + raise HTTPException(status_code=422, detail=str(exc)) from exc + row_index = body.row_index if body is not None else None + row_binding = next( + (binding for binding in (row_plan or {}).get("bindings", []) if binding["agent_node_id"] == node_id), + None, + ) + if row_index is not None and row_binding is None: + raise HTTPException(status_code=422, detail="no_row_driver") + dataset_row = None + if row_binding is not None: + registration = ( + await db.execute( + select(RegisteredDataset).where( + RegisteredDataset.id == uuid.UUID(row_plan["driver_dataset_id"]), + RegisteredDataset.owner_id == user.id, + ) + ) + ).scalar_one_or_none() + if registration is None: + raise HTTPException(status_code=422, detail="Registered row source is unavailable.") + try: + source = read_row_source( + dataset_id=str(registration.id), + raw_path=registration.raw_path, + raw_sha256=registration.raw_sha256, + ) + columns = row_binding["columns"] + if not columns or set(columns) - set(source.columns): + raise DatasetRowCsvError( + "invalid_columns", "row input columns must exist in the registered CSV header" + ) + dataset_row = project_row(source, row_index=0 if row_index is None else row_index, columns=columns) + except DatasetRowCsvError as exc: + raise HTTPException(status_code=422, detail=str(exc)) from exc run = await create_protocol_run( - db, protocol_id=protocol_id, owner_id=user.id, target_node_id=node_id, protocol_revision_id=revision.id + db, + protocol_id=protocol_id, + owner_id=user.id, + target_node_id=node_id, + protocol_revision_id=revision.id, + dataset_row=dataset_row, ) + await db.commit() await enqueue_protocol_run(run.id) return _protocol_run_response(run) @@ -556,26 +694,43 @@ async def preview_node_prompt_endpoint( resource. """ protocol = await _get_owned_protocol(db, protocol_id, user) + dataset_rows: list[dict[str, Any]] = [] try: text = await preview_node_prompt( normalize_protocol_graph(body.graph) if body.graph is not None else protocol.graph, node_id, owner_id=user.id, experiment_id=protocol.experiment_id, + row_index=body.row_index, + dataset_row_out=dataset_rows, ) - except ProtocolValidationError as exc: + except (ProtocolValidationError, DatasetRowInputError, DatasetRowCsvError) as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc - return PromptPreviewResponse(text=text) + return PromptPreviewResponse(text=text, dataset_row=dataset_rows[0] if dataset_rows else None) @router.post("/{protocol_id}/cell-runs", response_model=CellRunBatchResponse, status_code=201) async def create_cell_runs_endpoint( protocol_id: uuid.UUID, user: CurrentUser, db: DbSession, body: CellRunBatchRequest | None = None ) -> CellRunBatchResponse: - """Run every pending replicate plus any explicitly selected reruns.""" + """Run pending replicates or retry explicitly selected failed row slots.""" protocol = await _get_owned_protocol(db, protocol_id, user) revision = await _require_published_revision(db, protocol) try: + retry_row_result_ids = body.retry_row_result_ids if body is not None else None + if retry_row_result_ids is not None: + try: + is_row_mode = resolve_dataset_row_plan(revision.graph) is not None + except DatasetRowInputError: + is_row_mode = False + if ( + not retry_row_result_ids + or len(set(retry_row_result_ids)) != len(retry_row_result_ids) + or (body is not None and body.replicate_labels is not None) + or (body is not None and body.rerun_replicate_labels) + or not is_row_mode + ): + raise ProtocolValidationError("invalid_retry_selection") runs, skipped = await plan_cell_runs( db, protocol_id=protocol.id, @@ -585,11 +740,24 @@ async def create_cell_runs_endpoint( protocol_revision_id=revision.id, replicate_labels=set(body.replicate_labels) if body and body.replicate_labels is not None else None, rerun_replicate_labels=set(body.rerun_replicate_labels) if body is not None else None, + retry_row_result_ids=retry_row_result_ids, ) except ProtocolValidationError as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc - for run in runs: - await enqueue_protocol_run(run.id) + row_mode = resolve_dataset_row_plan(revision.graph) is not None + if row_mode: + await db.commit() + for run in runs: + try: + await enqueue_protocol_run(run.id) + except Exception as exc: + await fail_protocol_run( + db, run.id, error=f"enqueue_failure: {type(exc).__name__}: {exc}" + ) + await db.commit() + else: + for run in runs: + await enqueue_protocol_run(run.id) return CellRunBatchResponse( protocol_run_ids=[r.id for r in runs], replicate_labels=[r.replicate_label for r in runs if r.replicate_label is not None], @@ -597,7 +765,7 @@ async def create_cell_runs_endpoint( protocol_revision_id=revision.id, protocol_revision=revision.revision, consumption_mode=( - "per_row" if resolve_dataset_row_plan(revision.graph) is not None else "whole_dataset" + "per_row" if row_mode else "whole_dataset" ), row_result_ids=[r.row_result_id for r in runs if r.row_result_id is not None], ) diff --git a/src/asaree/mcp_servers/script_server.py b/src/asaree/mcp_servers/script_server.py index adddb8e7..6c0b5133 100644 --- a/src/asaree/mcp_servers/script_server.py +++ b/src/asaree/mcp_servers/script_server.py @@ -81,6 +81,7 @@ _META_KEY_TARGET_COLUMN = "motoro.ambient.target_column" _META_KEY_DATASET_NAMES = "motoro.ambient.dataset_names" _META_KEY_DATASET_MODE = "motoro.ambient.dataset_mode" +_META_KEY_ROW_INPUTS = "motoro.ambient.row_inputs" _RUN_CONTEXT_ENV = "ASAREE_RUN_CONTEXT" @@ -100,7 +101,7 @@ # (the product database URL and the internal API key among them — see .env), and # a wired script is user-authored code that has no business reading either. # Inheriting os.environ would hand every one of them over. What's left is what a -# script plausibly needs to run and find the workspace. +# script plausibly needs to run; its workspace is supplied as its cwd. _ENV_PASSTHROUGH = ("PATH", "HOME", "LANG", "LC_ALL", "TMPDIR", "ASAREE_DATASET_WORKSPACE_DIR") @@ -204,6 +205,46 @@ def _runtime_manifest(ctx: Context[Any, Any, Any] | None, workspace_id: str) -> raw_names = _ambient_value(ctx, _META_KEY_DATASET_NAMES) names = [str(name) for name in raw_names if isinstance(name, str)] if isinstance(raw_names, list) else [] dataset_mode = _ambient(ctx, _META_KEY_DATASET_MODE) + row_inputs = _ambient_value(ctx, _META_KEY_ROW_INPUTS) + if dataset_mode == "per_row" and isinstance(row_inputs, list): + # The list was resolved for this invoking Agent. It is authoritative, + # including [], and must never be supplemented from shared workspace + # state or a registration lookup. + training_inputs = [] + for item in row_inputs: + if not isinstance(item, dict): + continue + mode = item.get("mode") + if mode == "per_row": + training_inputs.append( + { + "name": str(item.get("name") or ""), + "path": str(item.get("path") or ""), + "target_column": str(item.get("target_column") or ""), + "mode": "per_row", + "slot": None, + "workspace_version": None, + "dataset_id": item.get("dataset_id"), + "raw_sha256": item.get("raw_sha256"), + "row_index": item.get("row_index"), + "columns": item.get("columns"), + } + ) + elif mode in {"workspace", "raw_unsplit"} and item.get("path"): + # Whole-context inputs are included only if they occur in the + # invoking Agent's authorized row_inputs list. + training_inputs.append( + { + "name": str(item.get("name") or ""), + "path": str(item["path"]), + "target_column": str(item.get("target_column") or ""), + "mode": mode, + "slot": item.get("slot"), + "workspace_version": item.get("workspace_version"), + } + ) + return {"schema_version": 1, "training_inputs": training_inputs} + # An explicitly wired unsplit dataset must not inherit a durable workspace # left by an older protocol revision for the same experiment/cell. locators = raw_training_data_locators(workspace_id) if workspace_id and dataset_mode != "raw_unsplit" else {} @@ -243,15 +284,15 @@ def _runtime_manifest(ctx: Context[Any, Any, Any] | None, workspace_id: str) -> def _legacy_unsplit_state(manifest: dict[str, Any]) -> dict[str, Any] | None: """A read-only workspace-shaped view for scripts written before the API. - This is never placed at a real workspace root. It exists only in an - isolated execution directory while a single unsplit input's script runs, + This is never placed at a real workspace root. It exists only in an + isolated execution directory while a single input's script runs, so workspace tools cannot mistake it for a seeded train/test lineage. """ inputs = manifest.get("training_inputs") if not isinstance(inputs, list) or len(inputs) != 1: return None item = inputs[0] - if not isinstance(item, dict) or item.get("mode") != "raw_unsplit" or not item.get("path"): + if not isinstance(item, dict) or item.get("mode") not in {"raw_unsplit", "per_row"} or not item.get("path"): return None return { "target_column": str(item.get("target_column") or ""), @@ -356,14 +397,18 @@ def run_wired_script( ) env = {k: os.environ[k] for k in _ENV_PASSTHROUGH if k in os.environ} + if not workspace_id: + # A standalone script has no ASAREE workspace to resolve. Do not expose + # the server's workspace root merely because it exists in the parent. + env.pop("ASAREE_DATASET_WORKSPACE_DIR", None) # Unbuffered so a script killed by the timeout has still flushed what it # printed -- the whole value of a partial result is that it survives. env["PYTHONUNBUFFERED"] = "1" result: dict[str, Any] = {"code_sha256": code_sha256, "script": script_file.name} cwd = _working_dir(workspace_id, script_file) manifest = _runtime_manifest(ctx, workspace_id) - # Preserve cwd for ordinary/helper-based scripts. The isolated legacy - # view is only needed when the authored source explicitly expects the old + # Preserve the invoking Agent's private workspace cwd for ordinary scripts. + # The isolated legacy view is only needed when the source expects the old # state-file contract. legacy_state = _legacy_unsplit_state(manifest) if "state.json" in source else None compatibility_dir: Path | None = None diff --git a/src/asaree/mcp_servers/workspace_server.py b/src/asaree/mcp_servers/workspace_server.py index 0c5a91b8..4cbb10c8 100644 --- a/src/asaree/mcp_servers/workspace_server.py +++ b/src/asaree/mcp_servers/workspace_server.py @@ -31,6 +31,7 @@ Stage, Workspace, WorkspaceError, + ambient_value_from_ctx, dataset_slot, make_workspace_id, provenance, @@ -387,6 +388,64 @@ async def open_workspace( lineage, target column and HEAD without you naming anything. The response echoes the slot to pass to later staging calls. """ + row_inputs = ambient_value_from_ctx("row_inputs", ctx) + dataset_row = ambient_value_from_ctx("dataset_row", ctx) + if isinstance(row_inputs, list): + if dataset_row and isinstance(dataset_row, dict): + context_item = next((item for item in row_inputs if isinstance(item, dict) and + item.get("mode") in {"workspace", "raw_unsplit"} and + (item.get("slot") == slot if slot else bool(name) and item.get("name") == name)), None) + if context_item is not None: + if experiment_id or cell_label or stage or target_column: + return json.dumps({"error": "row workspace override is not authorized"}) + if context_item.get("mode") == "raw_unsplit": + return json.dumps({"workspace_id": str(resolve_workspace_id_from_ctx("", ctx)), + "dataset_mode": "raw_unsplit", "dataset_name": context_item["name"], + "data_path": context_item["path"], + "target_column": context_item["target_column"]}) + return json.dumps({"workspace_id": str(resolve_workspace_id_from_ctx("", ctx)), + "dataset_mode": "workspace", "dataset_name": context_item["name"], + "slot": context_item.get("slot"), + "workspace_version": context_item.get("workspace_version")}) + bound_name = str(dataset_row.get("dataset_name") or "") + bound_id = str(dataset_row.get("dataset_id") or "") + if experiment_id or cell_label or stage or slot or target_column or (name and name != bound_name): + return json.dumps({"error": "row workspace is fixed to the authorized Agent view"}) + if name and name != bound_name: + return json.dumps({"error": "dataset is not authorized for this row view"}) + entries = [item for item in row_inputs if isinstance(item, dict) and item.get("mode") == "per_row"] + if len(entries) != 1: + return json.dumps({"error": "no authorized row driver is bound"}) + item = entries[0] + if item.get("dataset_id") != bound_id: + return json.dumps({"error": "row driver does not match authorized snapshot"}) + return json.dumps({ + "workspace_id": str(resolve_workspace_id_from_ctx("", ctx)), + "dataset_mode": "per_row", "dataset_id": bound_id, + "raw_sha256": str(dataset_row.get("raw_sha256") or ""), + "row_index": int(dataset_row.get("row_index")), + "columns": list(item.get("columns") or []), "data_path": str(item.get("path") or ""), + "target_column": str(item.get("target_column") or ""), + }) + # A row-mode Agent without a driver can open only an explicitly bound + # whole-context slot; it cannot ask the global registry to resolve a name. + if experiment_id or cell_label or target_column or stage: + return json.dumps({"error": "row workspace override is not authorized"}) + if slot or name: + item = next((item for item in row_inputs if isinstance(item, dict) and + (item.get("slot") == slot if slot else item.get("name") == name)), None) + if item is None or item.get("mode") not in {"workspace", "raw_unsplit"}: + return json.dumps({"error": "dataset or slot is not authorized for this Agent"}) + if item.get("mode") == "raw_unsplit": + return json.dumps({"workspace_id": str(resolve_workspace_id_from_ctx("", ctx)), + "dataset_mode": "raw_unsplit", "dataset_name": item["name"], + "data_path": item["path"], "target_column": item["target_column"]}) + if item.get("mode") == "workspace": + return json.dumps({"workspace_id": str(resolve_workspace_id_from_ctx("", ctx)), + "dataset_mode": "workspace", "dataset_name": item["name"], + "slot": item.get("slot"), "workspace_version": item.get("workspace_version")}) + return json.dumps({"error": "dataset view is not authorized for this Agent"}) + return json.dumps({"error": "no authorized row driver is bound"}) # Both halves of the workspace id, or neither: a half-specified pair would # have to be reconciled against the ambient id, and there is no sensible # answer when they disagree. @@ -512,6 +571,19 @@ def workspace_status(workspace_id: str = "", ctx: Context[Any, Any, Any] | None workspace_id: ``"{experiment_id}/{cell_label}"``. Optional — resolved from the ambient request ``_meta`` when omitted. """ + row_inputs = ambient_value_from_ctx("row_inputs", ctx) + if isinstance(row_inputs, list): + try: + bound = resolve_workspace_id_from_ctx("", ctx) + if workspace_id.strip() and workspace_id.strip() != bound: + return json.dumps({"error": "workspace override is not authorized for this Agent"}) + except WorkspaceError as e: + return json.dumps({"error": f"workspace: {e}"}) + return json.dumps({"workspace_id": bound, "dataset_mode": "per_row", + "inputs": [{"name": item.get("name"), "mode": item.get("mode"), + "slot": item.get("slot"), "path": item.get("path"), + "columns": item.get("columns"), "target_column": item.get("target_column")} + for item in row_inputs if isinstance(item, dict)]}) try: wid = resolve_workspace_id_from_ctx(workspace_id, ctx) ws = Workspace(wid) diff --git a/src/asaree/script_context.py b/src/asaree/script_context.py index 8c050724..2ae92331 100644 --- a/src/asaree/script_context.py +++ b/src/asaree/script_context.py @@ -31,6 +31,10 @@ class TrainingInput: mode: str slot: str | None = None workspace_version: str | None = None + dataset_id: str | None = None + raw_sha256: str | None = None + row_index: int | None = None + columns: tuple[str, ...] | None = None def _manifest() -> dict[str, Any]: @@ -71,6 +75,12 @@ def training_inputs() -> tuple[TrainingInput, ...]: workspace_version=( str(item["workspace_version"]) if item.get("workspace_version") is not None else None ), + dataset_id=str(item["dataset_id"]) if item.get("dataset_id") is not None else None, + raw_sha256=str(item["raw_sha256"]) if item.get("raw_sha256") is not None else None, + row_index=int(item["row_index"]) if item.get("row_index") is not None else None, + columns=tuple(str(column) for column in item["columns"]) + if isinstance(item.get("columns"), list) + else None, ) ) return tuple(resolved) diff --git a/src/asaree/services/agent_messenger.py b/src/asaree/services/agent_messenger.py index 4c2fc845..11c68628 100644 --- a/src/asaree/services/agent_messenger.py +++ b/src/asaree/services/agent_messenger.py @@ -124,6 +124,9 @@ def __init__( entry_agent_id: str, workspace_id: str | None = None, stage_plan: Any = None, + row_source: Any = None, + row_plan: dict[str, Any] | None = None, + dataset_row: dict[str, Any] | None = None, ) -> None: #: The *revision* graph -- capability: who each peer is and how it runs. #: Authorization reads the live draft graph instead, on every call. @@ -136,6 +139,9 @@ def __init__( #: consulted mid-conversation seeds into the same pipeline as everyone #: else. ``None`` means "whatever this cell already stages through". self._stage_plan = stage_plan + self._row_source = row_source + self._row_plan = row_plan or {} + self._dataset_row = dataset_row self._entry_agent_id = entry_agent_id self._sequence = 0 #: Canvas node ids of the agents whose turns are currently on the stack, @@ -526,6 +532,8 @@ async def _run_peer( self._owner_id, protocol_run_id=self._protocol_run_id, stage_plan=self._stage_plan, + row_source=self._row_source, + row_bindings=self._row_plan.get("bindings"), ) # A configured Output Parser still defines a consulted worker's reply # contract. Return its compact payload to the parent and retain the @@ -542,6 +550,7 @@ async def _run_peer( available_agents=await resolve_available_agents(self._graph, to_agent_id, owner_id=self._owner_id), agent_messenger=self, unsplit_dataset=dataset.unsplit_name, + row_input_context=(ambient_meta.get("row_inputs") or [None])[0], ) # The canvas shows a consulted peer as a node that ran, because it did. # A peer consulted twice keeps only its latest turn here; the full @@ -585,6 +594,9 @@ async def execute_conversation( ambient_meta: dict[str, Any] | None = None, stage_plan: Any = None, unsplit_dataset: str = "", + dataset_row: dict[str, Any] | None = None, + row_source: Any = None, + row_plan: dict[str, Any] | None = None, ) -> tuple[dict[str, Any], str]: """The conversation itself: seed the transcript, run the entry agent, map its outcome to a terminal conversation state, checkpoint. @@ -615,6 +627,9 @@ async def execute_conversation( entry_agent_id=entry_agent_id, workspace_id=workspace_id, stage_plan=stage_plan, + row_source=row_source, + row_plan=row_plan, + dataset_row=dataset_row, ) messenger.append( from_agent_id=USER_PARTICIPANT, @@ -635,6 +650,8 @@ async def execute_conversation( owner_id, protocol_run_id=protocol_run_id, stage_plan=stage_plan, + row_source=row_source, + row_bindings=(row_plan or {}).get("bindings"), ) unsplit_dataset = unsplit_dataset or entry_dataset.unsplit_name @@ -653,6 +670,7 @@ async def execute_conversation( available_agents=await resolve_available_agents(graph, entry_agent_id, owner_id=owner_id), agent_messenger=messenger, unsplit_dataset=unsplit_dataset, + row_input_context=(ambient_meta.get("row_inputs") or [None])[0], ) cancelled = error == _AGENT_CANCELLED @@ -777,6 +795,9 @@ async def execute_supervisor_architecture( experiment_id: uuid.UUID | None = None, effective_cell_label: str | None = None, stage_plan: Any = None, + dataset_row: dict[str, Any] | None = None, + row_source: Any = None, + row_plan: dict[str, Any] | None = None, ) -> tuple[dict[str, Any], str]: """Run one cell as a supervisor dispatching to workers. @@ -818,6 +839,9 @@ async def execute_supervisor_architecture( entry_agent_id=roles.supervisor, workspace_id=workspace_id, stage_plan=stage_plan, + row_source=row_source, + row_plan=row_plan, + dataset_row=dataset_row, ) started_at = time.monotonic() deadline = _MAX_SUPERVISOR_TURN_DURATION.total_seconds() * roles.execution_budget @@ -879,6 +903,8 @@ async def _turn( protocol_run_id=protocol_run_id, slot_prefix=slot_prefix, stage_plan=stage_plan, + row_source=row_source, + row_bindings=(row_plan or {}).get("bindings"), ) prompt = _build_user_input( nodes[node_id], @@ -905,6 +931,7 @@ async def _turn( workspace_id=workspace_id, ambient_meta=ambient_meta, unsplit_dataset=dataset.unsplit_name, + row_input_context=(ambient_meta.get("row_inputs") or [None])[0], ) run: dict[str, Any] = { "status": "cancelled" if error == _AGENT_CANCELLED else ("failed" if error else "completed"), diff --git a/src/asaree/services/csv_export.py b/src/asaree/services/csv_export.py index 289255d6..6f8c18b6 100644 --- a/src/asaree/services/csv_export.py +++ b/src/asaree/services/csv_export.py @@ -385,8 +385,12 @@ def result_rows_schema( metric_types: dict[str, str] | None = None, metric_aggregations: dict[str, str] | None = None, design_spec: dict[str, Any] | None = None, + *, + consumption_mode: str = "whole_dataset", ) -> dict[str, Any]: """Machine-readable companion metadata for a Results analysis CSV.""" + if consumption_mode == "per_row": + return _row_results_schema(rows, design_spec) factor_columns, metric_keys, script_columns = _result_csv_layout(rows, design_spec) metadata_fields = _result_metadata_fields(rows) reported_keys = { @@ -481,7 +485,189 @@ def result_rows_schema( } -def result_rows_to_csv(rows: Sequence[dict[str, Any]], design_spec: dict[str, Any] | None = None) -> str: +def _row_metric_layout( + rows: Sequence[dict[str, Any]], design_spec: dict[str, Any] | None +) -> tuple[list[tuple[str, str]], list[tuple[str, str]]]: + metrics: dict[str, str] = {} + for metric in (design_spec or {}).get("metrics", []): + if isinstance(metric, dict) and metric.get("kind") == "custom": + metric_id, name = metric.get("id"), metric.get("name") + if isinstance(metric_id, str) and isinstance(name, str): + metrics[metric_id] = name + for row in rows: + measurement = row.get("measurement") + observations = measurement.get("observations") if isinstance(measurement, dict) else [] + for observation in observations or []: + if isinstance(observation, dict) and isinstance(observation.get("metric_id"), str): + metrics.setdefault( + observation["metric_id"], + str(observation.get("metric_name") or observation["metric_id"]), + ) + row_factors = _factor_columns(rows, reserved=_ROW_FIXED_FIELDS, design_spec=design_spec) + reserved = {*_ROW_FIXED_FIELDS, *(column for column, _, _, _ in row_factors)} + output: list[tuple[str, str]] = [] + for metric_id, name in sorted(metrics.items(), key=lambda item: (item[1], item[0])): + base = name + suffix = 2 + while any(candidate in reserved for candidate in (base, base + "__status", base + "__producer")): + base = f"{name}_{suffix}" + suffix += 1 + column = base + reserved.update((column, column + "__status", column + "__producer")) + output.append((metric_id, column)) + provenance = [ + (metric_id, column + suffix) + for metric_id, column in output + for suffix in ("__status", "__producer") + ] + return output, provenance + + +_ROW_FIXED_FIELDS = [ + "row_result_id", + "cell_id", + "cell_label", + "replicate_result_id", + "replicate_label", + "replicate_number", + "dataset_id", + "raw_sha256", + "row_index", + "protocol_revision_id", + "design_revision_id", + "run_id", + "status", + "workspace_id", +] + + +def _row_results_schema( + rows: Sequence[dict[str, Any]], design_spec: dict[str, Any] | None +) -> dict[str, Any]: + metrics, _ = _row_metric_layout(rows, design_spec) + factor_columns = _factor_columns(rows, reserved=_ROW_FIXED_FIELDS, design_spec=design_spec) + metadata = [ + { + "name": field, + "role": "row_identity", + **({"value_type": "integer", "zero_based": True} if field == "row_index" else {}), + } + for field in _ROW_FIXED_FIELDS + ] + factor_schema = [ + {"name": col, "role": "factor", "source_factor": key, "encoding": kind} + for col, key, kind, _ in factor_columns + ] + metric_schema = [] + for metric_id, column in metrics: + metric_schema.append( + { + "name": column, + "role": "reported", + "metric_id": metric_id, + "value_type": "json", + "display_export_only": True, + "aggregation": None, + } + ) + metric_schema.append( + {"name": column + "__status", "role": "observation_status", "metric_id": metric_id} + ) + metric_schema.append( + { + "name": column + "__producer", + "role": "producer_provenance", + "metric_id": metric_id, + "value_type": "json", + } + ) + return { + "schema_version": 5, + "row_unit": "row_execution", + "consumption_mode": "per_row", + "coverage_semantics": ( + "One row per stable dataset row slot; pending, failed, cancelled, and unavailable slots are " + "included. Attempts are not separate rows." + ), + "row_identity": metadata, + "columns": [*metadata, *factor_schema, *metric_schema], + } + + +def _row_results_to_csv(rows: Sequence[dict[str, Any]], design_spec: dict[str, Any] | None) -> str: + metric_columns, _ = _row_metric_layout(rows, design_spec) + factors = _factor_columns(rows, reserved=_ROW_FIXED_FIELDS, design_spec=design_spec) + fields = [ + *_ROW_FIXED_FIELDS, + *(col for col, _, _, _ in factors), + *(name for _, name in metric_columns), + *(name + suffix for _, name in metric_columns for suffix in ("__status", "__producer")), + ] + output = io.StringIO() + writer = csv.DictWriter(output, fieldnames=fields, extrasaction="ignore") + writer.writeheader() + for source in rows: + dataset = source.get("dataset_row") or {} + attempt = source.get("latest_attempt") or {} + row = {key: source.get(key, "") for key in _ROW_FIXED_FIELDS} + row.update(dataset) + row["run_id"] = attempt.get("run_id") or "" + row["workspace_id"] = attempt.get("workspace_id") or source.get("workspace_id") or "" + values = source.get("metric_values") or {} + measurement = source.get("measurement") + observations = measurement.get("observations") if isinstance(measurement, dict) else [] + by_id = {obs.get("metric_id"): obs for obs in observations or [] if isinstance(obs, dict)} + by_name = {obs.get("metric_name"): obs for obs in observations or [] if isinstance(obs, dict)} + for col, factor_key, kind, labels in factors: + value = (source.get("factor_values") or {}).get(factor_key) + if kind == "categorical": + row[col] = (labels or {}).get(_factor_level_key(value), "") + elif kind == "boolean" and value is not None: + row[col] = int(value) + else: + row[col] = value if value is not None else "" + for metric_id, column in metric_columns: + declared_name = next( + ( + metric.get("name") + for metric in (design_spec or {}).get("metrics", []) + if isinstance(metric, dict) and metric.get("id") == metric_id + ), + None, + ) + obs = by_id.get(metric_id) or by_name.get(declared_name) + status = obs.get("status", "unavailable") if obs else "unavailable" + raw_value = ( + obs.get("value") + if obs and status == "measured" + else values.get(column, values.get(obs.get("metric_name") if obs else "")) + ) + row[column] = ( + json.dumps(raw_value, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + if status == "measured" + else "" + ) + row[column + "__status"] = status + row[column + "__producer"] = ( + json.dumps( + obs.get("producer") if obs else None, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ) + if obs + else "" + ) + writer.writerow(row) + return output.getvalue() + + +def result_rows_to_csv( + rows: Sequence[dict[str, Any]], + design_spec: dict[str, Any] | None = None, + *, + consumption_mode: str = "whole_dataset", +) -> str: """Export the enriched Results response as an analysis-ready CSV. Unlike :func:`replicates_to_csv`, this receives the read-only Results @@ -491,6 +677,8 @@ def result_rows_to_csv(rows: Sequence[dict[str, Any]], design_spec: dict[str, An become their short persisted level labels rather than raw prompt or configuration payloads. """ + if consumption_mode == "per_row": + return _row_results_to_csv(rows, design_spec) factor_columns, metric_keys, script_columns = _result_csv_layout(rows, design_spec) declared_custom_metric_names = { metric["name"] diff --git a/src/asaree/services/dataset_row_workspaces.py b/src/asaree/services/dataset_row_workspaces.py new file mode 100644 index 00000000..83256baf --- /dev/null +++ b/src/asaree/services/dataset_row_workspaces.py @@ -0,0 +1,179 @@ +"""Materialize private, run-scoped CSV snapshots for Dataset row inputs.""" + +from __future__ import annotations + +import hashlib +import json +import uuid +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +from asaree.config import get_settings +from asaree.services.dataset_row_csv import ( + DatasetRowCsvError, + RowSource, + materialize_row_csv, + project_row, +) + + +def row_attempt_workspace_id(run_id: uuid.UUID) -> str: + """Return the opaque workspace id for this immutable protocol attempt.""" + if not isinstance(run_id, uuid.UUID): + raise DatasetRowCsvError("invalid_run_id", "run id must be a UUID") + return f"_protocol_runs/{run_id}" + + +def _workspace_root() -> Path: + # Keep attempt data beside registered datasets while separating it from + # their immutable originals and the shared whole-dataset workspaces. + return Path(get_settings().dataset_storage_dir).resolve().parent + + +def _node_directory(root: Path, workspace_id: str, agent_node_id: str) -> Path: + key = hashlib.sha256(agent_node_id.encode("utf-8")).hexdigest() + attempt_root = (root / workspace_id).resolve() + directory = (attempt_root / "agents" / key).resolve() + if not directory.is_relative_to(attempt_root): + raise DatasetRowCsvError("invalid_agent_node_id", "Agent node directory escapes attempt root") + return directory + + +def prepare_agent_row_inputs( + *, + run_id: uuid.UUID, + agent_node_id: str, + source: RowSource, + bindings: Sequence[dict], +) -> dict[str, Any]: + """Create exact one-row CSV views for the requested Agent's driver edges.""" + workspace_id = row_attempt_workspace_id(run_id) + if not isinstance(agent_node_id, str) or not agent_node_id: + raise DatasetRowCsvError("invalid_agent_node_id", "Agent node id must be a nonempty string") + try: + source_id = str(uuid.UUID(source.dataset_id)) + except (ValueError, AttributeError, TypeError) as exc: + raise DatasetRowCsvError("source_identity_mismatch", "row source dataset id must be a UUID") from exc + + selected: list[dict] = [] + for binding in bindings: + if not isinstance(binding, dict): + raise DatasetRowCsvError("invalid_binding", "row input binding must be an object") + if binding.get("agent_node_id") != agent_node_id: + continue + try: + binding_id = str(uuid.UUID(binding.get("dataset_id"))) + except (ValueError, AttributeError, TypeError) as exc: + raise DatasetRowCsvError("source_identity_mismatch", "binding dataset id must be a UUID") from exc + if binding_id != source_id: + raise DatasetRowCsvError("source_identity_mismatch", "binding dataset does not match row source") + selected.append(binding) + + if not selected: + return {"workspace_id": workspace_id, "row_inputs": []} + + views: dict[str, dict] = {} + for binding in selected: + name = binding.get("dataset_name") + columns = binding.get("columns") + target = binding.get("target_column") + if not isinstance(name, str): + raise DatasetRowCsvError("invalid_binding", "dataset_name must be a string") + if target is None: + target = "" + if not isinstance(target, str): + raise DatasetRowCsvError("invalid_binding", "target_column must be a string or null") + view = project_row(source, row_index=binding.get("row_index"), columns=columns) + if target not in view["columns"]: + target = "" + identity = json.dumps( + {"dataset_id": binding_id, "row_index": view["row_index"], "columns": view["columns"], "target": target}, + sort_keys=True, + separators=(",", ":"), + ) + record = {"binding": binding, "view": view, "target": target, "name": name} + previous = views.get(identity) + if previous is not None: + if previous["name"] != name: + raise DatasetRowCsvError("conflicting_agent_views", "duplicate row views have different names") + continue + # One Agent gets one unambiguous view of a driver for this run. + if views: + raise DatasetRowCsvError("conflicting_agent_views", "Agent row bindings select different views") + views[identity] = record + + directory = _node_directory(_workspace_root(), workspace_id, agent_node_id) + row_inputs: list[dict] = [] + for record in views.values(): + view = record["view"] + token = hashlib.sha256( + json.dumps( + { + "dataset_id": source_id, + "raw_sha256": source.raw_sha256, + "row_index": view["row_index"], + "columns": view["columns"], + }, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + ).hexdigest() + destination = directory / f"{token}.csv" + manifest = directory / "row-view.json" + manifest_value = json.dumps( + { + "dataset_id": source_id, + "raw_sha256": source.raw_sha256, + "row_index": view["row_index"], + "columns": view["columns"], + "values": view["values"], + }, + sort_keys=True, + separators=(",", ":"), + ) + if manifest.exists(): + try: + if manifest.read_text(encoding="utf-8") != manifest_value: + raise DatasetRowCsvError( + "workspace_conflict", "Agent workspace already contains a different source view" + ) + except OSError as exc: + raise DatasetRowCsvError("source_unavailable", "Agent workspace cannot be verified") from exc + elif directory.exists() and any(directory.iterdir()): + raise DatasetRowCsvError("workspace_conflict", "Agent workspace contains unrecognized files") + # Let the CSV writer quote fields, then compare with an independently + # materialized candidate so existing same-run files are never trusted. + if destination.exists(): + temp_view = directory / f".{token}.verify.csv" + try: + materialize_row_csv(view, destination=temp_view) + expected_bytes = temp_view.read_bytes() + temp_view.unlink() + except OSError as exc: + raise DatasetRowCsvError("source_unavailable", "row snapshot cannot be verified") from exc + if destination.read_bytes() != expected_bytes: + raise DatasetRowCsvError( + "workspace_conflict", "existing row snapshot differs from requested source view" + ) + else: + materialize_row_csv(view, destination=destination) + if not manifest.exists(): + try: + manifest.write_text(manifest_value, encoding="utf-8") + except OSError as exc: + raise DatasetRowCsvError("source_unavailable", "Agent workspace metadata cannot be written") from exc + row_inputs.append( + { + "name": record["name"], + "dataset_id": source_id, + "raw_sha256": source.raw_sha256, + "row_index": view["row_index"], + "columns": view["columns"], + "values": view["values"], + "path": str(destination), + "mode": "per_row", + "target_column": record["target"], + } + ) + return {"workspace_id": workspace_id, "row_inputs": row_inputs} diff --git a/src/asaree/services/design_generation.py b/src/asaree/services/design_generation.py index 70f6f7d0..7a17733c 100644 --- a/src/asaree/services/design_generation.py +++ b/src/asaree/services/design_generation.py @@ -24,8 +24,11 @@ from dataclasses import dataclass from typing import Any +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession +from asaree.models.experiment import ResearchExperiment +from asaree.models.experiment_design_revision import ExperimentDesignRevision from asaree.models.factorial_replicate_result import FactorialReplicateResult from asaree.services.coordination import coordination_strategy_slug from asaree.services.design_revisions import get_current_revision, supersede_and_create @@ -385,6 +388,16 @@ async def generate_design_cells( practice) — it never affects which combinations/replicates are generated or their (deterministic) cell_label. """ + experiment = ( + await db.execute( + select(ResearchExperiment) + .where(ResearchExperiment.id == experiment_id) + .with_for_update() + .execution_options(populate_existing=True) + ) + ).scalar_one_or_none() + if experiment is None: + raise DesignValidationError("Experiment not found") # An empty factor declaration is a meaningful replacement when an # existing design's final factor was removed: it retires every current # cell into history and leaves an empty current revision. Do not call @@ -393,7 +406,17 @@ async def generate_design_cells( planned = _planned_replicates(factors, replicates) if factors else [] planned_labels = {label for label, _, _ in planned} - current = await get_current_revision(db, experiment_id) + current = ( + await db.execute( + select(ExperimentDesignRevision) + .where( + ExperimentDesignRevision.experiment_id == experiment_id, + ExperimentDesignRevision.superseded_at.is_(None), + ) + .with_for_update() + .execution_options(populate_existing=True) + ) + ).scalar_one_or_none() existing = ( { replicate.replicate_label: replicate diff --git a/src/asaree/services/experiment_run_results.py b/src/asaree/services/experiment_run_results.py index 7f4bd0f0..bd0d3765 100644 --- a/src/asaree/services/experiment_run_results.py +++ b/src/asaree/services/experiment_run_results.py @@ -13,17 +13,26 @@ from collections import defaultdict from datetime import datetime from decimal import Decimal -from math import isfinite +from math import isfinite, prod from typing import Any from motoro.runner import get_run from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession +from asaree.models.dataset import RegisteredDataset +from asaree.models.experiment_design_revision import ExperimentDesignRevision +from asaree.models.factorial_cell import FactorialCell +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult from asaree.models.protocol import Protocol from asaree.models.protocol_revision import ProtocolRevision from asaree.models.protocol_run import ProtocolRun +from asaree.services.dataset_row_csv import DatasetRowCsvError, read_row_source +from asaree.services.dataset_row_inputs import resolve_dataset_row_plan +from asaree.services.dataset_row_workspaces import row_attempt_workspace_id from asaree.services.factorial_cells import list_replicates +from asaree.services.factorial_row_results import list_row_results from asaree.services.measurement_migration import LegacyResultFacets, legacy_measurement_facets from asaree.services.metrics import normalize_metrics from asaree.services.protocol_runs import list_experiment_trials @@ -346,18 +355,384 @@ async def _agent_runs_by_id(run_ids: set[uuid.UUID]) -> dict[uuid.UUID, Any]: } +class RunResultsProjectionError(ValueError): + """A requested Results selector does not belong to the experiment.""" + + +def _reported_metric_ids(design_spec: dict[str, Any] | None) -> list[str]: + return [ + metric["id"] + for metric in normalize_metrics((design_spec or {}).get("metrics")) + if metric.get("kind") == "custom" and isinstance(metric.get("id"), str) + ] + + +async def _source_row_count( + db: AsyncSession, *, dataset_id: uuid.UUID | None, raw_sha256: str | None +) -> int | None: + if dataset_id is None or not isinstance(raw_sha256, str): + return None + dataset = await db.get(RegisteredDataset, dataset_id) + if dataset is None: + return None + try: + source = read_row_source( + dataset_id=str(dataset.id), + raw_path=dataset.raw_path, + raw_sha256=dataset.raw_sha256, + expected_sha256=raw_sha256, + ) + except DatasetRowCsvError: + return None + return len(source.rows) + + +async def _summarize_row_results( + db: AsyncSession, + *, + experiment_id: uuid.UUID, + design_spec: dict[str, Any] | None, + protocol_revision: ProtocolRevision, + design_revision: ExperimentDesignRevision | None, + row_plan: dict[str, Any], +) -> dict[str, Any]: + revision_id = design_revision.id if design_revision is not None else None + protocol_revision_id = protocol_revision.id + cells = list( + ( + await db.execute( + select(FactorialCell).where( + FactorialCell.experiment_id == experiment_id, + *([FactorialCell.design_revision_id == revision_id] if revision_id else []), + ) + ) + ).scalars().all() + ) if revision_id else [] + replicates = list( + ( + await db.execute( + select(FactorialReplicateResult) + .join(FactorialCell, FactorialReplicateResult.cell_id == FactorialCell.id) + .where( + FactorialCell.experiment_id == experiment_id, + FactorialCell.design_revision_id == revision_id, + ) + ) + ).scalars().all() + ) if revision_id else [] + slots = await list_row_results( + db, + experiment_id=experiment_id, + design_revision_id=revision_id, + protocol_revision_id=protocol_revision_id, + ) if revision_id else [] + attempts: dict[uuid.UUID, list[ProtocolRun]] = defaultdict(list) + if slots: + run_rows = ( + await db.execute( + select(ProtocolRun) + .join(FactorialRowResult, ProtocolRun.row_result_id == FactorialRowResult.id) + .join(FactorialReplicateResult, FactorialRowResult.replicate_result_id == FactorialReplicateResult.id) + .join(FactorialCell, FactorialReplicateResult.cell_id == FactorialCell.id) + .where( + ProtocolRun.row_result_id.in_([slot.id for slot in slots]), + FactorialCell.experiment_id == experiment_id, + FactorialCell.design_revision_id == revision_id, + FactorialRowResult.protocol_revision_id == protocol_revision_id, + ) + .order_by(ProtocolRun.created_at, ProtocolRun.id) + ) + ).scalars().all() + for run in run_rows: + if run.row_result_id is not None: + attempts[run.row_result_id].append(run) + + replicate_by_id = {replicate.id: replicate for replicate in replicates} + cell_by_id = {cell.id: cell for cell in cells} + declared_ids = _reported_metric_ids(design_spec) + coverage = { + metric_id: {"measured": 0, "unavailable": 0, "failed": 0, "cancelled": 0, "other": 0} + for metric_id in declared_ids + } + row_results: list[dict[str, Any]] = [] + counts = {key: 0 for key in ("pending", "running", "completed", "failed", "cancelled")} + scored = missing_reported = 0 + dataset_ids = {slot.dataset_id for slot in slots} + source_dataset_id = min(dataset_ids, key=str) if dataset_ids else None + source_sha = next((slot.raw_sha256 for slot in slots), None) + if source_dataset_id is None: + try: + source_dataset_id = uuid.UUID(str(row_plan.get("driver_dataset_id"))) + except (ValueError, TypeError, AttributeError): + source_dataset_id = None + if source_sha is None and source_dataset_id is not None: + source_registration = await db.get(RegisteredDataset, source_dataset_id) + source_sha = source_registration.raw_sha256 if source_registration is not None else None + row_count = await _source_row_count(db, dataset_id=source_dataset_id, raw_sha256=source_sha) if source_sha else None + + for slot in slots: + parent = replicate_by_id.get(slot.replicate_result_id) + cell = cell_by_id.get(parent.cell_id) if parent is not None else None + history = attempts.get(slot.id, []) + latest = history[-1] if history else None + current_run = next((run for run in history if run.id == slot.run_id), None) + status = current_run.status if current_run is not None else "pending" + if status in {"running", "finalizing"}: + counts["running"] += 1 + elif status == "completed": + counts["completed"] += 1 + elif status in {"failed", "limit_reached"}: + counts["failed"] += 1 + elif status == "cancelled": + counts["cancelled"] += 1 + else: + counts["pending"] += 1 + + attempt_measurement = None + if latest is not None and isinstance(latest.attempt_result, dict): + attempt_measurement = latest.attempt_result.get("measurement") + measurement = attempt_measurement + if not isinstance(measurement, dict) and isinstance(slot.artifacts, dict): + measurement = slot.artifacts.get("measurement") + observations = ( + attempt_measurement.get("observations") + if isinstance(attempt_measurement, dict) and isinstance(attempt_measurement.get("observations"), list) + else [] + ) + by_metric = { + observation.get("metric_id"): observation + for observation in observations + if isinstance(observation, dict) and isinstance(observation.get("metric_id"), str) + } + all_measured = bool(declared_ids) and all( + by_metric.get(metric_id, {}).get("status") == "measured" for metric_id in declared_ids + ) + truncated = any( + isinstance(node, dict) and isinstance(node.get("truncation"), dict) + for node in (latest.node_runs or {}).values() + ) if latest is not None else False + if status == "completed": + if declared_ids and not all_measured: + missing_reported += 1 + if declared_ids and all_measured and not truncated: + scored += 1 + for metric_id in declared_ids: + observation = by_metric.get(metric_id) + observation_status = observation.get("status") if observation is not None else "unavailable" + bucket = observation_status if observation_status in coverage[metric_id] else "other" + coverage[metric_id][bucket] += 1 + + row_results.append( + { + "row_result_id": str(slot.id), + "cell_id": str(cell.id) if cell else str(parent.cell_id) if parent else "", + "cell_label": cell.cell_label if cell else "", + "factor_values": (cell.factor_values or {}) if cell else {}, + "replicate_result_id": str(parent.id) if parent else str(slot.replicate_result_id), + "replicate_label": parent.replicate_label if parent else "", + "replicate_number": parent.replicate_number if parent else 0, + "dataset_row": { + "dataset_id": str(slot.dataset_id), + "raw_sha256": slot.raw_sha256, + "row_index": slot.row_index, + }, + "design_revision_id": str(revision_id) if revision_id else "", + "protocol_revision_id": str(protocol_revision_id), + "status": status, + "workspace_id": slot.workspace_id, + "metric_values": slot.metric_values, + "measurement": measurement, + "artifacts": slot.artifacts, + "latest_attempt": _row_attempt_payload(latest, current=latest.id == slot.run_id) if latest else None, + "attempts": [_row_attempt_payload(run, current=run.id == slot.run_id) for run in history], + } + ) + + # Before the first row plan, retain the design's forecast while planned stays + # the number of persisted slots (zero). + actual_cells = len(cells) + actual_parents = len(replicates) + forecast_parents = actual_parents + spec = design_revision.design_spec if design_revision is not None else design_spec + factors = (spec or {}).get("factors") if isinstance(spec, dict) else None + if actual_cells == 0 and isinstance(factors, list) and factors: + sizes = [ + len(factor.get("levels")) + for factor in factors + if isinstance(factor, dict) and isinstance(factor.get("levels"), list) + ] + if len(sizes) == len(factors): + forecast_parents = prod(sizes) * max(1, int((spec or {}).get("replicates") or 1)) + expected = row_count * forecast_parents if row_count is not None else None + row_summary = { + "cell_count": actual_cells, + "parent_replicate_count": actual_parents, + "row_count": row_count, + "expected": expected, + "planned": len(slots), + **counts, + "scored": scored, + "missing_reported": missing_reported, + "metric_coverage": coverage, + } + # Keep compatibility keys, but no cell rollups or outcome scorecard exists + # for row-mode runs. + return { + "consumption_mode": "per_row", + "row_results": row_results, + "row_summary": row_summary, + "overview": { + "total_replicates": actual_parents, + "completed_replicates": counts["completed"], + "running_replicates": counts["running"], + "queued_replicates": counts["pending"], + "failed_replicates": counts["failed"], + "not_started_replicates": 0, + }, + "metric_keys": [], + "metric_types": {}, + "metric_aggregations": {}, + "metric_directions": {}, + "primary_metric": None, + "primary_metric_direction": None, + "cells": [], + "replicates": [], + } + + +def _row_attempt_payload(run: ProtocolRun, *, current: bool) -> dict[str, Any]: + return { + "run_id": str(run.id), + "status": run.status, + "error": run.error, + "started_at": run.started_at, + "completed_at": run.completed_at, + "workspace_id": row_attempt_workspace_id(run.id), + "attempt_result": run.attempt_result, + "node_runs": run.node_runs, + "conversation": run.conversation, + "dataset_row": run.dataset_row, + "protocol_revision_id": str(run.protocol_revision_id) if run.protocol_revision_id else None, + "design_revision_id": str(run.design_revision_id) if run.design_revision_id else None, + "current": current, + } + + async def summarize_experiment_run_results( - db: AsyncSession, *, experiment_id: uuid.UUID, design_spec: dict[str, Any] | None = None + db: AsyncSession, + *, + experiment_id: uuid.UUID, + design_spec: dict[str, Any] | None = None, + protocol_id: uuid.UUID | None = None, + design_revision_id: uuid.UUID | None = None, + protocol_revision_id: uuid.UUID | None = None, ) -> dict[str, Any]: - """Build a current-design results scorecard plus cell/replicate detail. + """Build whole-dataset scorecards or scoped per-row Results projections. - The summary intentionally excludes obsolete results from metric and usage - aggregates. They remain visible in the returned replicate list so people - can inspect history without allowing an older canvas to influence today's - comparison. + Whole-dataset rollups preserve their existing aggregate behavior. Row mode + returns stable slots and their immutable attempt histories without combining + values across rows. """ - replicates = await list_replicates(db, experiment_id=experiment_id) - trials = await list_experiment_trials(db, experiment_id=experiment_id) + protocols = list( + (await db.execute(select(Protocol).where(Protocol.experiment_id == experiment_id))).scalars().all() + ) + if protocol_id is not None: + protocol = next((item for item in protocols if item.id == protocol_id), None) + if protocol is None: + raise RunResultsProjectionError("protocol_not_found") + elif len(protocols) > 1: + raise RunResultsProjectionError("ambiguous_protocol") + else: + protocol = protocols[0] if protocols else None + + if design_revision_id is not None: + revision = await db.scalar( + select(ExperimentDesignRevision).where( + ExperimentDesignRevision.id == design_revision_id, + ExperimentDesignRevision.experiment_id == experiment_id, + ) + ) + if revision is None: + raise RunResultsProjectionError("design_revision_not_found") + else: + revision = await db.scalar( + select(ExperimentDesignRevision).where( + ExperimentDesignRevision.experiment_id == experiment_id, + ExperimentDesignRevision.superseded_at.is_(None), + ) + ) + if design_revision_id is not None and revision is not None: + design_spec = revision.design_spec + + published = None + if protocol is not None: + if protocol_revision_id is not None: + published = await db.scalar( + select(ProtocolRevision).where( + ProtocolRevision.id == protocol_revision_id, + ProtocolRevision.protocol_id == protocol.id, + ) + ) + if published is None: + raise RunResultsProjectionError("protocol_revision_not_found") + elif protocol.published_revision_id is not None: + published = await db.get(ProtocolRevision, protocol.published_revision_id) + elif protocol_revision_id is not None: + raise RunResultsProjectionError("protocol_revision_not_found") + + graph = published.graph if published is not None and isinstance(published.graph, dict) else {} + row_plan = resolve_dataset_row_plan(graph) if published is not None else None + if row_plan is not None: + return await _summarize_row_results( + db, + experiment_id=experiment_id, + design_spec=design_spec, + protocol_revision=published, + design_revision=revision, + row_plan=row_plan, + ) + + selected_revision_id = revision.id if revision is not None else None + replicates = await list_replicates(db, experiment_id=experiment_id, revision_id=selected_revision_id) + trials = await list_experiment_trials( + db, experiment_id=experiment_id, revision_id=selected_revision_id + ) + current_trial_run_ids = {trial.run_id for trial in trials if trial.run_id is not None} + excluded_trial_run_ids: set[uuid.UUID] = set() + if current_trial_run_ids: + excluded_trial_run_ids = set( + ( + await db.execute( + select(ProtocolRun.id).where( + ProtocolRun.id.in_(current_trial_run_ids), + ProtocolRun.row_result_id.is_not(None), + ) + ) + ).scalars().all() + ) + if protocol_id is not None or protocol_revision_id is not None: + compatible = set( + ( + await db.execute( + select(ProtocolRun.id).where( + ProtocolRun.id.in_(current_trial_run_ids), + ProtocolRun.row_result_id.is_(None), + *([ProtocolRun.protocol_id == protocol_id] if protocol_id is not None else []), + *( + [ProtocolRun.protocol_revision_id == protocol_revision_id] + if protocol_revision_id is not None + else [] + ), + ) + ) + ).scalars().all() + ) + excluded_trial_run_ids.update(current_trial_run_ids - compatible) + for trial in trials: + if trial.run_id in excluded_trial_run_ids: + trial.run_id = None + trial.status = "not_started" + trial.error = None trials_by_label = {trial.replicate_label: trial for trial in trials} protocol_run_ids = {trial.run_id for trial in trials if trial.run_id is not None} # The replicate row points to its latest ProtocolRun, but every earlier @@ -368,7 +743,14 @@ async def summarize_experiment_run_results( select(ProtocolRun, Protocol.published_revision_id, ProtocolRevision.published_at) .join(Protocol, ProtocolRun.protocol_id == Protocol.id) .outerjoin(ProtocolRevision, Protocol.published_revision_id == ProtocolRevision.id) - .where(Protocol.experiment_id == experiment_id, ProtocolRun.replicate_label.is_not(None)) + .where( + Protocol.experiment_id == experiment_id, + ProtocolRun.replicate_label.is_not(None), + ProtocolRun.row_result_id.is_(None), + *([ProtocolRun.protocol_id == protocol.id] if protocol is not None else []), + *([ProtocolRun.protocol_revision_id == protocol_revision_id] if protocol_revision_id else []), + *([ProtocolRun.design_revision_id == design_revision_id] if design_revision_id else []), + ) ) ).all() history_by_label: defaultdict[str, list[tuple[ProtocolRun, bool]]] = defaultdict(list) @@ -692,6 +1074,9 @@ def historical_run_payload(protocol_run: ProtocolRun, *, obsolete: bool) -> dict } primary_metric, primary_metric_direction = _primary_metric(design_spec) return { + "consumption_mode": "whole_dataset", + "row_results": [], + "row_summary": None, "overview": overview, "metric_keys": sorted(metric_keys), "metric_types": {key: metric_types.get(key, "number") for key in sorted(metric_keys)}, diff --git a/src/asaree/services/factorial_analysis.py b/src/asaree/services/factorial_analysis.py index 045ee1c2..bc0d05e3 100644 --- a/src/asaree/services/factorial_analysis.py +++ b/src/asaree/services/factorial_analysis.py @@ -517,7 +517,12 @@ def objective(beta: np.ndarray) -> tuple[float, np.ndarray]: } -def analyze_experiment_design(design_spec: dict[str, Any] | None, replicates: Sequence[Any]) -> dict[str, Any]: +def analyze_experiment_design( + design_spec: dict[str, Any] | None, + replicates: Sequence[Any], + *, + consumption_mode: str = "whole_dataset", +) -> dict[str, Any]: """The Results tab's own entry point -- a thin wrapper around ``analyze_factorial`` that needs no caller-supplied parameters at all. ``condition_factors``/``positive_levels``/``reference_condition``/ @@ -536,6 +541,26 @@ def analyze_experiment_design(design_spec: dict[str, Any] | None, replicates: Se whether that's a missing declaration or ``analyze_factorial`` itself raising ``FactorialAnalysisError`` (e.g. not enough scored replicates). """ + if consumption_mode == "per_row": + return { + "available": False, + "reason": ( + "Per-row executions are available for inspection and export; " + "factorial analysis is not supported." + ), + "analysis": None, + "best_condition": None, + } + if consumption_mode != "whole_dataset": + raise FactorialAnalysisError(f"unsupported consumption mode: {consumption_mode}") + for replicate in replicates: + if isinstance(replicate, dict): + row_shaped = any(key in replicate for key in ("row_result_id", "row_index", "dataset_row")) + else: + row_shaped = hasattr(replicate, "row_result_id") or type(replicate).__name__ == "FactorialRowResult" + if row_shaped: + raise FactorialAnalysisError("row execution results cannot be analyzed as factorial replicates") + design_spec = design_spec or {} factors = design_spec.get("factors") or [] metrics = design_spec.get("metrics") or [] diff --git a/src/asaree/services/factorial_row_results.py b/src/asaree/services/factorial_row_results.py index bb82ade8..d2a2e4c7 100644 --- a/src/asaree/services/factorial_row_results.py +++ b/src/asaree/services/factorial_row_results.py @@ -4,6 +4,7 @@ import re import uuid +from datetime import UTC, datetime from sqlalchemy import exists, select, update from sqlalchemy.exc import IntegrityError @@ -64,7 +65,7 @@ async def project_row_attempt( update(FactorialRowResult) .where(FactorialRowResult.id == row_result_id, valid_scope) .values(**fields) - .execution_options(synchronize_session=False) + .execution_options(synchronize_session="fetch") ) return result.rowcount == 1 @@ -282,7 +283,14 @@ async def claim_row_attempt( from asaree.services.protocol_runs import create_protocol_run # The locked slot defines the identity, even when callers include it in kwargs. - return await create_protocol_run(db, **{**create_kwargs, "row_result_id": row_result_id}) + run = await create_protocol_run(db, **{**create_kwargs, "row_result_id": row_result_id}) + # PostgreSQL's ``now()`` is fixed at transaction start. A retry commonly + # happens in the same transaction as its failed predecessor, so give it + # its actual creation time and retain the documented created_at/id history + # ordering instead of letting a random UUID decide which attempt is latest. + run.created_at = datetime.now(UTC) + await db.flush() + return run async def list_row_attempts( diff --git a/src/asaree/services/protocol_execution.py b/src/asaree/services/protocol_execution.py index 286886e8..5c671314 100644 --- a/src/asaree/services/protocol_execution.py +++ b/src/asaree/services/protocol_execution.py @@ -56,6 +56,7 @@ from asaree.models.database import get_session from asaree.models.dataset import RegisteredDataset from asaree.models.experiment import ResearchExperiment +from asaree.models.experiment_design_revision import ExperimentDesignRevision from asaree.models.protocol import Protocol from asaree.models.protocol_revision import ProtocolRevision from asaree.models.protocol_run import ProtocolRun @@ -65,6 +66,7 @@ from asaree.services.dataset_row_csv import DatasetRowCsvError, project_row, read_row_source from asaree.services.dataset_row_inputs import DatasetRowInputError, resolve_dataset_row_plan from asaree.services.dataset_row_planning import enumerate_row_candidates +from asaree.services.dataset_row_workspaces import prepare_agent_row_inputs, row_attempt_workspace_id from asaree.services.dataset_workspaces import ( WorkspaceSeedError, fetch_owned_registration, @@ -82,7 +84,7 @@ from asaree.services.experiments import get_experiment from asaree.services.factor_bindings import validate_factor_bindings from asaree.services.factorial_cells import get_replicate, list_replicates, upsert_replicate -from asaree.services.factorial_row_results import ensure_row_result, list_row_attempts +from asaree.services.factorial_row_results import claim_row_attempt, ensure_row_result, get_row_result from asaree.services.protocol_revisions import get_revision from asaree.services.protocol_runs import ( TERMINAL_PROTOCOL_RUN_STATUSES, @@ -1554,6 +1556,9 @@ def _ambient_meta_for( *, script_workspace_id: str | None = None, slots: tuple[str, ...] = (), + row_mode: bool = False, + row_inputs: list[dict[str, Any]] | None = None, + dataset_row: dict[str, Any] | None = None, ) -> dict[str, Any]: """The node's Reference-route values, for Motoro's caller-ambient ``_meta``. @@ -1610,7 +1615,30 @@ def _ambient_meta_for( wire call is unchanged for a node with no references. """ meta: dict[str, Any] = {} - dataset_names = [str(c["dataset_name"]) for c in _resolve_dataset_configs(graph, node_id) if c.get("dataset_name")] + if row_mode: + # Row attempts never inherit locators from a shared cell workspace. + meta["dataset_mode"] = "per_row" + meta["row_inputs"] = list(row_inputs or []) + if dataset_row is not None: + meta["dataset_row"] = dataset_row + if len(row_inputs or []) == 1: + item = row_inputs[0] + meta["data_path"] = str(item["path"]) + if item.get("target_column"): + meta["target_column"] = str(item["target_column"]) + if len(row_inputs or []) > 1: + meta["data_slots"] = { + str(item["slot"]): {"name": item["name"], "data_path": item["path"], + "target_column": item.get("target_column", "")} + for item in row_inputs if item.get("slot") + } + # Scripts remain private to this Agent, but the cell workspace is not. + workspace_id = None + dataset_names = ( + [str(item["name"]) for item in (row_inputs or []) if item.get("name")] + if row_mode else + [str(c["dataset_name"]) for c in _resolve_dataset_configs(graph, node_id) if c.get("dataset_name")] + ) if dataset_names: meta["dataset_names"] = dataset_names if workspace_id: @@ -1825,6 +1853,10 @@ async def _node_run_context( protocol_run_id: uuid.UUID | None = None, slot_prefix: str | None = None, stage_plan: Any = None, + row_input_context: Mapping[str, Any] | None = None, + row_inputs: list[dict[str, Any]] | None = None, + row_source: Any = None, + row_bindings: list[dict[str, Any]] | None = None, ) -> tuple[dict[str, Any], NodeDataset]: """``(ambient_meta, dataset)`` for one node -- everything the node's References contribute, resolved together so the three call sites (gated @@ -1849,7 +1881,18 @@ async def _node_run_context( its ambient view narrowed to the slots just resolved for those connectors; this also ensures a worker sharing a cell workspace with sibling workers sees its own HEAD rather than everybody's.""" - dataset = await _resolve_node_dataset( + row_mode = row_source is not None or row_input_context is not None or row_inputs is not None + if row_source is not None: + prepared = prepare_agent_row_inputs( + run_id=protocol_run_id, + agent_node_id=node_id, + source=row_source, + bindings=row_bindings or [], + ) + workspace_id = prepared["workspace_id"] + row_inputs = prepared["row_inputs"] + row_input_context = row_inputs[0] if row_inputs else None + dataset = NodeDataset() if row_mode else await _resolve_node_dataset( graph, node_id, workspace_id, owner_id, slot_prefix=slot_prefix, stage_plan=stage_plan ) ambient_meta = _ambient_meta_for( @@ -1861,6 +1904,9 @@ async def _node_run_context( # just resolved. Nodes with none keep the whole-cell view so downstream # stages can consume the workspace produced upstream. slots=tuple(slot for _name, slot in dataset.seeded), + row_mode=row_mode, + row_inputs=row_inputs, + dataset_row=dict(row_input_context) if row_input_context is not None else None, ) if dataset.data_path: ambient_meta.pop("data_slots", None) @@ -1887,7 +1933,9 @@ def _resolve_model_config(graph: dict[str, Any], node_id: str) -> dict[str, Any] return (source.get("data") or {}).get("config") or {} -def _resolve_dataset_configs(graph: dict[str, Any], node_id: str) -> list[dict[str, Any]]: +def _resolve_dataset_configs( + graph: dict[str, Any], node_id: str, *, row_input_context: Mapping[str, Any] | None = None +) -> list[dict[str, Any]]: """Every Dataset node wired into this agent, as a list of ``{"dataset_id": ..., "dataset_name": ...}`` configs in canvas wiring order -- ``[]`` when none is connected. @@ -1941,6 +1989,9 @@ def _resolve_dataset_configs(graph: dict[str, Any], node_id: str) -> list[dict[s continue seen.add(key) configs.append(config) + if row_input_context is not None: + driver_id = str(row_input_context.get("dataset_id")) + configs = [config for config in configs if str(config.get("dataset_id")) != driver_id] return configs @@ -2596,7 +2647,9 @@ def _output_shape_block(contract: dict[str, Any] | None) -> str: ) -def _resolve_dataset_tool_config(graph: dict[str, Any], node_id: str, *, unsplit_dataset: str = "") -> dict[str, Any]: +def _resolve_dataset_tool_config( + graph: dict[str, Any], node_id: str, *, unsplit_dataset: str = "", row_mode: bool = False +) -> dict[str, Any]: """The Dataset connector's contribution to the tool allow-list: ASAREE's own ``asaree-workspace`` server, shaped like ``_resolve_tool_config``'s output so it merges with the rest. @@ -2625,14 +2678,19 @@ def _resolve_dataset_tool_config(graph: dict[str, Any], node_id: str, *, unsplit stay granted alongside them -- ``workspace_status`` reporting "no workspace here" is a better answer than a missing tool. """ - if not _resolve_dataset_configs(graph, node_id): + if not _resolve_dataset_configs(graph, node_id) and not row_mode: return {"server_names": [], "tool_names": []} server_names = [WORKSPACE_SERVER_NAME] - tool_names = [f"{WORKSPACE_SERVER_NAME}.{name}" for name in WORKSPACE_AGENT_TOOLS] - if unsplit_dataset: + if row_mode: + tool_names = [f"{WORKSPACE_SERVER_NAME}.open_workspace", f"{WORKSPACE_SERVER_NAME}.workspace_status"] + server_names.append(SCIKIT_LEARN_SERVER_NAME) + tool_names.append(f"{SCIKIT_LEARN_SERVER_NAME}.describe_dataset") + else: + tool_names = [f"{WORKSPACE_SERVER_NAME}.{name}" for name in WORKSPACE_AGENT_TOOLS] + if unsplit_dataset and not row_mode: server_names.append(SCIKIT_LEARN_SERVER_NAME) tool_names.extend(f"{SCIKIT_LEARN_SERVER_NAME}.{name}" for name in UNSPLIT_DATASET_AGENT_TOOLS) - if any(config.get("dictionary_available") for config in _resolve_dataset_configs(graph, node_id)): + if not row_mode and any(config.get("dictionary_available") for config in _resolve_dataset_configs(graph, node_id)): server_names.append(EDA_SERVER_NAME) tool_names.extend(f"{EDA_SERVER_NAME}.{name}" for name in DATASET_DICTIONARY_AGENT_TOOLS) return {"server_names": server_names, "tool_names": tool_names} @@ -3283,7 +3341,9 @@ def validate_edge_handoffs(*, graph: dict[str, Any]) -> None: raise ProtocolValidationError(f"{where} narrows {name!r} to item keys, but it is not a list field.") -def _resource_catalog(graph: dict[str, Any], node_id: str) -> str: +def _resource_catalog( + graph: dict[str, Any], node_id: str, *, row_input_context: Mapping[str, Any] | None = None +) -> str: """Compact semantic metadata for references wired into one agent. Paths, ids, and source bodies stay out of the prompt. This block gives the @@ -3292,7 +3352,7 @@ def _resource_catalog(graph: dict[str, Any], node_id: str) -> str: """ sections: list[str] = [] - datasets = _resolve_dataset_configs(graph, node_id) + datasets = _resolve_dataset_configs(graph, node_id, row_input_context=row_input_context) if datasets: lines = [] for config in datasets: @@ -3359,6 +3419,7 @@ def _build_user_input( script_bound: bool = False, seeded_datasets: tuple[tuple[str, str], ...] = (), unsplit_dataset: str = "", + row_input_context: Mapping[str, Any] | None = None, upstream_ids: list[str] | None = None, unresolved_out: list[str] | None = None, ) -> str: @@ -3442,11 +3503,33 @@ def _build_user_input( if upstream_context: parts.append(upstream_context) - resource_catalog = _resource_catalog(graph, node["id"]) + resource_catalog = _resource_catalog(graph, node["id"], row_input_context=row_input_context) if resource_catalog: parts.append(resource_catalog) - dataset_configs = _resolve_dataset_configs(graph, node["id"]) + dataset_configs = _resolve_dataset_configs(graph, node["id"], row_input_context=row_input_context) + row_driver_ids = set() + for binding in (resolve_dataset_row_plan(graph) or {}).get("bindings", []): + if binding.get("agent_node_id") == node["id"]: + source = next( + ( + candidate + for candidate in graph.get("nodes", []) + if candidate.get("id") == binding.get("dataset_node_id") + ), + {}, + ) + config = (source.get("data") or {}).get("config") or {} + row_driver_ids.add(str(config.get("dataset_id"))) + if row_input_context is not None and row_driver_ids and str(row_input_context.get("dataset_id")) in row_driver_ids: + row_payload = { + "dataset_id": str(row_input_context["dataset_id"]), + "raw_sha256": str(row_input_context["raw_sha256"]), + "row_index": int(row_input_context["row_index"]), + "columns": list(row_input_context["columns"]), + "values": {str(key): str(value) for key, value in row_input_context["values"].items()}, + } + parts.append("Dataset row input:\n" + json.dumps(row_payload, ensure_ascii=False, separators=(",", ":"))) if dataset_configs and experiment_id is not None and effective_cell_label is not None: dataset_names = [str(c["dataset_name"]) for c in dataset_configs] if len(seeded_datasets) == 1: @@ -3625,7 +3708,10 @@ def _build_system_prompt( _PREVIEW_CELL_LABEL = "preview" -async def _preview_node_dataset(graph: dict[str, Any], node_id: str, owner_id: uuid.UUID) -> NodeDataset: +async def _preview_node_dataset( + graph: dict[str, Any], node_id: str, owner_id: uuid.UUID, + *, row_input_context: Mapping[str, Any] | None = None, +) -> NodeDataset: """:func:`_resolve_node_dataset`'s answer, without doing any of the work. That function seeds the cell's workspace as a side effect, which a preview @@ -3639,7 +3725,8 @@ async def _preview_node_dataset(graph: dict[str, Any], node_id: str, owner_id: u canvas implies, which is what a workspace created fresh for the next cell will actually use. """ - names = [str(c["dataset_name"]) for c in _resolve_dataset_configs(graph, node_id) if c.get("dataset_name")] + configs = _resolve_dataset_configs(graph, node_id, row_input_context=row_input_context) + names = [str(c["dataset_name"]) for c in configs if c.get("dataset_name")] if not names: return NodeDataset() solo = len(names) == 1 @@ -3649,7 +3736,7 @@ async def _preview_node_dataset(graph: dict[str, Any], node_id: str, owner_id: u reg = await fetch_owned_registration(name, owner_id) if reg is None: continue - for config in _resolve_dataset_configs(graph, node_id): + for config in configs: if str(config.get("dataset_name") or "") == name: config.update( description=reg.get("description"), @@ -3700,6 +3787,8 @@ async def preview_node_prompt( *, owner_id: uuid.UUID, experiment_id: uuid.UUID | None = None, + row_index: int | None = None, + dataset_row_out: list[dict[str, Any]] | None = None, ) -> str: """The exact prompt this agent would be given, assembled from the draft canvas. @@ -3716,6 +3805,10 @@ async def preview_node_prompt( Creates nothing: no ``ProtocolRun``, no agent run, no workspace (see :func:`_preview_node_dataset`). + A row input is verified against the owned draft driver and projected only + to this Agent's columns. ``dataset_row_out`` collects that snapshot while + preserving the string return used by existing preview callers. + Raises :class:`ProtocolValidationError` for a node that is not an agent -- only an agent is given a prompt, and previewing a connector would be inventing one. @@ -3727,12 +3820,40 @@ async def preview_node_prompt( if node.get("type") not in ("agent", "sub_agent"): raise ProtocolValidationError(f"{_node_display_name(node)} is not an agent, so it is never given a prompt.") + row_plan = resolve_dataset_row_plan(graph) + row_binding = next( + (binding for binding in (row_plan or {}).get("bindings", []) if binding["agent_node_id"] == node_id), + None, + ) + if row_index is not None and row_binding is None: + raise ProtocolValidationError("no_row_driver") + dataset_row = None + if row_binding is not None: + async with get_session() as db: + registration = (await db.execute(select( + RegisteredDataset.id, RegisteredDataset.raw_path, RegisteredDataset.raw_sha256, + ).where( + RegisteredDataset.id == uuid.UUID(row_plan["driver_dataset_id"]), + RegisteredDataset.owner_id == owner_id, + ))).one_or_none() + if registration is None: + raise ProtocolValidationError("Registered row source is unavailable.") + source = read_row_source( + dataset_id=str(registration.id), raw_path=registration.raw_path, + raw_sha256=registration.raw_sha256, + ) + dataset_row = project_row( + source, row_index=0 if row_index is None else row_index, columns=row_binding["columns"], + ) + if dataset_row_out is not None: + dataset_row_out.append(dataset_row) + node_runs = { upstream_id: _preview_node_run(graph, nodes[upstream_id], upstream_id) for upstream_id in referenceable_node_ids(graph, node_id) if upstream_id in nodes } - dataset = await _preview_node_dataset(graph, node_id, owner_id) + dataset = await _preview_node_dataset(graph, node_id, owner_id, row_input_context=dataset_row) return _build_user_input( node, graph, @@ -3745,6 +3866,7 @@ async def preview_node_prompt( script_bound=experiment_id is not None, seeded_datasets=dataset.seeded, unsplit_dataset=dataset.unsplit_name, + row_input_context=dataset_row, ) @@ -3986,6 +4108,7 @@ async def _run_agent_node( available_agents: list[dict[str, Any]] | None = None, agent_messenger: Any = None, unsplit_dataset: str = "", + row_input_context: Mapping[str, Any] | None = None, ) -> tuple[str | None, str | None, uuid.UUID | None, dict[str, Any] | None]: """Create-or-sync the real agent and run it to completion. Returns ``(output_text, error, run_id, extraction)`` -- exactly one of @@ -4021,7 +4144,8 @@ async def _run_agent_node( # config.name directly risks two unrelated nodes silently overwriting # each other's agent definition on every run. config.name is folded # into the description instead, purely as a human label. - agent_name = f"protocol-{protocol_id}-{node['id']}" + row_attempt = bool(ambient_meta and ambient_meta.get("dataset_mode") == "per_row") + agent_name = f"protocol-{protocol_id}-{node['id']}" + (f"-row-{protocol_run_id}" if row_attempt else "") # Model/tool/execution-pattern are no longer fields on the agent's own # config -- resolved from its required Model connector, its (optional, # repeatable) Tool connectors, and its optional Architectural Pattern @@ -4042,9 +4166,29 @@ async def _run_agent_node( tool_config = _merge_tool_configs( _resolve_tool_config(graph, node["id"]), _resolve_knowledge_config(graph, node["id"]), - _resolve_dataset_tool_config(graph, node["id"], unsplit_dataset=unsplit_dataset), + _resolve_dataset_tool_config( + graph, node["id"], unsplit_dataset=unsplit_dataset, + row_mode=bool(ambient_meta and ambient_meta.get("dataset_mode") == "per_row"), + ), _resolve_script_tool_config(graph, node["id"]), ) + if ambient_meta and ambient_meta.get("dataset_mode") == "per_row": + allowed = { + f"{WORKSPACE_SERVER_NAME}.open_workspace", f"{WORKSPACE_SERVER_NAME}.workspace_status", + f"{SCIKIT_LEARN_SERVER_NAME}.describe_dataset", + } + unsafe_managed_servers = {WORKSPACE_SERVER_NAME, EDA_SERVER_NAME, SCIKIT_LEARN_SERVER_NAME} + tool_config["tool_names"] = [ + name for name in tool_config.get("tool_names", []) + if name not in {f"{server}.{tool}" for server in unsafe_managed_servers for tool in + (WORKSPACE_AGENT_TOOLS + DATASET_DICTIONARY_AGENT_TOOLS + UNSPLIT_DATASET_AGENT_TOOLS)} + or name in allowed + ] + tool_config["server_names"] = [ + server for server in tool_config.get("server_names", []) + if server not in unsafe_managed_servers + or any(name.startswith(server + ".") for name in tool_config["tool_names"]) + ] pattern_config_data = _resolve_pattern_config(graph, node["id"]) pattern_config = PatternConfig( execution_pattern=pattern_config_data.get("execution_pattern"), @@ -4102,6 +4246,14 @@ async def _run_agent_node( script_workspace_id=_script_workspace_id(workspace_id, protocol_run_id, str(node["id"])), ) ) + if row_input_context is not None: + row_payload = { + key: row_input_context[key] + for key in ("dataset_id", "raw_sha256", "row_index", "columns", "values") + } + user_input = user_input + "\n\nDataset row input:\n" + json.dumps( + row_payload, ensure_ascii=False, separators=(",", ":") + ) run = await create_run( agent_id=agent.id, user_input=user_input, @@ -4151,6 +4303,10 @@ async def _run_agent_node( # that completed, and they are usually seen together: hitting the ceiling # is the single most common reason the parser has nothing to read. node_fields = {**(_extraction_fields(envelope) or {}), **(_truncation_fields(finished) or {})} + # Handoffs retain their empty-string fallback, but measurement must + # distinguish an absent provider response from a reported empty string. + if finished.output is None: + node_fields["final_output_available"] = False return output_text, None, run.id, node_fields or None @@ -4162,6 +4318,7 @@ async def _run_critic( owner_id: uuid.UUID, worker_output: str, graph: dict[str, Any], + row_attempt: bool = False, ) -> tuple[dict[str, Any] | None, str | None, str | None]: """Create-or-sync the gate's own critic agent and run it once. Returns ``(verdict, error, critic_run_id)`` -- exactly one of verdict/error is @@ -4174,7 +4331,7 @@ async def _run_critic( whatever (if anything) is in the node's own config. Model is resolved from its required Model connector, same as an agent node.""" config = gate["data"]["config"] - agent_name = f"protocol-{protocol_id}-{gate['id']}" + agent_name = f"protocol-{protocol_id}-{gate['id']}" + (f"-row-{protocol_run_id}" if row_attempt else "") model_config_data = {k: v for k, v in _resolve_model_config(graph, gate["id"]).items() if v is not None} model_config = ModelConfig(**model_config_data) pattern_config = PatternConfig(execution_pattern="single_agent_baseline").model_dump() @@ -4281,6 +4438,9 @@ async def _run_gated_worker( experiment_id: uuid.UUID | None = None, effective_cell_label: str | None = None, stage_plan: Any = None, + row_source: Any = None, + row_plan: dict[str, Any] | None = None, + dataset_row: dict[str, Any] | None = None, ) -> tuple[dict[str, Any], dict[str, Any]]: """Generalizes the notebook's ``run_stage`` revision loop (cell 19): run worker -> if the gate is enabled, run critic on its output -> on @@ -4305,6 +4465,8 @@ async def _run_gated_worker( owner_id, protocol_run_id=protocol_run_id, stage_plan=stage_plan, + row_source=row_source, + row_bindings=(row_plan or {}).get("bindings"), ) base_instruction = _build_user_input( worker, @@ -4339,6 +4501,7 @@ async def _run_gated_worker( workspace_id=workspace_id, ambient_meta=worker_ambient, unsplit_dataset=worker_dataset.unsplit_name, + row_input_context=(worker_ambient.get("row_inputs") or [None])[0], ) run_id_str = str(run_id) if run_id else None if error == _AGENT_CANCELLED: @@ -4404,6 +4567,7 @@ async def _run_gated_worker( owner_id=owner_id, worker_output=complete_output, graph=graph, + row_attempt=row_source is not None, ) if verdict_error == _AGENT_CANCELLED: # The worker's own output is real and already complete -- only @@ -4461,6 +4625,7 @@ async def plan_cell_runs( protocol_revision_id: uuid.UUID | None = None, replicate_labels: set[str] | None = None, rerun_replicate_labels: set[str] | None = None, + retry_row_result_ids: Sequence[uuid.UUID] | None = None, ) -> tuple[list[ProtocolRun], int]: """ "Run all cells": creates one pending :class:`ProtocolRun` per not-yet-completed replicate row under *experiment_id*, each carrying its @@ -4484,10 +4649,54 @@ async def plan_cell_runs( ``create_protocol_run_endpoint`` already uses for a plain run.""" if experiment_id is None: raise ProtocolValidationError("This protocol has no linked experiment to run replicates for.") + # Serialize all production row planning against publication and design + # regeneration. Keep this global lock order in retry paths too. + row_mode_requested = resolve_dataset_row_plan(graph) is not None + if retry_row_result_ids is not None and ( + not retry_row_result_ids + or len(set(retry_row_result_ids)) != len(retry_row_result_ids) + or replicate_labels is not None + or rerun_replicate_labels + or not row_mode_requested + ): + raise ProtocolValidationError("invalid_retry_selection") + protocol = None + current_design = None + if row_mode_requested: + protocol = ( + await db.execute( + select(Protocol) + .where(Protocol.id == protocol_id) + .with_for_update() + .execution_options(populate_existing=True) + ) + ).scalar_one_or_none() + experiment = ( + await db.execute( + select(ResearchExperiment) + .where(ResearchExperiment.id == experiment_id) + .with_for_update() + .execution_options(populate_existing=True) + ) + ).scalar_one_or_none() + current_design = ( + await db.execute( + select(ExperimentDesignRevision) + .where( + ExperimentDesignRevision.experiment_id == experiment_id, + ExperimentDesignRevision.superseded_at.is_(None), + ) + .with_for_update() + .execution_options(populate_existing=True) + ) + ).scalar_one_or_none() + else: + experiment = await get_experiment(db, experiment_id) + if experiment is None or (row_mode_requested and (protocol is None or protocol.experiment_id != experiment_id)): + raise ProtocolValidationError("This protocol has no linked experiment to run replicates for.") # Strategy first, so its own message wins over a pipeline requirement that # may not apply to this canvas at all -- see is_conversation_strategy. - experiment = await get_experiment(db, experiment_id) - design_spec = experiment.design_spec if experiment is not None else None + design_spec = experiment.design_spec measurement_plan = ( experiment.locked_measurement_plan if experiment is not None and experiment.locked_at is not None @@ -4523,7 +4732,7 @@ async def plan_cell_runs( impact = await get_design_impact(db, experiment_id=experiment_id, design_spec=design_spec) if impact.regeneration_required: raise ProtocolValidationError( - "Design changed — review and regenerate before running all cells. " + "design_changed: review and regenerate before running all cells. " f"Current {impact.current_cell_count} cells/{impact.current_replicate_count} replicates, " f"proposed {impact.proposed_cell_count} cells/{impact.proposed_replicate_count} replicates." ) @@ -4531,8 +4740,17 @@ async def plan_cell_runs( # A row batch is pinned to an immutable published revision. Validate the # complete ownership/scope relationship and the caller-supplied snapshot # before the shared current-design parents can create any row slots. - revision = await db.get(ProtocolRevision, protocol_revision_id) if protocol_revision_id else None - protocol = await db.get(Protocol, protocol_id) + revision = ( + await db.get(ProtocolRevision, protocol_revision_id, populate_existing=True) + if protocol_revision_id + else None + ) + if ( + revision is not None + and resolve_dataset_row_plan(revision.graph, design_spec) is not None + and (protocol is None or protocol.published_revision_id != revision.id) + ): + raise ProtocolValidationError("revision_changed") if ( revision is not None and protocol is not None @@ -4549,6 +4767,42 @@ async def plan_cell_runs( if original_row_plan is not None: if rerun_replicate_labels: raise ProtocolValidationError("Row executions cannot be rerun through the cell batch; use row retry.") + retry_targets = None + if retry_row_result_ids is not None: + if current_design is None: + raise ProtocolValidationError("invalid_retry_target") + parents = await list_replicates(db, experiment_id=experiment_id) + parents_by_id = {parent.id: parent for parent in parents} + retry_targets = [] + for target_id in retry_row_result_ids: + slot = await get_row_result( + db, + experiment_id=experiment_id, + row_result_id=target_id, + design_revision_id=current_design.id, + protocol_revision_id=revision.id, + ) + parent = parents_by_id.get(slot.replicate_result_id) if slot is not None else None + latest = ( + await db.get(ProtocolRun, slot.run_id, populate_existing=True) + if slot is not None and slot.run_id is not None + else None + ) + if ( + slot is None + or parent is None + or slot.protocol_revision_id != revision.id + or latest is None + or latest.row_result_id != slot.id + or latest.protocol_id != protocol_id + or latest.owner_id != owner_id + or latest.replicate_result_id != parent.id + or latest.design_revision_id != current_design.id + or latest.protocol_revision_id != revision.id + or latest.status not in {"failed", "cancelled"} + ): + raise ProtocolValidationError("invalid_retry_target") + retry_targets.append((slot, latest, parent)) registration = ( await db.execute( select(RegisteredDataset).where( @@ -4570,7 +4824,77 @@ async def plan_cell_runs( ) except DatasetRowCsvError as exc: raise ProtocolValidationError(str(exc)) from exc + # Validate the immutable published plan against the verified source + # header even when the caller selected no parents (and thus would + # otherwise produce no candidates to project). + try: + source_columns = { + column + for binding in original_row_plan["bindings"] + for column in binding["columns"] + } + if not source_columns or source_columns - set(source.columns): + raise DatasetRowCsvError( + "invalid_columns", "row input columns must exist in the registered CSV header" + ) + except DatasetRowCsvError as exc: + raise ProtocolValidationError(str(exc)) from exc + if retry_row_result_ids is not None: + planned_retries: list[tuple[uuid.UUID, uuid.UUID, str, dict, dict, uuid.UUID]] = [] + for slot, latest, parent in retry_targets or []: + if slot.dataset_id != registration.id: + raise ProtocolValidationError("invalid_retry_target") + try: + pinned_source = read_row_source( + dataset_id=str(slot.dataset_id), + raw_path=registration.raw_path, + raw_sha256=registration.raw_sha256, + expected_sha256=slot.raw_sha256, + ) + patched_graph = apply_factor_bindings(revision.graph, parent.factor_values or {}) + patched_plan = resolve_dataset_row_plan(patched_graph, design_spec) + if patched_plan is None or patched_plan["driver_dataset_id"] != str(slot.dataset_id): + raise DatasetRowInputError("driver_mismatch", "factor substitution changed the row driver") + columns = list(dict.fromkeys( + column for binding in patched_plan["bindings"] for column in binding["columns"] + )) + view = project_row(pinned_source, row_index=slot.row_index, columns=columns) + except DatasetRowCsvError as exc: + raise ProtocolValidationError(str(exc)) from exc + except DatasetRowInputError as exc: + raise ProtocolValidationError(str(exc)) from exc + planned_retries.append( + (slot.id, latest.id, parent.replicate_label, parent.factor_values or {}, view, parent.id) + ) + + created_retries: list[ProtocolRun] = [] + for slot_id, expected_run_id, label, factor_values, view, parent_id in sorted( + planned_retries, key=lambda item: item[0].int + ): + run = await claim_row_attempt( + db, + row_result_id=slot_id, + expected_run_id=expected_run_id, + create_kwargs={ + "protocol_id": protocol_id, + "owner_id": owner_id, + "replicate_label": label, + "factor_values": factor_values, + "replicate_result_id": parent_id, + "design_revision_id": current_design.id, + "protocol_revision_id": revision.id, + "dataset_row": view, + }, + ) + if run is None: + raise ProtocolValidationError("invalid_retry_target") + created_retries.append(run) + retry_order = {slot_id: index for index, slot_id in enumerate(retry_row_result_ids)} + created_retries.sort(key=lambda run: retry_order[run.row_result_id]) + return created_retries, 0 parents = await list_replicates(db, experiment_id=experiment_id) + if current_design is None or any(parent.design_revision_id != current_design.id for parent in parents): + raise ProtocolValidationError("design_changed") labels = {parent.replicate_label for parent in parents} requested_labels = labels if replicate_labels is None else replicate_labels unknown = requested_labels - labels @@ -4613,6 +4937,7 @@ async def plan_cell_runs( created: list[ProtocolRun] = [] skipped = 0 + claims = [] for candidate, view in planned: row_result = await ensure_row_result( db, @@ -4624,29 +4949,29 @@ async def plan_cell_runs( raw_sha256=source.raw_sha256, row_index=candidate["dataset_row"]["row_index"], ) - attempts = await list_row_attempts( + claims.append((row_result.id, candidate, view)) + for row_result_id, candidate, view in sorted(claims, key=lambda item: item[0].int): + run = await claim_row_attempt( db, - experiment_id=experiment_id, - row_result_id=row_result.id, - design_revision_id=uuid.UUID(candidate["design_revision_id"]), - protocol_revision_id=revision.id, + row_result_id=row_result_id, + expected_run_id=None, + create_kwargs={ + "protocol_id": protocol_id, + "owner_id": owner_id, + "replicate_label": candidate["replicate_label"], + "factor_values": candidate["factor_values"], + "replicate_result_id": uuid.UUID(candidate["replicate_result_id"]), + "design_revision_id": uuid.UUID(candidate["design_revision_id"]), + "protocol_revision_id": revision.id, + "dataset_row": view, + }, ) - if attempts: + if run is None: skipped += 1 - continue - run = await create_protocol_run( - db, - protocol_id=protocol_id, - owner_id=owner_id, - replicate_label=candidate["replicate_label"], - factor_values=candidate["factor_values"], - replicate_result_id=uuid.UUID(candidate["replicate_result_id"]), - design_revision_id=uuid.UUID(candidate["design_revision_id"]), - protocol_revision_id=revision.id, - row_result_id=row_result.id, - dataset_row=view, - ) - created.append(run) + else: + created.append(run) + candidate_order = {row_id: index for index, (row_id, _, _) in enumerate(claims)} + created.sort(key=lambda run: candidate_order[run.row_result_id]) return created, skipped elif resolve_dataset_row_plan(revision.graph if revision is not None else graph, design_spec) is not None: raise ProtocolValidationError("Published protocol revision does not match the requested protocol and graph.") @@ -4726,7 +5051,9 @@ async def plan_single_replicate_run( experiment_id: uuid.UUID | None, owner_id: uuid.UUID, graph: dict[str, Any], - replicate_label: str, + replicate_label: str | None, + row_index: int | None = None, + snapshot_only: bool = False, protocol_revision_id: uuid.UUID | None = None, ) -> ProtocolRun: """Run one already-generated replicate for real, by name -- the single-run @@ -4780,9 +5107,71 @@ async def plan_single_replicate_run( if impact.regeneration_required: raise ProtocolValidationError("Design changed — review and regenerate before running a replicate.") - replicate = await get_replicate(db, experiment_id=experiment_id, replicate_label=replicate_label) + row_plan = resolve_dataset_row_plan(graph, design_spec) + if replicate_label is not None: + replicate = await get_replicate(db, experiment_id=experiment_id, replicate_label=replicate_label) + elif row_plan is not None and snapshot_only: + current_replicates = await list_replicates(db, experiment_id=experiment_id) + replicate = current_replicates[0] if current_replicates else None + else: + replicate = None if replicate is None: - raise ProtocolValidationError(f"No such replicate: {replicate_label!r}") + if replicate_label is not None: + raise ProtocolValidationError(f"No such replicate: {replicate_label!r}") + raise ProtocolValidationError("No current design replicate is available for row execution.") + if row_plan is not None: + if protocol_revision_id is None: + raise ProtocolValidationError("A published protocol revision is required for row execution.") + if isinstance(row_index, bool) or (row_index is not None and (not isinstance(row_index, int) or row_index < 0)): + raise ProtocolValidationError("row_index must be a nonnegative integer.") + protocol = await db.get(Protocol, protocol_id, populate_existing=True) + revision = await db.get(ProtocolRevision, protocol_revision_id, populate_existing=True) + from asaree.services.design_revisions import get_current_revision + + current = await get_current_revision(db, experiment_id) + if ( + protocol is None or protocol.owner_id != owner_id or protocol.experiment_id != experiment_id + or protocol.published_revision_id != protocol_revision_id or revision is None + or revision.protocol_id != protocol_id or revision.graph != graph + or current is None or replicate.design_revision_id != current.id + ): + raise ProtocolValidationError("revision_changed") + registration = (await db.execute(select(RegisteredDataset).where( + RegisteredDataset.id == uuid.UUID(row_plan["driver_dataset_id"]), + RegisteredDataset.owner_id == owner_id, + ))).scalar_one_or_none() + if registration is None or experiment.owner_id != owner_id: + raise ProtocolValidationError("Registered row source is unavailable.") + try: + source = read_row_source(dataset_id=str(registration.id), raw_path=registration.raw_path, + raw_sha256=registration.raw_sha256) + columns = list(dict.fromkeys(column for binding in row_plan["bindings"] for column in binding["columns"])) + if not columns or set(columns) - set(source.columns): + raise DatasetRowCsvError("invalid_columns", "row input columns must exist in the registered CSV header") + selected_index = 0 if row_index is None else row_index + view = project_row(source, row_index=selected_index, columns=columns) + except (DatasetRowCsvError, IndexError) as exc: + raise ProtocolValidationError(str(exc)) from exc + row_result = await ensure_row_result( + db, experiment_id=experiment_id, design_revision_id=replicate.design_revision_id, + protocol_revision_id=protocol_revision_id, replicate_result_id=replicate.id, + dataset_id=registration.id, raw_sha256=source.raw_sha256, row_index=selected_index, + ) + if row_result.run_id is not None: + raise ProtocolValidationError("row_already_attempted") + run = await claim_row_attempt( + db, row_result_id=row_result.id, expected_run_id=None, + create_kwargs={"protocol_id": protocol_id, "owner_id": owner_id, + "replicate_label": replicate.replicate_label, + "factor_values": replicate.factor_values or {}, + "replicate_result_id": replicate.id, + "design_revision_id": replicate.design_revision_id, + "protocol_revision_id": protocol_revision_id, "dataset_row": view, + "snapshot_only_row": snapshot_only}, + ) + if run is None: + raise ProtocolValidationError("row_already_attempted") + return run return await create_protocol_run( db, protocol_id=protocol_id, @@ -4875,6 +5264,9 @@ async def _run_single_node( graph: dict[str, Any], node_id: str, experiment_id: uuid.UUID | None = None, + row_source: Any = None, + row_plan: dict[str, Any] | None = None, + dataset_row: dict[str, Any] | None = None, ) -> None: """The canvas's per-node Play run: one Agent node, no upstream, no gated pair, no factor substitution, no coordination-strategy check -- none of @@ -4896,6 +5288,8 @@ async def _run_single_node( # synthetic per-run label (see _effective_cell_label). effective_cell_label = _effective_cell_label(None, protocol_run_id) workspace_id = _compute_workspace_id(experiment_id, None, protocol_run_id) + if dataset_row is not None: + workspace_id = row_attempt_workspace_id(protocol_run_id) async with get_session() as db: experiment = await get_experiment(db, experiment_id) if experiment_id else None single_design_spec = experiment.design_spec if experiment is not None else None @@ -4906,6 +5300,8 @@ async def _run_single_node( owner_id, protocol_run_id=protocol_run_id, stage_plan=stage_plan_spec(single_design_spec, graph=graph), + row_source=row_source, + row_bindings=(row_plan or {}).get("bindings"), ) user_input = _build_user_input( node, @@ -4916,6 +5312,7 @@ async def _run_single_node( script_bound="script_paths" in ambient_meta, seeded_datasets=node_dataset.seeded, unsplit_dataset=node_dataset.unsplit_name, + row_input_context=(ambient_meta.get("row_inputs") or [None])[0], ) output_text, error, run_id, extraction = await _run_agent_node( node, @@ -4927,6 +5324,7 @@ async def _run_single_node( workspace_id=workspace_id, ambient_meta=ambient_meta, unsplit_dataset=node_dataset.unsplit_name, + row_input_context=(ambient_meta.get("row_inputs") or [None])[0], ) node_run: dict[str, Any] = { "status": "failed" if error else "completed", @@ -4941,12 +5339,46 @@ async def _run_single_node( await set_status(db, protocol_run_id, status="failed", error=error) else: await set_status(db, protocol_run_id, status="finalizing") - if experiment_id is not None: - await finalize_attempt_measurement(db, protocol_run_id) - await set_status(db, protocol_run_id, status="completed") + if error: + return + # Finalization claims and commits its own transaction. Keep the terminal + # transition in a fresh session so node Play cannot be left finalizing + # after its reported observations have been persisted. + if experiment_id is not None: + async with get_session() as db: + await finalize_attempt_measurement(db, protocol_run_id) + async with get_session() as db: + await set_status(db, protocol_run_id, status="completed") async def run_protocol(protocol_run_id: uuid.UUID) -> None: + row_mode = False + row_context_mode = False + dataset_row = None + async with get_session() as db: + run = await db.execute( + select(ProtocolRun).where(ProtocolRun.id == protocol_run_id).with_for_update() + ) + run = run.scalar_one_or_none() + if run is None: + return + row_mode = run.row_result_id is not None and run.replicate_result_id is not None + attempt_result = run.attempt_result if isinstance(run.attempt_result, dict) else {} + row_context_mode = ( + run.dataset_row is not None + or run.row_result_id is not None + or isinstance(attempt_result.get("row_provenance"), dict) + ) + if row_context_mode: + if run.status != "pending": + return + if run.cancel_requested_at is not None: + await set_status(db, protocol_run_id, status="cancelled") + return + # Claim before registry hydration, source reads, workspace writes, or + # any provider work. A redelivered job sees running and exits above. + await set_status(db, protocol_run_id, status="running") + dataset_row = dict(run.dataset_row) if run.dataset_row is not None else None # The worker hydrates its MCP registry once, at startup (worker/settings.py), # so any server registered SINCE then -- an OKF bundle the user added # mid-session being the case this exists for -- isn't live in this process @@ -5004,6 +5436,61 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: if pinned is not None and pinned.design_spec is not None: pinned_spec = pinned.design_spec + row_source = None + row_plan = None + if row_context_mode: + try: + if dataset_row is None: + raise DatasetRowCsvError("source_unavailable", "row attempt has no dataset snapshot") + row_plan = resolve_dataset_row_plan(graph, design_spec) + if row_plan is None or row_plan["driver_dataset_id"] != dataset_row["dataset_id"]: + raise DatasetRowCsvError("source_identity_mismatch", "published row driver differs from attempt") + registration = ( + await db.execute( + select(RegisteredDataset).where( + RegisteredDataset.id == uuid.UUID(dataset_row["dataset_id"]), + RegisteredDataset.owner_id == owner_id, + ) + ) + ).scalar_one_or_none() + if registration is None: + raise DatasetRowCsvError("source_unavailable", "registered row source is unavailable") + row_source = read_row_source( + dataset_id=str(registration.id), + raw_path=registration.raw_path, + raw_sha256=registration.raw_sha256, + expected_sha256=dataset_row["raw_sha256"], + ) + projected = project_row( + row_source, row_index=dataset_row["row_index"], columns=dataset_row["columns"] + ) + if projected != dataset_row: + raise DatasetRowCsvError("source_hash_mismatch", "published row snapshot no longer matches source") + bindings = [] + nodes = {str(node.get("id")): node for node in graph.get("nodes") or []} + for binding in row_plan["bindings"]: + config = ((nodes.get(binding["dataset_node_id"]) or {}).get("data") or {}).get("config") or {} + bindings.append({ + **binding, + "dataset_id": row_plan["driver_dataset_id"], + "dataset_name": str(config.get("dataset_name") or registration.name), + "row_index": dataset_row["row_index"], + "target_column": registration.target_column or "", + }) + row_plan["bindings"] = bindings + if row_mode: + from asaree.services.factorial_row_results import project_row_attempt + + await project_row_attempt( + db, + row_result_id=run.row_result_id, + run_id=protocol_run_id, + fields={"workspace_id": row_attempt_workspace_id(protocol_run_id)}, + ) + except (DatasetRowCsvError, DatasetRowInputError, ValueError, TypeError) as exc: + await set_status(db, protocol_run_id, status="failed", error=f"Row source validation failed: {exc}") + return + if target_node_id: await _run_single_node( protocol_run_id, @@ -5012,6 +5499,9 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: graph=graph, node_id=target_node_id, experiment_id=experiment_id, + row_source=row_source, + row_plan=row_plan, + dataset_row=dataset_row, ) return @@ -5038,6 +5528,11 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: effective_cell_label = _effective_cell_label(replicate_label, protocol_run_id) workspace_id = _compute_workspace_id(experiment_id, replicate_label, protocol_run_id) + if row_context_mode: + # Every workspace-backed write in a row attempt (including direct + # tool steps and Script materialization) belongs to this attempt's + # inspectable row slot, never the parent replicate's shared workspace. + workspace_id = row_attempt_workspace_id(protocol_run_id) # Derived from the *pinned* graph, not the live canvas: a canvas edit # mid-run must not change which stages this run's later nodes are staging # through. @@ -5045,7 +5540,12 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: async with get_session() as db: await set_status(db, protocol_run_id, status="running") - if replicate_label and experiment_id and await is_current_replicate_attempt(db, protocol_run_id): + if ( + not row_context_mode + and replicate_label + and experiment_id + and await is_current_replicate_attempt(db, protocol_run_id) + ): # Pre-write, before any node executes: a crash/timeout mid-run # still leaves this cell's provenance recorded (mirrors the # notebook's own pre-scoring upsert_replicate call). workspace_id is @@ -5083,6 +5583,8 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: owner_id, protocol_run_id=protocol_run_id, stage_plan=stage_plan, + row_source=row_source, + row_bindings=(row_plan or {}).get("bindings"), ) node_run, conversation_status = await execute_conversation( protocol_run_id, @@ -5108,6 +5610,9 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: ambient_meta=ambient_meta, stage_plan=stage_plan, unsplit_dataset=entry_dataset.unsplit_name, + dataset_row=dataset_row, + row_source=row_source, + row_plan=row_plan, ) node_runs[entry_agent_id] = node_run cancelled = conversation_status == "cancelled" @@ -5136,6 +5641,8 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: owner_id, protocol_run_id=protocol_run_id, stage_plan=stage_plan, + row_source=row_source, + row_bindings=(row_plan or {}).get("bindings"), ) node_run, supervisor_status = await execute_supervisor_architecture( protocol_run_id, @@ -5158,6 +5665,9 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: experiment_id=experiment_id, effective_cell_label=effective_cell_label, stage_plan=stage_plan, + dataset_row=dataset_row, + row_source=row_source, + row_plan=row_plan, ) node_runs[roles.supervisor] = node_run cancelled = supervisor_status == "cancelled" @@ -5191,6 +5701,9 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: entry_agent_id=pipeline_parents[0], workspace_id=workspace_id, stage_plan=stage_plan, + row_source=row_source, + row_plan=row_plan, + dataset_row=dataset_row, ) for node in order: @@ -5249,6 +5762,9 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: experiment_id=experiment_id, effective_cell_label=effective_cell_label, stage_plan=stage_plan, + row_source=row_source, + row_plan=row_plan, + dataset_row=dataset_row, ) node_runs[node_id] = worker_run node_runs[gate["id"]] = gate_run @@ -5268,7 +5784,12 @@ async def record_step_provenance(step: dict[str, Any], *, step_node_id: str = no # crash mid-call still leaves the approved payload on record. async with get_session() as db: await update_node_run(db, protocol_run_id, step_node_id, {"tool_step": step}) - if replicate_label and experiment_id and await is_current_replicate_attempt(db, protocol_run_id): + if ( + not row_context_mode + and replicate_label + and experiment_id + and await is_current_replicate_attempt(db, protocol_run_id) + ): await upsert_replicate( db, experiment_id=experiment_id, @@ -5326,6 +5847,8 @@ async def record_step_provenance(step: dict[str, Any], *, step_node_id: str = no owner_id, protocol_run_id=protocol_run_id, stage_plan=stage_plan, + row_source=row_source, + row_bindings=(row_plan or {}).get("bindings"), ) user_input = _build_user_input( node, @@ -5371,6 +5894,7 @@ async def record_step_provenance(step: dict[str, Any], *, step_node_id: str = no available_agents=available_sub_agents, agent_messenger=pipeline_messenger, unsplit_dataset=node_dataset.unsplit_name, + row_input_context=(ambient_meta.get("row_inputs") or [None])[0], ) else: output_text, error, run_id, extraction = await _run_agent_node( @@ -5384,6 +5908,7 @@ async def record_step_provenance(step: dict[str, Any], *, step_node_id: str = no workspace_id=workspace_id, ambient_meta=ambient_meta, unsplit_dataset=node_dataset.unsplit_name, + row_input_context=(ambient_meta.get("row_inputs") or [None])[0], ) if error == _AGENT_CANCELLED: @@ -5458,7 +5983,8 @@ async def record_step_provenance(step: dict[str, Any], *, step_node_id: str = no # Metric observations are finalized below exclusively through # the attempt's explicit measurement-plan producer bindings. if ( - replicate_label + not row_context_mode + and replicate_label and experiment_id and result_node_id is not None and node_runs.get(result_node_id, {}).get("status") == "completed" @@ -5494,7 +6020,12 @@ async def record_step_provenance(step: dict[str, Any], *, step_node_id: str = no current = await get_protocol_run(db, protocol_run_id) if current is None or current.status in TERMINAL_PROTOCOL_RUN_STATUSES: return - if experiment_id and (replicate_label or current.is_test_run or current.target_node_id): + if experiment_id and ( + current.is_test_run + or current.row_result_id is not None + or current.dataset_row is not None + or (not row_context_mode and (replicate_label or current.target_node_id)) + ): await finalize_attempt_measurement(db, protocol_run_id) await db.refresh(current) if current.status in TERMINAL_PROTOCOL_RUN_STATUSES: diff --git a/src/asaree/services/protocol_revisions.py b/src/asaree/services/protocol_revisions.py index 22a062ab..b7d5fbce 100644 --- a/src/asaree/services/protocol_revisions.py +++ b/src/asaree/services/protocol_revisions.py @@ -34,6 +34,14 @@ async def publish_protocol( db: AsyncSession, protocol: Protocol, *, owner_id: uuid.UUID | None = None ) -> ProtocolRevision: """Freeze the protocol's current draft as its next production revision.""" + protocol = ( + await db.execute( + select(Protocol) + .where(Protocol.id == protocol.id) + .with_for_update() + .execution_options(populate_existing=True) + ) + ).scalar_one() plan = resolve_dataset_row_plan(protocol.graph) if plan is not None: if owner_id is None: diff --git a/src/asaree/services/protocol_runs.py b/src/asaree/services/protocol_runs.py index 803358e2..b104ac4d 100644 --- a/src/asaree/services/protocol_runs.py +++ b/src/asaree/services/protocol_runs.py @@ -91,6 +91,7 @@ async def create_protocol_run( protocol_revision_id: uuid.UUID | None = None, row_result_id: uuid.UUID | None = None, dataset_row: dict[str, Any] | None = None, + snapshot_only_row: bool = False, is_test_run: bool = False, ) -> ProtocolRun: """``replicate_label``/``factor_values``/``design_revision_id`` are set together @@ -154,9 +155,9 @@ async def create_protocol_run( status="pending", is_test_run=is_test_run, node_runs={}, - replicate_label=replicate_label, - factor_values=factor_values, - replicate_result_id=replicate_result_id, + replicate_label=None if snapshot_only_row else replicate_label, + factor_values=None if snapshot_only_row else factor_values, + replicate_result_id=None if snapshot_only_row else replicate_result_id, row_result_id=row_result_id, dataset_row=dataset_row, target_node_id=target_node_id, @@ -313,7 +314,12 @@ async def _validate_row_attempt( async def create_test_run( - db: AsyncSession, *, protocol_id: uuid.UUID, owner_id: uuid.UUID, protocol_revision_id: uuid.UUID + db: AsyncSession, + *, + protocol_id: uuid.UUID, + owner_id: uuid.UUID, + protocol_revision_id: uuid.UUID, + dataset_row: dict[str, Any] | None = None, ) -> ProtocolRun: """Create the experiment's next canvas validation attempt. @@ -338,6 +344,7 @@ async def create_test_run( protocol_id=protocol_id, owner_id=owner_id, protocol_revision_id=protocol_revision_id, + dataset_row=dataset_row, is_test_run=True, ) experiment.latest_test_run_id = run.id @@ -437,9 +444,18 @@ async def list_stale_protocol_runs( async def set_status( db: AsyncSession, protocol_run_id: uuid.UUID, *, status: str, error: str | None = None ) -> ProtocolRun | None: - run = await get_protocol_run(db, protocol_run_id) + run = ( + await db.execute( + select(ProtocolRun) + .where(ProtocolRun.id == protocol_run_id) + .with_for_update() + .execution_options(populate_existing=True) + ) + ).scalar_one_or_none() if run is None: return None + if run.status in TERMINAL_PROTOCOL_RUN_STATUSES: + return run now = datetime.now(UTC) _apply_status(run, status=status, error=error, now=now) attempt_result = dict(run.attempt_result or {}) @@ -517,10 +533,33 @@ async def update_attempt_result( an earlier attempt's values. Other callers may add independently named result facets without overwriting the stored node timeline. """ - run = await get_protocol_run(db, protocol_run_id) + run = ( + await db.execute( + select(ProtocolRun) + .where(ProtocolRun.id == protocol_run_id) + .with_for_update() + .execution_options(populate_existing=True) + ) + ).scalar_one_or_none() if run is None: return None result = dict(run.attempt_result or {}) + if "measurement" in result and any( + key in fields + for key in ( + "measurement", + "metric_values", + "row_provenance", + "metric_evaluation", + "evaluation_summary", + "evaluation_claimed_at", + "task_completed_at", + "evaluation_started_at", + "evaluation_completed_at", + "evaluation_state", + ) + ): + raise ValueError("immutable_attempt_result") result.update(fields) run.attempt_result = result await db.flush() @@ -535,6 +574,75 @@ def _required_metric_was_measured(evaluation: MeasurementEvaluation, metric_id: ) +@dataclass(frozen=True) +class MeasurementSubject: + """Stable identity and current result projection for one attempt.""" + + subject_id: str + replicate: FactorialReplicateResult | None = None + row: FactorialRowResult | None = None + is_current: bool = False + + +def measurement_subject_id(run: ProtocolRun) -> str: + """The MeasurementEvaluation ``replicate_id`` compatibility field's subject.""" + attempt_result = run.attempt_result if isinstance(run.attempt_result, Mapping) else {} + provenance = attempt_result.get("row_provenance") + row_id = run.row_result_id or (provenance.get("row_result_id") if isinstance(provenance, Mapping) else None) + return str(row_id or run.replicate_result_id or run.id) + + +async def resolve_measurement_subject( + db: AsyncSession, + run: ProtocolRun, + *, + experiment_id: uuid.UUID, +) -> MeasurementSubject: + """Resolve the row slot or whole-dataset replicate projection for an attempt.""" + if run.row_result_id is not None: + if run.replicate_result_id is None or run.replicate_label is None: + raise ValueError("measurement evaluation row is outside the run's design scope") + replicate = await get_replicate( + db, + experiment_id=experiment_id, + replicate_label=run.replicate_label, + revision_id=run.design_revision_id, + ) + row = await db.get(FactorialRowResult, run.row_result_id) + if ( + replicate is None + or replicate.id != run.replicate_result_id + or row is None + or row.replicate_result_id != replicate.id + or row.protocol_revision_id != run.protocol_revision_id + ): + raise ValueError("measurement evaluation row is outside the run's design scope") + return MeasurementSubject( + subject_id=str(row.id), + row=row, + is_current=row.run_id == run.id, + ) + + if run.replicate_result_id is None: + # Preview snapshot-only runs intentionally have no durable result slot. + return MeasurementSubject(subject_id=str(run.id)) + if run.replicate_label is None: + raise ValueError("measurement evaluation run is not scoped to an experiment replicate") + replicate = await get_replicate( + db, + experiment_id=experiment_id, + replicate_label=run.replicate_label, + revision_id=run.design_revision_id, + ) + if replicate is None or replicate.id != run.replicate_result_id: + raise ValueError("measurement evaluation replicate is outside the run's design scope") + return MeasurementSubject( + subject_id=str(replicate.id), + replicate=replicate, + is_current=replicate.run_id == run.id, + ) + + async def record_measurement_evaluation( db: AsyncSession, protocol_run_id: uuid.UUID, @@ -545,20 +653,27 @@ async def record_measurement_evaluation( """Freeze one engine result on its attempt and update only its current projection. A later attempt writes its own ``ProtocolRun.attempt_result`` and may replace - the replicate projection, but it never edits this attempt's document. + the current row or replicate projection, but it never edits this attempt's document. """ - run = await get_protocol_run(db, protocol_run_id) + run = ( + await db.execute( + select(ProtocolRun) + .where(ProtocolRun.id == protocol_run_id) + .with_for_update() + .execution_options(populate_existing=True) + ) + ).scalar_one_or_none() if run is None: return None if evaluation.attempt_id != str(run.id): raise ValueError("measurement evaluation attempt does not match the protocol run") - expected_subject_id = str(run.replicate_result_id or run.id) + expected_subject_id = measurement_subject_id(run) if evaluation.replicate_id != expected_subject_id: raise ValueError("measurement evaluation replicate does not match the protocol run") attempt_result = dict(run.attempt_result or {}) if "measurement" in attempt_result: - raise ValueError("measurement evaluation is immutable once recorded") + raise ValueError("immutable_attempt_result") document = evaluation.to_document() attempt_result["measurement"] = document attempt_result["evaluation_completed_at"] = datetime.now(UTC).isoformat() @@ -577,23 +692,71 @@ async def record_measurement_evaluation( # Canvas Test Runs and per-node Play runs retain the immutable measurement # document on the run itself; neither projects into a factorial replicate. - if run.is_test_run or run.target_node_id is not None: + if run.is_test_run or run.target_node_id is not None or ( + run.dataset_row is not None and run.row_result_id is None + and run.replicate_result_id is None + and not (attempt_result.get("row_provenance") or {}).get("row_result_id") + ): await db.flush() await db.refresh(run) return run protocol = await db.get(Protocol, run.protocol_id) - if protocol is None or protocol.experiment_id is None or run.replicate_label is None: + if protocol is None or protocol.experiment_id is None: + raise ValueError("measurement evaluation run is not scoped to an experiment") + if ( + run.row_result_id is None + and run.replicate_result_id is None + and run.dataset_row is None + ): raise ValueError("measurement evaluation run is not scoped to an experiment replicate") - replicate = await get_replicate( - db, - experiment_id=protocol.experiment_id, - replicate_label=run.replicate_label, - revision_id=run.design_revision_id, + attempt_provenance = attempt_result.get("row_provenance") + is_row_execution = ( + run.row_result_id is not None + or run.dataset_row is not None + or isinstance(attempt_provenance, Mapping) ) - if replicate is None or replicate.id != run.replicate_result_id: - raise ValueError("measurement evaluation replicate is outside the run's design scope") - if replicate.run_id == run.id: + if is_row_execution: + # A deleted or superseded row slot does not invalidate this run's own + # immutable facts. Projection authorization happens atomically below. + if not isinstance(attempt_provenance, Mapping): + raise ValueError("measurement evaluation row is outside the run's design scope") + provenance_row_id = attempt_provenance.get("row_result_id") + if not isinstance(provenance_row_id, str): + raise ValueError("measurement evaluation row is outside the run's design scope") + if run.row_result_id is not None and provenance_row_id != str(run.row_result_id): + raise ValueError("measurement evaluation row is outside the run's design scope") + subject = MeasurementSubject(subject_id=provenance_row_id) + else: + subject = await resolve_measurement_subject(db, run, experiment_id=protocol.experiment_id) + if not is_row_execution and subject.row is None and subject.replicate is None: + await db.flush() + await db.refresh(run) + return run + if is_row_execution and run.row_result_id is None: + # The row FK uses SET NULL when its slot is deleted. Its provenance + # still proves this attempt was row-scoped, so retain its facts without + # ever redirecting projection to the parent replicate. + await db.flush() + await db.refresh(run) + return run + if is_row_execution: + from asaree.services.factorial_row_results import project_row_attempt + + projection: dict[str, Any] = {"artifacts": {"measurement": evaluation.to_document()}} + # Keep observed values on the immutable attempt even when execution + # did not finish, but only a completed, non-truncated row execution is + # eligible for the latest scored projection. + if ( + run.status == "completed" + and measured_values + and required_metric_measured + and node_run_truncation(run.node_runs) is None + ): + projection["metric_values"] = measured_values + await project_row_attempt(db, row_result_id=run.row_result_id, run_id=run.id, fields=projection) + elif subject.replicate is not None and subject.is_current: + replicate = subject.replicate artifacts = dict(replicate.artifacts or {}) artifacts["measurement"] = evaluation.to_document() replicate.artifacts = artifacts diff --git a/src/asaree/services/reported_metrics.py b/src/asaree/services/reported_metrics.py index 19d9f9d6..db4a51e1 100644 --- a/src/asaree/services/reported_metrics.py +++ b/src/asaree/services/reported_metrics.py @@ -167,6 +167,7 @@ async def collect_reported_metrics( and isinstance(node_run, Mapping) and (node_run.get("status") == "completed" or "last_successful_output_text" in node_run) and successful_output is not None + and node_run.get("final_output_available", True) ) # Field projections read the Agent's typed output: its Output Parser # payload, else a JSON object in its final answer. @@ -234,7 +235,7 @@ async def collect_reported_metrics( ) ) return MeasurementEvaluation( - replicate_id=str(run.replicate_result_id or run.id), + replicate_id=str(getattr(run, "row_result_id", None) or run.replicate_result_id or run.id), attempt_id=attempt_id, observations=tuple(observations), artifacts=(), diff --git a/src/asaree/services/runtime_metrics.py b/src/asaree/services/runtime_metrics.py index fb4d5632..c2b3a0d8 100644 --- a/src/asaree/services/runtime_metrics.py +++ b/src/asaree/services/runtime_metrics.py @@ -52,6 +52,7 @@ from asaree.services.protocol_runs import ( TERMINAL_PROTOCOL_RUN_STATUSES, get_cancel_requested_at, + measurement_subject_id, record_measurement_evaluation, ) from asaree.services.reported_metrics import ( @@ -684,7 +685,13 @@ async def finalize_attempt_measurement( if ( run is None or (run.status not in TERMINAL_PROTOCOL_RUN_STATUSES and run.status != "finalizing") - or (run.replicate_result_id is None and not run.is_test_run and run.target_node_id is None) + or ( + run.replicate_result_id is None + and run.row_result_id is None + and run.dataset_row is None + and not run.is_test_run + and run.target_node_id is None + ) or (run.attempt_result or {}).get("measurement") is not None ): return False @@ -716,7 +723,7 @@ async def finalize_attempt_measurement( ) if interrupted_status is not None: attempt = CompletedReplicate( - replicate_id=str(run.replicate_result_id or run.id), + replicate_id=measurement_subject_id(run), attempt_id=str(run.id), ) interrupted = MeasurementEngine(()).complete_declared_metrics( @@ -745,6 +752,13 @@ async def finalize_attempt_measurement( interrupted_result = dict(run.attempt_result or {}) interrupted_result["evaluation_summary"] = {"cost_usd": 0.0} run.attempt_result = interrupted_result + revision = await db.get(ProtocolRevision, run.protocol_revision_id) if run.protocol_revision_id else None + graph = resolve_reported_metric_graph( + revision.graph if revision is not None else {}, + run.factor_values or {}, + ) + reported = await collect_reported_metrics(run, declared, graph) + interrupted = _merge_evaluations(interrupted, reported) await record_measurement_evaluation(db, run.id, interrupted) completed_result = dict(run.attempt_result or {}) completed_result["evaluation_state"] = "completed" @@ -889,7 +903,7 @@ async def finalize_attempt_measurement( evaluation = await engine.evaluate( runtime_plan, snapshot, - CompletedReplicate(replicate_id=str(run.replicate_result_id or run.id), attempt_id=str(run.id)), + CompletedReplicate(replicate_id=measurement_subject_id(run), attempt_id=str(run.id)), validate_plan=True, unavailable_outputs={ binding_id: reasons @@ -908,7 +922,7 @@ async def finalize_attempt_measurement( evaluation = engine.complete_declared_metrics( declared, - CompletedReplicate(replicate_id=str(run.replicate_result_id or run.id), attempt_id=str(run.id)), + CompletedReplicate(replicate_id=measurement_subject_id(run), attempt_id=str(run.id)), evaluation, status="unavailable", error=lambda _metric, binding: ( @@ -931,7 +945,11 @@ async def finalize_attempt_measurement( completed_result = dict(run.attempt_result or {}) completed_result["evaluation_state"] = "completed" run.attempt_result = completed_result - if any(observation.status == "cancelled" for observation in evaluation.observations): + await db.flush() + if ( + any(observation.status == "cancelled" for observation in evaluation.observations) + and run.status not in TERMINAL_PROTOCOL_RUN_STATUSES + ): run.status = "cancelled" run.completed_at = datetime.now(UTC) await db.flush() diff --git a/src/asaree/services/test_run_results.py b/src/asaree/services/test_run_results.py index b16930bf..a1efe57f 100644 --- a/src/asaree/services/test_run_results.py +++ b/src/asaree/services/test_run_results.py @@ -6,6 +6,7 @@ from datetime import datetime from decimal import Decimal from typing import Any, Literal +from uuid import UUID from motoro.runner import list_runs from sqlalchemy.ext.asyncio import AsyncSession @@ -33,6 +34,8 @@ class TestRunResourceSummary: @dataclass(frozen=True) class TestRunResultProjection: + dataset_row: dict[str, Any] | None + protocol_revision_id: UUID | None observations: list[dict[str, Any]] artifacts: list[dict[str, Any]] tested_published_revision: ProtocolRevision | None @@ -143,6 +146,8 @@ async def project_test_run_result( freshness_reasons.append("measurement_plan") return TestRunResultProjection( + dataset_row=run.dataset_row, + protocol_revision_id=run.protocol_revision_id, observations=observations, artifacts=artifacts, tested_published_revision=tested_revision, diff --git a/tests/test_dataset_row_acceptance.py b/tests/test_dataset_row_acceptance.py new file mode 100644 index 00000000..7f2f2caa --- /dev/null +++ b/tests/test_dataset_row_acceptance.py @@ -0,0 +1,433 @@ +"""Integrated five-cell, three-original-row, two-replicate production acceptance.""" + +from __future__ import annotations + +import csv +import hashlib +import json +import os +import uuid +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException +from motoro.models import database as core_database +from motoro.models.run import AgentRun, RunStatus +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from asaree.api import protocols as api +from asaree.api.experiments import export_run_results_csv_endpoint, get_experiment_run_results_endpoint +from asaree.models import database +from asaree.models.dataset import RegisteredDataset +from asaree.models.factorial_cell import FactorialCell +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol import Protocol +from asaree.models.protocol_run import ProtocolRun +from asaree.services import dataset_row_workspaces, run_tools +from asaree.services import protocol_execution as execution +from asaree.services.design_generation import generate_design_cells +from asaree.services.experiments import create_experiment, delete_experiment +from asaree.services.factorial_row_results import project_row_attempt +from asaree.services.protocol_revisions import publish_protocol +from asaree.services.protocols import create_protocol, delete_protocol +from asaree.services.users import get_user_by_email + + +@pytest.mark.asyncio +async def test_five_cells_three_rows_two_replicates_and_one_failed_row_retry(tmp_path, monkeypatch): + product_engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + core_engine = create_async_engine(os.environ["DATABASE_URL"]) + sessions = async_sessionmaker(product_engine, expire_on_commit=False) + core_sessions = async_sessionmaker(core_engine, expire_on_commit=False) + monkeypatch.setattr(database, "_get_session_factory", lambda: sessions) + monkeypatch.setattr(core_database, "_get_session_factory", lambda: core_sessions) + monkeypatch.setattr(execution, "get_session", database.get_session) + monkeypatch.setattr(dataset_row_workspaces, "_workspace_root", lambda: tmp_path / "datasets") + monkeypatch.setattr(execution, "WORKSPACE_ROOT", str(tmp_path / "workspaces")) + monkeypatch.setenv("ASAREE_DATASET_WORKSPACE_DIR", str(tmp_path / "workspaces")) + source_path = tmp_path / "original.csv" + raw = b"question,reference\nQUESTION-1,GOLD-1\nQUESTION-1,GOLD-1\nQUESTION-3,GOLD-3\n" + source_path.write_bytes(raw) + digest = hashlib.sha256(raw).hexdigest() + experiment_id = protocol_id = dataset_id = None + queued = [] + provider_calls = [] + private_paths = set() + failure_run = None + null_run = None + missing_run = None + attempted_failure = False + + async def enqueue(run_id): + async with sessions() as db: + assert await db.get(ProtocolRun, run_id) is not None, "enqueue follows commit" + queued.append(run_id) + + async def hydrate(): + return None + + registry = SimpleNamespace( + get_all_tools=lambda **kwargs: [ + { + "name": f"{execution.SCRIPT_SERVER_NAME}.run_wired_script", + "tool_name": "run_wired_script", + "server": execution.SCRIPT_SERVER_NAME, + "description": "Inspect authorized row", + "input_schema": {}, + } + ] + ) + monkeypatch.setattr(api, "enqueue_protocol_run", enqueue) + monkeypatch.setattr(execution, "hydrate_registry", hydrate) + monkeypatch.setattr(execution, "get_registry", lambda: registry) + monkeypatch.setattr(run_tools, "get_registry", lambda: registry) + + async def provider(*, run_id, available_tools, **kwargs): + nonlocal attempted_failure + from asaree.mcp_servers import script_server + + async with core_sessions() as db: + agent_run = await db.get(AgentRun, run_id) + meta = agent_run.run_metadata + node_id = meta["node_id"] + ambient = meta["ambient_meta"] + if ambient.get("dataset_mode") != "per_row": + assert "Dataset row input:" not in agent_run.input + agent_run.status = RunStatus.COMPLETED + agent_run.output = "whole dataset regression" + await db.commit() + return + values = ambient["row_inputs"][0]["values"] + row_index = ambient["row_inputs"][0]["row_index"] + protocol_run_id = uuid.UUID(meta["protocol_run_id"]) + input_text = agent_run.input + assert ambient["dataset_mode"] == "per_row" + assert values["question"] == ("QUESTION-3" if row_index == 2 else "QUESTION-1") + if node_id == "prediction": + assert "GOLD" not in input_text and "reference" not in input_text + assert "GOLD" not in json.dumps(ambient) + assert any(tool["tool_name"] == "run_wired_script" for tool in available_tools) + metadata = {f"motoro.ambient.{key}": value for key, value in ambient.items()} + metadata["motoro.workspace_id"] = f"_protocol_runs/{protocol_run_id}" + result = json.loads( + script_server.run_wired_script( + ctx=SimpleNamespace(request_context=SimpleNamespace(meta=SimpleNamespace(model_extra=metadata))) + ) + ) + assert result["exit_code"] == 0, result + output = json.loads(result["stdout"]) + assert output["rows"] == [{"question": values["question"]}] + assert "GOLD" not in json.dumps(output) + path = output["path"] + manifest = json.loads((Path(path).parent / "row-view.json").read_text()) + assert manifest["values"] == {"question": values["question"]} + assert manifest["row_index"] == row_index and manifest["raw_sha256"] == digest + assert "GOLD" not in json.dumps(manifest) + assert path not in private_paths, "every attempt has a fresh private input" + private_paths.add(path) + else: + assert node_id == "grading" + assert values["reference"] == ("GOLD-3" if row_index == 2 else "GOLD-1") + assert values["reference"] in input_text + provider_calls.append((protocol_run_id, node_id, row_index, input_text)) + if protocol_run_id == failure_run and node_id == "prediction" and not attempted_failure: + attempted_failure = True + raise RuntimeError("injected row failure") + async with core_sessions.begin() as db: + agent_run = await db.get(AgentRun, run_id) + agent_run.status = RunStatus.COMPLETED + agent_run.output = ( + None + if protocol_run_id == missing_run and node_id == "grading" + else "prediction" + if node_id == "prediction" + else "graded" + ) + + async def steps(run_id): + async with core_sessions() as db: + run = await db.get(AgentRun, run_id) + if run.run_metadata["node_id"] != "prediction": + return [] + protocol_run_id = uuid.UUID(run.run_metadata["protocol_run_id"]) + return [ + SimpleNamespace( + iteration=1, + sequence=1, + tool_call={ + "server": execution.SCRIPT_SERVER_NAME, + "tool": "run_wired_script", + "arguments": {"script": "inspect-row"}, + "success": True, + "result": None if protocol_run_id == null_run else {"opaque": "12"}, + }, + ) + ] + + monkeypatch.setattr(execution, "execute_run", provider) + monkeypatch.setattr("asaree.services.reported_metrics.get_run_steps", steps) + try: + async with sessions.begin() as db: + user = await get_user_by_email(db, "test@test.com") + assert user is not None + dataset = RegisteredDataset( + id=uuid.uuid4(), + name=f"acceptance-{uuid.uuid4()}", + owner_id=user.id, + raw_path=str(source_path), + raw_sha256=digest, + target_column="reference", + ) + db.add(dataset) + dataset_id = dataset.id + design = { + "factors": [{"name": "Prompt", "levels": [f"LEVEL-{i}" for i in range(5)]}], + "replicates": 2, + "metrics": [ + {"id": "grade", "name": "Grade", "kind": "custom", "valueType": "opaque"}, + {"id": "script", "name": "Script report", "kind": "custom", "valueType": "opaque"}, + ], + } + plan = { + "metrics": [{"id": "grade", "name": "Grade"}, {"id": "script", "name": "Script report"}], + "producers": [ + { + "id": "grade-output", + "kind": "reported", + "producer_id": "asaree.agent_output", + "outputs": {"value": "grade"}, + "config": {"agent_node_id": "grading"}, + }, + { + "id": "script-output", + "kind": "reported", + "producer_id": "asaree.python_script", + "outputs": {"value": "script"}, + "config": {"agent_node_id": "prediction", "script_node_id": "script"}, + }, + ], + "inputs": [], + } + experiment = await create_experiment( + db, name=f"acceptance-{uuid.uuid4()}", owner_id=user.id, design_spec=design, measurement_plan=plan + ) + experiment_id = experiment.id + graph = { + "nodes": [ + { + "id": "source", + "type": "dataset", + "data": {"config": {"dataset_id": str(dataset.id), "dataset_name": dataset.name}}, + }, + { + "id": "prediction", + "type": "agent", + "data": { + "config": {"name": "Prediction", "prompt": "Predict.", "system_prompt": "LEVEL-0"}, + "factor_bindings": {"config.system_prompt": "Prompt"}, + }, + }, + {"id": "grading", "type": "agent", "data": {"config": {"name": "Grading", "prompt": "Grade."}}}, + {"id": "model-p", "type": "model_openai", "data": {"config": {}}}, + {"id": "model-g", "type": "model_openai", "data": {"config": {}}}, + { + "id": "script", + "type": "script", + "data": { + "config": { + "name": "inspect-row", + "code": "import csv, json\nfrom asaree.script_context import training_input\n" + 'item=training_input()\n' + 'print(json.dumps({"path": str(item.path), ' + '"rows": list(csv.DictReader(item.path.open()))}))\n', + } + }, + }, + ], + "edges": [ + { + "id": "data-p", + "source": "source", + "target": "prediction", + "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["question"]}}, + }, + { + "id": "data-g", + "source": "source", + "target": "grading", + "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["question", "reference"]}}, + }, + {"id": "m-p", "source": "model-p", "target": "prediction", "targetHandle": "model"}, + {"id": "m-g", "source": "model-g", "target": "grading", "targetHandle": "model"}, + {"id": "s-p", "source": "script", "target": "prediction", "targetHandle": "tool"}, + {"id": "p-g", "source": "prediction", "target": "grading"}, + ], + } + protocol = await create_protocol( + db, name=f"acceptance-{uuid.uuid4()}", owner_id=user.id, experiment_id=experiment.id, graph=graph + ) + protocol_id = protocol.id + revision = await publish_protocol(db, protocol, owner_id=user.id) + protocol.published_revision_id = revision.id + parents = await generate_design_cells( + db, experiment_id=experiment.id, factors=design["factors"], replicates=2, design_spec=design + ) + assert len(parents) == 10 + assert ( + await db.scalar( + select(func.count()).select_from(FactorialCell).where(FactorialCell.experiment_id == experiment.id) + ) + == 5 + ) + user_id = user.id + async with sessions() as db: + user = await get_user_by_email(db, "test@test.com") + batch = await api.create_cell_runs_endpoint(protocol_id, user, db) + assert len(batch.protocol_run_ids) == len(batch.row_result_ids) == 30 + assert batch.consumption_mode == "per_row" + for run_id in batch.protocol_run_ids: + pinned = await db.get(ProtocolRun, run_id) + assert pinned.protocol_revision_id == revision.id + assert pinned.design_revision_id is not None + assert pinned.dataset_row["dataset_id"] == str(dataset_id) + assert pinned.dataset_row["columns"] == ["question", "reference"] + assert pinned.factor_values["Prompt"] in [f"LEVEL-{i}" for i in range(5)] + failure_run, null_run, missing_run = batch.protocol_run_ids[:3] + protocol = await db.get(Protocol, protocol_id) + edited = deepcopy(protocol.graph) + edited["nodes"][1]["data"]["config"]["prompt"] = "DRAFT-SENTINEL" + protocol.graph = edited + await db.commit() + for run_id in batch.protocol_run_ids: + await execution.run_protocol(run_id) + async with sessions() as db: + user = await get_user_by_email(db, "test@test.com") + result = await get_experiment_run_results_endpoint(experiment_id, user, db, protocol_id=protocol_id) + assert result.row_summary["planned"] == 30 + assert result.row_summary["completed"] == 29 and result.row_summary["failed"] == 1 + assert result.row_summary["missing_reported"] == 1 + assert result.cells == [] and result.replicates == [] and result.primary_metric is None + assert "DRAFT-SENTINEL" not in json.dumps(provider_calls, default=str) + rows = result.row_results + failed = next(row for row in rows if row["status"] == "failed") + null_row = next(row for row in rows if row["latest_attempt"]["run_id"] == str(null_run)) + observation = next( + item for item in null_row["measurement"]["observations"] if item["metric_id"] == "script" + ) + assert observation["status"] == "measured" and observation["value"] is None + retry = await api.create_cell_runs_endpoint( + protocol_id, user, db, api.CellRunBatchRequest(retry_row_result_ids=[failed["row_result_id"]]) + ) + assert len(retry.protocol_run_ids) == 1 + assert retry.row_result_ids == [uuid.UUID(failed["row_result_id"])] + await execution.run_protocol(retry.protocol_run_ids[0]) + async with sessions() as db: + user = await get_user_by_email(db, "test@test.com") + result = await get_experiment_run_results_endpoint(experiment_id, user, db, protocol_id=protocol_id) + assert result.row_summary["completed"] == result.row_summary["planned"] == 30 + assert result.row_summary["missing_reported"] == 1 + assert sum(len(row["attempts"]) for row in result.row_results) == 31 + retried = next(row for row in result.row_results if row["row_result_id"] == failed["row_result_id"]) + assert {item["status"] for item in retried["attempts"]} == {"failed", "completed"} + assert not await project_row_attempt( + db, row_result_id=uuid.UUID(failed["row_result_id"]), run_id=failure_run, + fields={"metric_values": {"stale": 999}}, + ) + await db.commit() + assert ( + await db.scalar( + select(func.count()) + .select_from(FactorialReplicateResult) + .where(FactorialReplicateResult.id.in_([parent.id for parent in parents])) + ) + == 10 + ) + resume = await api.create_cell_runs_endpoint(protocol_id, user, db) + assert resume.protocol_run_ids == [] and resume.skipped == 30 + csv_response = await export_run_results_csv_endpoint(experiment_id, user, db, protocol_id=protocol_id) + exported = csv_response.body.decode() + assert digest in exported and str(revision.id) in exported and "row_index" in exported + assert "null" in exported and len(list(csv.DictReader(exported.splitlines()))) == 30 + for row in result.row_results: + assert row["dataset_row"]["raw_sha256"] == digest + assert row["dataset_row"]["row_index"] in (0, 1, 2) + snapshots_before = await db.scalar(select(func.count()).select_from(FactorialRowResult)) + selected_test = await api.create_test_run_endpoint(protocol_id, user, db, api.TestRunRequest(row_index=2)) + selected_play = await api.run_single_node_endpoint( + protocol_id, "prediction", user, db, api.NodePlayRequest(row_index=2) + ) + assert selected_test.dataset_row["row_index"] == selected_play.dataset_row["row_index"] == 2 + preview = await api.preview_node_prompt_endpoint( + protocol_id, "prediction", api.PromptPreviewRequest(graph=graph, row_index=2), user, db + ) + assert preview.dataset_row["row_index"] == 2 and "GOLD" not in preview.text + assert await db.scalar(select(func.count()).select_from(FactorialRowResult)) == snapshots_before + await execution.run_protocol(selected_test.id) + await execution.run_protocol(selected_play.id) + assert len(private_paths) == 33 + assert source_path.read_bytes() == raw + async with sessions() as db: + user = await get_user_by_email(db, "test@test.com") + protocol = await db.get(Protocol, protocol_id) + later = await publish_protocol(db, protocol, owner_id=user.id) + protocol.published_revision_id = later.id + await db.commit() + assert later.id != revision.id + current = await get_experiment_run_results_endpoint(experiment_id, user, db, protocol_id=protocol_id) + assert current.row_results == [] + history = await get_experiment_run_results_endpoint( + experiment_id, user, db, protocol_id=protocol_id, protocol_revision_id=revision.id, + ) + assert len(history.row_results) == 30 and history.row_summary["completed"] == 30 + source_path.write_bytes(raw + b"CORRUPTED,SECRET\n") + with pytest.raises(HTTPException) as error: + await api.create_test_run_endpoint(protocol_id, user, db, api.TestRunRequest(row_index=0)) + assert error.value.status_code == 422 + history_after = await get_experiment_run_results_endpoint( + experiment_id, user, db, protocol_id=protocol_id, protocol_revision_id=revision.id, + ) + assert len(history_after.row_results) == 30 + source_path.write_bytes(raw) + whole_graph = deepcopy(graph) + for edge in whole_graph["edges"]: + edge.pop("data", None) + protocol.graph = whole_graph + whole_revision = await publish_protocol(db, protocol, owner_id=user.id) + protocol.published_revision_id = whole_revision.id + await db.commit() + async def enqueue_whole(run_id): + queued.append(run_id) + monkeypatch.setattr(api, "enqueue_protocol_run", enqueue_whole) + whole = await api.create_protocol_run_endpoint(protocol_id, user, db) + await db.commit() + assert whole.dataset_row is None + assert await db.scalar(select(func.count()).select_from(FactorialRowResult)) == snapshots_before + await execution.run_protocol(whole.id) + async with sessions() as db: + assert (await db.get(ProtocolRun, whole.id)).status == "completed" + finally: + if protocol_id is not None: + from motoro.models.agent import Agent + + async with core_sessions.begin() as db: + agents = ( + await db.scalars( + select(Agent).where(Agent.owner_id == user_id, Agent.name.like(f"protocol-{protocol_id}-%")) + ) + ).all() + for agent in agents: + await db.delete(agent) + async with sessions.begin() as db: + await delete_protocol(db, protocol_id) + await delete_experiment(db, experiment_id) + dataset = await db.get(RegisteredDataset, dataset_id) + if dataset is not None: + await db.delete(dataset) + await product_engine.dispose() + await core_engine.dispose() diff --git a/tests/test_dataset_row_analysis_exclusion.py b/tests/test_dataset_row_analysis_exclusion.py new file mode 100644 index 00000000..ec6eac8a --- /dev/null +++ b/tests/test_dataset_row_analysis_exclusion.py @@ -0,0 +1,121 @@ +"""Row executions remain inspectable but never enter factorial analysis.""" + +from __future__ import annotations + +import pytest +from fastapi import HTTPException + +from asaree.api.experiments import ( + export_replicates_csv_endpoint, + get_experiment_results_endpoint, +) +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.services.factorial_analysis import FactorialAnalysisError, analyze_experiment_design + +pytest_plugins = ("test_dataset_row_results_projection",) + + +@pytest.mark.asyncio +async def test_published_row_protocol_disables_results_without_calling_analyzer(row_results_setup, monkeypatch) -> None: + db, ctx = row_results_setup + ctx["experiment"].design_spec = { + "factors": [{"name": "prompt", "levels": ["a", "b", "c", "d", "e"]}], + "metrics": [{"name": "Report", "primary": True, "direction": "maximize"}], + } + ctx["design_revision"].design_spec = ctx["experiment"].design_spec + + def unexpected(*_args, **_kwargs): + raise AssertionError("numerical analyzer must not run for row mode") + + monkeypatch.setattr("asaree.services.factorial_analysis.analyze_factorial", unexpected) + result = await get_experiment_results_endpoint( + ctx["experiment"].id, ctx["user"], db, protocol_id=ctx["protocol"].id + ) + assert result.model_dump() == { + "available": False, + "reason": "Per-row executions are available for inspection and export; factorial analysis is not supported.", + "analysis": None, + "best_condition": None, + } + + +def test_row_result_input_is_rejected_under_whole_dataset() -> None: + row_result = FactorialRowResult(row_index=0) + with pytest.raises(FactorialAnalysisError, match="row execution"): + analyze_experiment_design( + {"factors": [{"name": "prompt", "levels": ["a", "b"]}], "metrics": []}, + [row_result], + ) + + +def test_per_row_custom_json_is_opaque_and_analysis_stays_unavailable() -> None: + opaque = {"score": 9, "passed": True} + result = analyze_experiment_design( + {"factors": [{"name": "prompt", "levels": ["a", "b"]}], "metrics": []}, + [{"row_result_id": "row-1", "metric_values": {"Report": opaque}}], + consumption_mode="per_row", + ) + assert result["available"] is False + assert result["analysis"] is None + assert opaque == {"score": 9, "passed": True} + + +@pytest.mark.asyncio +async def test_stale_parent_metric_does_not_rank_row_mode_and_legacy_export_is_blocked(row_results_setup) -> None: + db, ctx = row_results_setup + ctx["experiment"].design_spec = { + "factors": [{"name": "arm", "levels": ["left", "right"]}], + "metrics": [{"name": "Report", "primary": True, "direction": "maximize"}], + } + for parent in ctx["parents"]: + parent.metric_values = {"Report": 1000} + results = await get_experiment_results_endpoint( + ctx["experiment"].id, ctx["user"], db, protocol_id=ctx["protocol"].id + ) + assert results.available is False and results.best_condition is None + with pytest.raises(HTTPException) as caught: + await export_replicates_csv_endpoint( + ctx["experiment"].id, ctx["user"], db, protocol_id=ctx["protocol"].id + ) + assert caught.value.status_code == 422 + assert caught.value.detail["code"] == "use_row_results_export" + assert caught.value.detail["link_text"] == "run-results.csv" + + +@pytest.mark.asyncio +async def test_historical_whole_publication_keeps_legacy_analysis_and_export(row_results_setup) -> None: + db, ctx = row_results_setup + ctx["protocol"].graph = {"nodes": [{"id": "agent", "type": "agent", "data": {}}], "edges": []} + from asaree.services.protocol_revisions import publish_protocol + + whole = await publish_protocol(db, ctx["protocol"], owner_id=ctx["user"].id) + ctx["experiment"].design_spec = { + "factors": [{"name": "arm", "levels": ["left", "right"]}], + "metrics": [{"name": "Report", "primary": True, "direction": "maximize"}], + } + for parent in ctx["parents"]: + parent.metric_values = {"Report": 0.7} + for cell in ctx["cells"]: + db.add(FactorialReplicateResult( + cell_id=cell.id, + replicate_number=2, + replicate_label=f"historical-{cell.cell_label}", + metric_values={"Report": 0.71}, + )) + # Draft is switched back to row mode after publishing; the explicit whole + # publication remains authoritative for both legacy projections. + ctx["protocol"].graph = { + "nodes": [{"id": "dataset", "type": "dataset", "data": {"config": {"dataset_id": str(ctx["dataset"].id)}}}], + "edges": [], + } + analyzed = await get_experiment_results_endpoint( + ctx["experiment"].id, ctx["user"], db, + protocol_id=ctx["protocol"].id, protocol_revision_id=whole.id, + ) + assert analyzed.available is True + exported = await export_replicates_csv_endpoint( + ctx["experiment"].id, ctx["user"], db, + protocol_id=ctx["protocol"].id, protocol_revision_id=whole.id, + ) + assert exported.media_type == "text/csv" diff --git a/tests/test_dataset_row_batch_planning.py b/tests/test_dataset_row_batch_planning.py index ca2e9c5e..d222665d 100644 --- a/tests/test_dataset_row_batch_planning.py +++ b/tests/test_dataset_row_batch_planning.py @@ -6,6 +6,7 @@ import os import uuid from collections.abc import AsyncIterator +from copy import deepcopy from pathlib import Path import pytest @@ -170,3 +171,65 @@ async def enqueue(_run_id: uuid.UUID) -> None: result = await create_cell_runs_endpoint(protocol.id, ctx["user"], db) assert result.consumption_mode == "whole_dataset" assert result.row_result_ids == [] + + +@pytest.mark.asyncio +async def test_parent_selection_expands_every_source_row(row_batch, monkeypatch: pytest.MonkeyPatch) -> None: + db, ctx = row_batch + queued: list[uuid.UUID] = [] + + async def enqueue(run_id: uuid.UUID) -> None: + queued.append(run_id) + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + dataset_id = uuid.UUID( + ctx["revision"].graph["nodes"][0]["data"]["config"]["dataset_id"] + ) + dataset = await db.get(RegisteredDataset, dataset_id) + assert dataset is not None + raw = b"question,answer\nq1,a1\nq2,a2\nq3,a3\n" + Path(dataset.raw_path).write_bytes(raw) + dataset.raw_sha256 = hashlib.sha256(raw).hexdigest() + await upsert_replicate( + db, + experiment_id=ctx["protocol"].experiment_id, + replicate_label="second-parent", + fields={"factor_values": {}}, + ) + + result = await create_cell_runs_endpoint( + ctx["protocol"].id, + ctx["user"], + db, + CellRunBatchRequest(replicate_labels=["batch-parent", "second-parent"]), + ) + assert len(result.protocol_run_ids) == 6 + assert len(result.row_result_ids) == 6 + # Parents with identical cell labels and replicate numbers are ordered by + # their persistent UUID, not their display labels or insertion order. + assert sorted(result.replicate_labels) == ["batch-parent"] * 3 + ["second-parent"] * 3 + assert len(set(result.replicate_labels[:3])) == 1 + assert len(set(result.replicate_labels[3:])) == 1 + assert result.skipped == 0 + assert len(queued) == 6 + + +@pytest.mark.asyncio +async def test_new_publication_has_separate_row_slots(row_batch, monkeypatch: pytest.MonkeyPatch) -> None: + db, ctx = row_batch + + async def enqueue(_run_id: uuid.UUID) -> None: + return None + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + protocol = ctx["protocol"] + first = await create_cell_runs_endpoint(protocol.id, ctx["user"], db) + protocol.graph = deepcopy(protocol.graph) + protocol.graph["nodes"][1]["data"]["system_prompt"] = "published again" + next_revision = await publish_protocol(db, protocol, owner_id=ctx["user"].id) + protocol.published_revision_id = next_revision.id + + second = await create_cell_runs_endpoint(protocol.id, ctx["user"], db) + assert first.protocol_revision_id != second.protocol_revision_id + assert len(second.protocol_run_ids) == 30 + assert len(set(first.row_result_ids).intersection(second.row_result_ids)) == 0 diff --git a/tests/test_dataset_row_csv_export.py b/tests/test_dataset_row_csv_export.py new file mode 100644 index 00000000..2342b926 --- /dev/null +++ b/tests/test_dataset_row_csv_export.py @@ -0,0 +1,139 @@ +"""PostgreSQL Results-to-CSV export coverage for stable dataset row slots.""" + +from __future__ import annotations + +import csv +import io +import json +import uuid + +import pytest +from fastapi import HTTPException + +import asaree.api.protocols as protocol_api +from asaree.api.experiments import ( + export_run_results_csv_endpoint, + get_run_results_schema_endpoint, +) +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.protocol_run import ProtocolRun +from asaree.services.factorial_row_results import claim_row_attempt, list_row_results + +pytest_plugins = ("test_dataset_row_results_projection",) + + +@pytest.mark.asyncio +async def test_row_csv_and_schema_use_selected_results_projection(row_results_setup, monkeypatch) -> None: + db, ctx = row_results_setup + # Plan stable slots through the production planner. Each dataset position + # remains a separate CSV record even where source values are duplicates. + from asaree.api.protocols import CellRunBatchRequest, create_cell_runs_endpoint + + async def no_enqueue(_run_id: uuid.UUID) -> None: + return None + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", no_enqueue) + planned = await create_cell_runs_endpoint(ctx["protocol"].id, ctx["user"], db, CellRunBatchRequest()) + slots = await list_row_results( + db, + experiment_id=ctx["experiment"].id, + design_revision_id=ctx["design_revision"].id, + protocol_revision_id=ctx["publication"].id, + ) + assert slots + runs = [await db.get(ProtocolRun, run_id) for run_id in planned.protocol_run_ids] + run_by_slot = {run.row_result_id: run for run in runs if run is not None and run.row_result_id is not None} + failed_slot = slots[2] + failed_run = run_by_slot[failed_slot.id] + assert failed_run is not None + failed_replicate = await db.get(FactorialReplicateResult, failed_slot.replicate_result_id) + assert failed_replicate is not None + failed_run.status = "failed" + second_failed_run = run_by_slot[slots[6].id] + assert second_failed_run is not None + second_failed_run.status = "failed" + quoted_value = "quote,\n雪" + opaque_value = {"nested": [1, None], "text": "é"} + for slot, value in zip(slots[:5], [None, None, None, quoted_value, opaque_value], strict=True): + if slot is failed_slot: + continue + run = run_by_slot[slot.id] + assert run is not None + status = "unavailable" if slot is slots[1] else "measured" + run.status = "completed" + run.attempt_result = {"measurement": {"observations": [{ + "metric_id": "reported", "metric_name": "Report", "status": status, + "value": value, "producer": {"producer_id": "asaree.agent_output", "node_id": "agent"}, + }]}} + if status == "measured": + slot.metric_values = {"Report": value} + retry = await claim_row_attempt( + db, + row_result_id=failed_slot.id, + expected_run_id=failed_run.id, + create_kwargs={ + "protocol_id": ctx["protocol"].id, + "owner_id": ctx["user"].id, + "replicate_label": failed_replicate.replicate_label, + "factor_values": failed_run.factor_values, + "replicate_result_id": failed_slot.replicate_result_id, + "design_revision_id": ctx["design_revision"].id, + "protocol_revision_id": ctx["publication"].id, + "row_result_id": failed_slot.id, + "dataset_row": failed_run.dataset_row, + }, + ) + assert retry is not None + await db.flush() + + selectors = { + "protocol_id": ctx["protocol"].id, + "design_revision_id": ctx["design_revision"].id, + "protocol_revision_id": ctx["publication"].id, + } + response = await export_run_results_csv_endpoint(ctx["experiment"].id, ctx["user"], db, **selectors) + parsed = list(csv.DictReader(io.StringIO(response.body.decode("utf-8")))) + schema = await get_run_results_schema_endpoint(ctx["experiment"].id, ctx["user"], db, **selectors) + header = next(csv.reader(io.StringIO(response.body.decode("utf-8")))) + + assert len(parsed) == len(slots) + assert len({row["row_result_id"] for row in parsed}) == len(slots) + assert {row["row_index"] for row in parsed} == {str(slot.row_index) for slot in slots} + assert parsed[0]["Report"] == "null" + assert parsed[0]["Report__status"] == "measured" + assert '"producer_id":"asaree.agent_output"' in parsed[0]["Report__producer"] + assert parsed[1]["Report"] == "" and parsed[1]["Report__status"] == "unavailable" + assert parsed[2]["status"] == "pending" + assert any(row["status"] == "failed" for row in parsed) + retry_row = next(row for row in parsed if row["row_result_id"] == str(failed_slot.id)) + assert retry_row["run_id"] == str(retry.id) + assert json.loads(parsed[3]["Report"]) == quoted_value + assert json.loads(parsed[4]["Report"]) == opaque_value + assert schema["consumption_mode"] == "per_row" + assert schema["row_identity"][-6]["name"] == "row_index" + assert set(header) == {column["name"] for column in schema["columns"]} + + for endpoint in (export_run_results_csv_endpoint, get_run_results_schema_endpoint): + with pytest.raises(HTTPException) as foreign: + await endpoint( + ctx["experiment"].id, + ctx["user"], + db, + protocol_id=ctx["protocol"].id, + protocol_revision_id=uuid.uuid4(), + ) + assert foreign.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_whole_dataset_csv_route_keeps_legacy_layout(row_results_setup) -> None: + db, ctx = row_results_setup + ctx["protocol"].graph = {"nodes": [{"id": "agent", "type": "agent", "data": {}}], "edges": []} + from asaree.services.protocol_revisions import publish_protocol + + await publish_protocol(db, ctx["protocol"], owner_id=ctx["user"].id) + response = await export_run_results_csv_endpoint( + ctx["experiment"].id, ctx["user"], db, protocol_id=ctx["protocol"].id + ) + assert response.body.decode("utf-8").startswith("cell_label,replicate_number,") + assert "row_result_id" not in response.body.decode("utf-8").splitlines()[0] diff --git a/tests/test_dataset_row_executor.py b/tests/test_dataset_row_executor.py new file mode 100644 index 00000000..a72aa74a --- /dev/null +++ b/tests/test_dataset_row_executor.py @@ -0,0 +1,194 @@ +"""PostgreSQL integration coverage for a complete row-scoped protocol attempt.""" + +from __future__ import annotations + +import hashlib +import os +import uuid +from pathlib import Path + +import pytest +import pytest_asyncio +from motoro.runner import get_run +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +import asaree.services.protocol_execution as execution +from asaree.models.dataset import RegisteredDataset +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol_run import ProtocolRun +from asaree.services.experiments import create_experiment +from asaree.services.factorial_cells import upsert_replicate +from asaree.services.factorial_row_results import claim_row_attempt, ensure_row_result +from asaree.services.protocol_revisions import publish_protocol +from asaree.services.protocols import create_protocol +from asaree.services.users import get_user_by_email + + +@pytest_asyncio.fixture +async def row_protocol(tmp_path: Path): + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + sessions = async_sessionmaker(engine, expire_on_commit=False) + try: + async with sessions.begin() as db: + user = await get_user_by_email(db, "test@test.com") + assert user is not None, "run scripts/seed_row_test_users.py before row database tests" + source_path = tmp_path / "protocol-source.csv" + content = b"prompt,reference\nfirst,gold-a\nsecond,gold-b\n" + source_path.write_bytes(content) + dataset = RegisteredDataset( + id=uuid.uuid4(), name=f"row-executor-{uuid.uuid4()}", owner_id=user.id, + raw_path=str(source_path), raw_sha256=hashlib.sha256(content).hexdigest(), + target_column="reference", + ) + db.add(dataset) + experiment = await create_experiment( + db, name=f"row-executor-{uuid.uuid4()}", owner_id=user.id, + design_spec={"factors": [], "metrics": []}, + ) + replicate = await upsert_replicate( + db, experiment_id=experiment.id, replicate_label="row-executor-cell", + fields={"factor_values": {}}, + ) + graph = { + "nodes": [ + {"id": "source", "type": "dataset", "data": {"config": { + "dataset_id": str(dataset.id), "dataset_name": dataset.name, + }}}, + {"id": "first", "type": "agent", "data": {"config": {"name": "first"}}}, + {"id": "second", "type": "agent", "data": {"config": {"name": "second"}}}, + {"id": "model-first", "type": "model_openai", "data": {"config": {}}}, + {"id": "model-second", "type": "model_openai", "data": {"config": {}}}, + ], + "edges": [ + {"id": "driver-first", "source": "source", "target": "first", "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["prompt", "reference"]}}}, + {"id": "model-first-edge", "source": "model-first", "target": "first", "targetHandle": "model"}, + {"id": "first-second", "source": "first", "target": "second"}, + {"id": "model-second-edge", "source": "model-second", "target": "second", "targetHandle": "model"}, + ], + } + protocol = await create_protocol( + db, name=f"row-executor-{uuid.uuid4()}", owner_id=user.id, + experiment_id=experiment.id, graph=graph, + ) + revision = await publish_protocol(db, protocol, owner_id=user.id) + protocol.published_revision_id = revision.id + row = execution.project_row( + execution.read_row_source( + dataset_id=str(dataset.id), raw_path=str(source_path), raw_sha256=dataset.raw_sha256, + ), row_index=0, columns=["prompt", "reference"], + ) + # Production row creation validates and claims the stable row slot. + slot = await ensure_row_result( + db, experiment_id=experiment.id, design_revision_id=replicate.design_revision_id, + protocol_revision_id=revision.id, replicate_result_id=replicate.id, + dataset_id=dataset.id, raw_sha256=dataset.raw_sha256, row_index=0, + ) + run = await claim_row_attempt( + db, row_result_id=slot.id, expected_run_id=None, + create_kwargs={"protocol_id": protocol.id, "owner_id": user.id, + "replicate_label": replicate.replicate_label, "factor_values": {}, + "replicate_result_id": replicate.id, "design_revision_id": replicate.design_revision_id, + "protocol_revision_id": revision.id, "dataset_row": row}, + ) + assert run is not None + ids = (run.id,) + yield sessions, ids[0] + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_complete_protocol_uses_fresh_row_context_without_parent_writeback(row_protocol, monkeypatch): + sessions, run_id = row_protocol + calls = 0 + + async def execute_run_mock(**_kwargs): + nonlocal calls + calls += 1 + + async def no_hydration(): + return None + + monkeypatch.setattr(execution, "execute_run", execute_run_mock) + monkeypatch.setattr(execution, "hydrate_registry", no_hydration) + monkeypatch.setattr(execution, "get_registry", lambda: None) + + await execution.run_protocol(run_id) + + async with sessions() as db: + run = await db.get(ProtocolRun, run_id) + assert run is not None and run.status == "completed" + assert set(run.node_runs) >= {"first", "second"} + assert calls == 2 + replicate = await db.get(FactorialReplicateResult, run.replicate_result_id) + slot = await db.get(FactorialRowResult, run.row_result_id) + assert replicate is not None and replicate.workspace_id is None and replicate.metric_values is None + assert slot is not None and slot.workspace_id == f"_protocol_runs/{run_id}" + motoro_runs = [await get_run(uuid.UUID(run.node_runs[node]["run_id"])) for node in ("first", "second")] + assert all(item is not None for item in motoro_runs) + contexts = [item.run_metadata["ambient_meta"] for item in motoro_runs] + assert all(item["dataset_mode"] == "per_row" for item in contexts) + assert contexts[0]["row_inputs"][0]["values"] == {"prompt": "first", "reference": "gold-a"} + assert contexts[1]["row_inputs"] == [] + + +@pytest.mark.asyncio +async def test_duplicate_delivery_does_not_create_another_row_execution(row_protocol, monkeypatch): + _sessions, run_id = row_protocol + calls = 0 + + async def execute_run_mock(**_kwargs): + nonlocal calls + calls += 1 + + async def no_hydration(): + return None + + monkeypatch.setattr(execution, "execute_run", execute_run_mock) + monkeypatch.setattr(execution, "hydrate_registry", no_hydration) + monkeypatch.setattr(execution, "get_registry", lambda: None) + + await execution.run_protocol(run_id) + first_delivery_calls = calls + await execution.run_protocol(run_id) + + assert first_delivery_calls == 2 + assert calls == first_delivery_calls + + +@pytest.mark.asyncio +async def test_changed_original_source_fails_before_provider_execution(row_protocol, monkeypatch): + sessions, run_id = row_protocol + calls = 0 + + async def execute_run_mock(**_kwargs): + nonlocal calls + calls += 1 + + async def no_hydration(): + return None + + monkeypatch.setattr(execution, "execute_run", execute_run_mock) + monkeypatch.setattr(execution, "hydrate_registry", no_hydration) + monkeypatch.setattr(execution, "get_registry", lambda: None) + + async with sessions.begin() as db: + run = await db.get(ProtocolRun, run_id) + assert run is not None + registration = await db.get(RegisteredDataset, run.dataset_row["dataset_id"]) + assert registration is not None + changed_path = Path(registration.raw_path).with_name("changed-source.csv") + changed_bytes = b"prompt,reference\nchanged,gold-a\nsecond,gold-b\n" + changed_path.write_bytes(changed_bytes) + registration.raw_path = str(changed_path) + registration.raw_sha256 = hashlib.sha256(changed_bytes).hexdigest() + + await execution.run_protocol(run_id) + + async with sessions() as db: + run = await db.get(ProtocolRun, run_id) + assert run is not None and run.status == "failed" + assert "Row source validation failed" in (run.error or "") + assert calls == 0 diff --git a/tests/test_dataset_row_finalization.py b/tests/test_dataset_row_finalization.py new file mode 100644 index 00000000..5c26ba57 --- /dev/null +++ b/tests/test_dataset_row_finalization.py @@ -0,0 +1,415 @@ +"""PostgreSQL coverage for immutable row-attempt measurements and projections.""" + +from __future__ import annotations + +import hashlib +import os +import uuid +from collections.abc import AsyncIterator +from dataclasses import replace +from datetime import UTC, datetime +from types import SimpleNamespace + +import pytest +import pytest_asyncio +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from asaree.models.dataset import RegisteredDataset +from asaree.models.experiment import ResearchExperiment +from asaree.models.experiment_design_revision import ExperimentDesignRevision +from asaree.models.factorial_cell import FactorialCell +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol import Protocol +from asaree.models.protocol_revision import ProtocolRevision +from asaree.models.protocol_run import ProtocolRun +from asaree.services.dataset_row_csv import project_row, read_row_source +from asaree.services.measurement_engine import parse_measurement_plan +from asaree.services.protocol_runs import create_protocol_run, record_measurement_evaluation +from asaree.services.reported_metrics import collect_reported_metrics +from asaree.services.runtime_metrics import finalize_attempt_measurement +from asaree.services.users import get_user_by_email + +PLAN = { + "metrics": [ + {"id": "agent", "name": "Agent output"}, + {"id": "script", "name": "Script output"}, + {"id": "missing-agent", "name": "Missing Agent"}, + {"id": "missing-script", "name": "Missing Script"}, + ], + "producers": [ + { + "id": "agent-report", + "producer_id": "asaree.agent_output", + "kind": "reported", + "outputs": {"value": "agent"}, + "config": {"agent_node_id": "agent"}, + }, + { + "id": "script-report", + "producer_id": "asaree.python_script", + "kind": "reported", + "outputs": {"value": "script"}, + "config": {"agent_node_id": "agent", "script_node_id": "score"}, + }, + { + "id": "missing-agent-report", + "producer_id": "asaree.agent_output", + "kind": "reported", + "outputs": {"value": "missing-agent"}, + "config": {"agent_node_id": "missing-agent"}, + }, + { + "id": "missing-script-report", + "producer_id": "asaree.python_script", + "kind": "reported", + "outputs": {"value": "missing-script"}, + "config": {"agent_node_id": "agent", "script_node_id": "missing-score"}, + }, + ], + "inputs": [], +} +GRAPH = { + "nodes": [ + {"id": "agent", "type": "agent", "data": {"config": {}}}, + {"id": "missing-agent", "type": "agent", "data": {"config": {}}}, + {"id": "score", "type": "script", "data": {"config": {"name": "score", "code": "return 1"}}}, + { + "id": "missing-score", + "type": "script", + "data": {"config": {"name": "missing-score", "code": "return 1"}}, + }, + ], + "edges": [ + {"source": "score", "target": "agent", "targetHandle": "tool"}, + {"source": "missing-score", "target": "agent", "targetHandle": "tool"}, + ], +} + + +@pytest_asyncio.fixture +async def row_measurements(tmp_path) -> AsyncIterator[tuple[AsyncSession, dict]]: + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + try: + async with engine.connect() as connection: + transaction = await connection.begin() + sessions = async_sessionmaker(bind=connection, expire_on_commit=False) + async with sessions() as db: + user = await get_user_by_email(db, "test@test.com") + assert user is not None, "run scripts/seed_row_test_users.py before row database tests" + experiment_id, dataset_id, protocol_id = (uuid.uuid4() for _ in range(3)) + design_id, protocol_revision_id = uuid.uuid4(), uuid.uuid4() + cell_id, replicate_id = uuid.uuid4(), uuid.uuid4() + raw = b"input,value\nfirst,one\nsecond,two\nthird,three\n" + source_path = tmp_path / "observations.csv" + source_path.write_bytes(raw) + assert source_path.is_file() + digest = hashlib.sha256(raw).hexdigest() + source = read_row_source(dataset_id=str(dataset_id), raw_path=str(source_path), raw_sha256=digest) + workspace = tmp_path / "workspace" + workspace.mkdir() + (workspace / "trace.json").write_text('{"trace":true}', encoding="utf-8") + assert (workspace / "trace.json").is_file() + + experiment = ResearchExperiment( + id=experiment_id, + name=f"row-measurements-{uuid.uuid4()}", + owner_id=user.id, + measurement_plan=PLAN, + design_spec={"factors": [], "metrics": []}, + ) + dataset = RegisteredDataset( + id=dataset_id, + name=f"row-measurements-{uuid.uuid4()}", + owner_id=user.id, + raw_path=str(source_path), + raw_sha256=digest, + ) + protocol = Protocol( + id=protocol_id, + name=f"row-measurements-{uuid.uuid4()}", + owner_id=user.id, + experiment_id=experiment_id, + graph=GRAPH, + ) + db.add_all([experiment, dataset, protocol]) + await db.flush() + design = ExperimentDesignRevision(id=design_id, experiment_id=experiment_id, revision=1) + revision = ProtocolRevision( + id=protocol_revision_id, + protocol_id=protocol_id, + revision=1, + graph=GRAPH, + published_at=datetime.now(UTC), + ) + db.add_all([design, revision]) + await db.flush() + protocol.published_revision_id = protocol_revision_id + cell = FactorialCell( + id=cell_id, + experiment_id=experiment_id, + design_revision_id=design_id, + cell_label="row-cell", + factor_values={}, + ) + db.add(cell) + await db.flush() + replicate = FactorialReplicateResult( + id=replicate_id, + cell_id=cell_id, + replicate_number=1, + replicate_label="row-replicate", + run_id=uuid.uuid4(), + workspace_id="parent-workspace", + metric_values={"parent": "preserved"}, + artifacts={"parent": True}, + ) + db.add(replicate) + await db.flush() + rows = [ + FactorialRowResult( + replicate_result_id=replicate_id, + protocol_revision_id=protocol_revision_id, + dataset_id=dataset_id, + raw_sha256=digest, + row_index=index, + run_id=uuid.uuid4(), + workspace_id=f"old-workspace-{index}", + metric_values={"old": index}, + artifacts={"old": index}, + ) + for index in range(3) + ] + db.add_all(rows) + await db.flush() + yield db, { + "user_id": user.id, + "experiment": experiment, + "protocol": protocol, + "design_id": design_id, + "protocol_revision_id": protocol_revision_id, + "replicate": replicate, + "rows": rows, + "source": source, + "workspace": workspace, + } + if transaction.is_active: + await transaction.rollback() + finally: + await engine.dispose() + + +def _row_snapshot(ctx: dict, index: int) -> dict: + return project_row(ctx["source"], row_index=index, columns=["input", "value"]) + + +async def _create_attempt(db: AsyncSession, ctx: dict, index: int) -> ProtocolRun: + return await create_protocol_run( + db, + protocol_id=ctx["protocol"].id, + owner_id=ctx["user_id"], + replicate_label="row-replicate", + factor_values={}, + replicate_result_id=ctx["replicate"].id, + row_result_id=ctx["rows"][index].id, + dataset_row=_row_snapshot(ctx, index), + design_revision_id=ctx["design_id"], + protocol_revision_id=ctx["protocol_revision_id"], + ) + + +@pytest.mark.asyncio +async def test_finalize_row_attempts_freeze_opaque_reports_on_distinct_subjects(row_measurements, monkeypatch): + db, ctx = row_measurements + calls_by_run: dict[uuid.UUID, list[SimpleNamespace]] = {} + + async def trace(run_id): + return calls_by_run.get(run_id, []) + + monkeypatch.setattr("asaree.services.reported_metrics.get_run_steps", trace) + parent = ctx["replicate"] + parent_before = (parent.run_id, parent.workspace_id, parent.metric_values, parent.artifacts) + attempts = [await _create_attempt(db, ctx, index) for index in (0, 1)] + values = ["{\"row\":1}", "{\"row\":2}"] + for index, attempt in enumerate(attempts): + attempt.status = "completed" + attempt.node_runs = { + "agent": {"status": "completed", "output_text": values[index], "run_id": str(uuid.uuid4())}, + "missing-agent": {"status": "failed", "run_id": str(uuid.uuid4())}, + } + agent_run = uuid.UUID(attempt.node_runs["agent"]["run_id"]) + calls_by_run[agent_run] = [ + SimpleNamespace( + iteration=1, + sequence=1, + tool_call={ + "server": "asaree-script", + "tool": "run_wired_script", + "arguments": {"script": "score"}, + "result": {"row": index + 10}, + "success": True, + }, + ), + ] + if index == 1: + calls_by_run[agent_run].append(SimpleNamespace( + iteration=2, + sequence=1, + tool_call={ + "server": "asaree-script", + "tool": "run_wired_script", + "arguments": {"script": "score"}, + "result": None, + "success": False, + "error_type": "tool_reported", + }, + )) + assert await finalize_attempt_measurement(db, attempt.id) is True + + assert (parent.run_id, parent.workspace_id, parent.metric_values, parent.artifacts) == parent_before + for index, attempt in enumerate(attempts): + slot = ctx["rows"][index] + document = attempt.attempt_result["measurement"] + by_name = {item["metric_name"]: item for item in document["observations"]} + assert document["replicate_id"] == str(slot.id) + assert document["attempt_id"] == str(attempt.id) + assert by_name["Agent output"]["value"] == values[index] + assert by_name["Script output"]["value"] == ({"row": 10} if index == 0 else None) + assert by_name["Script output"]["status"] == "measured" + assert by_name["Script output"]["producer"]["evaluation"]["tool_call_success"] is (index == 0) + if index == 1: + assert by_name["Script output"]["producer"]["evaluation"]["tool_call_error_type"] == "tool_reported" + assert by_name["Missing Agent"]["status"] == "unavailable" + assert by_name["Missing Script"]["status"] == "unavailable" + assert slot.artifacts["measurement"] == document + assert slot.metric_values["Agent output"] == values[index] + assert slot.metric_values["Script output"] == ({"row": 10} if index == 0 else None) + assert slot.run_id == attempt.id + first_subject = attempts[0].attempt_result["measurement"]["replicate_id"] + second_subject = attempts[1].attempt_result["measurement"]["replicate_id"] + assert first_subject != second_subject + + +@pytest.mark.asyncio +async def test_row_attempt_cancellation_failure_truncation_and_preview_projection(row_measurements, monkeypatch): + db, ctx = row_measurements + + async def no_steps(_run_id): + return [] + + monkeypatch.setattr("asaree.services.reported_metrics.get_run_steps", no_steps) + + cancelled = await _create_attempt(db, ctx, 0) + cancelled.status = "cancelled" + cancelled.attempt_result["evaluation_state"] = "running" + cancelled.node_runs = {"agent": {"status": "completed", "output_text": "cancelled output"}} + assert await finalize_attempt_measurement(db, cancelled.id) is True + cancelled_document = cancelled.attempt_result["measurement"] + assert cancelled_document["observations"][0]["status"] == "measured" + assert cancelled_document["observations"][3]["status"] == "unavailable" + assert ctx["rows"][0].artifacts["measurement"] == cancelled_document + assert ctx["rows"][0].metric_values is None + + failed = await _create_attempt(db, ctx, 1) + failed.status = "failed" + failed.error = "enqueue failed" + assert await finalize_attempt_measurement(db, failed.id) is True + assert ctx["rows"][1].artifacts["measurement"]["attempt_id"] == str(failed.id) + assert ctx["rows"][1].metric_values is None + + truncated = await _create_attempt(db, ctx, 2) + truncated.status = "completed" + truncated.node_runs = { + "agent": {"status": "completed", "output_text": "unfinished", "truncation": {"reason": "limit"}} + } + assert await finalize_attempt_measurement(db, truncated.id) is True + assert truncated.attempt_result["measurement"]["observations"][0]["status"] == "measured" + assert "metric_values" in truncated.attempt_result + assert ctx["rows"][2].artifacts["measurement"]["attempt_id"] == str(truncated.id) + assert ctx["rows"][2].metric_values is None + + preview = ProtocolRun( + protocol_id=ctx["protocol"].id, + owner_id=ctx["user_id"], + status="completed", + dataset_row=_row_snapshot(ctx, 0), + attempt_result={"measurement_plan_snapshot": PLAN}, + node_runs={"agent": {"status": "completed", "output_text": "preview"}}, + ) + db.add(preview) + await db.flush() + assert await finalize_attempt_measurement(db, preview.id) is True + assert preview.attempt_result["measurement"]["replicate_id"] == str(preview.id) + assert all(row.run_id != preview.id for row in ctx["rows"]) + + +@pytest.mark.asyncio +async def test_record_measurement_validates_subject_and_keeps_history_immutable(row_measurements): + db, ctx = row_measurements + attempt = await _create_attempt(db, ctx, 0) + plan = parse_measurement_plan(PLAN) + from asaree.services.reported_metrics import collect_reported_metrics + + attempt.node_runs = {"agent": {"status": "completed", "output_text": "recorded"}} + evaluation = await collect_reported_metrics(attempt, plan, GRAPH) + assert evaluation.replicate_id == str(ctx["rows"][0].id) + assert await record_measurement_evaluation(db, attempt.id, evaluation) is not None + frozen = dict(attempt.attempt_result) + with pytest.raises(ValueError, match="immutable"): + await record_measurement_evaluation(db, attempt.id, evaluation) + assert attempt.attempt_result == frozen + + with pytest.raises(ValueError, match="attempt"): + await record_measurement_evaluation(db, attempt.id, replace(evaluation, attempt_id=str(uuid.uuid4()))) + replacement = await _create_attempt(db, ctx, 1) + wrong_subject = replace( + evaluation, + replicate_id=str(ctx["rows"][0].id), + attempt_id=str(replacement.id), + observations=tuple(replace(item, attempt_id=str(replacement.id)) for item in evaluation.observations), + ) + with pytest.raises(ValueError, match="replicate"): + await record_measurement_evaluation(db, replacement.id, wrong_subject) + + +@pytest.mark.asyncio +async def test_superseded_row_attempt_keeps_history_without_replacing_current_projection(row_measurements): + db, ctx = row_measurements + first = await _create_attempt(db, ctx, 0) + first.node_runs = {"agent": {"status": "completed", "output_text": "older"}} + evaluation = await collect_reported_metrics(first, parse_measurement_plan(PLAN), GRAPH) + replacement = await _create_attempt(db, ctx, 0) + + await record_measurement_evaluation(db, first.id, evaluation) + + assert first.attempt_result["measurement"]["observations"][0]["value"] == "older" + assert ctx["rows"][0].run_id == replacement.id + assert ctx["rows"][0].metric_values is None + assert ctx["rows"][0].artifacts is None + + +@pytest.mark.asyncio +async def test_whole_dataset_attempt_keeps_replicate_projection_compatibility(row_measurements, monkeypatch): + db, ctx = row_measurements + + async def no_steps(_run_id): + return [] + + monkeypatch.setattr("asaree.services.reported_metrics.get_run_steps", no_steps) + run = await create_protocol_run( + db, + protocol_id=ctx["protocol"].id, + owner_id=ctx["user_id"], + replicate_label="row-replicate", + factor_values={}, + replicate_result_id=ctx["replicate"].id, + design_revision_id=ctx["design_id"], + protocol_revision_id=ctx["protocol_revision_id"], + ) + run.status = "completed" + run.node_runs = {"agent": {"status": "completed", "output_text": "whole dataset"}} + + assert await finalize_attempt_measurement(db, run.id) is True + assert run.attempt_result["measurement"]["replicate_id"] == str(ctx["replicate"].id) + assert ctx["replicate"].artifacts["measurement"]["attempt_id"] == str(run.id) + assert ctx["replicate"].metric_values["Agent output"] == "whole dataset" diff --git a/tests/test_dataset_row_node_play.py b/tests/test_dataset_row_node_play.py new file mode 100644 index 00000000..a7ef256c --- /dev/null +++ b/tests/test_dataset_row_node_play.py @@ -0,0 +1,386 @@ +"""PostgreSQL coverage for row-selected per-node Play previews.""" + +from __future__ import annotations + +import hashlib +import json +import os +import uuid +from collections.abc import AsyncIterator +from pathlib import Path +from types import SimpleNamespace + +import pytest +import pytest_asyncio +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient +from motoro.models import database as core_database +from motoro.models.agent import Agent +from motoro.models.run import AgentRun, RunStatus +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +import asaree.api.protocols as protocol_api +import asaree.services.dataset_row_workspaces as row_workspaces +import asaree.services.protocol_execution as execution +import asaree.services.run_tools as run_tools +from asaree.deps import get_current_user +from asaree.models import database +from asaree.models.dataset import RegisteredDataset +from asaree.models.experiment import ResearchExperiment +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol import Protocol +from asaree.models.protocol_run import ProtocolRun +from asaree.models.user import User +from asaree.services.experiments import create_experiment, delete_experiment +from asaree.services.factorial_cells import upsert_replicate +from asaree.services.factorial_row_results import ensure_row_result +from asaree.services.protocol_revisions import publish_protocol +from asaree.services.protocols import create_protocol, delete_protocol +from asaree.services.users import get_user_by_email + + +@pytest_asyncio.fixture +async def node_play_context(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[dict]: + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + sessions = async_sessionmaker(engine, expire_on_commit=False) + core_engine = create_async_engine(os.environ["DATABASE_URL"]) + core_sessions = async_sessionmaker(core_engine, expire_on_commit=False) + monkeypatch.setattr(core_database, "_get_session_factory", lambda: core_sessions) + protocol_id = experiment_id = dataset_id = None + raw = b"question,reference\nfirst,gold-a\nsecond,gold-b\nthird,gold-c\n" + raw_path = tmp_path / "node-play.csv" + raw_path.write_bytes(raw) + dataset_name = f"node-play-{uuid.uuid4()}" + graph: dict = { + "nodes": [ + {"id": "source", "type": "dataset", "data": {"config": {"dataset_name": dataset_name}}}, + {"id": "agent", "type": "agent", "data": {"config": {"name": "prediction"}}}, + {"id": "script", "type": "script", "data": {"config": { + "name": "inspect-row", + "code": "import csv, json\nfrom asaree.script_context import training_input\n" + "item = training_input()\nprint(json.dumps(list(csv.DictReader(item.path.open()))))\n", + }}}, + {"id": "model", "type": "model_openai", "data": {"config": {}}}, + {"id": "upstream", "type": "agent", "data": {"config": {"name": "upstream"}}}, + {"id": "upstream-model", "type": "model_openai", "data": {"config": {}}}, + {"id": "unbound", "type": "agent", "data": {"config": {"name": "unbound"}}}, + {"id": "unbound-model", "type": "model_openai", "data": {"config": {}}}, + ], + "edges": [ + {"id": "dataset-agent", "source": "source", "target": "agent", "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["question"]}}}, + {"id": "script-agent", "source": "script", "target": "agent", "targetHandle": "tool"}, + {"id": "model-agent", "source": "model", "target": "agent", "targetHandle": "model"}, + {"id": "model-upstream", "source": "upstream-model", "target": "upstream", "targetHandle": "model"}, + {"id": "model-unbound", "source": "unbound-model", "target": "unbound", "targetHandle": "model"}, + {"id": "agent-upstream", "source": "agent", "target": "upstream"}, + ], + } + try: + async with sessions.begin() as db: + user = await get_user_by_email(db, "test@test.com") + assert user is not None, "run scripts/seed_row_test_users.py before row database tests" + dataset = RegisteredDataset( + id=uuid.uuid4(), name=dataset_name, owner_id=user.id, raw_path=str(raw_path), + raw_sha256=hashlib.sha256(raw).hexdigest(), + target_column="reference", + ) + db.add(dataset) + graph["nodes"][0]["data"]["config"]["dataset_id"] = str(dataset.id) + experiment = await create_experiment( + db, name=f"node-play-{uuid.uuid4()}", owner_id=user.id, + design_spec={"factors": [], "metrics": []}, + ) + experiment.measurement_plan = { + "metrics": [{"id": "prediction", "name": "Prediction"}], + "producers": [{"id": "prediction-output", "producer_id": "asaree.agent_output", + "kind": "reported", "outputs": {"value": "prediction"}, + "config": {"agent_node_id": "agent"}}], + "inputs": [], + } + protocol = await create_protocol( + db, name=f"node-play-{uuid.uuid4()}", owner_id=user.id, + experiment_id=experiment.id, graph=graph, + ) + revision = await publish_protocol(db, protocol, owner_id=user.id) + protocol.published_revision_id = revision.id + protocol_id, experiment_id, dataset_id = protocol.id, experiment.id, dataset.id + ids = user.id, protocol.id, revision.id, experiment.id, dataset.id + async with sessions() as db: + user = await db.get(User, ids[0]) + protocol = await db.get(Protocol, ids[1]) + assert user is not None and protocol is not None + # Worker sessions must commit status/node updates on successful exit. + monkeypatch.setattr(database, "_get_session_factory", lambda: sessions) + monkeypatch.setattr(execution, "get_session", database.get_session) + monkeypatch.setattr(row_workspaces, "_workspace_root", lambda: tmp_path / "dataset-workspaces") + monkeypatch.setattr(execution, "WORKSPACE_ROOT", str(tmp_path / "workspaces")) + monkeypatch.setenv("ASAREE_DATASET_WORKSPACE_DIR", str(tmp_path / "workspaces")) + yield {"sessions": sessions, "user": user, "protocol": protocol, + "core_sessions": core_sessions, + "revision_id": ids[2], "experiment_id": ids[3], "dataset_id": ids[4], + "graph": graph, "tmp_path": tmp_path} + finally: + if protocol_id is not None: + async with core_sessions.begin() as db: + agents = (await db.scalars(select(Agent).where( + Agent.owner_id == ids[0], Agent.name.like(f"protocol-{protocol_id}-%"), + ))).all() + for agent in agents: + await db.delete(agent) + if protocol_id is not None: + async with sessions.begin() as db: + await delete_protocol(db, protocol_id) + if experiment_id is not None: + await delete_experiment(db, experiment_id) + if dataset_id is not None: + dataset = await db.get(RegisteredDataset, dataset_id) + if dataset is not None: + await db.delete(dataset) + await engine.dispose() + await core_engine.dispose() + + +@pytest_asyncio.fixture +async def node_play_client(node_play_context): + ctx = node_play_context + app = FastAPI() + app.include_router(protocol_api.router) + + async def current_user(): + return ctx["user"] + + async def session(): + async with ctx["sessions"]() as db: + yield db + + app.dependency_overrides[get_current_user] = current_user + app.dependency_overrides[database.get_db] = session + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + yield client + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("selection", "row_index", "question"), [(2, 2, "third"), (None, 0, "first")]) +async def test_node_play_selects_one_row_and_runs_only_target_with_private_script_input( + node_play_context, node_play_client, monkeypatch: pytest.MonkeyPatch, selection, row_index, question, +) -> None: + ctx = node_play_context + queued: list[uuid.UUID] = [] + script_rows: list[list[dict]] = [] + + async def enqueue(run_id: uuid.UUID) -> None: + async with ctx["sessions"]() as db: + assert await db.get(ProtocolRun, run_id) is not None, "enqueue must follow commit" + queued.append(run_id) + + async def mock_provider(*, run_id, available_tools, **_kwargs): + from asaree.mcp_servers import script_server + + async with ctx["core_sessions"]() as db: + run = await db.get(AgentRun, run_id) + assert run is not None + assert run.run_metadata["node_id"] == "agent" + assert f'"row_index":{row_index}' in run.input and f'"question":"{question}"' in run.input + assert "gold-" not in run.input and "reference" not in run.input + assert "Dataset context:" not in run.input + ambient_meta = run.run_metadata["ambient_meta"] + assert ambient_meta["row_inputs"][0]["values"] == {"question": question} + assert "gold-" not in json.dumps(ambient_meta) + assert "run_wired_script" in " ".join(run.agent.tool_config_data["tool_names"]) + assert any(tool["tool_name"] == "run_wired_script" for tool in available_tools) + metadata = {f"motoro.ambient.{key}": value for key, value in ambient_meta.items()} + metadata["motoro.workspace_id"] = f"_protocol_runs/{queued[-1]}" + result = json.loads(script_server.run_wired_script(ctx=SimpleNamespace( + request_context=SimpleNamespace(meta=SimpleNamespace(model_extra=metadata)) + ))) + assert result["exit_code"] == 0, result + script_rows.append(json.loads(result["stdout"])) + async with ctx["core_sessions"].begin() as db: + run = await db.get(AgentRun, run_id) + run.output = "predicted" + run.status = RunStatus.COMPLETED + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + monkeypatch.setattr(execution, "hydrate_registry", lambda: _done()) + monkeypatch.setattr(execution, "execute_run", mock_provider) + # Supply descriptors locally; never connect a real MCP fleet. + registry = SimpleNamespace(get_all_tools=lambda **_kwargs: [{ + "name": f"{execution.SCRIPT_SERVER_NAME}.run_wired_script", "tool_name": "run_wired_script", + "server": execution.SCRIPT_SERVER_NAME, "description": "Run the wired Script", "input_schema": {}, + }]) + monkeypatch.setattr(run_tools, "get_registry", lambda: registry) + monkeypatch.setattr(execution, "get_registry", lambda: registry) + async with ctx["sessions"]() as db: + protocol = await db.get(Protocol, ctx["protocol"].id) + assert protocol is not None + # Full Test Run requires one sequential chain; node Play separately + # permits an unwired Agent without executing it. + unwired_ids = {"unbound", "unbound-model"} + protocol.graph = { + "nodes": [node for node in ctx["graph"]["nodes"] if node["id"] not in unwired_ids], + "edges": [edge for edge in ctx["graph"]["edges"] if edge["target"] not in unwired_ids], + } + revision = await publish_protocol(db, protocol, owner_id=ctx["user"].id) + protocol.published_revision_id = revision.id + await db.commit() + latest = await protocol_api.create_test_run_endpoint(ctx["protocol"].id, ctx["user"], db) + replicate = await upsert_replicate( + db, experiment_id=ctx["experiment_id"], replicate_label="production-slot", + fields={"factor_values": {}}, + ) + dataset = await db.get(RegisteredDataset, ctx["dataset_id"]) + slot = await ensure_row_result( + db, experiment_id=ctx["experiment_id"], design_revision_id=replicate.design_revision_id, + protocol_revision_id=revision.id, replicate_result_id=replicate.id, + dataset_id=dataset.id, raw_sha256=dataset.raw_sha256, row_index=0, + ) + slot_id = slot.id + slot_before = {column.name: getattr(slot, column.name) for column in slot.__table__.columns} + protocol.graph = ctx["graph"] + revision = await publish_protocol(db, protocol, owner_id=ctx["user"].id) + protocol.published_revision_id = revision.id + await db.commit() + request_kwargs = {} if selection is None else {"json": {"row_index": selection}} + response = await node_play_client.post(f"/protocols/{ctx['protocol'].id}/nodes/agent/run", **request_kwargs) + assert response.status_code == 201, response.text + selected = protocol_api.ProtocolRunResponse.model_validate(response.json()) + assert selected.dataset_row == { + "dataset_id": str(ctx["dataset_id"]), + "raw_sha256": hashlib.sha256((ctx["tmp_path"] / "node-play.csv").read_bytes()).hexdigest(), + "row_index": row_index, + "columns": ["question"], + "values": {"question": question}, + } + assert selected.row_result_id is None and selected.target_node_id == "agent" + assert selected.replicate_result_id is None and selected.replicate_label is None + assert queued == [latest.id, selected.id] + await execution.run_protocol(selected.id) + async with ctx["sessions"]() as db: + run = await db.get(ProtocolRun, selected.id) + experiment = await db.get(ResearchExperiment, ctx["experiment_id"]) + assert run is not None and run.status == "completed" + assert run.dataset_row == selected.dataset_row and run.row_result_id is None + assert run.node_runs.keys() == {"agent"} + assert experiment is not None and experiment.latest_test_run_id == latest.id + assert ( + await db.scalar( + select(func.count()) + .select_from(FactorialRowResult) + .where(FactorialRowResult.dataset_id == ctx["dataset_id"]) + ) + ) == 1 + slot = await db.get(FactorialRowResult, slot_id) + assert {column.name: getattr(slot, column.name) for column in slot.__table__.columns} == slot_before + assert run.attempt_result["measurement"]["observations"][0]["value"] == "predicted" + assert script_rows == [[{"question": question}]] + + +async def _done() -> None: + return None + + +@pytest.mark.asyncio +async def test_node_play_defaults_to_zero_and_rejects_invalid_or_ineligible_selection( + node_play_context, node_play_client, monkeypatch, +): + ctx = node_play_context + queued: list[uuid.UUID] = [] + + async def enqueue(run_id: uuid.UUID) -> None: + queued.append(run_id) + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + async with ctx["sessions"]() as db: + default = await protocol_api.run_single_node_endpoint(ctx["protocol"].id, "agent", ctx["user"], db) + assert default.dataset_row["row_index"] == 0 + for selection in (3, -1, True, 1.5, "2"): + invalid = await node_play_client.post( + f"/protocols/{ctx['protocol'].id}/nodes/agent/run", json={"row_index": selection}, + ) + assert invalid.status_code == 422, invalid.text + before = await db.scalar( + select(func.count()).select_from(ProtocolRun).where(ProtocolRun.protocol_id == ctx["protocol"].id) + ) + upstream = await node_play_client.post(f"/protocols/{ctx['protocol'].id}/nodes/upstream/run") + assert upstream.status_code == 422 and "upstream input" in upstream.json()["detail"] + assert ( + await db.scalar( + select(func.count()).select_from(ProtocolRun).where(ProtocolRun.protocol_id == ctx["protocol"].id) + ) + ) == before == 1 + assert queued == [default.id] + + +@pytest.mark.asyncio +async def test_row_selection_without_target_driver_or_in_whole_mode_is_rejected( + node_play_context, node_play_client, monkeypatch, +): + ctx = node_play_context + + async def enqueue(_run_id: uuid.UUID) -> None: + pytest.fail("invalid node Play must not enqueue") + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + no_binding = await node_play_client.post( + f"/protocols/{ctx['protocol'].id}/nodes/unbound/run", json={"row_index": 0}, + ) + assert no_binding.status_code == 422 and no_binding.json()["detail"] == "no_row_driver" + async with ctx["sessions"]() as db: + protocol = await db.get(Protocol, ctx["protocol"].id) + assert protocol is not None + protocol.graph = {"nodes": [ + {"id": "agent", "type": "agent", "data": {}}, + {"id": "model", "type": "model_openai", "data": {"config": {}}}, + ], "edges": [{"id": "model-agent", "source": "model", "target": "agent", "targetHandle": "model"}]} + revision = await publish_protocol(db, protocol, owner_id=ctx["user"].id) + protocol.published_revision_id = revision.id + await db.commit() + whole = await node_play_client.post( + f"/protocols/{ctx['protocol'].id}/nodes/agent/run", json={"row_index": 0}, + ) + assert whole.status_code == 422 and whole.json()["detail"] == "no_row_driver" + assert await db.scalar( + select(func.count()).select_from(ProtocolRun).where(ProtocolRun.protocol_id == protocol.id) + ) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("invalid_source", ["columns", "changed", "missing", "other_owner"]) +async def test_node_play_validates_source_before_creating_run( + node_play_context, node_play_client, monkeypatch, invalid_source, +): + ctx = node_play_context + + async def enqueue(_run_id): + pytest.fail("invalid source must not enqueue") + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + async with ctx["sessions"].begin() as db: + if invalid_source == "columns": + # Publication validates columns too. Change the registered source + # afterward so node Play must revalidate the published binding. + raw = b"replacement,reference\nfirst,gold-a\nsecond,gold-b\nthird,gold-c\n" + (ctx["tmp_path"] / "node-play.csv").write_bytes(raw) + dataset = await db.get(RegisteredDataset, ctx["dataset_id"]) + dataset.raw_sha256 = hashlib.sha256(raw).hexdigest() + elif invalid_source == "other_owner": + other = await get_user_by_email(db, "other@test.com") + assert other is not None, "run scripts/seed_row_test_users.py before row database tests" + dataset = await db.get(RegisteredDataset, ctx["dataset_id"]) + dataset.owner_id = other.id + elif invalid_source == "changed": + (ctx["tmp_path"] / "node-play.csv").write_text("question,reference\nchanged,private\n") + else: + (ctx["tmp_path"] / "node-play.csv").unlink() + response = await node_play_client.post( + f"/protocols/{ctx['protocol'].id}/nodes/agent/run", json={"row_index": 0}, + ) + assert response.status_code == 422, response.text + if invalid_source == "columns": + assert "dataset_row.invalid_columns" in response.json()["detail"] + async with ctx["sessions"]() as db: + assert await db.scalar( + select(func.count()).select_from(ProtocolRun).where(ProtocolRun.protocol_id == ctx["protocol"].id) + ) == 0 diff --git a/tests/test_dataset_row_prompt_preview.py b/tests/test_dataset_row_prompt_preview.py new file mode 100644 index 00000000..ec948989 --- /dev/null +++ b/tests/test_dataset_row_prompt_preview.py @@ -0,0 +1,133 @@ +"""Read-only draft row preview uses the same authorized runtime assembler.""" +import hashlib +import uuid +from contextlib import asynccontextmanager +from copy import deepcopy +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient + +from asaree.api import protocols as api +from asaree.deps import get_current_user +from asaree.models.database import get_db +from asaree.services import protocol_execution as execution +from asaree.services.dataset_row_csv import DatasetRowCsvError, project_row, read_row_source +from asaree.services.dataset_row_inputs import DatasetRowInputError + + +@pytest.fixture +def draft(tmp_path, monkeypatch): + source = tmp_path / 'original.csv' + source.write_text('question,reference\nfirst,GOLD-1\nsecond,GOLD-2\nthird,GOLD-3\n') + registration = SimpleNamespace(id=uuid.uuid4(), raw_path=str(source), + raw_sha256=hashlib.sha256(source.read_bytes()).hexdigest()) + graph = {'nodes': [ + {'id': 'dataset', 'type': 'dataset', 'data': {'config': {'dataset_id': str(registration.id)}}}, + {'id': 'agent', 'type': 'agent', 'data': {'config': {'prompt': 'Predict.'}}}, + ], 'edges': [{'id': 'edge', 'source': 'dataset', 'target': 'agent', 'targetHandle': 'dataset', + 'data': {'dataset_input': {'mode': 'per_row', 'columns': ['question']}}}]} + class Session: + async def execute(self, statement): + params = statement.compile().params + assert registration.id in params.values() + return SimpleNamespace(one_or_none=lambda: registration) + @asynccontextmanager + async def session(): + yield Session() + monkeypatch.setattr(execution, 'get_session', session) + async def forbidden(*args, **kwargs): + pytest.fail('preview must not execute, seed, enqueue or fetch full metadata') + for name in ('execute_run', 'fetch_owned_registration', 'create_protocol_run', 'seed_cell_workspace'): + monkeypatch.setattr(execution, name, forbidden) + return graph, registration, source + + +@pytest.mark.parametrize('index', [None, 2]) +async def test_default_and_selected_preview_equal_runtime_without_source_leak(draft, index): + graph, registration, source = draft + snapshots = [] + text = await execution.preview_node_prompt(graph, 'agent', owner_id=uuid.uuid4(), + row_index=index, dataset_row_out=snapshots) + original = read_row_source(dataset_id=str(registration.id), raw_path=str(source), + raw_sha256=registration.raw_sha256) + view = project_row(original, row_index=index or 0, columns=['question']) + assert snapshots == [view] + assert text == execution._build_user_input(graph['nodes'][1], graph, {}, row_input_context=view) + assert 'GOLD' not in text and 'reference' not in text + assert ('third' if index else 'first') in text + assert list(source.parent.iterdir()) == [source] + + +async def test_grader_explicit_columns_and_unwired_node(draft): + graph, _, _ = draft + graph['edges'][0]['data']['dataset_input']['columns'].append('reference') + text = await execution.preview_node_prompt(graph, 'agent', owner_id=uuid.uuid4(), row_index=2) + assert 'GOLD-3' in text and 'GOLD-1' not in text + graph['edges'] = [] + assert await execution.preview_node_prompt(graph, 'agent', owner_id=uuid.uuid4()) == 'Predict.' + with pytest.raises(execution.ProtocolValidationError, match='no_row_driver'): + await execution.preview_node_prompt(graph, 'agent', owner_id=uuid.uuid4(), row_index=2) + + +@pytest.mark.parametrize('index', [-1, 3, True, 1.5]) +async def test_invalid_bounds(draft, index): + graph, _, _ = draft + with pytest.raises(DatasetRowCsvError): + await execution.preview_node_prompt(graph, 'agent', owner_id=uuid.uuid4(), row_index=index) + + +async def test_invalid_columns_hash_empty_and_topology(draft): + graph, reg, source = draft + graph['edges'][0]['data']['dataset_input']['columns'] = ['missing'] + with pytest.raises(DatasetRowCsvError, match='invalid_columns'): + await execution.preview_node_prompt(graph, 'agent', owner_id=uuid.uuid4()) + graph['edges'][0]['data']['dataset_input']['columns'] = ['question'] + source.write_text('question,reference\n') + with pytest.raises(DatasetRowCsvError, match='source_hash_mismatch'): + await execution.preview_node_prompt(graph, 'agent', owner_id=uuid.uuid4()) + reg.raw_sha256 = hashlib.sha256(source.read_bytes()).hexdigest() + with pytest.raises(DatasetRowCsvError, match='empty_source'): + await execution.preview_node_prompt(graph, 'agent', owner_id=uuid.uuid4()) + graph['nodes'].append({'id': 'other', 'type': 'dataset', 'data': {'config': {'dataset_id': str(uuid.uuid4())}}}) + edge = deepcopy(graph['edges'][0]) + edge.update(id='other-edge', source='other') + graph['edges'].append(edge) + with pytest.raises(DatasetRowInputError, match='multiple_drivers'): + await execution.preview_node_prompt(graph, 'agent', owner_id=uuid.uuid4()) + graph['nodes'][-1]['data']['config']['dataset_id'] = str(reg.id) + graph['edges'][-1]['data'] = {} + with pytest.raises(DatasetRowInputError, match='mixed_driver_modes'): + await execution.preview_node_prompt(graph, 'agent', owner_id=uuid.uuid4()) + + +async def test_route_uses_supplied_draft_or_stored_draft_and_returns_identity(draft, monkeypatch): + graph, _, _ = draft + protocol_id = uuid.uuid4() + user = SimpleNamespace(id=uuid.uuid4()) + stored = deepcopy(graph) + stored['nodes'][1]['data']['config']['prompt'] = 'Stored draft.' + async def owned(*args): + return SimpleNamespace(graph=stored, experiment_id=None) + monkeypatch.setattr(api, '_get_owned_protocol', owned) + app = FastAPI() + app.include_router(api.router) + async def current_user(): + return user + + async def database_session(): + yield None + + app.dependency_overrides[get_current_user] = current_user + app.dependency_overrides[get_db] = database_session + async with AsyncClient(transport=ASGITransport(app=app), base_url='http://test') as client: + url = f'/protocols/{protocol_id}/nodes/agent/prompt-preview' + result = await client.post(url, json={'graph': graph, 'row_index': 2}) + assert result.status_code == 200, result.text + assert result.json()['dataset_row']['row_index'] == 2 + assert result.json()['text'].startswith('Predict.') + result = await client.post(url, json={}) + assert result.json()['text'].startswith('Stored draft.') + assert (await client.post(url, json={'row_index': True})).status_code == 422 + assert (await client.post(url, json={'row_index': 8})).status_code == 422 diff --git a/tests/test_dataset_row_prompts.py b/tests/test_dataset_row_prompts.py new file mode 100644 index 00000000..dee46b7f --- /dev/null +++ b/tests/test_dataset_row_prompts.py @@ -0,0 +1,142 @@ +from __future__ import annotations + +import hashlib +import json +import uuid +from pathlib import Path + +from asaree.services.dataset_row_csv import project_row, read_row_source +from asaree.services.protocol_execution import _build_user_input + +DATASET_ID = str(uuid.UUID("12345678-1234-5678-1234-567812345678")) + + +def _fixture(tmp_path: Path, *, columns: list[str]) -> tuple[dict, dict, dict]: + raw = ( + "question,reference,other\n" + "What is 2+2?,same-row-reference,first-row-secret\n" + "Next?,another-row-secret,second-row-secret\n" + ) + source_path = tmp_path / "source.csv" + source_path.write_text(raw, encoding="utf-8") + source = read_row_source( + dataset_id=DATASET_ID, + raw_path=str(source_path), + raw_sha256=hashlib.sha256(source_path.read_bytes()).hexdigest(), + ) + view = project_row(source, row_index=0, columns=columns) + agent = {"id": "agent", "type": "agent", "data": {"label": "Agent", "config": {"prompt": "Authored instructions."}}} + dataset = { + "id": "dataset", + "type": "dataset", + "data": { + "label": "Dataset", + "config": { + "dataset_id": DATASET_ID, + "dataset_name": "private-registration-name", + "description": "FULL DICTIONARY MUST NOT APPEAR", + "target_column": "reference", + "split_state": "test", + "dictionary_available": True, + }, + }, + } + graph = { + "nodes": [agent, dataset], + "edges": [ + { + "id": "driver", + "source": "dataset", + "target": "agent", + "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": columns}}, + } + ], + } + return agent, graph, view + + +def _assemble(agent: dict, graph: dict, view: dict) -> str: + return _build_user_input(agent, graph, {}, row_input_context=view) + + +def test_prediction_prompt_contains_only_projected_question_row(tmp_path: Path) -> None: + agent, graph, view = _fixture(tmp_path, columns=["question"]) + prompt = _assemble(agent, graph, view) + block = prompt.split("Dataset row input:\n", 1)[1] + row = json.loads(block) + assert row == view + assert "What is 2+2?" in prompt + assert "same-row-reference" not in prompt + assert "first-row-secret" not in prompt + assert "second-row-secret" not in prompt + assert "FULL DICTIONARY MUST NOT APPEAR" not in prompt + assert "private-registration-name" not in prompt + assert "target=reference" not in prompt + assert "split_state" not in prompt + assert "source.csv" not in prompt + + +def test_grading_prompt_contains_reference_for_the_same_projected_row(tmp_path: Path) -> None: + agent, graph, view = _fixture(tmp_path, columns=["question", "reference"]) + prompt = _assemble(agent, graph, view) + row = json.loads(prompt.split("Dataset row input:\n", 1)[1]) + assert row["values"] == {"question": "What is 2+2?", "reference": "same-row-reference"} + assert "another-row-secret" not in prompt + + +def test_row_json_preserves_empty_unicode_and_requested_column_order(tmp_path: Path) -> None: + raw = "question,reference,note\n雪,,café\n".encode() + source_path = tmp_path / "unicode.csv" + source_path.write_bytes(raw) + source = read_row_source( + dataset_id=DATASET_ID, raw_path=str(source_path), raw_sha256=hashlib.sha256(raw).hexdigest() + ) + view = project_row(source, row_index=0, columns=["note", "reference", "question"]) + agent = {"id": "agent", "type": "agent", "data": {"label": "Agent", "config": {"prompt": "Authored instructions."}}} + dataset = { + "id": "dataset", + "type": "dataset", + "data": {"label": "Dataset", "config": {"dataset_id": DATASET_ID, "dataset_name": "cohort"}}, + } + graph = { + "nodes": [agent, dataset], + "edges": [ + { + "id": "driver", + "source": "dataset", + "target": "agent", + "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": view["columns"]}}, + } + ], + } + prompt = _assemble(agent, graph, view) + encoded = prompt.split("Dataset row input:\n", 1)[1] + assert "café" in encoded and "雪" in encoded + row = json.loads(encoded) + assert row["columns"] == ["note", "reference", "question"] + assert row["values"] == {"note": "café", "reference": "", "question": "雪"} + + +def test_agent_without_driver_edge_gets_no_row_values_or_catalog_entry(tmp_path: Path) -> None: + agent, graph, view = _fixture(tmp_path, columns=["question"]) + graph["edges"] = [] + prompt = _build_user_input(agent, graph, {}, row_input_context=view) + assert "Dataset row input:" not in prompt + assert "Available datasets:" not in prompt + assert "What is 2+2?" not in prompt + + +def test_whole_context_and_legacy_prompt_remain_compatible() -> None: + agent = {"id": "agent", "type": "agent", "data": {"label": "Agent", "config": {"prompt": "Authored instructions."}}} + assert _build_user_input(agent, {"nodes": [agent], "edges": []}, {}) == "Authored instructions." + dataset = {"id": "dataset", "type": "dataset", "data": {"label": "Dataset", "config": {"dataset_name": "cohort"}}} + graph = { + "nodes": [agent, dataset], + "edges": [{"id": "d", "source": "dataset", "target": "agent", "targetHandle": "dataset"}], + } + prompt = _build_user_input(agent, graph, {}, experiment_id=uuid.uuid4(), effective_cell_label="cell") + assert "Dataset context:" in prompt and "open_workspace()" in prompt + assert "Available datasets:\n- cohort" in prompt + assert "Dataset row input:" not in prompt diff --git a/tests/test_dataset_row_results_projection.py b/tests/test_dataset_row_results_projection.py new file mode 100644 index 00000000..f98b25ac --- /dev/null +++ b/tests/test_dataset_row_results_projection.py @@ -0,0 +1,393 @@ +"""PostgreSQL Results projection coverage for immutable dataset row slots.""" + +from __future__ import annotations + +import hashlib +import os +import uuid +from collections.abc import AsyncIterator +from pathlib import Path + +import pytest +import pytest_asyncio +from fastapi import HTTPException +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +import asaree.api.protocols as protocol_api +from asaree.api.experiments import get_experiment_run_results_endpoint +from asaree.api.protocols import CellRunBatchRequest, create_cell_runs_endpoint +from asaree.models.dataset import RegisteredDataset +from asaree.models.factorial_cell import FactorialCell +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol_run import ProtocolRun +from asaree.services.design_revisions import get_or_create_current +from asaree.services.experiment_run_results import RunResultsProjectionError, summarize_experiment_run_results +from asaree.services.experiments import create_experiment +from asaree.services.factorial_row_results import claim_row_attempt +from asaree.services.protocol_revisions import publish_protocol +from asaree.services.protocols import create_protocol +from asaree.services.users import get_user_by_email + + +@pytest_asyncio.fixture +async def row_results_setup(tmp_path: Path) -> AsyncIterator[tuple[AsyncSession, dict]]: + """Use the seeded application user and a real hash-verified CSV source.""" + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + try: + async with engine.begin() as connection: + sessions = async_sessionmaker(bind=connection, expire_on_commit=False) + async with sessions() as db: + user = await get_user_by_email(db, "test@test.com") + assert user is not None, "run scripts/seed_row_test_users.py before row database tests" + raw = b"question,answer\n" + b"prompt,response\n" * 15 + source_path = tmp_path / "results.csv" + source_path.write_bytes(raw) + dataset = RegisteredDataset( + name=f"results-{uuid.uuid4()}", + owner_id=user.id, + raw_path=str(source_path), + raw_sha256=hashlib.sha256(raw).hexdigest(), + ) + db.add(dataset) + await db.flush() + design_spec = { + "factors": [], + "replicates": 1, + "metrics": [ + {"id": "reported", "name": "Report", "kind": "custom", "valueType": "opaque"} + ], + } + experiment = await create_experiment( + db, + name=f"row-results-{uuid.uuid4()}", + owner_id=user.id, + design_spec=design_spec, + measurement_plan={ + "metrics": [{"id": "reported", "name": "Report"}], + "producers": [ + { + "id": "report-output", + "producer_id": "asaree.agent_output", + "kind": "reported", + "outputs": {"value": "reported"}, + "config": {"agent_node_id": "agent"}, + } + ], + "inputs": [], + }, + ) + design_revision = await get_or_create_current(db, experiment.id, design_spec=design_spec) + parents = [] + cells = [] + for arm in ("left", "right"): + cell = FactorialCell( + experiment_id=experiment.id, + design_revision_id=design_revision.id, + cell_label=f"arm={arm}", + factor_values={"arm": arm}, + ) + db.add(cell) + await db.flush() + parent = FactorialReplicateResult( + cell_id=cell.id, + replicate_number=1, + replicate_label="duplicate-slot-label", + ) + db.add(parent) + await db.flush() + cells.append(cell) + parents.append(parent) + graph = { + "nodes": [ + { + "id": "dataset", + "type": "dataset", + "data": {"config": {"dataset_id": str(dataset.id), "dataset_name": dataset.name}}, + }, + {"id": "agent", "type": "agent", "data": {}}, + {"id": "model", "type": "model_openai", "data": {"config": {}}}, + ], + "edges": [ + { + "id": "dataset-agent", + "source": "dataset", + "target": "agent", + "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["question"]}}, + }, + {"id": "model-agent", "source": "model", "target": "agent", "targetHandle": "model"}, + ], + } + protocol = await create_protocol( + db, + name=f"row-results-protocol-{uuid.uuid4()}", + owner_id=user.id, + experiment_id=experiment.id, + graph=graph, + ) + publication = await publish_protocol(db, protocol, owner_id=user.id) + protocol.published_revision_id = publication.id + await db.flush() + context = { + "user": user, + "experiment": experiment, + "dataset": dataset, + "source_path": source_path, + "protocol": protocol, + "publication": publication, + "design_revision": design_revision, + "parents": parents, + "cells": cells, + "design_spec": design_spec, + } + yield db, context + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_row_projection_counts_attempts_and_preserves_slots(row_results_setup, monkeypatch) -> None: + db, ctx = row_results_setup + + async def no_enqueue(_run_id: uuid.UUID) -> None: + return None + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", no_enqueue) + planned = await create_cell_runs_endpoint( + ctx["protocol"].id, ctx["user"], db, CellRunBatchRequest() + ) + assert len(planned.row_result_ids) == 30 + runs = [await db.get(ProtocolRun, run_id) for run_id in planned.protocol_run_ids] + slots = [await db.get(FactorialRowResult, row_id) for row_id in planned.row_result_ids] + assert all(run is not None for run in runs) and all(slot is not None for slot in slots) + + for index, (run, slot) in enumerate(zip(runs, slots, strict=True)): + assert run is not None and slot is not None + workspace_id = f"_protocol_runs/{run.id}" + workspace = ctx["source_path"].parent / "workspaces" / str(run.id) + workspace.mkdir(parents=True) + (workspace / "row.txt").write_text(f"row attempt {index}", encoding="utf-8") + slot.workspace_id = workspace_id + if index == 1: + run.status = "running" + elif index == 2: + run.status = "finalizing" + elif index == 29: + run.status = "failed" + run.error = "failed attempt" + elif index >= 3: + run.status = "completed" + metric_status = "unavailable" if index == 3 else "measured" + value = None if index == 4 else {"opaque": [index, None]} if index == 5 else index + run.attempt_result = { + "measurement": { + "observations": [ + {"metric_id": "reported", "status": metric_status, "value": value} + ] + } + } + if metric_status == "measured": + slot.metric_values = {"Report": value} + + results = await summarize_experiment_run_results( + db, experiment_id=ctx["experiment"].id, design_spec=ctx["design_spec"] + ) + summary = results["row_summary"] + assert results["consumption_mode"] == "per_row" + assert results["cells"] == [] and results["replicates"] == [] + assert results["primary_metric"] is None and results["primary_metric_direction"] is None + assert len(results["row_results"]) == 30 + assert summary == { + "cell_count": 2, + "parent_replicate_count": 2, + "row_count": 15, + "expected": 30, + "planned": 30, + "pending": 1, + "running": 2, + "completed": 26, + "failed": 1, + "cancelled": 0, + "scored": 25, + "missing_reported": 1, + "metric_coverage": { + "reported": {"measured": 25, "unavailable": 5, "failed": 0, "cancelled": 0, "other": 0} + }, + } + duplicate_label_rows = [row for row in results["row_results"] if row["replicate_label"] == "duplicate-slot-label"] + assert len({row["replicate_result_id"] for row in duplicate_label_rows}) == 2 + # Planning claims slots in UUID order; Results orders by cell/replicate/row. + # Match persisted slot identities rather than assuming those orders agree. + rows_by_id = {row["row_result_id"]: row for row in results["row_results"]} + assert set(rows_by_id) == {str(slot.id) for slot in slots} + null_row = rows_by_id[str(slots[4].id)] + opaque_row = rows_by_id[str(slots[5].id)] + assert null_row["metric_values"] == {"Report": None} + assert opaque_row["metric_values"] == {"Report": {"opaque": [5, None]}} + assert null_row["latest_attempt"]["run_id"] == str(runs[4].id) + assert null_row["latest_attempt"]["workspace_id"] == null_row["workspace_id"] + + # Retry one failed row through the production claim API; history is immutable + # and only the run still attached to the stable slot is current. + failed_run = runs[29] + failed_slot = slots[29] + assert failed_run is not None and failed_slot is not None + retry = await claim_row_attempt( + db, + row_result_id=failed_slot.id, + expected_run_id=failed_run.id, + create_kwargs={ + "protocol_id": ctx["protocol"].id, + "owner_id": ctx["user"].id, + "replicate_label": "duplicate-slot-label", + "factor_values": {"arm": "right"}, + "replicate_result_id": failed_slot.replicate_result_id, + "design_revision_id": ctx["design_revision"].id, + "protocol_revision_id": ctx["publication"].id, + "row_result_id": failed_slot.id, + "dataset_row": failed_run.dataset_row, + }, + ) + assert retry is not None + retried = await summarize_experiment_run_results( + db, experiment_id=ctx["experiment"].id, design_spec=ctx["design_spec"] + ) + retried_slot = next(row for row in retried["row_results"] if row["row_result_id"] == str(failed_slot.id)) + assert [attempt["current"] for attempt in retried_slot["attempts"]] == [False, True] + assert retried_slot["latest_attempt"]["run_id"] == str(retry.id) + assert retried["row_summary"]["failed"] == 0 + assert retried["row_summary"]["pending"] == 2 + + # A removed source never makes persisted slot history disappear. + ctx["source_path"].unlink() + source_missing = await summarize_experiment_run_results( + db, experiment_id=ctx["experiment"].id, design_spec=ctx["design_spec"] + ) + assert source_missing["row_summary"]["row_count"] is None + assert source_missing["row_summary"]["expected"] is None + assert len(source_missing["row_results"]) == 30 + + no_metrics = await summarize_experiment_run_results( + db, experiment_id=ctx["experiment"].id, design_spec={"metrics": []} + ) + assert no_metrics["row_summary"]["scored"] == 0 + assert no_metrics["row_summary"]["metric_coverage"] == {} + + +@pytest.mark.asyncio +async def test_row_results_selectors_forecast_and_whole_scorecard(row_results_setup, monkeypatch) -> None: + db, ctx = row_results_setup + # A protocol revision and design revision are scoped to this experiment. + other_user = await get_user_by_email(db, "other@test.com") + assert other_user is not None, "run scripts/seed_row_test_users.py before row database tests" + other = await create_protocol( + db, + name=f"foreign-results-{uuid.uuid4()}", + owner_id=other_user.id, + graph={"nodes": [], "edges": []}, + ) + foreign_experiment = await create_experiment( + db, + name=f"foreign-results-experiment-{uuid.uuid4()}", + owner_id=other_user.id, + ) + other.experiment_id = foreign_experiment.id + foreign_publication = await publish_protocol(db, other) + with pytest.raises(RunResultsProjectionError, match="protocol_revision_not_found"): + await summarize_experiment_run_results( + db, + experiment_id=ctx["experiment"].id, + design_spec=ctx["design_spec"], + protocol_revision_id=foreign_publication.id, + ) + with pytest.raises(HTTPException) as foreign_publication_response: + await get_experiment_run_results_endpoint( + ctx["experiment"].id, + ctx["user"], + db, + protocol_id=ctx["protocol"].id, + protocol_revision_id=foreign_publication.id, + ) + assert foreign_publication_response.value.status_code == 404 + foreign_design = await get_or_create_current(db, foreign_experiment.id) + with pytest.raises(RunResultsProjectionError, match="design_revision_not_found"): + await summarize_experiment_run_results( + db, + experiment_id=ctx["experiment"].id, + design_spec=ctx["design_spec"], + design_revision_id=foreign_design.id, + ) + with pytest.raises(HTTPException) as foreign_design_response: + await get_experiment_run_results_endpoint( + ctx["experiment"].id, + ctx["user"], + db, + design_revision_id=foreign_design.id, + ) + assert foreign_design_response.value.status_code == 404 + + await create_protocol( + db, + name=f"ambiguous-{uuid.uuid4()}", + owner_id=ctx["user"].id, + experiment_id=ctx["experiment"].id, + ) + with pytest.raises(HTTPException) as ambiguous: + await get_experiment_run_results_endpoint(ctx["experiment"].id, ctx["user"], db) + assert ambiguous.value.status_code == 422 + assert ambiguous.value.detail == "ambiguous_protocol" + + # The current draft is irrelevant: the selected publication controls mode. + ctx["protocol"].graph = {"nodes": [], "edges": []} + per_row = await summarize_experiment_run_results( + db, + experiment_id=ctx["experiment"].id, + design_spec=ctx["design_spec"], + protocol_id=ctx["protocol"].id, + protocol_revision_id=ctx["publication"].id, + ) + assert per_row["consumption_mode"] == "per_row" + + async def no_enqueue(_run_id: uuid.UUID) -> None: + return None + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", no_enqueue) + await create_cell_runs_endpoint( + ctx["protocol"].id, ctx["user"], db, CellRunBatchRequest() + ) + + # Publishing a whole-dataset canvas exposes the legacy scorecard shape and + # row attempts stay out of its label-based history. + ctx["protocol"].graph = {"nodes": [{"id": "agent", "type": "agent", "data": {}}], "edges": []} + whole_publication = await publish_protocol(db, ctx["protocol"], owner_id=ctx["user"].id) + whole = await get_experiment_run_results_endpoint( + ctx["experiment"].id, ctx["user"], db, protocol_id=ctx["protocol"].id + ) + assert whole.consumption_mode == "whole_dataset" + assert whole.row_summary is None and whole.row_results == [] + assert whole.cells and whole.replicates + assert whole_publication.id == ctx["protocol"].published_revision_id + + +@pytest.mark.asyncio +async def test_preplanning_forecast_uses_verified_source(row_results_setup) -> None: + db, ctx = row_results_setup + # Remove the planned cells, but retain the declaration and verified source. + for cell in ctx["cells"]: + await db.delete(cell) + ctx["experiment"].design_spec = { + "factors": [{"name": "arm", "levels": ["left", "right"]}], + "replicates": 3, + "metrics": [], + } + ctx["design_revision"].design_spec = ctx["experiment"].design_spec + forecast = await summarize_experiment_run_results( + db, + experiment_id=ctx["experiment"].id, + design_spec=ctx["experiment"].design_spec, + ) + assert forecast["row_summary"]["planned"] == 0 + assert forecast["row_summary"]["cell_count"] == 0 + assert forecast["row_summary"]["parent_replicate_count"] == 0 + assert forecast["row_summary"]["expected"] == 90 diff --git a/tests/test_dataset_row_retry.py b/tests/test_dataset_row_retry.py new file mode 100644 index 00000000..06a28bd4 --- /dev/null +++ b/tests/test_dataset_row_retry.py @@ -0,0 +1,405 @@ +"""PostgreSQL integration coverage for explicitly retrying row slots.""" + +from __future__ import annotations + +import asyncio +import hashlib +import os +import uuid +from collections.abc import AsyncIterator +from copy import deepcopy +from pathlib import Path + +import pytest +import pytest_asyncio +from fastapi import HTTPException +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +import asaree.api.protocols as protocol_api +from asaree.api.protocols import CellRunBatchRequest, create_cell_runs_endpoint +from asaree.models.dataset import RegisteredDataset +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol import Protocol +from asaree.models.protocol_run import ProtocolRun +from asaree.services.experiments import create_experiment +from asaree.services.factorial_cells import upsert_replicate +from asaree.services.protocol_revisions import publish_protocol +from asaree.services.protocols import create_protocol +from asaree.services.users import get_user_by_email + + +@pytest_asyncio.fixture +async def retry_setup(tmp_path: Path) -> AsyncIterator[dict]: + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + sessions = async_sessionmaker(engine, expire_on_commit=False) + try: + async with sessions.begin() as db: + user = await get_user_by_email(db, "test@test.com") + assert user is not None, "run scripts/seed_row_test_users.py before row database tests" + raw = b"question,answer\n" + b"prompt,response\n" * 30 + raw_path = tmp_path / "retry.csv" + raw_path.write_bytes(raw) + dataset_id = uuid.uuid4() + dataset = RegisteredDataset( + id=dataset_id, + name=f"retry-{dataset_id}", + owner_id=user.id, + raw_path=str(raw_path), + raw_sha256=hashlib.sha256(raw).hexdigest(), + ) + db.add(dataset) + experiment = await create_experiment( + db, + name=f"row-retry-{uuid.uuid4()}", + owner_id=user.id, + design_spec={"factors": [], "metrics": []}, + ) + parent = await upsert_replicate( + db, + experiment_id=experiment.id, + replicate_label="retry-parent", + fields={"factor_values": {}}, + ) + graph = { + "nodes": [ + { + "id": "dataset", + "type": "dataset", + "data": {"config": {"dataset_id": str(dataset_id), "dataset_name": dataset.name}}, + }, + {"id": "agent", "type": "agent", "data": {}}, + {"id": "model", "type": "model_openai", "data": {"config": {}}}, + ], + "edges": [ + { + "id": "dataset-agent", + "source": "dataset", + "target": "agent", + "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["question"]}}, + }, + {"id": "model-agent", "source": "model", "target": "agent", "targetHandle": "model"}, + ], + } + protocol = await create_protocol( + db, + name=f"row-retry-{uuid.uuid4()}", + owner_id=user.id, + experiment_id=experiment.id, + graph=graph, + ) + revision = await publish_protocol(db, protocol, owner_id=user.id) + protocol.published_revision_id = revision.id + identifiers = user.id, protocol.id, revision.id, parent.id + async with sessions() as db: + user = await get_user_by_email(db, "test@test.com") + assert user is not None and user.id == identifiers[0] + protocol = await db.get(Protocol, identifiers[1]) + parent = await db.get(FactorialReplicateResult, identifiers[3]) + assert protocol is not None and parent is not None + yield { + "engine": engine, + "sessions": sessions, + "db": db, + "user": user, + "protocol": protocol, + "revision_id": identifiers[2], + "revision_graph": deepcopy(protocol.graph), + "parent": parent, + "dataset": dataset, + "raw_path": raw_path, + } + finally: + await engine.dispose() + + +async def _enqueue_rows(ctx: dict, monkeypatch: pytest.MonkeyPatch) -> list[uuid.UUID]: + queued: list[uuid.UUID] = [] + + async def enqueue(run_id: uuid.UUID) -> None: + queued.append(run_id) + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + result = await create_cell_runs_endpoint(ctx["protocol"].id, ctx["user"], ctx["db"]) + assert len(result.protocol_run_ids) == len(result.row_result_ids) == 30 + return result.row_result_ids + + +async def _latest(ctx: dict, slot_id: uuid.UUID) -> ProtocolRun: + slot = await ctx["db"].get(FactorialRowResult, slot_id) + assert slot is not None and slot.run_id is not None + run = await ctx["db"].get(ProtocolRun, slot.run_id) + assert run is not None + return run + + +@pytest.mark.asyncio +async def test_failed_and_cancelled_slots_retry_as_new_attempts(retry_setup, monkeypatch) -> None: + ctx = retry_setup + slot_ids = await _enqueue_rows(ctx, monkeypatch) + previous = [await _latest(ctx, slot_id) for slot_id in slot_ids[:2]] + previous[0].status = "failed" + previous[1].status = "cancelled" + previous[0].attempt_result = {"immutable": "failed facts"} + previous[1].attempt_result = {"immutable": "cancelled facts"} + parent_before = await ctx["db"].get(FactorialReplicateResult, ctx["parent"].id) + parent_run_id = parent_before.run_id if parent_before else None + + result = await create_cell_runs_endpoint( + ctx["protocol"].id, + ctx["user"], + ctx["db"], + CellRunBatchRequest(retry_row_result_ids=slot_ids[:2]), + ) + assert len(result.protocol_run_ids) == 2 + assert result.row_result_ids == slot_ids[:2] + assert result.replicate_labels == ["retry-parent", "retry-parent"] + assert result.skipped == 0 + attempts = [await ctx["db"].get(ProtocolRun, run_id) for run_id in result.protocol_run_ids] + assert all(run is not None for run in attempts) + for old, new in zip(previous, attempts, strict=True): + assert new is not None + assert new.row_result_id == old.row_result_id + assert new.replicate_result_id == old.replicate_result_id + assert new.replicate_label == old.replicate_label + assert new.replicate_label == "retry-parent" + assert new.factor_values == old.factor_values == {} + assert new.design_revision_id == old.design_revision_id + assert new.protocol_revision_id == old.protocol_revision_id + assert new.dataset_row == old.dataset_row + assert old.attempt_result in ({"immutable": "failed facts"}, {"immutable": "cancelled facts"}) + assert (await ctx["db"].get(FactorialReplicateResult, ctx["parent"].id)).run_id == parent_run_id + assert await ctx["db"].scalar( + select(func.count()).select_from(ProtocolRun).where(ProtocolRun.row_result_id.in_(slot_ids[:2])) + ) == 4 + + +@pytest.mark.asyncio +async def test_retry_one_of_thirty_adds_one_attempt_without_parent_mutation(retry_setup, monkeypatch) -> None: + ctx = retry_setup + slot_ids = await _enqueue_rows(ctx, monkeypatch) + previous = await _latest(ctx, slot_ids[0]) + previous.status = "failed" + previous.attempt_result = {"error": "original"} + row_before = await ctx["db"].get(FactorialRowResult, slot_ids[0]) + row_before.workspace_id = "previous-attempt-workspace" + parent = await ctx["db"].get(FactorialReplicateResult, ctx["parent"].id) + parent_snapshot = (parent.run_id, parent.metric_values, parent.artifacts, parent.replicate_number) + + result = await create_cell_runs_endpoint( + ctx["protocol"].id, + ctx["user"], + ctx["db"], + CellRunBatchRequest(retry_row_result_ids=[slot_ids[0]]), + ) + assert len(result.protocol_run_ids) == 1 + assert result.row_result_ids == [slot_ids[0]] + assert await ctx["db"].scalar( + select(func.count()).select_from(ProtocolRun).where( + ProtocolRun.protocol_id == ctx["protocol"].id, + ProtocolRun.is_test_run.is_(False), + ) + ) == 31 + parent = await ctx["db"].get(FactorialReplicateResult, ctx["parent"].id) + assert (parent.run_id, parent.metric_values, parent.artifacts, parent.replicate_number) == parent_snapshot + row = await ctx["db"].get(FactorialRowResult, slot_ids[0]) + assert row is not None and row.run_id == result.protocol_run_ids[0] + assert row.workspace_id is None + assert previous.attempt_result == {"error": "original"} + + +@pytest.mark.asyncio +async def test_success_unavailable_active_and_duplicate_targets_reject(retry_setup, monkeypatch) -> None: + ctx = retry_setup + slot_ids = await _enqueue_rows(ctx, monkeypatch) + statuses = ["completed", "completed", "pending", "running", "finalizing"] + for slot_id, status in zip(slot_ids[:5], statuses, strict=True): + run = await _latest(ctx, slot_id) + run.status = status + unavailable = await _latest(ctx, slot_ids[1]) + unavailable.attempt_result = {"measurement": {"observations": [{"status": "unavailable"}]}} + for selected in ([slot_ids[0]], [slot_ids[1]], [slot_ids[2]], [slot_ids[3]], [slot_ids[4]]): + with pytest.raises(HTTPException) as exc: + await create_cell_runs_endpoint( + ctx["protocol"].id, + ctx["user"], + ctx["db"], + CellRunBatchRequest(retry_row_result_ids=selected), + ) + assert exc.value.status_code == 422 + assert exc.value.detail == "invalid_retry_target" + with pytest.raises(HTTPException) as exc: + await create_cell_runs_endpoint( + ctx["protocol"].id, + ctx["user"], + ctx["db"], + CellRunBatchRequest(retry_row_result_ids=[slot_ids[0], slot_ids[0]]), + ) + assert exc.value.status_code == 422 and exc.value.detail == "invalid_retry_selection" + + for body in ( + CellRunBatchRequest(retry_row_result_ids=[]), + CellRunBatchRequest(retry_row_result_ids=[slot_ids[5]], replicate_labels=[]), + CellRunBatchRequest(retry_row_result_ids=[slot_ids[5]], rerun_replicate_labels=["retry-parent"]), + ): + with pytest.raises(HTTPException) as exc: + await create_cell_runs_endpoint(ctx["protocol"].id, ctx["user"], ctx["db"], body) + assert exc.value.status_code == 422 and exc.value.detail == "invalid_retry_selection" + + +@pytest.mark.asyncio +async def test_mixed_valid_and_invalid_targets_are_atomic(retry_setup, monkeypatch) -> None: + ctx = retry_setup + slot_ids = await _enqueue_rows(ctx, monkeypatch) + first, second = [await _latest(ctx, slot_id) for slot_id in slot_ids[:2]] + first.status = "failed" + second.status = "completed" + before = await ctx["db"].scalar(select(func.count()).select_from(ProtocolRun)) + with pytest.raises(HTTPException) as exc: + await create_cell_runs_endpoint( + ctx["protocol"].id, + ctx["user"], + ctx["db"], + CellRunBatchRequest(retry_row_result_ids=slot_ids[:2]), + ) + assert exc.value.status_code == 422 and exc.value.detail == "invalid_retry_target" + assert await ctx["db"].scalar(select(func.count()).select_from(ProtocolRun)) == before + assert (await ctx["db"].get(FactorialRowResult, slot_ids[0])).run_id == first.id + + +@pytest.mark.asyncio +async def test_foreign_historical_targets_and_changed_source_reject(retry_setup, monkeypatch) -> None: + ctx = retry_setup + slot_ids = await _enqueue_rows(ctx, monkeypatch) + previous = await _latest(ctx, slot_ids[0]) + previous.status = "failed" + + other = await get_user_by_email(ctx["db"], "other@test.com") + assert other is not None, "run scripts/seed_row_test_users.py before row database tests" + foreign_experiment = await create_experiment( + ctx["db"], + name=f"foreign-row-retry-{uuid.uuid4()}", + owner_id=other.id, + design_spec={"factors": [], "metrics": []}, + ) + await upsert_replicate( + ctx["db"], + experiment_id=foreign_experiment.id, + replicate_label="foreign-parent", + fields={"factor_values": {}}, + ) + foreign_dataset_id = uuid.uuid4() + ctx["db"].add( + RegisteredDataset( + id=foreign_dataset_id, + name=f"foreign-{foreign_dataset_id}", + owner_id=other.id, + raw_path=str(ctx["raw_path"]), + raw_sha256=ctx["dataset"].raw_sha256, + ) + ) + foreign_graph = deepcopy(ctx["revision_graph"]) + foreign_graph["nodes"][0]["data"]["config"]["dataset_id"] = str(foreign_dataset_id) + foreign_graph["nodes"][0]["data"]["config"]["dataset_name"] = f"foreign-{foreign_dataset_id}" + foreign_protocol = await create_protocol( + ctx["db"], + name=f"foreign-row-retry-{uuid.uuid4()}", + owner_id=other.id, + experiment_id=foreign_experiment.id, + graph=foreign_graph, + ) + foreign_revision = await publish_protocol(ctx["db"], foreign_protocol, owner_id=other.id) + foreign_protocol.published_revision_id = foreign_revision.id + async def no_enqueue(_run_id: uuid.UUID) -> None: + return None + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", no_enqueue) + foreign_batch = await create_cell_runs_endpoint(foreign_protocol.id, other, ctx["db"]) + foreign_slot_id = foreign_batch.row_result_ids[0] + + with pytest.raises(HTTPException) as exc: + await create_cell_runs_endpoint( + ctx["protocol"].id, + ctx["user"], + ctx["db"], + CellRunBatchRequest(retry_row_result_ids=[foreign_slot_id]), + ) + assert exc.value.status_code == 422 and exc.value.detail == "invalid_retry_target" + + protocol = ctx["protocol"] + protocol.graph = deepcopy(protocol.graph) + protocol.graph["nodes"][1]["data"]["system_prompt"] = "new published graph" + revision = await publish_protocol(ctx["db"], protocol, owner_id=ctx["user"].id) + protocol.published_revision_id = revision.id + with pytest.raises(HTTPException) as exc: + await create_cell_runs_endpoint( + protocol.id, + ctx["user"], + ctx["db"], + CellRunBatchRequest(retry_row_result_ids=[slot_ids[0]]), + ) + assert exc.value.status_code == 422 and exc.value.detail == "invalid_retry_target" + + # Republish original graph so the slot is current, then verify its pinned hash. + protocol.graph = deepcopy(ctx["revision_graph"]) + current_again = await publish_protocol(ctx["db"], protocol, owner_id=ctx["user"].id) + protocol.published_revision_id = current_again.id + with pytest.raises(HTTPException) as exc: + await create_cell_runs_endpoint( + protocol.id, + ctx["user"], + ctx["db"], + CellRunBatchRequest(retry_row_result_ids=[slot_ids[0]]), + ) + assert exc.value.status_code == 422 and exc.value.detail == "invalid_retry_target" + + +@pytest.mark.asyncio +async def test_changed_bytes_report_source_hash_mismatch(retry_setup, monkeypatch) -> None: + ctx = retry_setup + slot_ids = await _enqueue_rows(ctx, monkeypatch) + (await _latest(ctx, slot_ids[0])).status = "failed" + ctx["raw_path"].write_bytes(b"question,answer\nchanged,response\n") + with pytest.raises(HTTPException) as exc: + await create_cell_runs_endpoint( + ctx["protocol"].id, + ctx["user"], + ctx["db"], + CellRunBatchRequest(retry_row_result_ids=[slot_ids[0]]), + ) + assert exc.value.status_code == 422 and "source_hash_mismatch" in exc.value.detail + + +@pytest.mark.asyncio +async def test_simultaneous_retry_claims_at_most_one_attempt(retry_setup, monkeypatch) -> None: + ctx = retry_setup + slot_ids = await _enqueue_rows(ctx, monkeypatch) + (await _latest(ctx, slot_ids[0])).status = "cancelled" + await ctx["db"].commit() + sessions = ctx["sessions"] + + async def retry_once(): + async with sessions() as db: + user = await get_user_by_email(db, "test@test.com") + protocol = await db.get(Protocol, ctx["protocol"].id) + assert user is not None and protocol is not None + try: + result = await create_cell_runs_endpoint( + protocol.id, + user, + db, + CellRunBatchRequest(retry_row_result_ids=[slot_ids[0]]), + ) + return result.protocol_run_ids + except HTTPException as exc: + return exc.detail + + outcomes = await asyncio.gather(retry_once(), retry_once()) + assert sum(isinstance(outcome, list) and len(outcome) == 1 for outcome in outcomes) == 1 + assert sum(outcome == "invalid_retry_target" for outcome in outcomes) == 1 + async with sessions() as check: + assert await check.scalar( + select(func.count()).select_from(ProtocolRun).where(ProtocolRun.row_result_id == slot_ids[0]) + ) == 2 diff --git a/tests/test_dataset_row_script_inputs.py b/tests/test_dataset_row_script_inputs.py new file mode 100644 index 00000000..c851c2eb --- /dev/null +++ b/tests/test_dataset_row_script_inputs.py @@ -0,0 +1,224 @@ +"""Script subprocesses receive only the invoking Agent's authorized inputs.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +from asaree.mcp_servers import script_server as ss + + +class _FakeCtx: + def __init__(self, extra: dict[str, Any]) -> None: + self.request_context = type("_R", (), {"meta": type("_M", (), {"model_extra": extra})()})() + + +def _run(ctx: _FakeCtx) -> dict[str, Any]: + return json.loads(ss.run_wired_script(ctx=ctx)) + + +def _row_context(tmp_path: Path, path: Path, *, workspace_id: str, row_inputs: list[dict[str, Any]]) -> _FakeCtx: + script = tmp_path / "wired.py" + script.write_text( + "import csv, json\n" + "from asaree.script_context import training_inputs\n" + "inputs = training_inputs()\n" + "rows = list(csv.DictReader(inputs[0].path.open())) if inputs else []\n" + "item = inputs[0] if inputs else None\n" + "print(json.dumps({'inputs': [{'name': x.name, 'path': str(x.path), 'mode': x.mode, " + "'target': x.target_column, 'dataset_id': x.dataset_id, 'raw_sha256': x.raw_sha256, " + "'row_index': x.row_index, 'columns': x.columns} for x in inputs], 'rows': rows}))\n" + ) + return _FakeCtx( + { + "motoro.ambient.script_path": str(script), + "motoro.workspace_id": workspace_id, + "motoro.ambient.dataset_mode": "per_row", + "motoro.ambient.row_inputs": row_inputs, + "motoro.ambient.data_path": "/stale/full.csv", + "motoro.ambient.dataset_names": ["stale"], + } + ) + + +def test_prediction_script_reads_only_authorized_question_row( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("ASAREE_DATASET_WORKSPACE_DIR", str(tmp_path / "workspaces")) + row_dir = tmp_path / "workspaces" / "_protocol_runs" / "attempt-1" / "agents" / "agent-a" + row_dir.mkdir(parents=True) + row_csv = row_dir / "authorized.csv" + row_csv.write_text("question\nWhat is 2+2?\n") + monkeypatch.setattr(ss, "raw_training_data_locators", lambda _workspace_id: pytest.fail("stale lookup")) + row = { + "name": "prompt", + "path": str(row_csv), + "target_column": "", + "mode": "per_row", + "dataset_id": "dataset-1", + "raw_sha256": "a" * 64, + "row_index": 3, + "columns": ["question"], + } + + out = _run(_row_context(tmp_path, row_csv, workspace_id="_protocol_runs/attempt-1", row_inputs=[row])) + + payload = json.loads(out["stdout"]) + assert payload["rows"] == [{"question": "What is 2+2?"}] + assert len(payload["inputs"]) == 1 + assert payload["inputs"][0] == { + "name": "prompt", + "path": str(row_csv), + "mode": "per_row", + "target": "", + "dataset_id": "dataset-1", + "raw_sha256": "a" * 64, + "row_index": 3, + "columns": ["question"], + } + + +def test_grading_script_receives_same_row_reference_and_legacy_state_is_narrow( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("ASAREE_DATASET_WORKSPACE_DIR", str(tmp_path / "workspaces")) + row_dir = tmp_path / "workspaces" / "_protocol_runs" / "attempt-grade" / "agents" / "agent-b" + row_dir.mkdir(parents=True) + row_csv = row_dir / "grade.csv" + row_csv.write_text("question,reference\nQ,A\n") + row = { + "name": "grading", + "path": str(row_csv), + "target_column": "reference", + "mode": "per_row", + "dataset_id": "dataset-1", + "raw_sha256": "b" * 64, + "row_index": 3, + "columns": ["question", "reference"], + } + ctx = _row_context(tmp_path, row_csv, workspace_id="_protocol_runs/attempt-grade", row_inputs=[row]) + script = tmp_path / "grading.py" + script.write_text( + "import json\nfrom pathlib import Path\n" + "state = json.loads(Path('state.json').read_text())\n" + "from asaree.script_context import training_input\n" + "item = training_input()\n" + "print(json.dumps({'path': str(item.path), 'state': state}))\n" + ) + ctx.request_context.meta.model_extra["motoro.ambient.script_path"] = str(script) + + payload = json.loads(_run(ctx)["stdout"]) + + state = payload["state"] + assert payload["path"] == str(row_csv) + assert state == { + "target_column": "reference", + "head": "v0_raw", + "versions": [{"id": "v0_raw", "train": str(row_csv)}], + } + assert set(state["versions"][0]) == {"id", "train"} + + +def test_empty_authorized_row_inputs_do_not_inherit_stale_head( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("ASAREE_DATASET_WORKSPACE_DIR", str(tmp_path / "workspaces")) + workspace = tmp_path / "workspaces" / "_protocol_runs" / "attempt-empty" + workspace.mkdir(parents=True) + (workspace / "state.json").write_text('{"head":"stale"}') + monkeypatch.setattr(ss, "raw_training_data_locators", lambda _workspace_id: pytest.fail("stale lookup")) + out = _run( + _row_context( + tmp_path, + tmp_path / "unused.csv", + workspace_id="_protocol_runs/attempt-empty", + row_inputs=[], + ) + ) + assert out["exit_code"] == 0 + assert json.loads(out["stdout"]) == {"inputs": [], "rows": []} + + +def test_retry_uses_a_different_private_attempt_workspace(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + root = tmp_path / "workspaces" + monkeypatch.setenv("ASAREE_DATASET_WORKSPACE_DIR", str(root)) + for attempt in ("attempt-one", "attempt-two"): + (root / "_protocol_runs" / attempt).mkdir(parents=True) + + paths = [] + for attempt in ("attempt-one", "attempt-two"): + ctx = _row_context(tmp_path, tmp_path / "unused.csv", workspace_id=f"_protocol_runs/{attempt}", row_inputs=[]) + script = tmp_path / f"{attempt}.py" + script.write_text("from pathlib import Path\nprint(Path.cwd())") + ctx.request_context.meta.model_extra["motoro.ambient.script_path"] = str(script) + paths.append(Path(_run(ctx)["stdout"].strip())) + + assert paths == [root / "_protocol_runs" / "attempt-one", root / "_protocol_runs" / "attempt-two"] + assert paths[0] != paths[1] + + +def test_whole_context_inputs_keep_existing_manifest_shape(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + raw = tmp_path / "whole.csv" + raw.write_text("question\nhello\n") + monkeypatch.setattr(ss, "raw_training_data_locators", lambda _workspace_id: {}) + manifest = ss._runtime_manifest( + _FakeCtx( + { + "motoro.ambient.dataset_mode": "raw_unsplit", + "motoro.ambient.dataset_names": ["whole"], + "motoro.ambient.data_path": str(raw), + "motoro.ambient.target_column": "", + } + ), + "", + ) + assert manifest == { + "schema_version": 1, + "training_inputs": [ + {"name": "whole", "path": str(raw), "target_column": "", "mode": "raw_unsplit", "slot": None, + "workspace_version": None} + ], + } + + +def test_whole_mode_subprocess_and_result_envelope_are_preserved( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + raw = tmp_path / "whole.csv" + raw.write_text("question\nhello\n") + script = tmp_path / "whole.py" + script.write_text( + "from asaree.script_context import training_input\n" + "item = training_input()\n" + "print(item.mode, item.path)\n" + ) + monkeypatch.setattr(ss, "raw_training_data_locators", lambda _workspace_id: {}) + ctx = _FakeCtx( + { + "motoro.ambient.script_path": str(script), + "motoro.ambient.dataset_mode": "raw_unsplit", + "motoro.ambient.dataset_names": ["whole"], + "motoro.ambient.data_path": str(raw), + "motoro.ambient.target_column": "", + } + ) + out = _run(ctx) + assert out["exit_code"] == 0 + assert out["stdout"] == f"raw_unsplit {raw}\n" + assert len(out["code_sha256"]) == 64 + assert "stderr" in out + + failed_script = tmp_path / "failed.py" + failed_script.write_text("raise ValueError('expected failure')") + ctx.request_context.meta.model_extra["motoro.ambient.script_path"] = str(failed_script) + failed = _run(ctx) + assert failed["exit_code"] == 1 + assert "expected failure" in failed["stderr"] + assert "error" in failed + + failed_script.write_text("print('x' * 20000)") + truncated = _run(ctx) + assert "truncated" in truncated["stdout"] diff --git a/tests/test_dataset_row_single_run.py b/tests/test_dataset_row_single_run.py new file mode 100644 index 00000000..f8c78b50 --- /dev/null +++ b/tests/test_dataset_row_single_run.py @@ -0,0 +1,204 @@ +"""PostgreSQL integration coverage for production single-row execution.""" + +from __future__ import annotations + +import hashlib +import os +import uuid +from collections.abc import AsyncIterator +from pathlib import Path + +import pytest +import pytest_asyncio +from pydantic import ValidationError +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +import asaree.api.protocols as protocol_api +from asaree.api.protocols import CreateProtocolRunRequest, create_protocol_run_endpoint +from asaree.models.dataset import RegisteredDataset +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol_run import ProtocolRun +from asaree.services.experiments import create_experiment +from asaree.services.factorial_cells import upsert_replicate +from asaree.services.protocol_revisions import publish_protocol +from asaree.services.protocols import create_protocol +from asaree.services.users import get_user_by_email + + +@pytest_asyncio.fixture +async def row_single(tmp_path: Path) -> AsyncIterator[tuple[AsyncSession, dict]]: + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + sessions = async_sessionmaker(engine, expire_on_commit=False) + try: + async with sessions.begin() as db: + user = await get_user_by_email(db, "test@test.com") + assert user is not None, "run scripts/seed_row_test_users.py before row database tests" + raw = b"question,answer\nq0,a0\nq1,a1\nq2,a2\n" + raw_path = tmp_path / "single.csv" + raw_path.write_bytes(raw) + dataset = RegisteredDataset( + id=uuid.uuid4(), name=f"single-{uuid.uuid4()}", owner_id=user.id, + raw_path=str(raw_path), raw_sha256=hashlib.sha256(raw).hexdigest(), + ) + db.add(dataset) + experiment = await create_experiment( + db, name=f"row-single-{uuid.uuid4()}", owner_id=user.id, + design_spec={"factors": [], "metrics": []}, + ) + parent = await upsert_replicate( + db, experiment_id=experiment.id, replicate_label="single-parent", + fields={"factor_values": {}}, + ) + graph = { + "nodes": [ + {"id": "dataset", "type": "dataset", "data": {"config": { + "dataset_id": str(dataset.id), "dataset_name": dataset.name, + }}}, + {"id": "agent", "type": "agent", "data": {}}, + {"id": "model", "type": "model_openai", "data": {"config": {}}}, + ], + "edges": [ + {"id": "dataset-agent", "source": "dataset", "target": "agent", + "targetHandle": "dataset", "data": {"dataset_input": { + "mode": "per_row", "columns": ["question"], + }}}, + {"id": "model-agent", "source": "model", "target": "agent", "targetHandle": "model"}, + ], + } + protocol = await create_protocol( + db, name=f"row-single-{uuid.uuid4()}", owner_id=user.id, + experiment_id=experiment.id, graph=graph, + ) + revision = await publish_protocol(db, protocol, owner_id=user.id) + protocol.published_revision_id = revision.id + await db.flush() + user_id, protocol_id, revision_id, parent_id = user.id, protocol.id, revision.id, parent.id + async with sessions() as db: + from asaree.models.protocol import Protocol + from asaree.models.protocol_revision import ProtocolRevision + from asaree.models.user import User + + user = await db.get(User, user_id) + protocol = await db.get(Protocol, protocol_id) + revision = await db.get(ProtocolRevision, revision_id) + assert user is not None and protocol is not None and revision is not None + yield db, { + "user": user, + "protocol": protocol, + "revision": revision, + "parent_id": parent_id, + } + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_single_run_selects_row_and_default_row_without_factorial_writes(row_single, monkeypatch): + db, ctx = row_single + queued: list[uuid.UUID] = [] + + async def enqueue(run_id: uuid.UUID) -> None: + queued.append(run_id) + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + selected = await create_protocol_run_endpoint( + ctx["protocol"].id, ctx["user"], db, + CreateProtocolRunRequest(replicate_label="single-parent", row_index=2), + ) + default = await create_protocol_run_endpoint(ctx["protocol"].id, ctx["user"], db) + assert selected.dataset_row["row_index"] == 2 + assert default.dataset_row["row_index"] == 0 + assert selected.row_result_id != default.row_result_id + assert selected.replicate_result_id is not None + assert selected.replicate_label == "single-parent" + assert selected.factor_values == {} + assert default.replicate_result_id is None + assert default.replicate_label is None + assert default.factor_values is None + for run in (selected, default): + assert run.protocol_revision_id == ctx["revision"].id + assert run.design_revision_id is not None + assert len(queued) == 2 + assert await db.scalar(select(func.count()).select_from(ProtocolRun).where( + ProtocolRun.protocol_id == ctx["protocol"].id, + )) == 2 + parent = await db.get(FactorialReplicateResult, ctx["parent_id"]) + assert parent is not None and parent.run_id is None and parent.metric_values is None + assert await db.scalar(select(func.count()).select_from(FactorialRowResult).where( + FactorialRowResult.replicate_result_id == ctx["parent_id"], + )) == 2 + + +@pytest.mark.asyncio +async def test_attempted_slot_and_out_of_range_row_are_rejected_before_enqueue(row_single, monkeypatch): + from fastapi import HTTPException + + db, ctx = row_single + queued: list[uuid.UUID] = [] + + async def enqueue(run_id: uuid.UUID) -> None: + queued.append(run_id) + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + request = CreateProtocolRunRequest(replicate_label="single-parent", row_index=1) + await create_protocol_run_endpoint(ctx["protocol"].id, ctx["user"], db, request) + with pytest.raises(HTTPException) as attempted: + await create_protocol_run_endpoint(ctx["protocol"].id, ctx["user"], db, request) + assert attempted.value.status_code == 422 + assert attempted.value.detail == "row_already_attempted" + before = await db.scalar(select(func.count()).select_from(ProtocolRun).where( + ProtocolRun.protocol_id == ctx["protocol"].id, + )) + with pytest.raises(HTTPException): + await create_protocol_run_endpoint( + ctx["protocol"].id, ctx["user"], db, + CreateProtocolRunRequest(replicate_label="single-parent", row_index=3), + ) + assert await db.scalar(select(func.count()).select_from(ProtocolRun).where( + ProtocolRun.protocol_id == ctx["protocol"].id, + )) == before == 1 + assert len(queued) == 1 + + +@pytest.mark.asyncio +async def test_row_index_without_driver_is_rejected_before_run_creation(row_single, monkeypatch): + from fastapi import HTTPException + + db, ctx = row_single + queued: list[uuid.UUID] = [] + + async def enqueue(run_id: uuid.UUID) -> None: + queued.append(run_id) + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + protocol = ctx["protocol"] + protocol.graph = { + "nodes": [ + {"id": "agent", "type": "agent", "data": {}}, + {"id": "model", "type": "model_openai", "data": {"config": {}}}, + ], + "edges": [{"id": "model-agent", "source": "model", "target": "agent", "targetHandle": "model"}], + } + revision = await publish_protocol(db, protocol, owner_id=ctx["user"].id) + protocol.published_revision_id = revision.id + before = await db.scalar(select(func.count()).select_from(ProtocolRun).where( + ProtocolRun.protocol_id == protocol.id, + )) + with pytest.raises(HTTPException) as no_driver: + await create_protocol_run_endpoint( + protocol.id, ctx["user"], db, CreateProtocolRunRequest(row_index=0), + ) + assert no_driver.value.status_code == 422 + assert no_driver.value.detail == "no_row_driver" + assert await db.scalar(select(func.count()).select_from(ProtocolRun).where( + ProtocolRun.protocol_id == protocol.id, + )) == before == 0 + assert queued == [] + + +def test_row_index_is_strict_and_nonnegative(): + for value in (-1, True, "2"): + with pytest.raises(ValidationError): + CreateProtocolRunRequest(row_index=value) diff --git a/tests/test_dataset_row_stale_writes.py b/tests/test_dataset_row_stale_writes.py new file mode 100644 index 00000000..5df5add2 --- /dev/null +++ b/tests/test_dataset_row_stale_writes.py @@ -0,0 +1,301 @@ +"""PostgreSQL concurrency coverage for latest row-attempt projections.""" + +from __future__ import annotations + +import asyncio +import hashlib +import os +import uuid +from collections.abc import AsyncIterator +from datetime import UTC, datetime +from pathlib import Path + +import pytest +import pytest_asyncio +from sqlalchemy import delete +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from asaree.models.dataset import RegisteredDataset +from asaree.models.experiment import ResearchExperiment +from asaree.models.experiment_design_revision import ExperimentDesignRevision +from asaree.models.factorial_cell import FactorialCell +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol import Protocol +from asaree.models.protocol_revision import ProtocolRevision +from asaree.models.protocol_run import ProtocolRun +from asaree.services.factorial_row_results import claim_row_attempt, project_row_attempt +from asaree.services.measurement_engine import ( + MeasurementEvaluation, + MetricObservation, + ProducerProvenance, +) +from asaree.services.protocol_runs import create_protocol_run, record_measurement_evaluation, set_status +from asaree.services.users import get_user_by_email + + +@pytest_asyncio.fixture +async def row_slot(tmp_path: Path) -> AsyncIterator[dict]: + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + sessions = async_sessionmaker(engine, expire_on_commit=False) + try: + async with sessions() as db: + user = await get_user_by_email(db, "test@test.com") + assert user is not None, "run scripts/seed_row_test_users.py before row database tests" + experiment_id, dataset_id, protocol_id = (uuid.uuid4() for _ in range(3)) + design_id, revision_id = uuid.uuid4(), uuid.uuid4() + cell_id, replicate_id, row_id = (uuid.uuid4() for _ in range(3)) + raw = b"input,expected\nquestion,response\n" + csv_path = tmp_path / "stale-writes.csv" + csv_path.write_bytes(raw) + digest = hashlib.sha256(raw).hexdigest() + workspace = tmp_path / "attempt-workspace" + workspace.mkdir() + (workspace / "trace.json").write_text('{"attempt":"A"}', encoding="utf-8") + experiment = ResearchExperiment( + id=experiment_id, + name=f"stale-row-{experiment_id}", + owner_id=user.id, + design_spec={"factors": [], "metrics": []}, + ) + dataset = RegisteredDataset( + id=dataset_id, + name=f"stale-row-{dataset_id}", + owner_id=user.id, + raw_path=str(csv_path), + raw_sha256=digest, + ) + protocol = Protocol( + id=protocol_id, + name=f"stale-row-{protocol_id}", + owner_id=user.id, + experiment_id=experiment_id, + graph={"nodes": [], "edges": []}, + ) + db.add_all([experiment, dataset, protocol]) + await db.flush() + design = ExperimentDesignRevision( + id=design_id, + experiment_id=experiment_id, + revision=1, + design_spec={"factors": [], "metrics": []}, + ) + revision = ProtocolRevision( + id=revision_id, + protocol_id=protocol_id, + revision=1, + graph={"nodes": [], "edges": []}, + published_at=datetime.now(UTC), + ) + cell = FactorialCell( + id=cell_id, + experiment_id=experiment_id, + design_revision_id=design_id, + cell_label="row-cell", + factor_values={}, + ) + db.add_all([design, revision, cell]) + await db.flush() + protocol.published_revision_id = revision_id + replicate = FactorialReplicateResult( + id=replicate_id, + cell_id=cell_id, + replicate_number=1, + replicate_label="row-replicate", + run_id=None, + metric_values=None, + artifacts=None, + ) + row = FactorialRowResult( + id=row_id, + replicate_result_id=replicate_id, + protocol_revision_id=revision_id, + dataset_id=dataset_id, + raw_sha256=digest, + row_index=0, + artifacts={"previous": True}, + ) + db.add_all([replicate, row]) + await db.flush() + await db.commit() + yield { + "sessions": sessions, + "user_id": user.id, + "protocol_id": protocol_id, + "design_id": design_id, + "revision_id": revision_id, + "replicate_id": replicate_id, + "row_id": row_id, + "dataset_id": dataset_id, + "digest": digest, + "row_snapshot": { + "dataset_id": str(dataset_id), + "raw_sha256": digest, + "row_index": 0, + "columns": ["input", "expected"], + "values": {"input": "question", "expected": "response"}, + }, + "workspace": str(workspace), + } + finally: + await engine.dispose() + + +async def _create_run(db, ctx: dict, *, row_id: uuid.UUID) -> ProtocolRun: + return await create_protocol_run( + db, + protocol_id=ctx["protocol_id"], + owner_id=ctx["user_id"], + replicate_label="row-replicate", + factor_values={}, + replicate_result_id=ctx["replicate_id"], + row_result_id=row_id, + dataset_row=ctx["row_snapshot"], + design_revision_id=ctx["design_id"], + protocol_revision_id=ctx["revision_id"], + ) + + +@pytest.mark.asyncio +async def test_stale_row_projection_cannot_overwrite_replacement_attempt(row_slot): + sessions = row_slot["sessions"] + async with sessions.begin() as db: + run_a = await _create_run(db, row_slot, row_id=row_slot["row_id"]) + run_a_id = run_a.id + + work_finished = asyncio.Event() + allow_stale_write = asyncio.Event() + + async def stale_worker() -> None: + # Represents A finishing its work, then pausing before projection. + work_finished.set() + await allow_stale_write.wait() + async with sessions.begin() as stale_db: + assert await project_row_attempt( + stale_db, + row_result_id=row_slot["row_id"], + run_id=run_a_id, + fields={ + "workspace_id": "attempt-A", + "metric_values": {"score": 1}, + "artifacts": {"output_text": "stale A", "late": True}, + }, + ) is False + + task = asyncio.create_task(stale_worker()) + await work_finished.wait() + async with sessions.begin() as db: + await set_status(db, run_a_id, status="failed", error="A failed") + run_b = await claim_row_attempt( + db, + row_result_id=row_slot["row_id"], + expected_run_id=run_a_id, + create_kwargs={ + "protocol_id": row_slot["protocol_id"], + "owner_id": row_slot["user_id"], + "replicate_label": "row-replicate", + "factor_values": {}, + "replicate_result_id": row_slot["replicate_id"], + "design_revision_id": row_slot["design_id"], + "protocol_revision_id": row_slot["revision_id"], + "dataset_row": row_slot["row_snapshot"], + }, + ) + assert run_b is not None + run_b_id = run_b.id + await set_status(db, run_b_id, status="completed") + assert await project_row_attempt( + db, + row_result_id=row_slot["row_id"], + run_id=run_b_id, + fields={ + "workspace_id": "attempt-B", + "metric_values": {"score": 2}, + "artifacts": {"output_text": "current B"}, + }, + ) is True + + allow_stale_write.set() + await task + async with sessions() as db: + slot = await db.get(FactorialRowResult, row_slot["row_id"]) + old = await db.get(ProtocolRun, run_a_id) + assert slot is not None and old is not None + assert slot.run_id == run_b_id + assert slot.workspace_id == "attempt-B" + assert slot.metric_values == {"score": 2} + assert slot.artifacts["output_text"] == "current B" + assert old.error == "A failed" + + +@pytest.mark.asyncio +async def test_deleted_row_slot_late_projection_never_targets_parent(row_slot): + sessions = row_slot["sessions"] + async with sessions.begin() as db: + run = await _create_run(db, row_slot, row_id=row_slot["row_id"]) + run_id = run.id + parent_before = await db.get(FactorialReplicateResult, row_slot["replicate_id"]) + assert parent_before is not None + before = (parent_before.workspace_id, parent_before.metric_values, parent_before.artifacts) + await db.execute(delete(FactorialRowResult).where(FactorialRowResult.id == row_slot["row_id"])) + assert await project_row_attempt( + db, + row_result_id=row_slot["row_id"], + run_id=run_id, + fields={"workspace_id": "late", "metric_values": {"score": 7}, "artifacts": {"late": True}}, + ) is False + parent_after = await db.get(FactorialReplicateResult, row_slot["replicate_id"]) + assert parent_after is not None + assert (parent_after.workspace_id, parent_after.metric_values, parent_after.artifacts) == before + + +@pytest.mark.asyncio +async def test_duplicate_row_finalization_preserves_measured_null_and_result(row_slot): + sessions = row_slot["sessions"] + async with sessions.begin() as db: + run = await _create_run(db, row_slot, row_id=row_slot["row_id"]) + run_id = run.id + run.status = "completed" + provenance = ProducerProvenance( + binding_id="report", + producer_id="asaree.agent_output", + kind="reported", + version="1", + ) + + def evaluation(value: str | None) -> MeasurementEvaluation: + return MeasurementEvaluation( + replicate_id=str(row_slot["row_id"]), + attempt_id=str(run_id), + observations=( + MetricObservation( + metric_id="answer", + metric_name="Answer", + value_type=None, + status="measured", + value=value, + error=None, + attempt_id=str(run_id), + producer=provenance, + input_provenance={}, + ), + ), + artifacts=(), + ) + + await record_measurement_evaluation(db, run_id, evaluation(None)) + persisted = await db.get(ProtocolRun, run_id) + assert persisted is not None + original = dict(persisted.attempt_result) + with pytest.raises(ValueError, match="immutable_attempt_result"): + await record_measurement_evaluation(db, run_id, evaluation("replacement")) + persisted = await db.get(ProtocolRun, run_id) + assert persisted is not None + assert persisted.attempt_result == original + with pytest.raises(ValueError, match="immutable_attempt_result"): + from asaree.services.protocol_runs import update_attempt_result + + await update_attempt_result(db, run_id, fields={"metric_values": {"Answer": "replacement"}}) + slot = await db.get(FactorialRowResult, row_slot["row_id"]) + assert slot is not None + assert slot.metric_values == {"Answer": None} diff --git a/tests/test_dataset_row_test_run.py b/tests/test_dataset_row_test_run.py new file mode 100644 index 00000000..62bbef7e --- /dev/null +++ b/tests/test_dataset_row_test_run.py @@ -0,0 +1,308 @@ +"""PostgreSQL coverage for row-selected canvas Test Runs.""" + +from __future__ import annotations + +import hashlib +import os +import uuid +from collections.abc import AsyncIterator +from pathlib import Path + +import pytest +import pytest_asyncio +from fastapi import HTTPException +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +import asaree.api.protocols as protocol_api +import asaree.services.protocol_execution as execution +from asaree.models.dataset import RegisteredDataset +from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.factorial_row_result import FactorialRowResult +from asaree.models.protocol import Protocol +from asaree.models.protocol_run import ProtocolRun +from asaree.models.user import User +from asaree.services.experiments import create_experiment, delete_experiment +from asaree.services.factorial_cells import upsert_replicate +from asaree.services.factorial_row_results import ensure_row_result +from asaree.services.protocol_revisions import publish_protocol +from asaree.services.protocol_runs import create_test_run +from asaree.services.protocols import create_protocol, delete_protocol +from asaree.services.users import get_user_by_email + + +@pytest_asyncio.fixture +async def test_run_context(tmp_path: Path) -> AsyncIterator[dict]: + engine = create_async_engine(os.environ["ASAREE_PRODUCT_DATABASE_URL"]) + sessions = async_sessionmaker(engine, expire_on_commit=False) + protocol_id = experiment_id = dataset_id = None + try: + async with sessions.begin() as db: + user = await get_user_by_email(db, "test@test.com") + assert user is not None, "run scripts/seed_row_test_users.py before row database tests" + raw = b"prompt,reference\nfirst,gold-a\nsecond,gold-b\nthird,gold-c\n" + raw_path = tmp_path / "test-run.csv" + raw_path.write_bytes(raw) + dataset = RegisteredDataset( + id=uuid.uuid4(), name=f"test-run-{uuid.uuid4()}", owner_id=user.id, + raw_path=str(raw_path), raw_sha256=hashlib.sha256(raw).hexdigest(), + ) + db.add(dataset) + experiment = await create_experiment( + db, name=f"test-run-{uuid.uuid4()}", owner_id=user.id, + design_spec={"factors": [], "metrics": []}, + ) + graph = { + "nodes": [ + {"id": "source", "type": "dataset", "data": {"config": { + "dataset_id": str(dataset.id), "dataset_name": dataset.name, + }}}, + {"id": "first", "type": "agent", "data": {"config": {"name": "first"}}}, + {"id": "second", "type": "agent", "data": {"config": {"name": "second"}}}, + {"id": "model-first", "type": "model_openai", "data": {"config": {}}}, + {"id": "model-second", "type": "model_openai", "data": {"config": {}}}, + ], + "edges": [ + {"id": "driver-first", "source": "source", "target": "first", "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["prompt", "reference"]}}}, + {"id": "first-second", "source": "first", "target": "second"}, + {"id": "model-first-edge", "source": "model-first", "target": "first", "targetHandle": "model"}, + {"id": "model-second-edge", "source": "model-second", "target": "second", "targetHandle": "model"}, + ], + } + protocol = await create_protocol( + db, name=f"test-run-{uuid.uuid4()}", owner_id=user.id, + experiment_id=experiment.id, graph=graph, + ) + revision = await publish_protocol(db, protocol, owner_id=user.id) + protocol.published_revision_id = revision.id + protocol_id, experiment_id, dataset_id = protocol.id, experiment.id, dataset.id + context_ids = (user.id, protocol.id, revision.id, experiment.id, dataset.id) + async with sessions() as db: + user_id, protocol_id, revision_id, experiment_id, dataset_id = context_ids + user = await db.get(User, user_id) + protocol = await db.get(Protocol, protocol_id) + assert user is not None and protocol is not None + yield { + "sessions": sessions, + "user": user, + "protocol": protocol, + "revision_id": revision_id, + "experiment_id": experiment_id, + "dataset_id": dataset_id, + } + finally: + if protocol_id is not None: + async with sessions.begin() as db: + await delete_protocol(db, protocol_id) + if experiment_id is not None: + await delete_experiment(db, experiment_id) + if dataset_id is not None: + dataset = await db.get(RegisteredDataset, dataset_id) + if dataset is not None: + await db.delete(dataset) + await engine.dispose() + + +@pytest.fixture +def valid_test_run_endpoint(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(protocol_api, "validate_coordination_strategy", lambda *_args, **_kwargs: None) + monkeypatch.setattr(protocol_api, "validate_stage_plan", lambda *_args, **_kwargs: None) + monkeypatch.setattr(protocol_api, "validate_prompt_references", lambda *_args, **_kwargs: None) + monkeypatch.setattr(protocol_api, "topological_order", lambda *_args, **_kwargs: ["first", "second"]) + + +@pytest.mark.asyncio +async def test_row_test_run_defaults_to_zero_and_selected_index_two( + test_run_context, valid_test_run_endpoint, monkeypatch +): + queued: list[uuid.UUID] = [] + + async def enqueue(run_id: uuid.UUID) -> None: + queued.append(run_id) + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + ctx = test_run_context + async with ctx["sessions"]() as db: + zero = await protocol_api.create_test_run_endpoint(ctx["protocol"].id, ctx["user"], db) + selected = await protocol_api.create_test_run_endpoint( + ctx["protocol"].id, ctx["user"], db, protocol_api.TestRunRequest(row_index=2) + ) + assert zero.dataset_row["row_index"] == 0 + assert selected.dataset_row["row_index"] == 2 + assert selected.dataset_row["values"] == {"prompt": "third", "reference": "gold-c"} + assert zero.protocol_revision_id == selected.protocol_revision_id == ctx["revision_id"] + assert len(queued) == 2 + assert await db.get(ProtocolRun, zero.id) is None + assert await db.scalar( + select(func.count()) + .select_from(FactorialRowResult) + .join(FactorialReplicateResult) + .where(FactorialReplicateResult.experiment_id == ctx["experiment_id"]) + ) == 0 + + +@pytest.mark.asyncio +async def test_invalid_selection_preserves_latest_and_production_row_slots( + test_run_context, valid_test_run_endpoint, monkeypatch +): + async def enqueue(_run_id: uuid.UUID) -> None: + return None + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + ctx = test_run_context + async with ctx["sessions"]() as db: + replicate = await upsert_replicate( + db, experiment_id=ctx["experiment_id"], replicate_label="production-slot", + fields={"factor_values": {}}, + ) + dataset = await db.get(RegisteredDataset, ctx["dataset_id"]) + assert dataset is not None + slot = await ensure_row_result( + db, + experiment_id=ctx["experiment_id"], + design_revision_id=replicate.design_revision_id, + protocol_revision_id=ctx["revision_id"], + replicate_result_id=replicate.id, + dataset_id=dataset.id, + raw_sha256=dataset.raw_sha256, + row_index=0, + ) + latest = await protocol_api.create_test_run_endpoint(ctx["protocol"].id, ctx["user"], db) + await protocol_api.fail_protocol_run(db, latest.id, error="mocked worker failure") + replacement = await protocol_api.create_test_run_endpoint(ctx["protocol"].id, ctx["user"], db) + run_count = await db.scalar(select(func.count()).select_from(ProtocolRun)) + with pytest.raises(HTTPException) as invalid: + await protocol_api.create_test_run_endpoint( + ctx["protocol"].id, ctx["user"], db, protocol_api.TestRunRequest(row_index=3) + ) + assert invalid.value.status_code == 422 + assert await protocol_api.get_latest_test_run_endpoint(ctx["protocol"].id, ctx["user"], db) == replacement + assert await db.scalar(select(func.count()).select_from(ProtocolRun)) == run_count + assert await db.get(FactorialReplicateResult, replicate.id) is not None + assert await db.get(FactorialRowResult, slot.id) is not None + + +@pytest.mark.asyncio +async def test_latest_polling_returns_persisted_row_and_revision_after_draft_changes( + test_run_context, valid_test_run_endpoint, monkeypatch +): + async def enqueue(_run_id: uuid.UUID) -> None: + return None + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + ctx = test_run_context + async with ctx["sessions"]() as db: + created = await protocol_api.create_test_run_endpoint( + ctx["protocol"].id, ctx["user"], db, protocol_api.TestRunRequest(row_index=2) + ) + run_id = created.id + protocol = await db.get(Protocol, ctx["protocol"].id) + assert protocol is not None + protocol.graph = {"nodes": [], "edges": []} + await db.commit() + async with ctx["sessions"]() as db: + latest = await protocol_api.get_latest_test_run_endpoint(ctx["protocol"].id, ctx["user"], db) + assert latest.id == run_id + assert latest.protocol_revision_id == ctx["revision_id"] + assert latest.dataset_row["row_index"] == 2 + assert latest.dataset_row["raw_sha256"] + + +@pytest.mark.asyncio +async def test_two_agent_test_run_executes_full_published_protocol_without_factorial_writes( + test_run_context, valid_test_run_endpoint, monkeypatch +): + calls: list[str] = [] + queued: list[uuid.UUID] = [] + + async def enqueue(run_id: uuid.UUID) -> None: + queued.append(run_id) + + async def mocked_provider_execution(**kwargs) -> None: + calls.append(str(kwargs["run_id"])) + + async def no_hydration() -> None: + return None + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + monkeypatch.setattr(execution, "execute_run", mocked_provider_execution) + monkeypatch.setattr(execution, "hydrate_registry", no_hydration) + monkeypatch.setattr(execution, "get_registry", lambda: None) + ctx = test_run_context + async with ctx["sessions"]() as db: + created = await protocol_api.create_test_run_endpoint( + ctx["protocol"].id, ctx["user"], db, protocol_api.TestRunRequest(row_index=1) + ) + run_id = created.id + assert queued == [run_id] + await execution.run_protocol(run_id) + async with ctx["sessions"]() as db: + run = await db.get(ProtocolRun, run_id) + assert run is not None and run.status == "completed" and run.is_test_run + assert run.protocol_revision_id == ctx["revision_id"] + assert run.dataset_row["row_index"] == 1 + assert run.row_result_id is None + assert run.replicate_result_id is None + assert set(run.node_runs) >= {"first", "second"} + assert len(calls) == 2 + assert await db.scalar( + select(func.count()).select_from(FactorialReplicateResult).where( + FactorialReplicateResult.experiment_id == ctx["experiment_id"] + ) + ) == 0 + assert await db.scalar( + select(func.count()) + .select_from(FactorialRowResult) + .join(FactorialReplicateResult) + .where(FactorialReplicateResult.experiment_id == ctx["experiment_id"]) + ) == 0 + + +@pytest.mark.asyncio +async def test_whole_mode_test_run_remains_compatible(test_run_context, valid_test_run_endpoint, monkeypatch): + async def enqueue(_run_id: uuid.UUID) -> None: + return None + + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", enqueue) + ctx = test_run_context + async with ctx["sessions"]() as db: + protocol = await db.get(Protocol, ctx["protocol"].id) + assert protocol is not None + protocol.graph = { + "nodes": [ + {"id": "agent", "type": "agent", "data": {}}, + {"id": "model", "type": "model_openai", "data": {"config": {}}}, + ], + "edges": [{"id": "model-agent", "source": "model", "target": "agent", "targetHandle": "model"}], + } + revision = await publish_protocol(db, protocol, owner_id=ctx["user"].id) + protocol.published_revision_id = revision.id + with pytest.raises(HTTPException) as no_driver: + await protocol_api.create_test_run_endpoint( + protocol.id, ctx["user"], db, protocol_api.TestRunRequest(row_index=0) + ) + assert no_driver.value.status_code == 422 and no_driver.value.detail == "no_row_driver" + run = await protocol_api.create_test_run_endpoint(protocol.id, ctx["user"], db) + assert run.dataset_row is None + assert run.protocol_revision_id == revision.id + + +@pytest.mark.asyncio +async def test_create_test_run_persists_snapshot_without_row_slot(test_run_context): + ctx = test_run_context + async with ctx["sessions"]() as db: + snapshot = { + "dataset_id": str(ctx["dataset_id"]), + "raw_sha256": "a" * 64, + "row_index": 2, + "columns": ["prompt"], + "values": {"prompt": "third"}, + } + run = await create_test_run( + db, protocol_id=ctx["protocol"].id, owner_id=ctx["user"].id, + protocol_revision_id=ctx["revision_id"], dataset_row=snapshot, + ) + assert run.dataset_row == snapshot + assert run.is_test_run + assert run.row_result_id is None and run.replicate_result_id is None diff --git a/tests/test_dataset_row_tool_bindings.py b/tests/test_dataset_row_tool_bindings.py new file mode 100644 index 00000000..d2b9f27f --- /dev/null +++ b/tests/test_dataset_row_tool_bindings.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import json +import uuid +from types import SimpleNamespace + +import pytest +from mcp.server.fastmcp.server import Settings as FastMCPSettings + +# FastMCP otherwise reads the application .env while this isolated test module +# imports the server. These tests use only synthetic ambient context. +FastMCPSettings.model_config["env_file"] = None + +from asaree.mcp_servers import workspace_server # noqa: E402 +from asaree.services.protocol_execution import ( # noqa: E402 + _ambient_meta_for, + _node_run_context, + _resolve_dataset_tool_config, +) + + +def _ctx(meta: dict): + return SimpleNamespace(request_context=SimpleNamespace(meta=SimpleNamespace(model_extra=meta))) + + +def _graph(*, driver: bool = True): + agent = {"id": "agent", "type": "agent", "data": {"config": {}}} + model = {"id": "model", "type": "llm", "data": {"config": {}}} + nodes = [agent, model] + edges = [{"source": "model", "target": "agent", "targetHandle": "model"}] + if driver: + nodes.append({"id": "dataset", "type": "dataset", "data": {"config": { + "dataset_id": "11111111-1111-4111-8111-111111111111", "dataset_name": "driver", + "target_column": "secret", "description": "FULL DICTIONARY", "dictionary_available": True, + }}}) + edges.append({"source": "dataset", "target": "agent", "targetHandle": "dataset", + "data": {"dataset_input": {"mode": "per_row", "columns": ["question"]}}}) + return {"nodes": nodes, "edges": edges} + + +@pytest.mark.asyncio +async def test_row_node_context_is_private_and_never_inherits_cell_head(monkeypatch, tmp_path): + async def forbid_registration(*_args, **_kwargs): + pytest.fail("a supplied row view must never look up the full driver") + + monkeypatch.setattr("asaree.services.protocol_execution.fetch_owned_registration", forbid_registration) + monkeypatch.setattr( + "asaree.services.protocol_execution.head_data_locator", + lambda _wid: ("/shared/train.csv", "secret"), + ) + inputs = [{"name": "driver", "dataset_id": "11111111-1111-4111-8111-111111111111", + "raw_sha256": "a" * 64, "row_index": 0, "columns": ["question"], + "values": {"question": "only row"}, "path": str(tmp_path / "row.csv"), + "mode": "per_row", "target_column": ""}] + meta, dataset = await _node_run_context( + _graph(), "agent", "shared-cell", uuid.uuid4(), protocol_run_id=uuid.uuid4(), + row_input_context={"dataset_id": inputs[0]["dataset_id"], "raw_sha256": "a" * 64, + "row_index": 0, "columns": ["question"], "values": {"question": "only row"}}, + row_inputs=inputs, + ) + assert dataset.seeded == () + assert meta["dataset_mode"] == "per_row" + assert meta["row_inputs"] == inputs + assert meta["data_path"] == inputs[0]["path"] + assert "/shared/" not in json.dumps(meta) + assert _resolve_dataset_tool_config(_graph(), "agent", row_mode=True)["tool_names"] == [ + "asaree-workspace.open_workspace", "asaree-workspace.workspace_status", "scikit-learn-mcp.describe_dataset" + ] + + +def test_row_without_driver_has_authoritative_empty_inputs_and_no_locator(monkeypatch): + monkeypatch.setattr("asaree.services.protocol_execution.head_data_locator", lambda _wid: ("/shared/train.csv", "y")) + meta = _ambient_meta_for(_graph(driver=False), "agent", "shared", row_mode=True, row_inputs=[]) + assert meta["dataset_mode"] == "per_row" + assert meta["row_inputs"] == [] + assert "data_path" not in meta and "data_slots" not in meta + + +@pytest.mark.asyncio +async def test_row_open_and_status_are_bound_and_reject_overrides(tmp_path): + row_file = tmp_path / "view.csv" + row_file.write_text("question\nhello\n", encoding="utf-8") + dataset_id = "11111111-1111-4111-8111-111111111111" + row = {"dataset_id": dataset_id, "raw_sha256": "b" * 64, "row_index": 2} + entry = {"name": "driver", "dataset_id": dataset_id, "mode": "per_row", "path": str(row_file), + "columns": ["question"], "target_column": ""} + ctx = _ctx({"motoro.workspace_id": "_protocol_runs/run-a", "motoro.ambient.dataset_row": row, + "motoro.ambient.row_inputs": [entry]}) + result = json.loads(await workspace_server.open_workspace(ctx=ctx)) + assert result == {"workspace_id": "_protocol_runs/run-a", "dataset_mode": "per_row", + "dataset_id": dataset_id, "raw_sha256": "b" * 64, "row_index": 2, + "columns": ["question"], "data_path": str(row_file), "target_column": ""} + assert json.loads(await workspace_server.open_workspace(name="other", ctx=ctx)) == { + "error": "row workspace is fixed to the authorized Agent view"} + assert json.loads(workspace_server.workspace_status("_protocol_runs/run-b", ctx)) == { + "error": "workspace override is not authorized for this Agent"} + status = json.loads(workspace_server.workspace_status(ctx=ctx)) + assert status["inputs"][0]["path"] == str(row_file) + assert "values" not in json.dumps(status) + + +def test_row_tool_config_and_ambient_do_not_leak_target_or_dictionary(): + config = _resolve_dataset_tool_config(_graph(), "agent", row_mode=True) + assert all("dictionary" not in tool and "split" not in tool for tool in config["tool_names"]) + assert "asaree-workspace.train_test_split" not in config["tool_names"] + assert "asaree-eda.get_data_dictionary" not in config["tool_names"] + assert "FULL DICTIONARY" not in json.dumps(_ambient_meta_for( + _graph(), "agent", None, row_mode=True, row_inputs=[] + )) diff --git a/tests/test_dataset_row_workspace_lifecycle.py b/tests/test_dataset_row_workspace_lifecycle.py new file mode 100644 index 00000000..5eaa5b8b --- /dev/null +++ b/tests/test_dataset_row_workspace_lifecycle.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import hashlib +import uuid +from pathlib import Path + +import pytest + +from asaree.services import dataset_row_workspaces as workspaces +from asaree.services.dataset_row_csv import DatasetRowCsvError, read_row_source + + +@pytest.fixture +def row_setup(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + root = tmp_path / "workspace-root" + monkeypatch.setattr(workspaces, "_workspace_root", lambda: root) + original = tmp_path / "registered.csv" + original.write_text("question,answer,extra\nfirst,one,x\nsecond,two,y\n", encoding="utf-8") + data = original.read_bytes() + dataset_id = str(uuid.uuid4()) + source = read_row_source( + dataset_id=dataset_id, + raw_path=str(original), + raw_sha256=hashlib.sha256(data).hexdigest(), + ) + return root, original, source + + +def _binding(dataset_id: str, agent: str, row: int, columns: list[str], target: str | None = None) -> dict: + return { + "dataset_id": dataset_id, + "dataset_name": "Owned registration", + "agent_node_id": agent, + "columns": columns, + "row_index": row, + "target_column": target, + } + + +def test_rows_and_retry_get_fresh_attempt_files_without_inherited_scratch(row_setup) -> None: + root, _, source = row_setup + agent = "agent" + runs = [uuid.uuid4(), uuid.uuid4(), uuid.uuid4()] + first = workspaces.prepare_agent_row_inputs( + run_id=runs[0], agent_node_id=agent, source=source, + bindings=[_binding(source.dataset_id, agent, 0, ["question"])], + )["row_inputs"][0] + second = workspaces.prepare_agent_row_inputs( + run_id=runs[1], agent_node_id=agent, source=source, + bindings=[_binding(source.dataset_id, agent, 1, ["question"])], + )["row_inputs"][0] + retry = workspaces.prepare_agent_row_inputs( + run_id=runs[2], agent_node_id=agent, source=source, + bindings=[_binding(source.dataset_id, agent, 0, ["question"])], + )["row_inputs"][0] + assert len({first["path"], second["path"], retry["path"]}) == 3 + assert all(not (Path(item["path"]).parent / "scratch.py").exists() for item in (first, second, retry)) + assert (root / "_protocol_runs" / str(runs[0]) / "agents").is_dir() + + +def test_agents_get_separate_directories_and_column_views(row_setup) -> None: + _, _, source = row_setup + run_id = uuid.uuid4() + result_a = workspaces.prepare_agent_row_inputs( + run_id=run_id, agent_node_id="agent-a", source=source, + bindings=[_binding(source.dataset_id, "agent-a", 0, ["question"])], + )["row_inputs"][0] + result_b = workspaces.prepare_agent_row_inputs( + run_id=run_id, agent_node_id="agent-b", source=source, + bindings=[_binding(source.dataset_id, "agent-b", 0, ["answer"])], + )["row_inputs"][0] + assert result_a["columns"] == ["question"] + assert result_b["columns"] == ["answer"] + assert Path(result_a["path"]).parent != Path(result_b["path"]).parent + + +def test_repeated_same_run_view_is_idempotent_and_unselected_target_is_valid(row_setup) -> None: + _, original, source = row_setup + run_id = uuid.uuid4() + binding = _binding(source.dataset_id, "node/../agent", 1, ["question", "extra"]) + first = workspaces.prepare_agent_row_inputs( + run_id=run_id, agent_node_id="node/../agent", source=source, bindings=[binding] + ) + Path(first["row_inputs"][0]["path"]).parent.joinpath("scratch.txt").write_text("private") + second = workspaces.prepare_agent_row_inputs( + run_id=run_id, agent_node_id="node/../agent", source=source, bindings=[binding] + ) + assert first == second + assert second["row_inputs"][0]["target_column"] == "" + assert Path(second["row_inputs"][0]["path"]).is_relative_to(Path(first["row_inputs"][0]["path"]).parents[2]) + assert original.read_text(encoding="utf-8").startswith("question,answer,extra") + + +def test_empty_agent_edge_and_selected_registration_target(row_setup) -> None: + _, _, source = row_setup + empty = workspaces.prepare_agent_row_inputs( + run_id=uuid.uuid4(), agent_node_id="agent", source=source, + bindings=[_binding(source.dataset_id, "other", 0, ["question"])], + ) + assert empty["row_inputs"] == [] + result = workspaces.prepare_agent_row_inputs( + run_id=uuid.uuid4(), agent_node_id="agent", source=source, + bindings=[_binding(source.dataset_id, "agent", 0, ["question"], "answer")], + )["row_inputs"][0] + assert result["target_column"] == "" + selected_target = workspaces.prepare_agent_row_inputs( + run_id=uuid.uuid4(), agent_node_id="agent", source=source, + bindings=[_binding(source.dataset_id, "agent", 0, ["question", "answer"], "answer")], + )["row_inputs"][0] + assert selected_target["target_column"] == "answer" + + +def test_source_identity_and_corrupted_existing_snapshot_fail_closed(row_setup) -> None: + _, _, source = row_setup + run_id = uuid.uuid4() + binding = _binding(source.dataset_id, "agent", 0, ["question"]) + result = workspaces.prepare_agent_row_inputs( + run_id=run_id, agent_node_id="agent", source=source, bindings=[binding] + )["row_inputs"][0] + Path(result["path"]).write_text("question\ncorrupted\n", encoding="utf-8") + with pytest.raises(DatasetRowCsvError, match="workspace_conflict"): + workspaces.prepare_agent_row_inputs( + run_id=run_id, agent_node_id="agent", source=source, bindings=[binding] + ) + wrong = _binding(str(uuid.uuid4()), "agent", 0, ["question"]) + with pytest.raises(DatasetRowCsvError, match="source_identity_mismatch"): + workspaces.prepare_agent_row_inputs( + run_id=uuid.uuid4(), agent_node_id="agent", source=source, bindings=[wrong] + ) diff --git a/tests/test_protocol_runs.py b/tests/test_protocol_runs.py index b8ce8f52..5e615ebf 100644 --- a/tests/test_protocol_runs.py +++ b/tests/test_protocol_runs.py @@ -166,6 +166,7 @@ def test_test_run_response_accepts_projected_resource_usage() -> None: status="pending", error=None, protocol_revision_id=None, + dataset_row=None, created_at=now, updated_at=now, observations=[], diff --git a/tests/test_reported_metrics.py b/tests/test_reported_metrics.py index 8c95ca56..7618d2a5 100644 --- a/tests/test_reported_metrics.py +++ b/tests/test_reported_metrics.py @@ -332,6 +332,18 @@ async def test_mcp_metric_captures_the_complete_tool_result_without_projection( assert result.observations[0].value_type is None +@pytest.mark.asyncio +async def test_absent_provider_output_is_unavailable_despite_empty_handoff() -> None: + run = SimpleNamespace( + id=uuid4(), replicate_result_id=None, + node_runs={"agent": {"status": "completed", "output_text": "", "final_output_available": False}}, + ) + result = await collect_reported_metrics( + run, _plan("asaree.agent_output", {"agent_node_id": "agent"}), {"nodes": [], "edges": []}, + ) + assert result.observations[0].status == "unavailable" + + @pytest.mark.asyncio async def test_agent_output_report_is_unavailable_when_agent_did_not_complete() -> None: run = SimpleNamespace( diff --git a/workspace-core/src/asaree_workspace_core/__init__.py b/workspace-core/src/asaree_workspace_core/__init__.py index e002531b..7b4fe781 100644 --- a/workspace-core/src/asaree_workspace_core/__init__.py +++ b/workspace-core/src/asaree_workspace_core/__init__.py @@ -25,6 +25,7 @@ META_KEY_DATASET_NAMES, META_KEY_OWNER_ID, META_KEY_WORKSPACE_ID, + ambient_value_from_ctx, dataset_names_from_meta, meta_mapping_from_ctx, owner_id_from_meta, @@ -70,6 +71,7 @@ "META_KEY_DATASET_NAMES", "META_KEY_OWNER_ID", "META_KEY_WORKSPACE_ID", + "ambient_value_from_ctx", "dataset_names_from_meta", "meta_mapping_from_ctx", "owner_id_from_meta", diff --git a/workspace-core/src/asaree_workspace_core/context.py b/workspace-core/src/asaree_workspace_core/context.py index 71da8df6..4e51a21d 100644 --- a/workspace-core/src/asaree_workspace_core/context.py +++ b/workspace-core/src/asaree_workspace_core/context.py @@ -118,6 +118,8 @@ def resolve_owner_id_from_ctx(ctx: Any, *, required: bool = False) -> str: # workspace_id above cannot fully serve — can resolve it without the model # retyping a name it read out of its prompt. META_KEY_DATASET_NAMES = "motoro.ambient.dataset_names" +META_KEY_ROW_INPUTS = "motoro.ambient.row_inputs" +META_KEY_DATASET_ROW = "motoro.ambient.dataset_row" def dataset_names_from_meta(meta: Mapping[str, Any] | None) -> list[str]: @@ -134,6 +136,11 @@ def dataset_names_from_meta(meta: Mapping[str, Any] | None) -> list[str]: return [v for v in value if isinstance(v, str) and v] +def ambient_value_from_ctx(key: str, ctx: Any) -> Any: + """Read an ASAREE ambient value from a FastMCP context.""" + return meta_mapping_from_ctx(ctx).get(f"motoro.ambient.{key}") if meta_mapping_from_ctx(ctx) else None + + def resolve_dataset_name(explicit: str, meta: Mapping[str, Any] | None) -> str: """Resolve which registered dataset a call means (issue #1455 follow-on). From 206e7bb706cdc0ab4adb87ff8e717516d3b10f03 Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Mon, 5 Oct 2026 17:46:31 -0700 Subject: [PATCH 09/16] feat: publish canvas and design as unified experiment versions --- CONTEXT.md | 6 +- .../protocol/DatasetExecutionSummary.tsx | 60 ++++++++++++++ .../protocol/DatasetRowHistory.test.tsx | 16 ++-- .../src/components/protocol/DesignTab.tsx | 26 +++++- .../protocol/ExperimentSidePanel.tsx | 34 ++++++-- .../protocol/ExperimentVersionView.tsx | 60 ++++++++++++++ .../components/protocol/ProtocolCanvas.tsx | 54 ++----------- .../src/components/protocol/ResultsTab.tsx | 28 +++---- .../protocol/RunConfirmDialog.test.tsx | 4 +- .../components/protocol/RunConfirmDialog.tsx | 12 +-- frontend/src/components/protocol/RunsTab.tsx | 16 +++- .../protocol/edges/DatasetInputPanel.test.tsx | 13 +++ .../protocol/edges/DatasetInputPanel.tsx | 9 ++- .../components/protocol/protocolNodeTypes.ts | 27 +++++++ .../src/pages/ProtocolCanvasPage.test.tsx | 46 +++++++++++ frontend/src/pages/ProtocolCanvasPage.tsx | 80 +++++++++++++------ frontend/src/types/protocols.ts | 8 ++ src/asaree/api/experiments.py | 39 +++++---- src/asaree/api/protocols.py | 23 ++++-- ...2c3d4e5f60_publish_experiment_snapshots.py | 25 ++++++ src/asaree/models/protocol_revision.py | 11 +++ src/asaree/services/design_generation.py | 17 +++- src/asaree/services/design_revisions.py | 3 + src/asaree/services/experiment_run_results.py | 46 ++++++++--- src/asaree/services/experiment_versions.py | 50 ++++++++++++ src/asaree/services/factorial_cells.py | 15 ++++ src/asaree/services/protocol_execution.py | 32 +++++++- src/asaree/services/protocol_revisions.py | 27 ++++++- src/asaree/services/protocol_runs.py | 28 +++++-- tests/test_dataset_row_analysis_exclusion.py | 7 +- tests/test_design_generation.py | 61 ++++++++++++++ tests/test_experiment_versions.py | 73 +++++++++++++++++ 32 files changed, 786 insertions(+), 170 deletions(-) create mode 100644 frontend/src/components/protocol/DatasetExecutionSummary.tsx create mode 100644 frontend/src/components/protocol/ExperimentVersionView.tsx create mode 100644 frontend/src/components/protocol/protocolNodeTypes.ts create mode 100644 frontend/src/pages/ProtocolCanvasPage.test.tsx create mode 100644 src/asaree/migrations/versions/1b2c3d4e5f60_publish_experiment_snapshots.py create mode 100644 src/asaree/services/experiment_versions.py create mode 100644 tests/test_experiment_versions.py diff --git a/CONTEXT.md b/CONTEXT.md index f3eeae8e..2f5271a2 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -12,6 +12,10 @@ factor-combination cell and compares the recorded outcomes. A factorial study comparing agent protocol conditions. Owns the design declaration, the protocol canvas, and all runs. _Avoid_: Project, study, trial +**Experiment version**: +A published snapshot of an experiment's canvas, design settings, and measurement configuration. Runs and results belong to the version that produced them; subsequent edits form the experiment's draft. +_Avoid_: Canvas version, independent design/publication versions + **Factor**: A declared experimental variable with discrete levels, bound to at least one canvas field before cells can run. _Avoid_: Parameter, variable @@ -108,7 +112,7 @@ The draft editable graph. Production runs never read it directly. _Avoid_: Workflow, pipeline **Published revision**: -The immutable protocol snapshot a run executes; created by publishing the canvas. +The immutable published snapshot a run executes; for an experiment it includes the canvas and experiment settings together. _Avoid_: Protocol version **Gated pair**: diff --git a/frontend/src/components/protocol/DatasetExecutionSummary.tsx b/frontend/src/components/protocol/DatasetExecutionSummary.tsx new file mode 100644 index 00000000..e91b7827 --- /dev/null +++ b/frontend/src/components/protocol/DatasetExecutionSummary.tsx @@ -0,0 +1,60 @@ +import { useQuery } from '@tanstack/react-query' +import { datasetsApi, experimentsApi, protocolsApi } from '@/api/client' +import { Card } from '@/components/ui/card' +import { datasetRowTopologyError, rowBindingForNode } from '@/lib/datasetRows' +import type { Protocol, ProtocolGraph } from '@/types/protocols' + +export function DatasetExecutionSummary({ graph, protocol, experimentId, cellCount, replicateCount }: { + graph?: ProtocolGraph + protocol?: Protocol + experimentId: string + cellCount: number + replicateCount: number +}) { + const binding = graph && rowBindingForNode(graph) + const datasetId = binding?.config?.dataset_id + const source = useQuery({ + queryKey: ['datasets', datasetId, 'row-schema'], + queryFn: () => datasetsApi.getRowSchema(datasetId!), + enabled: !!datasetId, + retry: false, + }) + const publication = useQuery({ + queryKey: ['protocols', protocol?.id, 'row-published-graph', protocol?.published_revision_id], + queryFn: () => protocolsApi.getRevision(protocol!.id, protocol!.published_revision_id!), + enabled: !!protocol?.published_revision_id, + }) + const results = useQuery({ + queryKey: ['experiments', experimentId, 'run-results', protocol?.id, protocol?.published_revision_id], + queryFn: () => experimentsApi.getRunResults(experimentId, { protocol_id: protocol!.id }), + enabled: !!protocol?.id && !!protocol.published_revision_id, + refetchInterval: 5000, + }) + const error = graph && datasetRowTopologyError(graph) + const publishedBinding = publication.data && rowBindingForNode(publication.data.graph) + const summary = results.data?.row_summary + const rows = source.data?.row_count + const columnsMissing = binding?.input.mode === 'per_row' && source.data && binding.input.columns.some(column => !source.data.columns.includes(column)) + + return +

Dataset execution

+
+

Experiment draft · proposed design

+ {!graph ?

Loading canvas inputs…

: binding ? <> +

Per row · {binding.config?.dataset_name ?? datasetId ?? 'Unregistered Dataset'} · {rows === undefined ? 'row count unavailable' : `${rows} original rows`}

+

{cellCount} cells × {replicateCount} replicates per cell × {rows ?? '?'} rows = {rows === undefined ? '?' : cellCount * replicateCount * rows} executions

+

Each row runs the complete protocol. Agents receive their selected columns.

+ {(error || source.error || columnsMissing || rows === 0) &&

{error || source.error?.message || (columnsMissing ? 'Selected columns are absent from the original CSV.' : 'The original Dataset has no rows.')}

} + :

Whole dataset · one protocol execution per replicate.

} +
+
+

Production runs · {protocol?.published_revision_id ? `experiment version ${protocol.published_revision}` : 'not published'}

+ {!protocol?.published_revision_id ?

Publish the experiment before running.

: publication.isError || results.isError ?

Published execution forecast unavailable.

: !publication.data || !results.data ?

Loading published inputs…

: results.data.consumption_mode === 'per_row' ? <> +

Per row · {publishedBinding?.config?.dataset_name ?? publishedBinding?.config?.dataset_id ?? 'Dataset'} · {summary?.row_count ?? '?'} original rows

+ {summary?.parent_replicate_count ?

{summary.cell_count} cells / {summary.parent_replicate_count} generated replicates × {summary.row_count ?? '?'} rows = {summary.expected ?? '?'} executions

:

No generated replicates. Generate cells before running a batch.

} +

Uses the design saved with this experiment version. Review the forecast in Runs before launching.

+ :

Whole dataset · one protocol execution per replicate.

} +

Draft edits take effect after publishing; design changes require generation before running.

+
+ +} diff --git a/frontend/src/components/protocol/DatasetRowHistory.test.tsx b/frontend/src/components/protocol/DatasetRowHistory.test.tsx index 11d492ac..86660bfb 100644 --- a/frontend/src/components/protocol/DatasetRowHistory.test.tsx +++ b/frontend/src/components/protocol/DatasetRowHistory.test.tsx @@ -1,7 +1,7 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query' import { fireEvent, render, screen, waitFor } from '@testing-library/react' import { afterEach, expect, it, vi } from 'vitest' -import { experimentsApi, protocolsApi } from '@/api/client' +import { experimentsApi } from '@/api/client' import type { Experiment, ExperimentRunResults } from '@/types/experiments' import { ResultsInspectorPanel, ResultsTab } from './ResultsTab' afterEach(() => vi.restoreAllMocks()) @@ -20,15 +20,9 @@ it('preserves failed attempt output while inspecting the latest row', async () = expect(experimentsApi.getRunResults).toHaveBeenCalledWith('e', scope) }) -it('history selectors request their exact design and publication scope', async () => { - vi.spyOn(experimentsApi, 'listDesignRevisions').mockResolvedValue([{ id: 'design-old', revision: 1 }] as never) - vi.spyOn(protocolsApi, 'listRevisions').mockResolvedValue([{ id: 'publication-old', revision: 1 }] as never) +it('a selected experiment version requests one publication scope without independent design controls', async () => { const results = vi.spyOn(experimentsApi, 'getRunResults').mockResolvedValue({ consumption_mode: 'per_row', row_results: [], row_summary: null } as unknown as ExperimentRunResults) - render() - await screen.findByRole('combobox', { name: 'Results design revision' }) - await waitFor(() => expect(screen.getAllByText('Revision 1')).toHaveLength(2)) - fireEvent.change(screen.getByRole('combobox', { name: 'Results design revision' }), { target: { value: 'design-old' } }) - fireEvent.change(await screen.findByRole('combobox', { name: 'Results published revision' }), { target: { value: 'publication-old' } }) - await waitFor(() => expect(results).toHaveBeenLastCalledWith('e', { protocol_id: 'p', design_revision_id: 'design-old', protocol_revision_id: 'publication-old' })) - expect(await screen.findByText('Read-only historical scope')).toBeInTheDocument() + render() + await waitFor(() => expect(results).toHaveBeenLastCalledWith('e', { protocol_id: 'p', protocol_revision_id: 'publication-old' })) + expect(screen.queryByRole('combobox', { name: 'Results design revision' })).not.toBeInTheDocument() }) diff --git a/frontend/src/components/protocol/DesignTab.tsx b/frontend/src/components/protocol/DesignTab.tsx index e1a3d4d6..7a1abb93 100644 --- a/frontend/src/components/protocol/DesignTab.tsx +++ b/frontend/src/components/protocol/DesignTab.tsx @@ -25,10 +25,11 @@ import { LEVEL_TYPE_LABELS, levelTypeOf } from './factorLevels' import { FactorEditorDialog } from './FactorEditorDialog' import { InfoTooltip } from './InfoTooltip' import { MetricsEditor } from './MetricsEditor' +import { DatasetExecutionSummary } from './DatasetExecutionSummary' export { MetricsEditor } from './MetricsEditor' import { normalizeDesignMetrics } from '@/lib/metricCatalog' import type { ProtocolCanvasHandle } from './ProtocolCanvas' -import type { ProtocolEdge, ProtocolGraph, ProtocolNode } from '@/types/protocols' +import type { Protocol, ProtocolEdge, ProtocolGraph, ProtocolNode } from '@/types/protocols' import { COORDINATION_STRATEGY_CATALOG, type CoordinationStrategySlug, @@ -332,13 +333,17 @@ function FactorsEditor({ export function DesignTab({ experiment, protocolId, + protocol, canvasRef, onDesignUpdatePendingChange, + onDraftBusyChange, }: { experiment: Experiment protocolId: string | undefined + protocol?: Protocol canvasRef: RefObject onDesignUpdatePendingChange: (pending: boolean) => void + onDraftBusyChange?: (busy: boolean) => void }) { const queryClient = useQueryClient() const graphQuery = useProtocolGraph(protocolId) @@ -581,6 +586,21 @@ export function DesignTab({ onDesignUpdatePendingChange(needsDesignUpdate) return () => onDesignUpdatePendingChange(false) }, [needsDesignUpdate, onDesignUpdatePendingChange]) + useEffect(() => { + onDraftBusyChange?.(needsDesignUpdate || metadataDraftChanged || isAutosavingMetadata || generateMutation.isPending) + }, [needsDesignUpdate, metadataDraftChanged, isAutosavingMetadata, generateMutation.isPending, onDraftBusyChange]) + useEffect(() => { + if (!protocolId) return + let cancelled = false + protocolsApi.get(protocolId).then(fresh => { + if (cancelled) return + queryClient.setQueryData(['protocols', 'for-experiment', experiment.id], previous => previous + ? (previous.published_revision ?? 0) > (fresh.published_revision ?? 0) ? previous + : { ...previous, published_revision_id: fresh.published_revision_id, published_revision: fresh.published_revision, has_unpublished_changes: fresh.has_unpublished_changes } + : fresh) + }).catch(() => {}) + return () => { cancelled = true } + }, [experiment.updated_at, experiment.id, protocolId, queryClient]) // Keep the latest delayed write in a ref so unmounting (navigating away or // switching to Runs) can flush it immediately. Once a snapshot has been @@ -703,7 +723,7 @@ export function DesignTab({

Nothing on the canvas is deleted, and switching back restores this state. Cells already generated stay - readable under the design revision that produced them, but they'll need regenerating before the new + readable under the experiment version that produced them, but they'll need regenerating before the new strategy's results are comparable.

@@ -784,6 +804,8 @@ export function DesignTab({
+ +
) : (
- {activeTab === 'design' ? ( + {versionControls} +
- ) : activeTab === 'runs' ? ( +
+ {viewingHistory && !version ?

The selected version is unavailable. Return to draft to continue.

: activeTab === 'design' && version ? : activeTab === 'runs' && version ? : activeTab === 'runs' ? ( { onResultSelection({ type: 'row', rowResultId, scope: { protocol_id: protocol?.id } }); openPanel('results') }} onViewResult={viewReplicateResult} /> - ) : ( + ) : activeTab === 'results' ? ( onResultSelection({ ...selection, scope: { protocol_id: protocol?.id, protocol_revision_id: version?.id } })} /> - )} + ) : null}
)} diff --git a/frontend/src/components/protocol/ExperimentVersionView.tsx b/frontend/src/components/protocol/ExperimentVersionView.tsx new file mode 100644 index 00000000..fafec42c --- /dev/null +++ b/frontend/src/components/protocol/ExperimentVersionView.tsx @@ -0,0 +1,60 @@ +import { useState } from 'react' +import { useQuery } from '@tanstack/react-query' +import { Background, ReactFlow, type NodeTypes, type NodeProps } from '@xyflow/react' +import { Card } from '@/components/ui/card' +import { Dialog, DialogContent, DialogHeader, DialogTitle } from '@/components/ui/dialog' +import { NODE_TYPES } from './protocolNodeTypes' +import { ProtocolCanvasActionsProvider } from './ProtocolCanvasContext' +import type { ProtocolRevision } from '@/types/protocols' +import { experimentsApi } from '@/api/client' +import { datasetRowBindings } from '@/lib/datasetRows' +import { Button } from '@/components/ui/button' +import type { ResultsSelection } from './ResultsTab' + +// Historical graphs have no editing callbacks or autosave effects. Keep them +// outside ProtocolCanvas so browsing can never write into the draft cache. +const READ_ONLY_NODE_TYPES = Object.fromEntries(Object.entries(NODE_TYPES as NodeTypes).map(([kind, Renderer]) => [kind, (props: NodeProps) =>
])) as NodeTypes +const noop = () => {} +const READ_ONLY_ACTIONS = { experimentLocked: true, requestConnectorAdd: noop, requestMainEdgeAdd: noop, requestEdgeInsert: noop, requestRunNode: noop, requestMakeFactor: noop, requestEditFactor: noop, metricsForNode: () => [], convertLegacyOutputContract: noop } + +export function ExperimentVersionCanvas({ version }: { version: ProtocolRevision }) { + const [selectedId, setSelectedId] = useState(null) + const selected = version.graph.nodes.find(node => node.id === selectedId) + return ({ ...edge, sourceHandle: edge.sourceHandle ?? undefined, targetHandle: edge.targetHandle ?? undefined }))} + onNodeClick={(_, node) => setSelectedId(node.id)} + nodesDraggable={false} nodesConnectable={false} edgesReconnectable={false} deleteKeyCode={null} fitView + proOptions={{ hideAttribution: true }} + > + { if (!open) setSelectedId(null) }}>{String(selected?.data.label ?? selected?.id)} · Version {version.revision}

Saved settings · read-only

{JSON.stringify(selected?.data, null, 2)}
 + +} + +export function ExperimentVersionDesign({ version }: { version: ProtocolRevision }) { + const snapshot = version.experiment_snapshot + if (!snapshot) return

This legacy publication saved only its canvas. Its original design settings were not captured.

+ const spec = snapshot.design_spec + return
+

Experiment version {version.revision} · read-only

+ {snapshot.hypothesis &&

Hypothesis

{snapshot.hypothesis}

} +

Factors

{spec?.factors?.length ? spec.factors.map(factor =>

{factor.name}

{factor.levels.map((level, index) => factor.level_labels?.[index] ?? (typeof level === 'string' ? level : JSON.stringify(level))).join(' · ')}

) :

No factors declared.

}

{spec?.replicates ?? 1} replicates per cell

 +

Execution settings

Design: {snapshot.design_type ?? 'Not specified'}

Coordination: {spec?.coordination_strategy?.slug ?? 'sequential'}

Randomization seed: {spec?.randomization_seed ?? 'Not set'}

 +

Dataset execution

{datasetRowBindings(version.graph).map(binding =>

{binding.config?.dataset_name ?? 'Dataset'} → {String(version.graph.nodes.find(node => node.id === binding.edge.target)?.data.label ?? binding.edge.target)}

{binding.input.mode === 'per_row' ? `One execution per row · columns: ${binding.input.columns.join(', ')}` : 'Whole dataset · one execution per replicate'}

)}{!datasetRowBindings(version.graph).length &&

No dataset inputs connected.

} +

Metrics

{(snapshot.measurement_plan?.metrics ?? spec?.metrics ?? []).map(metric =>

{metric.name}{metric.primary ? ' · Primary' : ''}

{metric.direction ?? 'maximize'} · {metric.aggregation ?? 'mean'}

)}{!(snapshot.measurement_plan?.metrics?.length || spec?.metrics?.length) &&

No metrics declared.

} + {!!snapshot.measurement_plan?.producers.length &&

Measurement producers

{snapshot.measurement_plan.producers.map(producer =>

{producer.producer_id} · {producer.kind}

{Object.entries(producer.outputs).map(([key, value]) => `${key} → ${value}`).join(' · ')}

)}} + {snapshot.task_brief &&

Task brief

{typeof snapshot.task_brief === 'string' ? snapshot.task_brief : JSON.stringify(snapshot.task_brief, null, 2)}

} +
+} + +export function ExperimentVersionRuns({ experimentId, protocolId, version, onSelectResult }: { experimentId: string; protocolId?: string; version: ProtocolRevision; onSelectResult: (selection: ResultsSelection) => void }) { + const scope = { protocol_id: protocolId, protocol_revision_id: version.id } + const query = useQuery({ queryKey: ['experiments', experimentId, 'run-results', protocolId, version.id], queryFn: () => experimentsApi.getRunResults(experimentId, scope), refetchInterval: 5000 }) + if (query.isLoading) return

Loading version runs…

+ if (!query.data) return

Could not load version runs.

+ const rows = query.data.consumption_mode === 'per_row' + ? (query.data.row_results ?? []).map(row => ({ key: row.row_result_id, label: `${row.replicate_label} · row ${row.dataset_row.row_index}`, status: row.status, runId: row.latest_attempt?.run_id, selection: { type: 'row' as const, rowResultId: row.row_result_id, scope } })) + : query.data.replicates.map(row => ({ key: row.replicate_label, label: row.replicate_label, status: row.status, runId: row.run_id, selection: { type: 'replicate' as const, replicateLabel: row.replicate_label, scope } })) + return

Runs for experiment version {version.revision}. Return to draft to launch new runs.

{!rows.length &&

No runs planned for this version.

}
{rows.map(row => )}
ReplicateStatusDetails
{row.label}{row.status.replaceAll('_', ' ')}
+} diff --git a/frontend/src/components/protocol/ProtocolCanvas.tsx b/frontend/src/components/protocol/ProtocolCanvas.tsx index 0a7f2cf5..ae618bac 100644 --- a/frontend/src/components/protocol/ProtocolCanvas.tsx +++ b/frontend/src/components/protocol/ProtocolCanvas.tsx @@ -129,21 +129,7 @@ import { SkillBrowserPanel } from './SkillBrowserPanel' import { SKILL_BROWSE, nodeDataForSkill } from './skillCatalog' import { SkillNodeInspector } from './SkillNodeInspector' import { InteractEdge } from './edges/InteractEdge' -import { AgentNode } from './nodes/AgentNode' -import { CriticGateNode } from './nodes/CriticGateNode' -import { ToolStepNode } from './nodes/ToolStepNode' -import { DatasetNode } from './nodes/DatasetNode' -import { ModelNode } from './nodes/ModelNode' -import { McpClientToolNode } from './nodes/McpClientToolNode' -import { McpToolNode } from './nodes/McpToolNode' -import { MemoryNode } from './nodes/MemoryNode' -import { OutputParserNode } from './nodes/OutputParserNode' -import { ReasonActPatternNode } from './nodes/ReasonActPatternNode' -import { ScriptNode } from './nodes/ScriptNode' -import { SingleAgentBaselinePatternNode } from './nodes/SingleAgentBaselinePatternNode' -import { OkfBundleNode } from './nodes/OkfBundleNode' -import { OkfDocumentNode } from './nodes/OkfDocumentNode' -import { SkillNode } from './nodes/SkillNode' +import { NODE_TYPES } from "./protocolNodeTypes" import { ProtocolCanvasMenu } from './ProtocolCanvasMenu' import { MODEL_NODE_TYPES, @@ -160,38 +146,6 @@ import { // every other slot's source below it. const TOP_EDGE_SLOTS = new Set(['architectural_pattern', 'skill', 'dataset', 'knowledge']) -const NODE_TYPES = { - agent: AgentNode, - sub_agent: AgentNode, - // Both MCP-tool types render through the same component (as the five LLM - // provider types do) -- they carry identical data and differ only in - // whether their server was picked in the browser or in a dropdown. - mcp_tool: McpToolNode, - mcp_scikit_learn: McpToolNode, - // The exception to "both MCP-tool types render the same": a client tool has - // its own icon and hue, since where its server came from is the one thing - // that distinguishes it. Same data, same inspector. - mcp_client_tool: McpClientToolNode, - critic_gate: CriticGateNode, - tool_step: ToolStepNode, - // All five LLM provider types render through the same component -- it - // derives icon/accent/placeholder from data.config.provider, not from - // which of these five keys it was registered under. - model_anthropic: ModelNode, - model_openai: ModelNode, - model_azure_foundry: ModelNode, - model_openrouter: ModelNode, - model_local: ModelNode, - memory: MemoryNode, - output_parser: OutputParserNode, - dataset: DatasetNode, - skill: SkillNode, - okf_bundle: OkfBundleNode, - okf_document: OkfDocumentNode, - script: ScriptNode, - pattern_reason_act: ReasonActPatternNode, - pattern_single_agent_baseline: SingleAgentBaselinePatternNode, -} // Every edge (plain or connector) renders through InteractEdge -- no edge // ever has an explicit `type`, so overriding xyflow's own built-in // "default" key covers all of them, matching how none of them are wired @@ -628,8 +582,12 @@ export const ProtocolCanvas = forwardRef protocolsApi.publish(protocolId), + mutationFn: async () => { + await protocolsApi.update(protocolId, { graph: toPersistedGraph(nodes, edges) }) + return protocolsApi.publish(protocolId) + }, onSuccess: async (published) => { + queryClient.invalidateQueries({ queryKey: ['protocols', protocolId, 'revisions'] }) if (published.experiment_id) { queryClient.setQueryData(protocolForExperimentQueryKey(published.experiment_id), published) } diff --git a/frontend/src/components/protocol/ResultsTab.tsx b/frontend/src/components/protocol/ResultsTab.tsx index 7d672a34..4edd5558 100644 --- a/frontend/src/components/protocol/ResultsTab.tsx +++ b/frontend/src/components/protocol/ResultsTab.tsx @@ -1,7 +1,7 @@ import { useState } from 'react' import { useQuery } from '@tanstack/react-query' import { AlertTriangle, ChevronDown, ChevronRight, CircleDollarSign, Clock3, Coins, Cpu, Download, ExternalLink, Trophy, X } from 'lucide-react' -import { experimentsApi, protocolsApi, type ResultsScope } from '@/api/client' +import { experimentsApi, type ResultsScope } from '@/api/client' import { Badge } from '@/components/ui/badge' import { Button } from '@/components/ui/button' import { Card } from '@/components/ui/card' @@ -416,8 +416,8 @@ function latestObsoleteRun(replicate: ResultReplicate | null): ObsoleteRun | nul } export type ResultsSelection = - | { type: 'cell'; cellLabel: string } - | { type: 'replicate'; replicateLabel: string } + | { type: 'cell'; cellLabel: string; scope?: ResultsScope } + | { type: 'replicate'; replicateLabel: string; scope?: ResultsScope } | { type: 'row'; rowResultId: string; scope: ResultsScope } function resultsMetricKey(metric: DesignMetric): string { @@ -452,8 +452,8 @@ export function ResultsInspectorPanel({ }) { const [expandedObsoleteRuns, setExpandedObsoleteRuns] = useState>(() => new Set()) const resultsQuery = useQuery({ - queryKey: ['experiments', experimentId, 'run-results', selection?.type === 'row' ? selection.scope : null], - queryFn: () => experimentsApi.getRunResults(experimentId, selection?.type === 'row' ? selection.scope : undefined), + queryKey: ['experiments', experimentId, 'run-results', selection?.scope], + queryFn: () => experimentsApi.getRunResults(experimentId, selection?.scope), enabled: selection !== null, refetchInterval: 5000, }) @@ -504,7 +504,7 @@ export function ResultsInspectorPanel({ Current - Obsolete{obsoleteRuns.length > 0 ? ` (${obsoleteRuns.length})` : ''} + {obsoleteRuns.length > 0 && Earlier versions ({obsoleteRuns.length})} Prior attempts{supersededRuns.length > 0 ? ` (${supersededRuns.length})` : ''} @@ -609,26 +609,22 @@ export function ResultsTab({ experimentName, experiment, onSelectResult, + versionId, + versionControls, }: { protocolId?: string experimentId: string experimentName: string experiment: Experiment onSelectResult: (selection: ResultsSelection) => void + versionId?: string + versionControls?: React.ReactNode }) { const [metricPreference, setMetricPreference] = useState(null) const [expandedResultCells, setExpandedResultCells] = useState>(() => new Set()) const [downloading, setDownloading] = useState(false) - const [designRevisionId, setDesignRevisionId] = useState('') - const [protocolRevisionId, setProtocolRevisionId] = useState('') - const designs = useQuery({ queryKey: ['experiments', experimentId, 'design-revisions'], queryFn: () => experimentsApi.listDesignRevisions(experimentId), enabled: !!protocolId }) - const publications = useQuery({ queryKey: ['protocols', protocolId, 'revisions'], queryFn: () => protocolsApi.listRevisions(protocolId!), enabled: !!protocolId }) - const scope: ResultsScope = { protocol_id: protocolId, design_revision_id: designRevisionId || undefined, protocol_revision_id: protocolRevisionId || undefined } - const scopeControls = protocolId ?
- - - {(designRevisionId || protocolRevisionId) &&

Read-only historical scope

} -
: null + const scope: ResultsScope = { protocol_id: protocolId, protocol_revision_id: versionId || undefined } + const scopeControls = versionControls const resultsQuery = useQuery({ queryKey: ['experiments', experimentId, 'run-results', scope], queryFn: () => experimentsApi.getRunResults(experimentId, scope), refetchInterval: 5000 }) if (resultsQuery.isLoading) return
if (resultsQuery.isError || !resultsQuery.data) return <>{scopeControls}

{resultsQuery.error?.message ?? 'Could not load this experiment’s results.'}

diff --git a/frontend/src/components/protocol/RunConfirmDialog.test.tsx b/frontend/src/components/protocol/RunConfirmDialog.test.tsx index cee5526f..0636cde4 100644 --- a/frontend/src/components/protocol/RunConfirmDialog.test.tsx +++ b/frontend/src/components/protocol/RunConfirmDialog.test.tsx @@ -17,7 +17,7 @@ describe('RunConfirmDialog', () => { />, ) - expect(screen.getByText('Canvas version:')).toBeInTheDocument() + expect(screen.getByText('Experiment version:')).toBeInTheDocument() expect(screen.getByText('Published v7')).toBeInTheDocument() }) @@ -96,7 +96,7 @@ describe('RunConfirmDialog', () => { />, ) - expect(screen.getByText('Canvas has not been published')).toBeInTheDocument() + expect(screen.getByText('Experiment has not been published')).toBeInTheDocument() expect(screen.getByRole('button', { name: 'Publish & Run replicate' })).toBeInTheDocument() expect(screen.queryByRole('button', { name: /published v/ })).not.toBeInTheDocument() }) diff --git a/frontend/src/components/protocol/RunConfirmDialog.tsx b/frontend/src/components/protocol/RunConfirmDialog.tsx index 4764da04..9328f09f 100644 --- a/frontend/src/components/protocol/RunConfirmDialog.tsx +++ b/frontend/src/components/protocol/RunConfirmDialog.tsx @@ -108,7 +108,7 @@ export function RunConfirmDialog({
{scope.type === 'all-cells' && (

- The published canvas will run each pending replicate across {scope.cellCount}{' '} + The published experiment will run each pending replicate across {scope.cellCount}{' '} {scope.cellCount === 1 ? 'cell' : 'cells'}. {scope.replicateCount > scope.pendingReplicateCount ? ` ${scope.replicateCount - scope.pendingReplicateCount} previously completed replicates will be skipped.` @@ -117,7 +117,7 @@ export function RunConfirmDialog({ )} {scope.type === 'selected-cells' && (

- The published canvas will run {scope.pendingReplicateCount + scope.rerunReplicateCount} replicate + The published experiment will run {scope.pendingReplicateCount + scope.rerunReplicateCount} replicate {scope.pendingReplicateCount + scope.rerunReplicateCount === 1 ? '' : 's'} across {scope.cellCount}{' '} {scope.cellCount === 1 ? 'cell' : 'cells'}. {scope.replicateCount - scope.pendingReplicateCount - scope.rerunReplicateCount > 0 @@ -132,7 +132,7 @@ export function RunConfirmDialog({

-
Canvas version:
+
Experiment version:
{publishedRevision === null ? 'not published' : `Published v${publishedRevision}`}
@@ -170,11 +170,11 @@ export function RunConfirmDialog({ {needsPublishChoice && (
-

{neverPublished ? 'Canvas has not been published' : 'Canvas has unpublished changes'}

+

{neverPublished ? 'Experiment has not been published' : 'Experiment has unpublished changes'}

{neverPublished - ? 'Runs always use a published canvas. Publish the canvas you are viewing to run it.' - : `This run will use published canvas v${publishedRevision}. Publish the latest canvas first to run the changes you are viewing.`} + ? 'Publish the experiment to freeze its canvas and design settings before running.' + : `This run will use experiment version ${publishedRevision}. Publish the draft experiment first to run your changes.`}

{publishError &&

{publishError}

}
diff --git a/frontend/src/components/protocol/RunsTab.tsx b/frontend/src/components/protocol/RunsTab.tsx index f494e539..fe3f2dc6 100644 --- a/frontend/src/components/protocol/RunsTab.tsx +++ b/frontend/src/components/protocol/RunsTab.tsx @@ -9,7 +9,7 @@ import { Skeleton } from '@/components/ui/skeleton' import { ApiError, experimentsApi, protocolsApi } from '@/api/client' import { displayFactorLevel, factorValueKey, groupReplicatesIntoCells, type ExperimentalCell } from '@/lib/experiment' import { factorBindingDiscrepancies, type FactorBindingDiscrepancy } from '@/lib/factorBindings' -import { protocolForExperimentQueryKey } from '@/lib/protocolGraph' +import { protocolForExperimentQueryKey, protocolGraphQueryKey, toPersistedGraph } from '@/lib/protocolGraph' import type { Experiment, ResultCell, ResultReplicate, Trial } from '@/types/experiments' import type { Protocol } from '@/types/protocols' import { DatasetRowRuns } from './DatasetRowRuns' @@ -257,8 +257,13 @@ export function RunAllCellsButton({ } const publishAndRunMutation = useMutation({ - mutationFn: () => protocolsApi.publish(protocol!.id), + mutationFn: async () => { + const live = queryClient.getQueryData<{ nodes: Node[]; edges: Edge[] }>(protocolGraphQueryKey(protocol!.id)) + if (live) await protocolsApi.update(protocol!.id, { graph: toPersistedGraph(live.nodes, live.edges) }) + return protocolsApi.publish(protocol!.id) + }, onSuccess: (published) => { + queryClient.invalidateQueries({ queryKey: ['protocols', protocol!.id, 'revisions'] }) queryClient.setQueryData(protocolForExperimentQueryKey(experimentId), published) beginRun() }, @@ -650,8 +655,13 @@ function RunReplicateButton({ }, }) const publishAndRunMutation = useMutation({ - mutationFn: () => protocolsApi.publish(protocol!.id), + mutationFn: async () => { + const live = queryClient.getQueryData<{ nodes: Node[]; edges: Edge[] }>(protocolGraphQueryKey(protocol!.id)) + if (live) await protocolsApi.update(protocol!.id, { graph: toPersistedGraph(live.nodes, live.edges) }) + return protocolsApi.publish(protocol!.id) + }, onSuccess: (published) => { + queryClient.invalidateQueries({ queryKey: ['protocols', protocol!.id, 'revisions'] }) queryClient.setQueryData(protocolForExperimentQueryKey(experimentId), published) runMutation.mutate() }, diff --git a/frontend/src/components/protocol/edges/DatasetInputPanel.test.tsx b/frontend/src/components/protocol/edges/DatasetInputPanel.test.tsx index 0110cc6d..499bcd18 100644 --- a/frontend/src/components/protocol/edges/DatasetInputPanel.test.tsx +++ b/frontend/src/components/protocol/edges/DatasetInputPanel.test.tsx @@ -41,3 +41,16 @@ it('reports source errors without saving drafts', async () => { await waitFor(() => expect(screen.getByRole('alert')).toHaveTextContent('Source unavailable')) expect(state.setEdges).not.toHaveBeenCalled() }) +it('acknowledges Apply and explains that production needs publication in both modes', async () => { + mount() + await screen.findByText(/3 original rows/) + fireEvent.click(screen.getByLabelText('Per row')) + fireEvent.click(screen.getByRole('checkbox', { name: 'question' })) + fireEvent.click(screen.getByRole('button', { name: 'Apply' })) + expect(screen.getByRole('status')).toHaveTextContent('Per row applied to the draft') + expect(screen.getByRole('status')).toHaveTextContent('Publish the canvas to update production runs and results') + fireEvent.click(screen.getByLabelText('Whole dataset')) + expect(screen.queryByRole('status')).not.toBeInTheDocument() + fireEvent.click(screen.getByRole('button', { name: 'Apply' })) + expect(screen.getByRole('status')).toHaveTextContent('Whole dataset applied to the draft') +}) diff --git a/frontend/src/components/protocol/edges/DatasetInputPanel.tsx b/frontend/src/components/protocol/edges/DatasetInputPanel.tsx index 0ed9e18f..e9021050 100644 --- a/frontend/src/components/protocol/edges/DatasetInputPanel.tsx +++ b/frontend/src/components/protocol/edges/DatasetInputPanel.tsx @@ -15,6 +15,7 @@ export function DatasetInputPanel({ edgeId, source, disabled = false }: { edgeId const saved = edge?.data?.dataset_input const [mode, setMode] = useState(saved?.mode ?? 'whole_dataset') const [columns, setColumns] = useState(saved?.mode === 'per_row' ? saved.columns : []) + const [applied, setApplied] = useState(null) const id = (nodes.find(node => node.id === source)?.data.config as { dataset_id?: string })?.dataset_id const schema = useQuery({ queryKey: ['datasets', id, 'row-schema'], queryFn: () => datasetsApi.getRowSchema(id!), enabled: !!id }) const input: DatasetInput = mode === 'per_row' ? { mode, columns } : { mode } @@ -29,15 +30,17 @@ export function DatasetInputPanel({ edgeId, source, disabled = false }: { edgeId if (mode === 'per_row') data.dataset_input = { mode, columns: schema.data!.columns.filter(column => columns.includes(column)) } return { ...item, data } })) + setApplied(mode) } return

Dataset input

- {(['whole_dataset', 'per_row'] as const).map(value => )} + {(['whole_dataset', 'per_row'] as const).map(value => )}
- {schema.data &&

{schema.data.row_count} original rows. Each row runs one complete protocol; Agents share one driver and select their own columns.

} - {mode === 'per_row' &&
{[...new Set([...(schema.data?.columns ?? []), ...columns])].map(column => )}
} + {schema.data &&

{schema.data.row_count} original rows. {mode === 'per_row' ? 'Each row runs one complete protocol; Agents share one driver and select their own columns.' : 'Each replicate runs the protocol with the whole dataset.'}

} + {mode === 'per_row' &&
{[...new Set([...(schema.data?.columns ?? []), ...columns])].map(column => )}
} {error &&

{error}

} + {applied ?

{applied === 'per_row' ? 'Per row' : 'Whole dataset'} applied to the draft. Publish the canvas to update production runs and results.

:

Apply changes the draft connection. Publish the canvas to use it for production runs and results.

}
} diff --git a/frontend/src/components/protocol/protocolNodeTypes.ts b/frontend/src/components/protocol/protocolNodeTypes.ts new file mode 100644 index 00000000..a277709b --- /dev/null +++ b/frontend/src/components/protocol/protocolNodeTypes.ts @@ -0,0 +1,27 @@ +import { AgentNode } from './nodes/AgentNode' +import { CriticGateNode } from './nodes/CriticGateNode' +import { ToolStepNode } from './nodes/ToolStepNode' +import { DatasetNode } from './nodes/DatasetNode' +import { ModelNode } from './nodes/ModelNode' +import { McpClientToolNode } from './nodes/McpClientToolNode' +import { McpToolNode } from './nodes/McpToolNode' +import { MemoryNode } from './nodes/MemoryNode' +import { OutputParserNode } from './nodes/OutputParserNode' +import { ReasonActPatternNode } from './nodes/ReasonActPatternNode' +import { ScriptNode } from './nodes/ScriptNode' +import { SingleAgentBaselinePatternNode } from './nodes/SingleAgentBaselinePatternNode' +import { OkfBundleNode } from './nodes/OkfBundleNode' +import { OkfDocumentNode } from './nodes/OkfDocumentNode' +import { SkillNode } from './nodes/SkillNode' + +export const NODE_TYPES = { + agent: AgentNode, sub_agent: AgentNode, + mcp_tool: McpToolNode, mcp_scikit_learn: McpToolNode, mcp_client_tool: McpClientToolNode, + critic_gate: CriticGateNode, tool_step: ToolStepNode, + model_anthropic: ModelNode, model_openai: ModelNode, model_azure_foundry: ModelNode, + model_openrouter: ModelNode, model_local: ModelNode, + memory: MemoryNode, output_parser: OutputParserNode, dataset: DatasetNode, + skill: SkillNode, okf_bundle: OkfBundleNode, okf_document: OkfDocumentNode, + script: ScriptNode, pattern_reason_act: ReasonActPatternNode, + pattern_single_agent_baseline: SingleAgentBaselinePatternNode, +} diff --git a/frontend/src/pages/ProtocolCanvasPage.test.tsx b/frontend/src/pages/ProtocolCanvasPage.test.tsx new file mode 100644 index 00000000..3a0b6585 --- /dev/null +++ b/frontend/src/pages/ProtocolCanvasPage.test.tsx @@ -0,0 +1,46 @@ +import { useState } from 'react' +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { fireEvent, render, screen } from '@testing-library/react' +import { MemoryRouter, Route, Routes } from 'react-router-dom' +import { afterEach, expect, it, vi } from 'vitest' +import { experimentsApi, protocolsApi } from '@/api/client' +import type { Experiment } from '@/types/experiments' +import type { Protocol, ProtocolRevision } from '@/types/protocols' +import { ProtocolCanvasPage } from './ProtocolCanvasPage' + +vi.mock('@/components/AppHeader', () => ({ AppHeader: () => null })) +vi.mock('@/components/protocol/ProtocolCanvas', () => ({ ProtocolCanvas: () => { + const [draft, setDraft] = useState('Unsaved draft') + return setDraft(event.target.value)} /> +} })) +vi.mock('@/components/protocol/ExperimentVersionView', () => ({ ExperimentVersionCanvas: ({ version }: { version: ProtocolRevision }) =>

Saved canvas {version.revision}

})) +vi.mock('@/components/protocol/ResultsTab', () => ({ ResultsInspectorPanel: () => null })) +vi.mock('@/components/protocol/ExperimentSidePanel', () => ({ ExperimentSidePanel: (props: { versionControls: React.ReactNode; viewingHistory: boolean; version?: ProtocolRevision; draftExperiment: Experiment; resultsExperiment: Experiment }) =>
{props.versionControls}

{props.viewingHistory ? `Historical design: ${props.version?.experiment_snapshot?.hypothesis}` : `Draft design: ${props.draftExperiment?.hypothesis}`}

Results design: {props.resultsExperiment?.hypothesis}

})) +afterEach(() => vi.restoreAllMocks()) + +it('switches canvas and design together and preserves local draft edits on return', async () => { + const experiment = { id: 'e', name: 'Experiment', hypothesis: 'Current', design_spec: {}, design_type: 'factorial' } as Experiment + const protocol = { id: 'p', name: 'Protocol', description: null, created_at: '2026-01-01', updated_at: '2026-01-01', experiment_id: 'e', graph: { nodes: [], edges: [] }, published_revision_id: 'v2', published_revision: 2, has_unpublished_changes: true } as Protocol + const versions = [{ id: 'v1', revision: 1, graph: { nodes: [], edges: [] }, published_at: '2026-01-01', experiment_snapshot: { hypothesis: 'Original', design_spec: {} } }, { id: 'v2', revision: 2, graph: { nodes: [], edges: [] }, published_at: '2026-01-02', experiment_snapshot: { hypothesis: 'Published', design_spec: {} } }] as unknown as ProtocolRevision[] + vi.spyOn(experimentsApi, 'get').mockResolvedValue(experiment) + vi.spyOn(experimentsApi, 'listReplicates').mockResolvedValue([]) + vi.spyOn(experimentsApi, 'getDesignImpact').mockResolvedValue({} as Awaited>) + vi.spyOn(experimentsApi, 'getRunResults').mockRejectedValue(new Error('No results')) + vi.spyOn(experimentsApi, 'listTrials').mockResolvedValue([]) + vi.spyOn(protocolsApi, 'list').mockResolvedValue([protocol]) + vi.spyOn(protocolsApi, 'listRevisions').mockResolvedValue(versions) + const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }) + render(} />) + const draft = await screen.findByRole('textbox', { name: 'Draft canvas' }) + await screen.findByText('Results design: Published') + fireEvent.change(draft, { target: { value: 'Local edit' } }) + fireEvent.change(screen.getByRole('combobox', { name: 'Experiment version' }), { target: { value: 'v1' } }) + expect(screen.getByText('Saved canvas 1')).toBeInTheDocument() + expect(screen.getByText('Historical design: Original')).toBeInTheDocument() + expect(screen.getByText('Results design: Original')).toBeInTheDocument() + expect(screen.queryByRole('button', { name: 'Publish experiment' })).not.toBeInTheDocument() + fireEvent.click(screen.getByRole('button', { name: 'Return to draft' })) + expect(screen.getByRole('textbox', { name: 'Draft canvas' })).toHaveValue('Local edit') + expect(screen.getByText('Draft design: Current')).toBeInTheDocument() + expect(screen.getByText('Results design: Published')).toBeInTheDocument() +}) diff --git a/frontend/src/pages/ProtocolCanvasPage.tsx b/frontend/src/pages/ProtocolCanvasPage.tsx index 25be6f52..24100c18 100644 --- a/frontend/src/pages/ProtocolCanvasPage.tsx +++ b/frontend/src/pages/ProtocolCanvasPage.tsx @@ -1,10 +1,11 @@ -import { useRef, useState } from 'react' +import { useEffect, useRef, useState } from 'react' import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query' import { ReactFlowProvider } from '@xyflow/react' import { AlertTriangle, Lock, Target, Trophy, type LucideIcon } from 'lucide-react' import { Link, useParams } from 'react-router-dom' import { AppHeader } from '@/components/AppHeader' import { ExperimentSidePanel } from '@/components/protocol/ExperimentSidePanel' +import { ExperimentVersionCanvas } from '@/components/protocol/ExperimentVersionView' import { ProtocolCanvas, type ProtocolCanvasHandle } from '@/components/protocol/ProtocolCanvas' import { ResultsInspectorPanel, type ResultsSelection } from '@/components/protocol/ResultsTab' import { Button } from '@/components/ui/button' @@ -19,9 +20,12 @@ import { applyExperimentRenameToProtocolCache, generatedProtocolName, protocolForExperimentQueryKey, + protocolGraphQueryKey, + toPersistedGraph, } from '@/lib/protocolGraph' import type { Experiment, Replicate } from '@/types/experiments' import type { Protocol } from '@/types/protocols' +import type { Node, Edge } from '@xyflow/react' // Click-to-rename, the pattern for anything created with a placeholder name: // no gate before creating, edit the name in place once you're looking at what @@ -139,7 +143,7 @@ function TopBarStats({ experiment, cells, obsoleteRunCount = 0, rowSummary }: { ) } -function ProtocolPublicationControl({ protocol, experimentId }: { protocol: Protocol; experimentId: string }) { +function ProtocolPublicationControl({ protocol, experimentId, draftBusy }: { protocol: Protocol; experimentId: string; draftBusy: boolean }) { const queryClient = useQueryClient() const [confirmOpen, setConfirmOpen] = useState(false) const trialsQuery = useQuery({ @@ -152,10 +156,16 @@ function ProtocolPublicationControl({ protocol, experimentId }: { protocol: Prot // become stale against, so do not overstate the impact in this warning. const affectedReplicateCount = (trialsQuery.data ?? []).filter((trial) => !!trial.run_id && !trial.obsolete).length const publishMutation = useMutation({ - mutationFn: () => protocolsApi.publish(protocol.id), + mutationFn: async () => { + const live = queryClient.getQueryData<{ nodes: Node[]; edges: Edge[] }>(protocolGraphQueryKey(protocol.id)) + if (live) await protocolsApi.update(protocol.id, { graph: toPersistedGraph(live.nodes, live.edges) }) + return protocolsApi.publish(protocol.id) + }, onSuccess: (published) => { queryClient.setQueryData(protocolForExperimentQueryKey(experimentId), published) queryClient.invalidateQueries({ queryKey: ['experiments', experimentId, 'runs'] }) + queryClient.invalidateQueries({ queryKey: ['experiments', experimentId, 'run-results'] }) + queryClient.invalidateQueries({ queryKey: ['protocols', protocol.id, 'revisions'] }) setConfirmOpen(false) }, }) @@ -163,7 +173,7 @@ function ProtocolPublicationControl({ protocol, experimentId }: { protocol: Prot ? protocol.has_unpublished_changes ? `Draft changes · production uses v${protocol.published_revision}` : `Published v${protocol.published_revision}` - : 'No published canvas' + : 'Experiment draft · not published' const error = publishMutation.error instanceof ApiError && typeof publishMutation.error.detail === 'string' ? publishMutation.error.detail : null function requestPublish() { if (affectedReplicateCount > 0) setConfirmOpen(true) @@ -172,33 +182,33 @@ function ProtocolPublicationControl({ protocol, experimentId }: { protocol: Prot return ( <>
- + {status} {error && {error}}
- Publish a new canvas version? + Publish a new experiment version?

- Publishing this draft changes the canvas version used by future runs. {affectedReplicateCount}{' '} - previously run replicate{affectedReplicateCount === 1 ? '' : 's'} will be marked obsolete because{affectedReplicateCount === 1 ? ' it was' : ' they were'} run against the current version. + Publish the canvas and design settings together as a new experiment version. Future runs will use it.

-

The results remain available for comparison; they are not deleted.

+

Existing runs and results remain available under their earlier experiment version.

 @@ -220,6 +230,9 @@ export function ProtocolCanvasPage() { // needs to be imperative rather than a plain prop. const canvasRef = useRef(null) const [resultSelection, setResultSelection] = useState(null) + const [versionId, setVersionId] = useState('') + const [draftBusy, setDraftBusy] = useState(false) + useEffect(() => { setVersionId(''); setResultSelection(null) }, [experimentId]) const experimentQuery = useQuery({ queryKey: ['experiments', experimentId], @@ -254,11 +267,6 @@ export function ProtocolCanvasPage() { enabled: !!experimentId, }) - const replicatesQuery = useQuery({ - queryKey: ['experiments', experimentId, 'replicates'], - queryFn: () => experimentsApi.listReplicates(experimentId!), - enabled: !!experimentId, - }) const runResultsQuery = useQuery({ queryKey: ['experiments', experimentId, 'run-results', protocolQuery.data?.id, protocolQuery.data?.published_revision_id], queryFn: () => experimentsApi.getRunResults(experimentId!, { protocol_id: protocolQuery.data?.id }), @@ -273,6 +281,22 @@ export function ProtocolCanvasPage() { const unboundFactors = experimentQuery.data && protocolQuery.data ? unboundFactorNames(experimentQuery.data.design_spec, protocolQuery.data.graph) : [] + const versionsQuery = useQuery({ queryKey: ['protocols', protocolQuery.data?.id, 'revisions'], queryFn: () => protocolsApi.listRevisions(protocolQuery.data!.id), enabled: !!protocolQuery.data?.id }) + const version = versionsQuery.data?.find(item => item.id === versionId) + const displayedExperiment = experimentQuery.data && version?.experiment_snapshot + ? { ...experimentQuery.data, ...version.experiment_snapshot } : experimentQuery.data + const resultsVersion = version ?? versionsQuery.data?.find(item => item.id === protocolQuery.data?.published_revision_id) + const resultsExperiment = experimentQuery.data && resultsVersion?.experiment_snapshot + ? { ...experimentQuery.data, ...resultsVersion.experiment_snapshot } : displayedExperiment + const resultReplicates: Replicate[] | undefined = runResultsQuery.data?.replicates.map(row => ({ + id: row.replicate_label, cell_id: row.cell_label, cell_label: row.cell_label, + replicate_label: row.replicate_label, replicate_number: row.replicate_number, + design_revision_id: resultsVersion?.design_revision_id ?? '', + factor_values: row.factor_values, metric_values: Object.keys(row.metric_values).length ? row.metric_values : null, + run_id: row.run_id, workspace_id: null, artifacts: null, + created_at: row.updated_at, updated_at: row.updated_at, + })) + const versionControls =
{version &&

Canvas, design, and results show this saved version. Read-only.

}{versionsQuery.isError &&

Could not load experiment versions.

}
return (
@@ -285,20 +309,27 @@ export function ProtocolCanvasPage() { ← Experiments - {experimentQuery.data && } + {experimentQuery.data && (versionId ? {experimentQuery.data.name} : )} {experimentQuery.data?.locked_at && Locked} - {experimentQuery.data && } + {!versionId && resultsExperiment && }
- {protocolQuery.data && experimentId && ( + {protocolQuery.data && experimentId && !versionId && ( <> - + )}
+ {versionId &&
Viewing {version?.experiment_snapshot ? 'experiment version' : 'legacy canvas'} {version?.revision ?? '…'} · read-only
}
Could not load this experiment's protocol.

) : ( +
- {experimentId && experimentQuery.data && setResultSelection(null)} />} +
+ {versionId && {version ? :

{versionsQuery.isError ? 'Could not load this experiment version. Return to draft to continue.' : 'Loading experiment version…'}

}} + {experimentId && resultsExperiment && setResultSelection(null)} />} )}
diff --git a/frontend/src/types/protocols.ts b/frontend/src/types/protocols.ts index ef8c96d1..1b47a4e9 100644 --- a/frontend/src/types/protocols.ts +++ b/frontend/src/types/protocols.ts @@ -243,6 +243,14 @@ export interface ProtocolRevision { revision: number graph: ProtocolGraph published_at: string + experiment_snapshot?: { + hypothesis: string | null + design_type: string + design_spec: import('./experiments').Experiment['design_spec'] + measurement_plan: import('./experiments').MeasurementPlan | null + task_brief: Record | null + } | null + design_revision_id?: string | null } // One "run all cells" trigger fans out into these -- one ProtocolRun per diff --git a/src/asaree/api/experiments.py b/src/asaree/api/experiments.py index 8b0a7842..f54d07aa 100644 --- a/src/asaree/api/experiments.py +++ b/src/asaree/api/experiments.py @@ -12,6 +12,7 @@ import uuid from copy import deepcopy from datetime import UTC, datetime +from types import SimpleNamespace from typing import Any from fastapi import APIRouter, HTTPException, Response @@ -988,7 +989,7 @@ async def analyze_factorial_endpoint( status_code=422, detail=f"{body.primary_metric!r} is not the experiment's declared primary metric.", ) - replicates = await list_replicates(db, experiment_id=experiment_id, revision_id=design_revision_id) + replicates = selection["replicates"] try: analysis = ( analyze_binary_factorial @@ -1053,13 +1054,16 @@ async def _factorial_analysis_selection( except RunResultsProjectionError as exc: status = 422 if str(exc) == "ambiguous_protocol" else 404 raise HTTPException(status_code=status, detail=str(exc)) from exc - design_spec = experiment.design_spec - if design_revision_id is not None: - revision = await get_revision(db, design_revision_id) - if revision is None or revision.experiment_id != experiment.id: - raise HTTPException(status_code=404, detail="design_revision_not_found") - design_spec = revision.design_spec - return {"consumption_mode": projection["consumption_mode"], "design_spec": design_spec} + replicates = [SimpleNamespace( + replicate_label=row["replicate_label"], cell_label=row["cell_label"], + factor_values=row["factor_values"], metric_values=row["metric_values"], + run_id=row.get("run_id"), workspace_id=row.get("workspace_id"), + artifacts={"metric_evaluation": row.get("metric_evaluation")}, + ) for row in projection["replicates"]] + return { + "consumption_mode": projection["consumption_mode"], + "design_spec": projection["selected_design_spec"], "replicates": replicates, + } class RunResultsResponse(BaseModel): @@ -1096,7 +1100,7 @@ async def get_experiment_results_endpoint( ) if selection["consumption_mode"] == "per_row": return ResultsResponse(**_row_analysis_unavailable()) - replicates = await list_replicates(db, experiment_id=experiment_id, revision_id=design_revision_id) + replicates = selection["replicates"] result = analyze_experiment_design(selection["design_spec"], replicates, consumption_mode="whole_dataset") return ResultsResponse(**result) @@ -1150,7 +1154,8 @@ async def export_run_results_csv_endpoint( raise HTTPException(status_code=404, detail=str(exc)) from exc row_mode = results["consumption_mode"] == "per_row" csv_text = result_rows_to_csv( - results["row_results"] if row_mode else results["replicates"], experiment.design_spec, + results["row_results"] if row_mode else results["replicates"], + results.get("selected_design_spec", experiment.design_spec), consumption_mode=results["consumption_mode"], ) filename = _UNSAFE_FILENAME_CHAR.sub("_", experiment.name.strip()) or "experiment" @@ -1183,7 +1188,8 @@ async def get_run_results_schema_endpoint( raise HTTPException(status_code=404, detail=str(exc)) from exc return result_rows_schema( results["row_results"] if results["consumption_mode"] == "per_row" else results["replicates"], - results["metric_types"], results["metric_aggregations"], experiment.design_spec, + results["metric_types"], results["metric_aggregations"], + results.get("selected_design_spec", experiment.design_spec), consumption_mode=results["consumption_mode"], ) @@ -1214,7 +1220,12 @@ async def upsert_replicate_endpoint( ) except ValueError as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc - replicate = await upsert_replicate(db, experiment_id=experiment_id, replicate_label=replicate_label, fields=fields) + try: + replicate = await upsert_replicate( + db, experiment_id=experiment_id, replicate_label=replicate_label, fields=fields + ) + except ValueError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc return ReplicateResponse.model_validate(replicate) @@ -1262,8 +1273,8 @@ async def export_replicates_csv_endpoint( "href": f"/experiments/{experiment_id}/run-results.csv", "link_text": "run-results.csv", }) - replicates = await list_replicates(db, experiment_id=experiment_id, revision_id=design_revision_id) - csv_text = replicates_to_csv(replicates_that_ran(replicates), design_spec=experiment.design_spec) + replicates = selection["replicates"] + csv_text = replicates_to_csv(replicates_that_ran(replicates), design_spec=selection["design_spec"]) filename = _UNSAFE_FILENAME_CHAR.sub("_", experiment.name.strip()) or "experiment" return Response( content=csv_text, diff --git a/src/asaree/api/protocols.py b/src/asaree/api/protocols.py index 1447f719..56fc297a 100644 --- a/src/asaree/api/protocols.py +++ b/src/asaree/api/protocols.py @@ -26,6 +26,7 @@ blocking_measurement_plan_issues, validate_experiment_measurement_plan, ) +from asaree.services.experiment_versions import publication_matches_experiment, version_design_spec from asaree.services.experiments import get_experiment from asaree.services.factor_bindings import validate_factor_bindings from asaree.services.protocol_execution import ( @@ -238,6 +239,8 @@ class ProtocolRevisionResponse(BaseModel): revision: int graph: dict[str, Any] published_at: datetime + experiment_snapshot: dict[str, Any] | None = None + design_revision_id: uuid.UUID | None = None class CreateProtocolRunRequest(BaseModel): @@ -320,7 +323,9 @@ async def _protocol_response(db: DbSession, protocol: Any) -> ProtocolResponse: graph=normalize_protocol_graph(protocol.graph), published_revision_id=published.id if published else None, published_revision=published.revision if published else None, - has_unpublished_changes=not is_draft_published(protocol, published), + has_unpublished_changes=not ( + is_draft_published(protocol, published) and await publication_matches_experiment(db, protocol, published) + ), created_at=protocol.created_at, updated_at=protocol.updated_at, ) @@ -424,7 +429,10 @@ async def publish_protocol_endpoint(protocol_id: uuid.UUID, user: CurrentUser, d raise ProtocolValidationError("; ".join(issue.message for issue in blocking_issues)) except (ProtocolValidationError, ValueError) as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc - if published is not None and is_draft_published(protocol, published): + if ( + published is not None and is_draft_published(protocol, published) + and await publication_matches_experiment(db, protocol, published) + ): try: await publish_protocol(db, protocol, owner_id=user.id) except ValueError as exc: @@ -470,6 +478,8 @@ async def get_protocol_revision_endpoint( revision=revision.revision, graph=normalize_protocol_graph(revision.graph), published_at=revision.published_at, + experiment_snapshot=revision.experiment_snapshot, + design_revision_id=revision.design_revision_id, ) @@ -541,16 +551,17 @@ async def create_test_run_endpoint( revision = await _require_published_revision(db, protocol) try: experiment = await get_experiment(db, protocol.experiment_id) if protocol.experiment_id else None - validate_coordination_strategy(experiment.design_spec if experiment is not None else None, graph=revision.graph) - validate_stage_plan(experiment.design_spec if experiment is not None else None) + published_design = version_design_spec(revision, experiment.design_spec if experiment is not None else None) + validate_coordination_strategy(published_design, graph=revision.graph) + validate_stage_plan(published_design) validate_prompt_references(graph=revision.graph) topological_order( revision.graph, - require_acyclic=not is_conversation_strategy(experiment.design_spec if experiment else None), + require_acyclic=not is_conversation_strategy(published_design), ) dataset_row = None try: - row_plan = resolve_dataset_row_plan(revision.graph, experiment.design_spec if experiment else None) + row_plan = resolve_dataset_row_plan(revision.graph, published_design) except DatasetRowInputError as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc row_index = body.row_index if body is not None else None diff --git a/src/asaree/migrations/versions/1b2c3d4e5f60_publish_experiment_snapshots.py b/src/asaree/migrations/versions/1b2c3d4e5f60_publish_experiment_snapshots.py new file mode 100644 index 00000000..550b2b39 --- /dev/null +++ b/src/asaree/migrations/versions/1b2c3d4e5f60_publish_experiment_snapshots.py @@ -0,0 +1,25 @@ +"""Freeze experiment settings and generated design with published canvases.""" + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +revision = "1b2c3d4e5f60" +down_revision = "0a1b2c3d4e5f" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.add_column("protocol_revisions", sa.Column("experiment_snapshot", postgresql.JSONB(), nullable=True)) + op.add_column("protocol_revisions", sa.Column("design_revision_id", postgresql.UUID(as_uuid=True), nullable=True)) + op.create_foreign_key( + "fk_protocol_revisions_design", "protocol_revisions", "experiment_design_revisions", + ["design_revision_id"], ["id"], deferrable=True, initially="DEFERRED", + ) + + +def downgrade() -> None: + op.drop_constraint("fk_protocol_revisions_design", "protocol_revisions", type_="foreignkey") + op.drop_column("protocol_revisions", "design_revision_id") + op.drop_column("protocol_revisions", "experiment_snapshot") diff --git a/src/asaree/models/protocol_revision.py b/src/asaree/models/protocol_revision.py index a5621f8f..5abc2715 100644 --- a/src/asaree/models/protocol_revision.py +++ b/src/asaree/models/protocol_revision.py @@ -31,6 +31,17 @@ class ProtocolRevision(Base, TimestampMixin): ) revision: Mapped[int] = mapped_column(Integer, nullable=False) graph: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False) + # An experiment publication freezes the declaration alongside its canvas. + # Null means a legacy canvas-only publication; never invent its settings. + experiment_snapshot: Mapped[dict[str, Any] | None] = mapped_column(JSONB, nullable=True) + design_revision_id: Mapped[uuid.UUID | None] = mapped_column( + # Defer the check so deleting an experiment can cascade through both + # its designs and publications in either order. Direct deletion of a + # published design is refused by the service and by this FK at commit. + UUID(as_uuid=True), ForeignKey( + "experiment_design_revisions.id", deferrable=True, initially="DEFERRED" + ), nullable=True + ) published_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False) diff --git a/src/asaree/services/design_generation.py b/src/asaree/services/design_generation.py index 7a17733c..52b24d72 100644 --- a/src/asaree/services/design_generation.py +++ b/src/asaree/services/design_generation.py @@ -30,6 +30,7 @@ from asaree.models.experiment import ResearchExperiment from asaree.models.experiment_design_revision import ExperimentDesignRevision from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.protocol_revision import ProtocolRevision from asaree.services.coordination import coordination_strategy_slug from asaree.services.design_revisions import get_current_revision, supersede_and_create from asaree.services.factorial_cells import list_replicates, upsert_replicate @@ -364,7 +365,7 @@ async def generate_design_cells( """Compute the design and materialize ``replicates`` child rows per combination (default 1), under the experiment's current design revision. - A new revision is opened exactly when the new design would *drop* a cell + An unpublished design opens a revision when the new design would *drop* a cell the current revision has — not on every change. Re-clicking generate, changing ``randomization_seed``, widening a factor's levels or raising ``replicates`` all leave every existing cell part of the new design, so @@ -372,6 +373,11 @@ async def generate_design_cells( duplicate revisions pile up, and (as before) the cells you already scored are untouched, right down to their row ids. + Once a publication references the design, its declaration and planned + slots are frozen. Any subsequent change creates a separate draft design; + results remain with the experiment version that produced them. Clicking + Generate again with the same declaration reuses the existing design. + Dropping a cell is the case that needs a revision (see services.design_revisions). This is the fix for cells outliving the design that created them: shrinking a design from 6 cells to 2 used to leave all 6 @@ -450,7 +456,12 @@ async def generate_design_cells( and coordination_strategy_slug(current.design_spec) != coordination_strategy_slug(design_spec) ) - if current is not None and not dropped and not strategy_changed: + published_design_changes = current is not None and ( + current.design_spec != design_spec or set(existing) != planned_labels or strategy_changed + ) and await db.scalar( + select(ProtocolRevision.id).where(ProtocolRevision.design_revision_id == current.id).limit(1) + ) is not None + if current is not None and not dropped and not strategy_changed and not published_design_changes: # Keep the revision and merge into its rows -- factor_values is # re-merged because a level's *value* can change without changing its # slugified label. @@ -464,7 +475,7 @@ async def generate_design_cells( revision_id = revision.id carry_over = ( {} - if strategy_changed + if strategy_changed or published_design_changes else { (_cell_key(previous.factor_values, previous.replicate_label), previous.replicate_number): previous for previous in existing.values() diff --git a/src/asaree/services/design_revisions.py b/src/asaree/services/design_revisions.py index b5a86835..694a6254 100644 --- a/src/asaree/services/design_revisions.py +++ b/src/asaree/services/design_revisions.py @@ -32,6 +32,7 @@ from asaree.models.experiment_design_revision import ExperimentDesignRevision from asaree.models.factorial_cell import FactorialCell from asaree.models.factorial_replicate_result import FactorialReplicateResult +from asaree.models.protocol_revision import ProtocolRevision class DesignRevisionError(ValueError): @@ -139,6 +140,8 @@ async def delete_revision(db: AsyncSession, revision_id: uuid.UUID) -> None: "This is the experiment's current design. Regenerate the design to replace it, " "or delete the experiment to remove it entirely." ) + if await db.scalar(select(ProtocolRevision.id).where(ProtocolRevision.design_revision_id == revision.id).limit(1)): + raise DesignRevisionError("This design belongs to a published experiment version and must be retained.") await db.delete(revision) await db.flush() diff --git a/src/asaree/services/experiment_run_results.py b/src/asaree/services/experiment_run_results.py index bd0d3765..1ab4faa6 100644 --- a/src/asaree/services/experiment_run_results.py +++ b/src/asaree/services/experiment_run_results.py @@ -553,7 +553,7 @@ async def _summarize_row_results( actual_cells = len(cells) actual_parents = len(replicates) forecast_parents = actual_parents - spec = design_revision.design_spec if design_revision is not None else design_spec + spec = design_spec factors = (spec or {}).get("factors") if isinstance(spec, dict) else None if actual_cells == 0 and isinstance(factors, list) and factors: sizes = [ @@ -680,10 +680,20 @@ async def summarize_experiment_run_results( elif protocol_revision_id is not None: raise RunResultsProjectionError("protocol_revision_not_found") + versioned = published is not None and published.experiment_snapshot is not None + if versioned: + if design_revision_id is not None and design_revision_id != published.design_revision_id: + raise RunResultsProjectionError("design_does_not_belong_to_experiment_version") + revision = ( + await db.get(ExperimentDesignRevision, published.design_revision_id) + if published.design_revision_id else None + ) + design_spec = published.experiment_snapshot.get("design_spec") + protocol_revision_id = published.id graph = published.graph if published is not None and isinstance(published.graph, dict) else {} row_plan = resolve_dataset_row_plan(graph) if published is not None else None if row_plan is not None: - return await _summarize_row_results( + projection = await _summarize_row_results( db, experiment_id=experiment_id, design_spec=design_spec, @@ -691,12 +701,17 @@ async def summarize_experiment_run_results( design_revision=revision, row_plan=row_plan, ) + projection["selected_design_spec"] = design_spec + return projection selected_revision_id = revision.id if revision is not None else None - replicates = await list_replicates(db, experiment_id=experiment_id, revision_id=selected_revision_id) - trials = await list_experiment_trials( - db, experiment_id=experiment_id, revision_id=selected_revision_id + replicates = ( + [] if versioned and revision is None else + await list_replicates(db, experiment_id=experiment_id, revision_id=selected_revision_id) ) + trials = ([] if versioned and revision is None else await list_experiment_trials( + db, experiment_id=experiment_id, revision_id=selected_revision_id + )) current_trial_run_ids = {trial.run_id for trial in trials if trial.run_id is not None} excluded_trial_run_ids: set[uuid.UUID] = set() if current_trial_run_ids: @@ -755,7 +770,7 @@ async def summarize_experiment_run_results( ).all() history_by_label: defaultdict[str, list[tuple[ProtocolRun, bool]]] = defaultdict(list) for historical_run, current_revision_id, current_published_at in history_rows: - obsolete = current_revision_id is not None and ( + obsolete = not versioned and current_revision_id is not None and ( ( historical_run.protocol_revision_id is not None and historical_run.protocol_revision_id != current_revision_id @@ -768,6 +783,14 @@ async def summarize_experiment_run_results( ) history_by_label[historical_run.replicate_label or ""].append((historical_run, obsolete)) protocol_run_ids.add(historical_run.id) + if versioned: + for trial in trials: + matching = history_by_label.get(trial.replicate_label, []) + latest = max((item[0] for item in matching), key=lambda run: (run.created_at, str(run.id)), default=None) + trial.run_id = latest.id if latest else None + trial.status = latest.status if latest else "not_started" + trial.error = latest.error if latest else None + trial.obsolete = False protocol_runs_by_id: dict[uuid.UUID, ProtocolRun] = {} if protocol_run_ids: result = await db.execute(select(ProtocolRun).where(ProtocolRun.id.in_(protocol_run_ids))) @@ -933,20 +956,20 @@ def historical_run_payload(protocol_run: ProtocolRun, *, obsolete: bool) -> dict metric_values = ( snapshot_metrics if "metric_values" in stored_attempt - else _normalize_metric_values(replicate.metric_values) + else {} if versioned else _normalize_metric_values(replicate.metric_values) ) metric_evaluation = snapshot_evaluation or ( (replicate.artifacts or {}).get("metric_evaluation") - if isinstance((replicate.artifacts or {}).get("metric_evaluation"), dict) + if not versioned and isinstance((replicate.artifacts or {}).get("metric_evaluation"), dict) else None ) - if not metric_observations and not evaluation_artifacts: + if not versioned and not metric_observations and not evaluation_artifacts: metric_observations, evaluation_artifacts = measurement_facets( (replicate.artifacts or {}).get("measurement") ) facets = _merge_legacy_facets( metric_values, - replicate.artifacts, + stored_attempt.get("artifacts") if versioned else replicate.artifacts, metric_observations, evaluation_artifacts, metrics=(design_spec or {}).get("metrics"), @@ -959,7 +982,7 @@ def historical_run_payload(protocol_run: ProtocolRun, *, obsolete: bool) -> dict metric_values, metric_evaluation = {}, None metric_observations, evaluation_artifacts = [], [] legacy_values = [] - if latest_run is None and (replicate.metric_values or replicate.artifacts): + if latest_run is None and not versioned and (replicate.metric_values or replicate.artifacts): metric_values = _normalize_metric_values(replicate.metric_values) facets = _merge_legacy_facets( metric_values, @@ -1075,6 +1098,7 @@ def historical_run_payload(protocol_run: ProtocolRun, *, obsolete: bool) -> dict primary_metric, primary_metric_direction = _primary_metric(design_spec) return { "consumption_mode": "whole_dataset", + "selected_design_spec": design_spec, "row_results": [], "row_summary": None, "overview": overview, diff --git a/src/asaree/services/experiment_versions.py b/src/asaree/services/experiment_versions.py new file mode 100644 index 00000000..dd24f2d1 --- /dev/null +++ b/src/asaree/services/experiment_versions.py @@ -0,0 +1,50 @@ +"""One publication freezes the experiment's canvas, settings, and design scope.""" + +from __future__ import annotations + +import copy +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from asaree.models.experiment import ResearchExperiment +from asaree.models.protocol import Protocol +from asaree.models.protocol_revision import ProtocolRevision +from asaree.services.design_revisions import get_current_revision + + +def experiment_settings(experiment: ResearchExperiment) -> dict[str, Any]: + return copy.deepcopy({ + "hypothesis": experiment.hypothesis, + "design_type": experiment.design_type, + "design_spec": experiment.design_spec, + "measurement_plan": experiment.measurement_plan, + "task_brief": experiment.task_brief, + }) + + +async def publication_matches_experiment( + db: AsyncSession, protocol: Protocol, publication: ProtocolRevision | None +) -> bool: + if publication is None: + return False + if protocol.experiment_id is None: + return True + experiment = await db.get(ResearchExperiment, protocol.experiment_id) + if experiment is None: + return False + current = await get_current_revision(db, experiment.id) + return ( + publication.experiment_snapshot == experiment_settings(experiment) + and publication.design_revision_id == (current.id if current else None) + ) + + +def version_design_spec(publication: ProtocolRevision | None, fallback: dict | None) -> dict | None: + snapshot = publication.experiment_snapshot if publication is not None else None + return snapshot.get("design_spec") if isinstance(snapshot, dict) else fallback + + +def version_measurement_plan(publication: ProtocolRevision | None, fallback: dict | None) -> dict | None: + snapshot = publication.experiment_snapshot if publication is not None else None + return snapshot.get("measurement_plan") if isinstance(snapshot, dict) else fallback diff --git a/src/asaree/services/factorial_cells.py b/src/asaree/services/factorial_cells.py index 901114dc..261eba9a 100644 --- a/src/asaree/services/factorial_cells.py +++ b/src/asaree/services/factorial_cells.py @@ -169,6 +169,11 @@ async def upsert_replicate( db, experiment_id=experiment_id, cell_label=cell_label, revision_id=revision_id ) if cell is None: + from asaree.models.protocol_revision import ProtocolRevision + if await db.scalar(select(ProtocolRevision.id).where( + ProtocolRevision.design_revision_id == revision_id + ).limit(1)): + raise ValueError("Generate a draft design before adding cells to a published experiment version.") cell = FactorialCell( experiment_id=experiment_id, design_revision_id=revision_id, @@ -180,12 +185,22 @@ async def upsert_replicate( if "factor_values" in fields: merged_factors = dict(cell.factor_values or {}) merged_factors.update(supplied_factors) + from asaree.models.protocol_revision import ProtocolRevision + if merged_factors != (cell.factor_values or {}) and await db.scalar(select(ProtocolRevision.id).where( + ProtocolRevision.design_revision_id == revision_id + ).limit(1)): + raise ValueError("Published experiment cell factors are immutable. Generate a draft design first.") cell.factor_values = merged_factors replicate = await get_replicate( db, experiment_id=experiment_id, replicate_label=replicate_label, revision_id=revision_id ) if replicate is None: + from asaree.models.protocol_revision import ProtocolRevision + if await db.scalar(select(ProtocolRevision.id).where( + ProtocolRevision.design_revision_id == revision_id + ).limit(1)): + raise ValueError("Generate a draft design before adding replicates to a published experiment version.") replicate = FactorialReplicateResult( cell=cell, replicate_number=replicate_number, diff --git a/src/asaree/services/protocol_execution.py b/src/asaree/services/protocol_execution.py index 5c671314..a5eae9ed 100644 --- a/src/asaree/services/protocol_execution.py +++ b/src/asaree/services/protocol_execution.py @@ -4696,7 +4696,9 @@ async def plan_cell_runs( raise ProtocolValidationError("This protocol has no linked experiment to run replicates for.") # Strategy first, so its own message wins over a pipeline requirement that # may not apply to this canvas at all -- see is_conversation_strategy. - design_spec = experiment.design_spec + publication = await db.get(ProtocolRevision, protocol_revision_id) if protocol_revision_id else None + from asaree.services.experiment_versions import version_design_spec, version_measurement_plan + design_spec = version_design_spec(publication, experiment.design_spec) measurement_plan = ( experiment.locked_measurement_plan if experiment is not None and experiment.locked_at is not None @@ -4704,6 +4706,12 @@ async def plan_cell_runs( if experiment is not None else None ) + measurement_plan = version_measurement_plan(publication, measurement_plan) + if publication is not None and publication.experiment_snapshot is not None: + from asaree.services.design_revisions import get_current_revision + generated = await get_current_revision(db, experiment_id) + if publication.design_revision_id != (generated.id if generated else None): + raise ProtocolValidationError("Publish the generated draft design before running this experiment.") validate_coordination_strategy(design_spec, graph=graph) validate_stage_plan(design_spec) validate_prompt_references(graph=graph) @@ -5071,6 +5079,9 @@ async def plan_single_replicate_run( # Same order and same reason as plan_cell_runs above. experiment = await get_experiment(db, experiment_id) design_spec = experiment.design_spec if experiment is not None else None + publication = await db.get(ProtocolRevision, protocol_revision_id) if protocol_revision_id else None + from asaree.services.experiment_versions import version_design_spec, version_measurement_plan + design_spec = version_design_spec(publication, design_spec) measurement_plan = ( experiment.locked_measurement_plan if experiment is not None and experiment.locked_at is not None @@ -5078,6 +5089,12 @@ async def plan_single_replicate_run( if experiment is not None else None ) + measurement_plan = version_measurement_plan(publication, measurement_plan) + if publication is not None and publication.experiment_snapshot is not None: + from asaree.services.design_revisions import get_current_revision + generated = await get_current_revision(db, experiment_id) + if publication.design_revision_id != (generated.id if generated else None): + raise ProtocolValidationError("Publish the generated draft design before running this experiment.") validate_coordination_strategy(design_spec, graph=graph) validate_stage_plan(design_spec) validate_prompt_references(graph=graph) @@ -5292,7 +5309,13 @@ async def _run_single_node( workspace_id = row_attempt_workspace_id(protocol_run_id) async with get_session() as db: experiment = await get_experiment(db, experiment_id) if experiment_id else None - single_design_spec = experiment.design_spec if experiment is not None else None + single_run = await db.get(ProtocolRun, protocol_run_id) + single_publication = ( + await db.get(ProtocolRevision, single_run.protocol_revision_id) + if single_run is not None and single_run.protocol_revision_id else None + ) + from asaree.services.experiment_versions import version_design_spec + single_design_spec = version_design_spec(single_publication, experiment.design_spec if experiment else None) ambient_meta, node_dataset = await _node_run_context( graph, node["id"], @@ -5352,6 +5375,7 @@ async def _run_single_node( async def run_protocol(protocol_run_id: uuid.UUID) -> None: + revision = None row_mode = False row_context_mode = False dataset_row = None @@ -5424,6 +5448,8 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: target_node_id = run.target_node_id experiment = await get_experiment(db, experiment_id) if experiment_id else None design_spec = experiment.design_spec if experiment is not None else None + from asaree.services.experiment_versions import version_design_spec + design_spec = version_design_spec(revision, design_spec) # The stage plan comes from the PINNED revision's snapshot, not from the # live design_spec: a plan edit made while this replicate was queued # would otherwise stage a cell through a pipeline its own design never @@ -5431,7 +5457,7 @@ async def run_protocol(protocol_run_id: uuid.UUID) -> None: # pre-existing behaviour; the plan is singled out because it is the one # design field that writes durable, versioned artifacts to disk. pinned_spec = design_spec - if design_revision_id is not None: + if design_revision_id is not None and (revision is None or revision.experiment_snapshot is None): pinned = await get_design_revision(db, design_revision_id) if pinned is not None and pinned.design_spec is not None: pinned_spec = pinned.design_spec diff --git a/src/asaree/services/protocol_revisions.py b/src/asaree/services/protocol_revisions.py index b7d5fbce..2dc7b56b 100644 --- a/src/asaree/services/protocol_revisions.py +++ b/src/asaree/services/protocol_revisions.py @@ -10,6 +10,8 @@ from sqlalchemy.ext.asyncio import AsyncSession from asaree.models.dataset import RegisteredDataset +from asaree.models.experiment import ResearchExperiment +from asaree.models.factorial_cell import FactorialCell from asaree.models.factorial_replicate_result import FactorialReplicateResult from asaree.models.protocol import Protocol from asaree.models.protocol_revision import ProtocolRevision @@ -17,6 +19,8 @@ from asaree.services.dataset_row_csv import DatasetRowCsvError, read_row_source from asaree.services.dataset_row_inputs import DatasetRowInputError, resolve_dataset_row_plan from asaree.services.datasets import get_dataset_by_name +from asaree.services.design_revisions import get_current_revision +from asaree.services.experiment_versions import experiment_settings, publication_matches_experiment from asaree.services.protocol_graph_schema import functional_protocol_graph_hash @@ -33,7 +37,7 @@ async def get_published_revision(db: AsyncSession, protocol: Protocol) -> Protoc async def publish_protocol( db: AsyncSession, protocol: Protocol, *, owner_id: uuid.UUID | None = None ) -> ProtocolRevision: - """Freeze the protocol's current draft as its next production revision.""" + """Freeze canvas, experiment settings, and generated design as one version.""" protocol = ( await db.execute( select(Protocol) @@ -42,6 +46,12 @@ async def publish_protocol( .execution_options(populate_existing=True) ) ).scalar_one() + experiment = ( + await db.scalar(select(ResearchExperiment).where(ResearchExperiment.id == protocol.experiment_id) + .with_for_update().execution_options(populate_existing=True)) + if protocol.experiment_id is not None else None + ) + design = await get_current_revision(db, experiment.id) if experiment is not None else None plan = resolve_dataset_row_plan(protocol.graph) if plan is not None: if owner_id is None: @@ -129,8 +139,18 @@ async def publish_protocol( if missing: raise DatasetRowCsvError("invalid_columns", "selected columns are not present in the registered CSV") published = await get_published_revision(db, protocol) - if published is not None and is_draft_published(protocol, published): + unchanged = ( + is_draft_published(protocol, published) and await publication_matches_experiment(db, protocol, published) + ) + if published is not None and unchanged: return published + if experiment is not None and experiment.locked_at is not None: + raise ValueError("Experiment is locked. Unlock it before publishing a changed experiment.") + if experiment is not None: + from asaree.services.design_generation import get_design_impact + impact = await get_design_impact(db, experiment_id=experiment.id, design_spec=experiment.design_spec) + if impact.regeneration_required: + raise ValueError("Generate cells for the draft design before publishing the experiment.") highest = ( await db.execute( @@ -141,6 +161,8 @@ async def publish_protocol( protocol_id=protocol.id, revision=(highest or 0) + 1, graph=copy.deepcopy(protocol.graph), + experiment_snapshot=experiment_settings(experiment) if experiment is not None else None, + design_revision_id=design.id if design is not None else None, published_at=datetime.now(UTC), ) db.add(revision) @@ -155,6 +177,7 @@ async def publish_protocol( await db.execute( select(ProtocolRun, FactorialReplicateResult) .join(FactorialReplicateResult, FactorialReplicateResult.run_id == ProtocolRun.id) + .join(FactorialCell, FactorialReplicateResult.cell_id == FactorialCell.id) .where(ProtocolRun.protocol_id == protocol.id) ) ).all() diff --git a/src/asaree/services/protocol_runs.py b/src/asaree/services/protocol_runs.py index b104ac4d..2ced07e3 100644 --- a/src/asaree/services/protocol_runs.py +++ b/src/asaree/services/protocol_runs.py @@ -26,6 +26,7 @@ from asaree.models.protocol_revision import ProtocolRevision from asaree.models.protocol_run import ProtocolRun from asaree.services.design_revisions import get_current_revision +from asaree.services.experiment_versions import version_design_spec, version_measurement_plan from asaree.services.factorial_cells import get_replicate, list_replicates from asaree.services.measurement_engine import MeasurementEvaluation, normalize_measurement_plan from asaree.services.measurement_migration import normalize_experiment_measurement_plan @@ -129,6 +130,7 @@ async def create_protocol_run( measurement_plan_snapshot: dict[str, Any] | None = None reference_values: dict[str, Any] = {} protocol = await db.get(Protocol, protocol_id) + publication = await db.get(ProtocolRevision, protocol_revision_id) if protocol_revision_id else None if protocol is not None and protocol.experiment_id is not None: experiment = await db.get(ResearchExperiment, protocol.experiment_id) if experiment is not None: @@ -138,6 +140,8 @@ async def create_protocol_run( effective_design_spec = ( experiment.locked_design_spec if experiment.locked_at is not None else experiment.design_spec ) + effective_design_spec = version_design_spec(publication, effective_design_spec) + effective_plan = version_measurement_plan(publication, effective_plan) attempt_plan = normalize_experiment_measurement_plan( effective_plan, (effective_design_spec or {}).get("metrics"), @@ -145,6 +149,8 @@ async def create_protocol_run( if attempt_plan["metrics"]: measurement_plan_snapshot = normalize_measurement_plan(attempt_plan) task_brief = experiment.task_brief if isinstance(experiment.task_brief, dict) else {} + if publication is not None and isinstance(publication.experiment_snapshot, dict): + task_brief = publication.experiment_snapshot.get("task_brief") or {} declared_references = task_brief.get("reference_values") if isinstance(declared_references, dict): reference_values = dict(declared_references) @@ -880,15 +886,20 @@ async def list_experiment_trials( # the result stale by comparing the run's creation time to the current # published revision's timestamp. current_revisions: dict[uuid.UUID, tuple[uuid.UUID | None, datetime | None]] = {} + versioned_protocols: set[uuid.UUID] = set() if protocol_ids: result = await db.execute( - select(Protocol.id, Protocol.published_revision_id, ProtocolRevision.published_at) + select(Protocol.id, Protocol.published_revision_id, ProtocolRevision.published_at, + ProtocolRevision.experiment_snapshot) .outerjoin(ProtocolRevision, Protocol.published_revision_id == ProtocolRevision.id) .where(Protocol.id.in_(protocol_ids)) ) + rows = result.all() current_revisions = { - protocol_id: (revision_id, published_at) for protocol_id, revision_id, published_at in result.all() + protocol_id: (revision_id, published_at) + for protocol_id, revision_id, published_at, _ in rows } + versioned_protocols = {protocol_id for protocol_id, _, _, snapshot in rows if snapshot is not None} trials = [] for replicate in replicates: @@ -908,11 +919,16 @@ async def list_experiment_trials( or (run.protocol_revision_id is None and published_at is not None and run.created_at < published_at) ) ) + prior_version = bool(run and run.protocol_id in versioned_protocols and obsolete) + if prior_version: + # Earlier executions are inspected under their experiment version. + run = None + obsolete = False if run is not None: status = run.status error = run.error updated_at = run.updated_at - elif replicate.metric_values: + elif replicate.metric_values and not prior_version: status, error, updated_at = "completed", None, replicate.updated_at else: status, error, updated_at = "not_started", None, replicate.updated_at @@ -920,11 +936,11 @@ async def list_experiment_trials( ExperimentTrial( replicate_label=replicate.replicate_label, factor_values=replicate.factor_values or {}, - metric_values=replicate.metric_values or {}, + metric_values={} if prior_version else replicate.metric_values or {}, status=status, - run_id=replicate.run_id, + run_id=run.id if run is not None else None, obsolete=obsolete, - truncated=bool((replicate.artifacts or {}).get("truncation")), + truncated=not prior_version and bool((replicate.artifacts or {}).get("truncation")), error=error, updated_at=updated_at, ) diff --git a/tests/test_dataset_row_analysis_exclusion.py b/tests/test_dataset_row_analysis_exclusion.py index ec6eac8a..db33d5ac 100644 --- a/tests/test_dataset_row_analysis_exclusion.py +++ b/tests/test_dataset_row_analysis_exclusion.py @@ -84,7 +84,7 @@ async def test_stale_parent_metric_does_not_rank_row_mode_and_legacy_export_is_b @pytest.mark.asyncio -async def test_historical_whole_publication_keeps_legacy_analysis_and_export(row_results_setup) -> None: +async def test_historical_whole_publication_ignores_later_draft_scores(row_results_setup) -> None: db, ctx = row_results_setup ctx["protocol"].graph = {"nodes": [{"id": "agent", "type": "agent", "data": {}}], "edges": []} from asaree.services.protocol_revisions import publish_protocol @@ -113,9 +113,12 @@ async def test_historical_whole_publication_keeps_legacy_analysis_and_export(row ctx["experiment"].id, ctx["user"], db, protocol_id=ctx["protocol"].id, protocol_revision_id=whole.id, ) - assert analyzed.available is True + # This publication captured no factorial declaration. Later draft edits + # and parent scores cannot manufacture historical analysis or results. + assert analyzed.available is False exported = await export_replicates_csv_endpoint( ctx["experiment"].id, ctx["user"], db, protocol_id=ctx["protocol"].id, protocol_revision_id=whole.id, ) assert exported.media_type == "text/csv" + assert "0.71" not in exported.body.decode() diff --git a/tests/test_design_generation.py b/tests/test_design_generation.py index f8f5dc9a..6dccc2b9 100644 --- a/tests/test_design_generation.py +++ b/tests/test_design_generation.py @@ -59,6 +59,67 @@ def test_generate_design_is_the_cross_product() -> None: assert {"tier": "large", "effort": "high"} in combos +@pytest.mark.asyncio +async def test_publishing_freezes_design_and_expansion_preserves_version_results(owner_id): + from asaree.models.protocol_run import ProtocolRun + from asaree.services.experiment_run_results import summarize_experiment_run_results + from asaree.services.protocol_revisions import publish_protocol + from asaree.services.protocols import create_protocol + + async with get_session() as db: + spec = {"factors": [{"name": "arm", "levels": ["a", "b"]}], "replicates": 1} + experiment = await create_experiment( + db, name=f"versions-{uuid.uuid4()}", owner_id=owner_id, design_spec=spec, + ) + spec = experiment.design_spec + protocol = await create_protocol( + db, name="Version test", owner_id=owner_id, experiment_id=experiment.id, + graph={"nodes": [], "edges": []}, + ) + parents = await generate_design_cells( + db, experiment_id=experiment.id, factors=spec["factors"], design_spec=spec, + ) + first = await publish_protocol(db, protocol) + assert (await publish_protocol(db, protocol)).id == first.id + first_parent_ids = [parent.id for parent in parents] + attempt = ProtocolRun( + protocol_id=protocol.id, owner_id=owner_id, status="completed", + protocol_revision_id=first.id, design_revision_id=first.design_revision_id, + replicate_label=parents[0].replicate_label, replicate_result_id=parents[0].id, + attempt_result={"metric_values": {"total_tokens": 42}}, + ) + db.add(attempt) + await db.flush() + parents[0].run_id = attempt.id + parents[0].metric_values = {"total_tokens": 42} + experiment.hypothesis = "Updated hypothesis" + second = await publish_protocol(db, protocol) + assert second.id != first.id + assert second.design_revision_id == first.design_revision_id + assert first.experiment_snapshot["hypothesis"] is None + old_results = await summarize_experiment_run_results( + db, experiment_id=experiment.id, design_spec=spec, + protocol_id=protocol.id, protocol_revision_id=first.id, + ) + assert old_results["replicates"][0]["run_id"] == str(attempt.id) + assert old_results["replicates"][0]["metric_values"]["total_tokens"] == 42 + current_results = await summarize_experiment_run_results( + db, experiment_id=experiment.id, design_spec=spec, protocol_id=protocol.id, + ) + assert all(row["run_id"] is None for row in current_results["replicates"]) + expanded = {**spec, "factors": [{"name": "arm", "levels": ["a", "b", "c"]}]} + experiment.design_spec = expanded + new_parents = await generate_design_cells( + db, experiment_id=experiment.id, factors=expanded["factors"], design_spec=expanded, + ) + assert len(new_parents) == 3 + assert set(first_parent_ids).isdisjoint(parent.id for parent in new_parents) + assert len(await list_replicates(db, experiment_id=experiment.id, revision_id=first.design_revision_id)) == 2 + third = await publish_protocol(db, protocol) + assert third.design_revision_id != first.design_revision_id + assert first.experiment_snapshot["design_spec"]["factors"][0]["levels"] == ["a", "b"] + + def test_generate_design_rejects_empty_factors() -> None: with pytest.raises(DesignValidationError, match="non-empty list"): generate_design([]) diff --git a/tests/test_experiment_versions.py b/tests/test_experiment_versions.py new file mode 100644 index 00000000..b3a5326c --- /dev/null +++ b/tests/test_experiment_versions.py @@ -0,0 +1,73 @@ +"""Frozen experiment settings and version-aligned analysis projections.""" + +import uuid +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from asaree.api.experiments import _factorial_analysis_selection +from asaree.services.experiment_versions import ( + experiment_settings, + publication_matches_experiment, + version_design_spec, + version_measurement_plan, +) + + +def test_snapshot_does_not_share_nested_draft_settings(): + experiment = SimpleNamespace( + hypothesis="Original", design_type="factorial", + design_spec={"factors": [{"name": "arm", "levels": ["a", "b"]}]}, + measurement_plan={"metrics": [{"name": "Quality"}]}, + task_brief={"reference_values": {"answer": "old"}}, + ) + snapshot = experiment_settings(experiment) + experiment.design_spec["factors"][0]["levels"].append("c") + experiment.measurement_plan["metrics"][0]["name"] = "New metric" + experiment.task_brief["reference_values"]["answer"] = "new" + publication = SimpleNamespace(experiment_snapshot=snapshot) + assert version_design_spec(publication, experiment.design_spec)["factors"][0]["levels"] == ["a", "b"] + assert version_measurement_plan(publication, experiment.measurement_plan)["metrics"][0]["name"] == "Quality" + assert snapshot["task_brief"]["reference_values"] == {"answer": "old"} + + +def test_explicit_empty_version_settings_never_fall_back_to_draft(): + publication = SimpleNamespace(experiment_snapshot={"design_spec": None, "measurement_plan": None}) + assert version_design_spec(publication, {"factors": ["new"]}) is None + assert version_measurement_plan(publication, {"metrics": ["new"]}) is None + legacy = SimpleNamespace(experiment_snapshot=None) + assert version_design_spec(legacy, {"replicates": 2}) == {"replicates": 2} + + +@pytest.mark.asyncio +async def test_design_only_change_marks_experiment_unpublished(monkeypatch): + experiment = SimpleNamespace(id=uuid.uuid4(), hypothesis="old", design_type="factorial", + design_spec={}, measurement_plan=None, task_brief=None) + publication = SimpleNamespace(experiment_snapshot=experiment_settings(experiment), design_revision_id=None) + protocol = SimpleNamespace(experiment_id=experiment.id) + db = SimpleNamespace(get=AsyncMock(return_value=experiment)) + monkeypatch.setattr("asaree.services.experiment_versions.get_current_revision", AsyncMock(return_value=None)) + assert await publication_matches_experiment(db, protocol, publication) + experiment.hypothesis = "new" + assert not await publication_matches_experiment(db, protocol, publication) + + +@pytest.mark.asyncio +async def test_analysis_reads_selected_attempt_metrics_and_design(monkeypatch): + version_id = uuid.uuid4() + projection = { + "consumption_mode": "whole_dataset", "selected_design_spec": {"metrics": [{"name": "old"}]}, + "replicates": [{"replicate_label": "a", "cell_label": "a", "factor_values": {"arm": "a"}, + "metric_values": {"old": 0.75}, "run_id": "old-run"}], + } + summarize = AsyncMock(return_value=projection) + monkeypatch.setattr("asaree.api.experiments.summarize_experiment_run_results", summarize) + experiment = SimpleNamespace(id=uuid.uuid4(), design_spec={"metrics": [{"name": "new"}]}) + selected = await _factorial_analysis_selection( + None, experiment, protocol_id=None, design_revision_id=None, protocol_revision_id=version_id, + ) + assert selected["design_spec"] == projection["selected_design_spec"] + assert selected["replicates"][0].metric_values == {"old": 0.75} + assert selected["replicates"][0].run_id == "old-run" + assert summarize.call_args.kwargs["protocol_revision_id"] == version_id From 7e1788383309b58401dfbbe0b08ba7b874dcfdca Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Mon, 5 Oct 2026 17:57:08 -0700 Subject: [PATCH 10/16] fix: clarify experiment version selection and group publication controls --- .../protocol/ExperimentSidePanel.tsx | 3 -- .../src/components/protocol/ResultsTab.tsx | 8 +--- .../src/pages/ProtocolCanvasPage.test.tsx | 20 ++++++++-- frontend/src/pages/ProtocolCanvasPage.tsx | 39 +++++++++++++++---- 4 files changed, 50 insertions(+), 20 deletions(-) diff --git a/frontend/src/components/protocol/ExperimentSidePanel.tsx b/frontend/src/components/protocol/ExperimentSidePanel.tsx index 0f7a9921..0c8c3dd9 100644 --- a/frontend/src/components/protocol/ExperimentSidePanel.tsx +++ b/frontend/src/components/protocol/ExperimentSidePanel.tsx @@ -72,7 +72,6 @@ export function ExperimentSidePanel({ unboundFactors, onResultSelection, version, - versionControls, viewingHistory = false, onDraftBusyChange, }: { @@ -88,7 +87,6 @@ export function ExperimentSidePanel({ unboundFactors: string[] onResultSelection: (selection: ResultsSelection | null) => void version?: ProtocolRevision - versionControls?: React.ReactNode viewingHistory?: boolean onDraftBusyChange?: (busy: boolean) => void }) { @@ -255,7 +253,6 @@ export function ExperimentSidePanel({
) : (
- {versionControls}
void versionId?: string - versionControls?: React.ReactNode }) { const [metricPreference, setMetricPreference] = useState(null) const [expandedResultCells, setExpandedResultCells] = useState>(() => new Set()) const [downloading, setDownloading] = useState(false) const scope: ResultsScope = { protocol_id: protocolId, protocol_revision_id: versionId || undefined } - const scopeControls = versionControls const resultsQuery = useQuery({ queryKey: ['experiments', experimentId, 'run-results', scope], queryFn: () => experimentsApi.getRunResults(experimentId, scope), refetchInterval: 5000 }) if (resultsQuery.isLoading) return
- if (resultsQuery.isError || !resultsQuery.data) return <>{scopeControls}

{resultsQuery.error?.message ?? 'Could not load this experiment’s results.'}

+ if (resultsQuery.isError || !resultsQuery.data) return

{resultsQuery.error?.message ?? 'Could not load this experiment’s results.'}

- if (resultsQuery.data.consumption_mode === 'per_row') return <>{scopeControls} onSelectResult({ type: 'row', rowResultId, scope })} /> + if (resultsQuery.data.consumption_mode === 'per_row') return onSelectResult({ type: 'row', rowResultId, scope })} /> const { overview, cells, replicates, metric_types: metricTypes, metric_aggregations: metricAggregations, metric_directions: metricDirections, primary_metric: primaryMetric, primary_metric_direction: primaryMetricDirection } = resultsQuery.data const metricKeys = orderedResultsMetricKeys(resultsQuery.data.metric_keys, experiment) @@ -683,7 +680,6 @@ export function ResultsTab({ return (
- {scopeControls}
{metricKeys.length > 0 && (
diff --git a/frontend/src/pages/ProtocolCanvasPage.test.tsx b/frontend/src/pages/ProtocolCanvasPage.test.tsx index 3a0b6585..0daa5924 100644 --- a/frontend/src/pages/ProtocolCanvasPage.test.tsx +++ b/frontend/src/pages/ProtocolCanvasPage.test.tsx @@ -1,6 +1,6 @@ import { useState } from 'react' import { QueryClient, QueryClientProvider } from '@tanstack/react-query' -import { fireEvent, render, screen } from '@testing-library/react' +import { act, fireEvent, render, screen } from '@testing-library/react' import { MemoryRouter, Route, Routes } from 'react-router-dom' import { afterEach, expect, it, vi } from 'vitest' import { experimentsApi, protocolsApi } from '@/api/client' @@ -15,7 +15,7 @@ vi.mock('@/components/protocol/ProtocolCanvas', () => ({ ProtocolCanvas: () => { } })) vi.mock('@/components/protocol/ExperimentVersionView', () => ({ ExperimentVersionCanvas: ({ version }: { version: ProtocolRevision }) =>

Saved canvas {version.revision}

})) vi.mock('@/components/protocol/ResultsTab', () => ({ ResultsInspectorPanel: () => null })) -vi.mock('@/components/protocol/ExperimentSidePanel', () => ({ ExperimentSidePanel: (props: { versionControls: React.ReactNode; viewingHistory: boolean; version?: ProtocolRevision; draftExperiment: Experiment; resultsExperiment: Experiment }) =>
{props.versionControls}

{props.viewingHistory ? `Historical design: ${props.version?.experiment_snapshot?.hypothesis}` : `Draft design: ${props.draftExperiment?.hypothesis}`}

Results design: {props.resultsExperiment?.hypothesis}

})) +vi.mock('@/components/protocol/ExperimentSidePanel', () => ({ ExperimentSidePanel: (props: { viewingHistory: boolean; version?: ProtocolRevision; draftExperiment: Experiment; resultsExperiment: Experiment }) =>

{props.viewingHistory ? `Historical design: ${props.version?.experiment_snapshot?.hypothesis}` : `Draft design: ${props.draftExperiment?.hypothesis}`}

Results design: {props.resultsExperiment?.hypothesis}

})) afterEach(() => vi.restoreAllMocks()) it('switches canvas and design together and preserves local draft edits on return', async () => { @@ -28,19 +28,31 @@ it('switches canvas and design together and preserves local draft edits on retur vi.spyOn(experimentsApi, 'getRunResults').mockRejectedValue(new Error('No results')) vi.spyOn(experimentsApi, 'listTrials').mockResolvedValue([]) vi.spyOn(protocolsApi, 'list').mockResolvedValue([protocol]) - vi.spyOn(protocolsApi, 'listRevisions').mockResolvedValue(versions) + vi.spyOn(protocolsApi, 'listRevisions').mockResolvedValue([...versions, { id: 'legacy', protocol_id: 'p', revision: 0, graph: { nodes: [], edges: [] }, published_at: '2025-12-01', experiment_snapshot: null }]) const client = new QueryClient({ defaultOptions: { queries: { retry: false } } }) render(} />) const draft = await screen.findByRole('textbox', { name: 'Draft canvas' }) await screen.findByText('Results design: Published') + const selector = screen.getByRole('combobox', { name: 'Experiment version' }) + expect(selector.closest('[aria-label="Experiment version controls"]')).toBeInTheDocument() + expect(screen.getByLabelText('Experiment panels')).not.toContainElement(selector) + expect(screen.getAllByRole('option')).toHaveLength(3) + expect(screen.queryByRole('option', { name: /Legacy|Version 0/ })).not.toBeInTheDocument() + expect(screen.getByRole('option', { name: 'Unpublished changes' })).toBeInTheDocument() fireEvent.change(draft, { target: { value: 'Local edit' } }) fireEvent.change(screen.getByRole('combobox', { name: 'Experiment version' }), { target: { value: 'v1' } }) expect(screen.getByText('Saved canvas 1')).toBeInTheDocument() expect(screen.getByText('Historical design: Original')).toBeInTheDocument() expect(screen.getByText('Results design: Original')).toBeInTheDocument() expect(screen.queryByRole('button', { name: 'Publish experiment' })).not.toBeInTheDocument() - fireEvent.click(screen.getByRole('button', { name: 'Return to draft' })) + fireEvent.click(screen.getByRole('button', { name: 'Return to experiment' })) expect(screen.getByRole('textbox', { name: 'Draft canvas' })).toHaveValue('Local edit') expect(screen.getByText('Draft design: Current')).toBeInTheDocument() expect(screen.getByText('Results design: Published')).toBeInTheDocument() + await act(async () => { + client.setQueryData(['protocols', 'for-experiment', 'e'], { ...protocol, has_unpublished_changes: false }) + }) + expect(await screen.findByRole('option', { name: 'Version 2' })).toBeInTheDocument() + expect(screen.getAllByRole('option')).toHaveLength(2) + expect(screen.queryByRole('option', { name: 'Unpublished changes' })).not.toBeInTheDocument() }) diff --git a/frontend/src/pages/ProtocolCanvasPage.tsx b/frontend/src/pages/ProtocolCanvasPage.tsx index 24100c18..26b29881 100644 --- a/frontend/src/pages/ProtocolCanvasPage.tsx +++ b/frontend/src/pages/ProtocolCanvasPage.tsx @@ -282,7 +282,15 @@ export function ProtocolCanvasPage() { ? unboundFactorNames(experimentQuery.data.design_spec, protocolQuery.data.graph) : [] const versionsQuery = useQuery({ queryKey: ['protocols', protocolQuery.data?.id, 'revisions'], queryFn: () => protocolsApi.listRevisions(protocolQuery.data!.id), enabled: !!protocolQuery.data?.id }) - const version = versionsQuery.data?.find(item => item.id === versionId) + const experimentVersions = versionsQuery.data?.filter(item => item.experiment_snapshot != null) + const hasChanges = !!protocolQuery.data?.has_unpublished_changes || draftBusy + const publishedVersion = protocolQuery.data?.published_revision + const editableVersionLabel = !publishedVersion ? 'Not published' + : hasChanges ? 'Unpublished changes' : `Version ${publishedVersion}` + const historicalOptions = experimentVersions?.filter(item => + hasChanges || item.id !== protocolQuery.data?.published_revision_id || item.id === versionId, + ) + const version = experimentVersions?.find(item => item.id === versionId) const displayedExperiment = experimentQuery.data && version?.experiment_snapshot ? { ...experimentQuery.data, ...version.experiment_snapshot } : experimentQuery.data const resultsVersion = version ?? versionsQuery.data?.find(item => item.id === protocolQuery.data?.published_revision_id) @@ -296,7 +304,6 @@ export function ProtocolCanvasPage() { run_id: row.run_id, workspace_id: null, artifacts: null, created_at: row.updated_at, updated_at: row.updated_at, })) - const versionControls =
{version &&

Canvas, design, and results show this saved version. Read-only.

}{versionsQuery.isError &&

Could not load experiment versions.

}
return (
@@ -312,14 +319,33 @@ export function ProtocolCanvasPage() { {experimentQuery.data && (versionId ? {experimentQuery.data.name} : )} {experimentQuery.data?.locked_at && Locked} {!versionId && resultsExperiment && } -
+
+
+ +

{versionId ? 'Canvas, Design, Runs, and Results show this saved version · read-only.' + : !publishedVersion ? 'Publish the experiment to create its first version.' + : hasChanges ? `Changes are not published. Runs and Results use version ${publishedVersion}.` + : 'Canvas, Design, Runs, and Results use this published version.'}

+ {versionsQuery.isError &&

Could not load experiment versions.

} {protocolQuery.data && experimentId && !versionId && ( - <> +
- +
)} + {versionId && }
- {versionId &&
Viewing {version?.experiment_snapshot ? 'experiment version' : 'legacy canvas'} {version?.revision ?? '…'} · read-only
}
Date: Mon, 5 Oct 2026 19:20:41 -0700 Subject: [PATCH 11/16] Unify dataset execution panels and support per-row reruns --- .../components/protocol/CellPresentation.tsx | 66 +++++++ .../components/protocol/DatasetRowCells.tsx | 83 ++++++++ .../protocol/DatasetRowResults.test.tsx | 33 ++++ .../components/protocol/DatasetRowResults.tsx | 61 +++--- .../protocol/DatasetRowRuns.test.tsx | 65 ++++++- .../components/protocol/DatasetRowRuns.tsx | 46 +++-- .../protocol/ExperimentSidePanel.tsx | 2 +- .../protocol/ExperimentVersionView.tsx | 7 +- .../src/components/protocol/ResultsPanel.tsx | 70 +++++++ .../components/protocol/ResultsTab.test.tsx | 14 +- .../src/components/protocol/ResultsTab.tsx | 125 ++++-------- frontend/src/components/protocol/RunsTab.tsx | 180 +++--------------- .../src/components/protocol/StopRowRuns.tsx | 18 ++ .../protocol/edges/DatasetInputPanel.test.tsx | 21 +- .../protocol/edges/DatasetInputPanel.tsx | 11 +- frontend/src/lib/experiment.ts | 11 ++ frontend/src/types/experiments.ts | 9 + src/asaree/api/experiments.py | 1 + src/asaree/services/experiment_run_results.py | 13 ++ src/asaree/services/factorial_row_results.py | 3 +- src/asaree/services/protocol_execution.py | 15 +- tests/test_dataset_row_claims.py | 81 ++++++++ tests/test_dataset_row_results_projection.py | 3 + 23 files changed, 633 insertions(+), 305 deletions(-) create mode 100644 frontend/src/components/protocol/CellPresentation.tsx create mode 100644 frontend/src/components/protocol/DatasetRowCells.tsx create mode 100644 frontend/src/components/protocol/ResultsPanel.tsx create mode 100644 frontend/src/components/protocol/StopRowRuns.tsx diff --git a/frontend/src/components/protocol/CellPresentation.tsx b/frontend/src/components/protocol/CellPresentation.tsx new file mode 100644 index 00000000..3312c0eb --- /dev/null +++ b/frontend/src/components/protocol/CellPresentation.tsx @@ -0,0 +1,66 @@ +import type { ReactNode } from 'react' +import { Badge } from '@/components/ui/badge' +import { Button } from '@/components/ui/button' +import { ChevronDown, ExternalLink } from 'lucide-react' + +/** Both execution modes use this layout; callers supply their progress and actions. */ +export function CellCard({ label, conditions, replicateCount, expanded, onToggle, progress, badges, actions, metric, children, results = false, listId, toggleLabel }: { + label: string; conditions: string[]; replicateCount: number; expanded: boolean; onToggle: () => void + progress?: ReactNode; badges?: ReactNode; actions?: ReactNode; metric?: ReactNode; children: ReactNode + results?: boolean; listId: string; toggleLabel?: string +}) { + const summary = (results ? conditions : conditions.slice(0, 2)).join(' · ') || (results ? 'No varying factors' : 'Cell') + return +
+ + {actions &&
{actions}
} +
+ {expanded &&
{children}
} + +} + +export function ReplicateRow({ number, detail, badges, status, actions, onView, viewLabel = 'View result', viewAriaLabel, results = false, error }: { + number: number; detail?: ReactNode; badges?: ReactNode; status?: ReactNode; actions?: ReactNode + onView?: () => void; viewLabel?: string; viewAriaLabel?: string; results?: boolean; error?: string | null +}) { + const content = <>

Replicate {number}

{badges}
{detail &&

{detail}

}
+
{status}{results ? onView && : <>{onView && }{actions}}
+ if (results) return
    • + return
    {content}
    {error &&

    {error}

    }     +} + +export function CellFrame({ children, results = false, expanded = false }: { children: ReactNode; results?: boolean; expanded?: boolean }) { + return
    {children}
    +} + +export function ReplicateFrame({ children }: { children: ReactNode }) { + return
  • {children}
    • +} + +export function ReplicateStatus({ status }: { status: string }) { + const classes = status === 'completed' ? 'bg-[color:var(--chart-3)]/10 text-[color:var(--chart-3)]' : status === 'failed' ? 'bg-destructive/10 text-destructive' : ['running', 'queued', 'pending', 'finalizing'].includes(status) ? 'bg-[color:var(--primary)]/10 text-[color:var(--primary)]' : 'bg-muted text-muted-foreground' + const label = status === 'pending' ? 'Queued' : status === 'not_started' ? 'Not started' : status.replaceAll('_', ' ').replace(/^./, letter => letter.toUpperCase()) + return {label} +} + +export function CellRunAction({ children, onClick, disabled, label }: { children: ReactNode; onClick: () => void; disabled?: boolean; label?: string }) { + return +} + +export function ReplicateRunAction({ children, onClick, disabled, label, title }: { children: ReactNode; onClick: () => void; disabled?: boolean; label?: string; title?: string }) { + return +} diff --git a/frontend/src/components/protocol/DatasetRowCells.tsx b/frontend/src/components/protocol/DatasetRowCells.tsx new file mode 100644 index 00000000..b84fe552 --- /dev/null +++ b/frontend/src/components/protocol/DatasetRowCells.tsx @@ -0,0 +1,83 @@ +import { useState } from 'react' +import { Badge } from '@/components/ui/badge' +import { Button } from '@/components/ui/button' +import { displayFactorCondition } from '@/lib/experiment' +import type { DesignSpec, ExperimentRunResults, RowCell } from '@/types/experiments' +import { CellCard, ReplicateRow, CellRunAction, ReplicateRunAction, ReplicateStatus } from './CellPresentation' + +// Generated cells exist before any execution slots. Older API responses can +// still provide the cells represented by their rows, without fabricating slots. +function datasetRowCells(results: ExperimentRunResults): RowCell[] { + if (results.row_cells) return results.row_cells + const cells = new Map() + const replicates = new Map>() + for (const row of results.row_results ?? []) { + if (!row.cell_id) continue + cells.set(row.cell_id, { cell_id: row.cell_id, cell_label: row.cell_label, factor_values: row.factor_values ?? {}, replicate_count: 0 }) + const identities = replicates.get(row.cell_id) ?? new Set() + identities.add(row.replicate_result_id ?? row.replicate_number) + replicates.set(row.cell_id, identities) + } + return [...cells.values()].map(cell => ({ ...cell, replicate_count: replicates.get(cell.cell_id)?.size ?? 0 })) +} + +export function DatasetRowCells({ results, selectedCellId, onSelect, designSpec, onSelectReplicate, onRun, onStop, runDisabled = false, headerActions, resultView = false }: { + results: ExperimentRunResults; selectedCellId: string; onSelect: (cellId: string) => void; designSpec?: DesignSpec | null + onSelectReplicate?: (cellId: string, replicateId: string) => void + onRun?: (labels: string[], cellCount: number, rerunLabels: string[]) => void + onStop?: (runIds: string[]) => React.ReactNode + headerActions?: React.ReactNode + runDisabled?: boolean + resultView?: boolean +}) { + const cells = datasetRowCells(results) + const [page, setPage] = useState(0) + const [expandedCells, setExpandedCells] = useState>(new Set()) + const currentPage = Math.min(page, Math.max(0, Math.ceil(cells.length / 20) - 1)) + if (!cells.length) return null + return
    +

    Cells

    {headerActions}{cells.length}
    +
    {cells.slice(currentPage * 20, (currentPage + 1) * 20).map(cell => { + const rows = (results.row_results ?? []).filter(row => row.cell_id === cell.cell_id) + const rowCount = results.row_summary?.row_count + const expected = rowCount == null ? null : rowCount * cell.replicate_count + const completed = rows.filter(row => row.status === 'completed').length + const failed = rows.filter(row => row.status === 'failed').length + const expanded = expandedCells.has(cell.cell_id) + const replicates = cell.replicates ?? [...new Map(rows.map(row => [row.replicate_result_id, { + replicate_result_id: row.replicate_result_id, replicate_label: row.replicate_label, + replicate_number: row.replicate_number, + }])).values()] + const orderedReplicates = [...replicates].sort((a, b) => a.replicate_number - b.replicate_number) + const labels = orderedReplicates.map(replicate => replicate.replicate_label).filter(Boolean) + const rerunLabels = orderedReplicates.filter(replicate => rows.some(row => row.replicate_result_id === replicate.replicate_result_id && ['completed', 'failed', 'cancelled', 'limit_reached'].includes(row.status))).map(replicate => replicate.replicate_label) + const activeRuns = rows.filter(row => ['pending', 'running', 'finalizing'].includes(row.status) && row.latest_attempt).map(row => row.latest_attempt!.run_id) + const conditions = Object.entries(cell.factor_values).map(([name, value]) => displayFactorCondition(name, value, designSpec, cell.cell_label)) + return { setExpandedCells(current => { const next = new Set(current); if (next.has(cell.cell_id)) next.delete(cell.cell_id); else next.add(cell.cell_id); return next }); onSelect(expanded ? '' : cell.cell_id) }} + progress={<>{completed}/{expected ?? '?'} row executions complete · {rows.length} planned} + badges={failed > 0 && {failed} failed} + actions={<>{resultView && }{onRun && onRun(labels, 1, rerunLabels)} label={`Run all replicates in ${cell.cell_label}`}>{rerunLabels.length === labels.length && labels.length ? 'Re-run all replicates' : 'Run all replicates'}}{onStop?.(activeRuns)}}> + {resultView &&

    Individual replicates

    Select one for full output

    } + {!orderedReplicates.length &&

    Replicate details unavailable. Refresh after updating the server.

    } +
      {orderedReplicates.map(replicate => { + const replicateRows = rows.filter(row => row.replicate_result_id === replicate.replicate_result_id) + const hasAttempt = replicateRows.some(row => Boolean(row.latest_attempt?.run_id)) + const done = replicateRows.filter(row => row.status === 'completed').length + const active = replicateRows.filter(row => ['pending', 'running', 'finalizing'].includes(row.status)).length + const errors = replicateRows.filter(row => ['failed', 'cancelled', 'limit_reached'].includes(row.status)).length + const finished = rowCount != null && done === rowCount + return {done}/{rowCount ?? '?'} row executions complete{active > 0 && ` · ${active} active`}{errors > 0 && ` · ${errors} failed`}} + status={} + onView={hasAttempt && onSelectReplicate ? () => onSelectReplicate(cell.cell_id, replicate.replicate_result_id) : undefined} + viewAriaLabel={resultView ? `View ${cell.cell_label} replicate ${replicate.replicate_number}` : undefined} + viewLabel="View results" + actions={<>{onRun && 0 || !replicate.replicate_label} onClick={() => onRun([replicate.replicate_label], 1, done > 0 || errors > 0 ? [replicate.replicate_label] : [])} label={`Run ${cell.cell_label} replicate ${replicate.replicate_number}`}>{done > 0 || errors > 0 ? 'Re-run' : 'Run'}}{onStop?.(replicateRows.filter(row => ['pending', 'running', 'finalizing'].includes(row.status) && row.latest_attempt).map(row => row.latest_attempt!.run_id))}} /> + })}
    +     + })}
    + {cells.length > 20 &&
    {currentPage + 1}/{Math.ceil(cells.length / 20)}
    } +
    +} diff --git a/frontend/src/components/protocol/DatasetRowResults.test.tsx b/frontend/src/components/protocol/DatasetRowResults.test.tsx index ec195a9b..9e74fe1a 100644 --- a/frontend/src/components/protocol/DatasetRowResults.test.tsx +++ b/frontend/src/components/protocol/DatasetRowResults.test.tsx @@ -42,3 +42,36 @@ it('paginates dense rows, maximizes with Escape, and exports the selected scope' expect(csv).toHaveBeenCalledWith('e', scope) await screen.findByRole('button', { name: 'Download CSV' }) }) + +it('shows generated cells with no outcomes and filters executions by cell identity', () => { + const results = { + row_cells: [ + { cell_id: 'a', cell_label: 'Prompt:short', factor_values: { Prompt: 'short' }, replicate_count: 1 }, + { cell_id: 'b', cell_label: 'Prompt:long', factor_values: { Prompt: 'long' }, replicate_count: 1 }, + ], + row_results: [ + { row_result_id: 'slot-a', cell_id: 'a', cell_label: 'Prompt:short', factor_values: { Prompt: 'short' }, dataset_row: { row_index: 0 }, replicate_number: 1, status: 'completed', latest_attempt: { run_id: 'run-a' } }, + ], + } as unknown as ExperimentRunResults + render() + fireEvent.click(screen.getByRole('button', { name: 'View cell Prompt:long' })) + expect(screen.queryByRole('button', { name: 'Inspect row' })).not.toBeInTheDocument() + expect(screen.getByText('No row executions planned for this cell.')).toBeInTheDocument() + fireEvent.click(screen.getByRole('button', { name: 'All cells' })) + expect(screen.getByRole('button', { name: 'Inspect row' })).toBeInTheDocument() +}) + +it('filters row outcomes to the selected replicate', () => { + const results = { + row_cells: [{ cell_id: 'a', cell_label: 'a', factor_values: {}, replicate_count: 2, replicates: [1, 2].map(number => ({ replicate_result_id: `parent-${number}`, replicate_label: `a-${number}`, replicate_number: number })) }], + row_results: [1, 2].map(number => ({ row_result_id: `slot-${number}`, cell_id: 'a', cell_label: 'a', factor_values: {}, replicate_result_id: `parent-${number}`, replicate_number: number, dataset_row: { row_index: 0 }, status: 'completed', latest_attempt: { run_id: `run-${number}` } })), + } as unknown as ExperimentRunResults + const inspect = vi.fn() + render() + fireEvent.click(screen.getByRole('button', { name: 'View cell a' })) + expect(screen.getAllByRole('button', { name: 'View results' })).toHaveLength(1) + fireEvent.click(screen.getByRole('button', { name: 'View a replicate 2' })) + expect(screen.getAllByRole('button', { name: 'Inspect row' })).toHaveLength(1) + fireEvent.click(screen.getByRole('button', { name: 'Inspect row' })) + expect(inspect).toHaveBeenCalledWith('slot-2') +}) diff --git a/frontend/src/components/protocol/DatasetRowResults.tsx b/frontend/src/components/protocol/DatasetRowResults.tsx index a16ea0ee..5878eb63 100644 --- a/frontend/src/components/protocol/DatasetRowResults.tsx +++ b/frontend/src/components/protocol/DatasetRowResults.tsx @@ -1,57 +1,46 @@ -import { useEffect, useState } from 'react' -import { experimentsApi, type ResultsScope } from '@/api/client' +import { useState } from 'react' +import type { ResultsScope } from '@/api/client' import { Button } from '@/components/ui/button' import { reportedMetricText } from '@/lib/reportedMetrics' +import { displayFactorLevel } from '@/lib/experiment' import type { Experiment, ExperimentRunResults, RowResult } from '@/types/experiments' import { RunStepTrace } from './NodeRunOutputPanel' +import { DatasetRowCells } from './DatasetRowCells' +import { ChevronRight, AlertTriangle, CircleCheck, CircleHelp } from 'lucide-react' +import { ResultsPanel, ResultsScorecard } from './ResultsPanel' export function DatasetRowResults({ results, experiment, scope, onInspect }: { results: ExperimentRunResults; experiment: Experiment; scope: ResultsScope; onInspect: (rowId: string) => void }) { const [page, setPage] = useState(0) - const [maximized, setMaximized] = useState(false) - useEffect(() => { - if (!maximized) return - const previous = document.body.style.overflow - document.body.style.overflow = 'hidden' - const close = (event: KeyboardEvent) => { if (event.key === 'Escape') setMaximized(false) } - document.addEventListener('keydown', close) - return () => { document.body.style.overflow = previous; document.removeEventListener('keydown', close) } - }, [maximized]) - const [error, setError] = useState(null) - const [downloading, setDownloading] = useState(false) - const rows = results.row_results ?? [] + const [selectedCellId, setSelectedCellId] = useState('') + const [selectedReplicateId, setSelectedReplicateId] = useState('') + const rows = (results.row_results ?? []).filter(row => (!selectedCellId || row.cell_id === selectedCellId) && (!selectedReplicateId || row.replicate_result_id === selectedReplicateId)) const currentPage = Math.min(page, Math.max(0, Math.ceil(rows.length / 20) - 1)) const visibleRows = rows.slice(currentPage * 20, (currentPage + 1) * 20) const factors = [...new Set(rows.flatMap(row => Object.keys(row.factor_values)))] const declared = experiment.design_spec?.metrics?.filter(metric => metric.kind === 'custom') ?? [] const observed = rows.flatMap(row => row.measurement?.observations ?? []).filter(item => item.producer?.kind !== 'runtime') const metrics = [...new Map([...observed.map(item => ({ id: item.metric_id, name: item.metric_name })), ...declared].map(metric => [metric.id, metric])).values()] - async function download() { - setDownloading(true) - try { - const blob = await experimentsApi.downloadRunResultsCsv(experiment.id, scope) - const url = URL.createObjectURL(blob) - const link = document.createElement('a') - link.href = url - link.download = 'row-results.csv' - link.click() - URL.revokeObjectURL(url) - } catch (error) { setError(error instanceof Error ? error.message : 'CSV unavailable') } - finally { setDownloading(false) } - } - return
    -

    Per-row outcomes are reported individually. Aggregation, ranking and factorial analysis are unavailable.

    - {results.row_summary && <>

    {results.row_summary.planned} planned · {results.row_summary.completed} completed · {results.row_summary.failed} failed · {results.row_summary.missing_reported} missing reported

    {JSON.stringify(results.row_summary.metric_coverage, null, 2)}
    } - - - {error &&

    {error}

    } + const summary = results.row_summary + return
    Per-row results

    Individual row outcomes

    Per-row outcomes are reported individually. Aggregation, ranking and factorial analysis are unavailable.

    } + scorecards={summary && <> + + + + + } + cells={ { setSelectedCellId(cellId); setSelectedReplicateId(''); setPage(0) }} onSelectReplicate={(cellId, replicateId) => { setSelectedCellId(cellId); setSelectedReplicateId(replicateId); setPage(0) }} />} + details={
    +

    {selectedCellId ? 'Cell row outcomes' : 'Row outcomes'}

    {[...factors, 'Source row', 'Replicate', 'Execution status', 'Reported coverage', ...metrics.map(metric => metric.name), 'Output'].map((heading, index) => )}{visibleRows.map(row => { const observations = row.measurement?.observations ?? [] const measured = metrics.filter(metric => observations.some(item => item.metric_id === metric.id && item.status === 'measured')).length - return {factors.map(factor => )}{metrics.map(metric => )} + return {factors.map(factor => )}{metrics.map(metric => )} })}
    {heading}
    {String(row.factor_values[factor] ?? '')}{row.dataset_row.row_index + 1}{row.replicate_number}{row.status}{measured}/{metrics.length}{reportedMetricText(observations.find(item => item.metric_id === metric.id))}
    {displayFactorLevel(experiment.design_spec, factor, row.factor_values[factor], row.cell_label)}{row.dataset_row.row_index + 1}{row.replicate_number}{row.status}{measured}/{metrics.length}{reportedMetricText(observations.find(item => item.metric_id === metric.id))}
    {rows.length > 20 &&
    Page {currentPage + 1} of {Math.ceil(rows.length / 20)}
    } - {!rows.length &&

    No row executions planned in this scope.

    } -
    + {!rows.length &&

    {selectedCellId ? 'No row executions planned for this cell.' : 'No row executions planned in this scope.'}

    } +
    } + /> } export function DatasetRowDetail({ row, onClose }: { row: RowResult; onClose: () => void }) { diff --git a/frontend/src/components/protocol/DatasetRowRuns.test.tsx b/frontend/src/components/protocol/DatasetRowRuns.test.tsx index 29459249..309c8e23 100644 --- a/frontend/src/components/protocol/DatasetRowRuns.test.tsx +++ b/frontend/src/components/protocol/DatasetRowRuns.test.tsx @@ -5,6 +5,7 @@ import { experimentsApi, protocolsApi } from '@/api/client' import type { ExperimentRunResults } from '@/types/experiments' import type { Protocol } from '@/types/protocols' import { RunsTab } from './RunsTab' +import { DatasetRowRuns } from './DatasetRowRuns' export const rowResults = { consumption_mode: 'per_row', cells: [], replicates: [], overview: {}, @@ -26,13 +27,73 @@ it('renders authoritative forecast and distinct row identities', async () => { mountRows(inspect) expect(await screen.findByText('5 cells / 10 replicates / 3 rows / 30 executions')).toBeInTheDocument() expect(screen.getByText(/26 completed.*1 failed.*1 missing reported/)).toBeInTheDocument() - fireEvent.click(screen.getAllByRole('button', { name: 'Output' })[1]) + fireEvent.click(screen.getAllByRole('button', { name: 'View results' })[1]) expect(inspect).toHaveBeenCalledWith('slot-1') expect(experimentsApi.getRunResults).toHaveBeenCalledWith('e', { protocol_id: 'p' }) }) it('cancels exactly the active row run', async () => { const cancel = vi.spyOn(protocolsApi, 'cancelRun').mockResolvedValue({} as never) mountRows() - fireEvent.click(await screen.findByRole('button', { name: 'Cancel' })) + fireEvent.click(await screen.findByRole('button', { name: 'Stop' })) await waitFor(() => expect(cancel).toHaveBeenCalledWith('p', 'run-2')) }) + +it('shows generated cells before row executions are planned', () => { + const results = { ...rowResults, row_results: [], row_cells: [ + { cell_id: 'cell-a', cell_label: 'Prompt:short', factor_values: { Prompt: 'short' }, replicate_count: 2, replicates: [ + { replicate_result_id: 'parent-1', replicate_label: 'short', replicate_number: 1 }, + { replicate_result_id: 'parent-2', replicate_label: 'short__rep2', replicate_number: 2 }, + ] }, + { cell_id: 'cell-b', cell_label: 'Prompt:long', factor_values: { Prompt: 'long' }, replicate_count: 2 }, + ] } + render() + expect(screen.getByRole('button', { name: 'View cell Prompt:short' })).toBeInTheDocument() + expect(screen.getByRole('button', { name: 'View cell Prompt:long' })).toBeInTheDocument() + expect(screen.getAllByText('0/6')).toHaveLength(2) + fireEvent.click(screen.getByRole('button', { name: 'View cell Prompt:short' })) + expect(screen.getByText('Replicate 1')).toBeInTheDocument() + expect(screen.getByText('Replicate 2')).toBeInTheDocument() + expect(screen.queryByRole('button', { name: 'View results' })).not.toBeInTheDocument() + expect(screen.queryByRole('button', { name: 'View Prompt:short replicate 1' })).not.toBeInTheDocument() + expect(screen.getByRole('button', { name: 'Run Prompt:short replicate 1' })).toBeEnabled() + expect(screen.getByRole('button', { name: 'Run all replicates in Prompt:short' })).toBeEnabled() +}) + +it('runs only the selected replicate across source rows after confirmation', async () => { + const run = vi.spyOn(protocolsApi, 'runCells').mockResolvedValue({} as never) + const results = { ...rowResults, row_results: [], row_cells: [{ cell_id: 'a', cell_label: 'short', factor_values: {}, replicate_count: 1, replicates: [{ replicate_result_id: 'parent', replicate_label: 'short', replicate_number: 1 }] }] } + render() + fireEvent.click(screen.getByRole('button', { name: 'View cell short' })) + fireEvent.click(screen.getByRole('button', { name: 'Run short replicate 1' })) + expect(run).not.toHaveBeenCalled() + expect(screen.getByText('1 replicates × 3 rows = 3 executions')).toBeInTheDocument() + fireEvent.click(screen.getByRole('button', { name: 'Run selected replicates' })) + await waitFor(() => expect(run).toHaveBeenCalledWith('p', { replicateLabels: ['short'], rerunReplicateLabels: [] })) +}) + +it('keeps multiple cells expanded and explicitly reruns completed rows', async () => { + const run = vi.spyOn(protocolsApi, 'runCells').mockResolvedValue({} as never) + const cells = ['a', 'b'].map(id => ({ cell_id: id, cell_label: id, factor_values: {}, replicate_count: 1, replicates: [{ replicate_result_id: `parent-${id}`, replicate_label: id, replicate_number: 1 }] })) + const results: ExperimentRunResults = { ...rowResults, row_cells: cells, row_results: rowResults.row_results!.map((row, index) => ({ ...row, cell_id: 'a', replicate_result_id: 'parent-a', replicate_label: 'a', status: index === 2 ? 'running' : 'completed' })) } + render() + fireEvent.click(screen.getByRole('button', { name: 'View cell a' })) + fireEvent.click(screen.getByRole('button', { name: 'View cell b' })) + expect(screen.getByRole('list', { name: 'Replicates for a' })).toBeInTheDocument() + expect(screen.getByRole('list', { name: 'Replicates for b' })).toBeInTheDocument() + expect(screen.getByRole('button', { name: 'Run a replicate 1' })).toBeDisabled() + fireEvent.click(screen.getByRole('button', { name: 'Run all replicates in a' })) + expect(screen.getByText(/2 executions will start/)).toBeInTheDocument() + expect(screen.getByText(/previous attempts remain available/)).toBeInTheDocument() + fireEvent.click(screen.getByRole('button', { name: 'Run selected replicates' })) + await waitFor(() => expect(run).toHaveBeenCalledWith('p', { replicateLabels: ['a'], rerunReplicateLabels: ['a'] })) +}) + +it('stops all active row runs without stopping finished rows', async () => { + const cancel = vi.spyOn(protocolsApi, 'cancelRun').mockResolvedValue({} as never) + const results: ExperimentRunResults = { ...rowResults, row_cells: [{ cell_id: 'a', cell_label: 'a', factor_values: {}, replicate_count: 1 }], row_results: rowResults.row_results!.map((row, index) => ({ ...row, cell_id: 'a', status: index === 0 ? 'completed' : 'running' })) } + render() + fireEvent.click(screen.getAllByRole('button', { name: 'Stop all' })[0]) + await waitFor(() => expect(cancel).toHaveBeenCalledTimes(2)) + expect(cancel).toHaveBeenCalledWith('p', 'run-1') + expect(cancel).toHaveBeenCalledWith('p', 'run-2') +}) diff --git a/frontend/src/components/protocol/DatasetRowRuns.tsx b/frontend/src/components/protocol/DatasetRowRuns.tsx index f9826de8..9d2c3d88 100644 --- a/frontend/src/components/protocol/DatasetRowRuns.tsx +++ b/frontend/src/components/protocol/DatasetRowRuns.tsx @@ -2,29 +2,38 @@ import { useState } from 'react' import { useMutation, useQueryClient } from '@tanstack/react-query' import { protocolsApi } from '@/api/client' import { Button } from '@/components/ui/button' -import type { ExperimentRunResults, RowSummary } from '@/types/experiments' +import type { DesignSpec, ExperimentRunResults, RowSummary } from '@/types/experiments' import type { Protocol } from '@/types/protocols' import { RunConfirmDialog } from './RunConfirmDialog' +import { DatasetRowCells } from './DatasetRowCells' +import { StopRowRuns } from './StopRowRuns' +import { ReplicateStatus } from './CellPresentation' +import { displayFactorCondition } from '@/lib/experiment' -export function rowExecutionForecast(summary: RowSummary): string { +function rowExecutionForecast(summary: RowSummary): string { return `${summary.cell_count} cells / ${summary.parent_replicate_count} replicates / ${summary.row_count ?? 'unknown'} rows / ${summary.expected ?? 'unknown'} executions` } -export function DatasetRowRuns({ results, protocol, blocked, onInspect }: { - results: ExperimentRunResults; protocol: Protocol; blocked: boolean; onInspect?: (rowId: string) => void +export function DatasetRowRuns({ results, protocol, blocked, onInspect, designSpec }: { + results: ExperimentRunResults; protocol: Protocol; blocked: boolean; onInspect?: (rowId: string) => void; designSpec?: DesignSpec | null }) { const queryClient = useQueryClient() - const [confirm, setConfirm] = useState(false) + const [confirm, setConfirm] = useState<{ labels?: string[]; cellCount: number; rerunLabels: string[] } | null>(null) const [retrying, setRetrying] = useState(null) const [page, setPage] = useState(0) + const [selectedCellId, setSelectedCellId] = useState('') + const [selectedReplicateId, setSelectedReplicateId] = useState('') const run = useMutation({ - mutationFn: () => protocolsApi.runCells(protocol.id), - onSuccess: () => { setConfirm(false); queryClient.invalidateQueries({ queryKey: ['experiments', protocol.experiment_id] }) }, + mutationFn: () => protocolsApi.runCells(protocol.id, { replicateLabels: confirm?.labels, rerunReplicateLabels: confirm?.rerunLabels }), + onSuccess: () => { setConfirm(null); queryClient.invalidateQueries({ queryKey: ['experiments', protocol.experiment_id] }) }, }) const summary = results.row_summary - const rows = results.row_results ?? [] + const rows = (results.row_results ?? []).filter(row => (!selectedCellId || row.cell_id === selectedCellId) && (!selectedReplicateId || row.replicate_result_id === selectedReplicateId)) const currentPage = Math.min(page, Math.max(0, Math.ceil(rows.length / 20) - 1)) const visibleRows = rows.slice(currentPage * 20, (currentPage + 1) * 20) + const selectedReplicateCount = confirm?.labels?.length ?? summary?.parent_replicate_count ?? 0 + const selectedRows = (results.row_results ?? []).filter(row => !confirm?.labels || confirm.labels.includes(row.replicate_label)) + const runnableExecutions = Math.max(0, selectedReplicateCount * (summary?.row_count ?? 0) - selectedRows.filter(row => ['pending', 'running', 'finalizing'].includes(row.status) || !confirm?.rerunLabels.includes(row.replicate_label)).length) const cancel = useMutation({ mutationFn: (runId: string) => protocolsApi.cancelRun(protocol.id, runId), onSuccess: () => queryClient.invalidateQueries({ queryKey: ['experiments', protocol.experiment_id] }) }) const retry = useMutation({ mutationFn: (rowId: string) => protocolsApi.runCells(protocol.id, { retry_row_result_ids: [rowId] }), @@ -35,14 +44,23 @@ export function DatasetRowRuns({ results, protocol, blocked, onInspect }: { return

    {summary ? rowExecutionForecast(summary) : 'Execution forecast unavailable: original source metadata unavailable.'}

    {summary &&

    {summary.planned} planned · {summary.pending} pending · {summary.running} running · {summary.completed} completed · {summary.failed} failed · {summary.cancelled} cancelled · {summary.missing_reported} missing reported

    } - {[run.error, cancel.error, retry.error].filter(Boolean).map((error, index) =>

    {error?.message}

    )} -
    {['Cell', 'Source row', 'Replicate', 'Status', 'Publication', 'Actions'].map(label => )}{visibleRows.map(row =>
    {label}
    {row.cell_label}{row.dataset_row.row_index + 1}{row.replicate_number}{row.status}{row.protocol_revision_id} - - {row.latest_attempt && ['pending', 'running', 'finalizing'].includes(row.status) && } - + { setSelectedCellId(cellId); setSelectedReplicateId(''); setPage(0) }} onSelectReplicate={(cellId, replicateId) => { setSelectedCellId(cellId); setSelectedReplicateId(replicateId); setPage(0) }} onRun={(labels, cellCount, rerunLabels) => setConfirm({ labels, cellCount, rerunLabels })} runDisabled={blocked || !summary?.row_count || run.isPending} + onStop={runIds => } + headerActions={<> ['pending', 'running', 'finalizing'].includes(row.status) && row.latest_attempt).map(row => row.latest_attempt!.run_id)} />} + /> +

    {selectedCellId ? 'Cell row executions' : 'Row executions'}

    + {!rows.length &&

    {selectedCellId ? 'No row executions planned for this cell.' : 'No row executions planned. Run all cells to create them.'}

    } +
    {['Cell', 'Source row', 'Replicate', 'Status', 'Actions'].map(label => )}{visibleRows.map(row => )}
    {label}
    {Object.entries(row.factor_values ?? {}).map(([name, value]) => displayFactorCondition(name, value, designSpec, row.cell_label)).join(' · ') || 'No varying factors'}{row.dataset_row.row_index + 1}{row.replicate_number} + + {row.latest_attempt && ['pending', 'running', 'finalizing'].includes(row.status) && } + {['failed', 'cancelled'].includes(row.status) && }
    {rows.length > 20 &&
    Page {currentPage + 1} of {Math.ceil(rows.length / 20)}
    } - {confirm && setConfirm(false)} onConfirm={() => run.mutate()} hasUnpublishedChanges={protocol.has_unpublished_changes} publishedRevision={protocol.published_revision} confirmDisabled={blocked || run.isPending} additionalContent={

    {summary && rowExecutionForecast(summary)}

    } />} + {confirm && setConfirm(null)} onConfirm={() => run.mutate()} hasUnpublishedChanges={protocol.has_unpublished_changes} publishedRevision={protocol.published_revision} confirmDisabled={blocked || run.isPending || !runnableExecutions} isConfirming={run.isPending} confirmError={run.error?.message} additionalContent={

    {confirm.labels ? `${confirm.labels.length} replicates × ${summary?.row_count ?? '?'} rows = ${(summary?.row_count ?? 0) * confirm.labels.length} executions` : summary && rowExecutionForecast(summary)}

    {runnableExecutions} executions will start. Active rows are skipped.{confirm.rerunLabels.length > 0 && ' Finished rows in the selected replicates will run again; previous attempts remain available.'}

    } />} } diff --git a/frontend/src/components/protocol/ExperimentSidePanel.tsx b/frontend/src/components/protocol/ExperimentSidePanel.tsx index 0c8c3dd9..bc26c2e7 100644 --- a/frontend/src/components/protocol/ExperimentSidePanel.tsx +++ b/frontend/src/components/protocol/ExperimentSidePanel.tsx @@ -263,7 +263,7 @@ export function ExperimentSidePanel({ onDraftBusyChange={onDraftBusyChange} /> - {viewingHistory && !version ?

    The selected version is unavailable. Return to draft to continue.

    : activeTab === 'design' && version ? : activeTab === 'runs' && version ? : activeTab === 'runs' ? ( + {viewingHistory && !version ?

    The selected version is unavailable. Return to draft to continue.

    : activeTab === 'design' && version ? : activeTab === 'runs' && version ? : activeTab === 'runs' ? ( void }) { + const [selectedCellId, setSelectedCellId] = useState('') + const [selectedReplicateId, setSelectedReplicateId] = useState('') const scope = { protocol_id: protocolId, protocol_revision_id: version.id } const query = useQuery({ queryKey: ['experiments', experimentId, 'run-results', protocolId, version.id], queryFn: () => experimentsApi.getRunResults(experimentId, scope), refetchInterval: 5000 }) if (query.isLoading) return

    Loading version runs…

    if (!query.data) return

    Could not load version runs.

    const rows = query.data.consumption_mode === 'per_row' - ? (query.data.row_results ?? []).map(row => ({ key: row.row_result_id, label: `${row.replicate_label} · row ${row.dataset_row.row_index}`, status: row.status, runId: row.latest_attempt?.run_id, selection: { type: 'row' as const, rowResultId: row.row_result_id, scope } })) + ? (query.data.row_results ?? []).filter(row => (!selectedCellId || row.cell_id === selectedCellId) && (!selectedReplicateId || row.replicate_result_id === selectedReplicateId)).map(row => ({ key: row.row_result_id, label: `${row.replicate_label} · row ${row.dataset_row.row_index + 1}`, status: row.status, runId: row.latest_attempt?.run_id, selection: { type: 'row' as const, rowResultId: row.row_result_id, scope } })) : query.data.replicates.map(row => ({ key: row.replicate_label, label: row.replicate_label, status: row.status, runId: row.run_id, selection: { type: 'replicate' as const, replicateLabel: row.replicate_label, scope } })) - return

    Runs for experiment version {version.revision}. Return to draft to launch new runs.

    {!rows.length &&

    No runs planned for this version.

    }
    {rows.map(row => )}
    ReplicateStatusDetails
    {row.label}{row.status.replaceAll('_', ' ')}
    + return

    Runs for experiment version {version.revision}. Select the current experiment to launch new runs.

    {query.data.consumption_mode === 'per_row' && { setSelectedCellId(cellId); setSelectedReplicateId('') }} onSelectReplicate={(cellId, replicateId) => { setSelectedCellId(cellId); setSelectedReplicateId(replicateId) }} />}{!rows.length &&

    No runs planned for this version.

    }
    {rows.map(row => )}
    ReplicateStatusDetails
    {row.label}{row.status.replaceAll('_', ' ')}
    } diff --git a/frontend/src/components/protocol/ResultsPanel.tsx b/frontend/src/components/protocol/ResultsPanel.tsx new file mode 100644 index 00000000..364747bf --- /dev/null +++ b/frontend/src/components/protocol/ResultsPanel.tsx @@ -0,0 +1,70 @@ +import { useEffect, useState, type ReactNode } from 'react' +import { Download, Maximize2, Minimize2, type LucideIcon } from 'lucide-react' +import { experimentsApi, type ResultsScope } from '@/api/client' +import { Button } from '@/components/ui/button' +import { Card } from '@/components/ui/card' +import { sanitizeFilename } from '@/lib/utils' +import { InfoTooltip } from './InfoTooltip' + +export function ResultsScorecard({ label, help, value, note, icon: Icon }: { label: string; help: string; value: string; note?: string; icon: LucideIcon }) { + return
    +
    {label}{help}
    +

    {value}

    + {note &&

    {note}

    } +
    +} + +/** One panel structure for both modes; data-specific content occupies these slots. */ +export function ResultsPanel({ experimentId, experimentName, scope, controls, summary, scorecards, notice, comparison, cells, details }: { + experimentId: string; experimentName: string; scope: ResultsScope + controls?: ReactNode; summary: ReactNode; scorecards?: ReactNode; notice?: ReactNode + comparison?: ReactNode; cells: ReactNode; details?: ReactNode +}) { + const [maximized, setMaximized] = useState(false) + const [downloading, setDownloading] = useState(false) + const [error, setError] = useState(null) + useEffect(() => { + if (!maximized) return + const previous = document.body.style.overflow + document.body.style.overflow = 'hidden' + const close = (event: KeyboardEvent) => { if (event.key === 'Escape') setMaximized(false) } + document.addEventListener('keydown', close) + return () => { document.body.style.overflow = previous; document.removeEventListener('keydown', close) } + }, [maximized]) + async function download() { + setDownloading(true) + setError(null) + try { + const blob = await experimentsApi.downloadRunResultsCsv(experimentId, scope) + const url = URL.createObjectURL(blob) + try { + const link = document.createElement('a') + link.href = url + link.download = `${sanitizeFilename(experimentName, 'experiment')}-results.csv` + link.click() + } finally { URL.revokeObjectURL(url) } + } catch (error) { setError(error instanceof Error ? error.message : 'CSV unavailable') } + finally { setDownloading(false) } + } + return
    +
    + {controls} + + {summary} +
    + + +
    +     + {scorecards &&
    {scorecards}
    } + {notice} + {error &&

    {error}

    } +
    + {comparison} +
    +

    Cell resultsEach card groups replicates that share the same experimental factor levels. Expand one to compare individual runs.

    Expand a condition to inspect its replicates, metrics, usage, outputs, and node activity.

    + {cells} +
    + {details} +
    +} diff --git a/frontend/src/components/protocol/ResultsTab.test.tsx b/frontend/src/components/protocol/ResultsTab.test.tsx index 95ea5898..1680284d 100644 --- a/frontend/src/components/protocol/ResultsTab.test.tsx +++ b/frontend/src/components/protocol/ResultsTab.test.tsx @@ -1,5 +1,5 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query' -import { render, screen, within } from '@testing-library/react' +import { fireEvent, render, screen, within } from '@testing-library/react' import userEvent from '@testing-library/user-event' import { describe, expect, it, vi } from 'vitest' import { experimentsApi } from '@/api/client' @@ -78,6 +78,18 @@ function orderedMetricsFixture() { } describe('results measurement states', () => { + it('uses the shared maximize controls for whole-dataset results', async () => { + const { orderedExperiment, results } = orderedMetricsFixture() + vi.mocked(experimentsApi.getRunResults).mockResolvedValue(results) + renderWithQuery() + fireEvent.click(await screen.findByRole('button', { name: 'Maximize Results' })) + expect(document.body.style.overflow).toBe('hidden') + expect(screen.getByRole('button', { name: 'Restore Results' })).toBeInTheDocument() + fireEvent.keyDown(document, { key: 'Escape' }) + expect(document.body.style.overflow).not.toBe('hidden') + expect(screen.getByRole('button', { name: 'Download CSV' })).toBeInTheDocument() + }) + it('shows explicit units for cost and duration in an individual replicate result', async () => { const replicate = { replicate_label: 'replicate-1', replicate_number: 1, cell_label: 'model_a', factor_values: { model: 'a' }, diff --git a/frontend/src/components/protocol/ResultsTab.tsx b/frontend/src/components/protocol/ResultsTab.tsx index ac72da7a..31d0a8fa 100644 --- a/frontend/src/components/protocol/ResultsTab.tsx +++ b/frontend/src/components/protocol/ResultsTab.tsx @@ -1,6 +1,6 @@ import { useState } from 'react' import { useQuery } from '@tanstack/react-query' -import { AlertTriangle, ChevronDown, ChevronRight, CircleDollarSign, Clock3, Coins, Cpu, Download, ExternalLink, Trophy, X } from 'lucide-react' +import { AlertTriangle, ChevronDown, ChevronRight, CircleDollarSign, Clock3, Coins, Cpu, Trophy, X } from 'lucide-react' import { experimentsApi, type ResultsScope } from '@/api/client' import { Badge } from '@/components/ui/badge' import { Button } from '@/components/ui/button' @@ -10,10 +10,11 @@ import { Skeleton } from '@/components/ui/skeleton' import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs' import { displayFactorLevel, formatMetricLabel, formatMetricValue } from '@/lib/experiment' import { OBSERVATION_LABELS } from '@/lib/measurementPlan' -import { sanitizeFilename } from '@/lib/utils' import type { DesignMetric, EvaluationArtifact, Experiment, HistoricalRun, MetricObservation, ObsoleteRun, ResultCell, ResultNodeRun, ResultReplicate, SupersededRun } from '@/types/experiments' import { DatasetRowDetail, DatasetRowResults } from './DatasetRowResults' +import { ResultsPanel, ResultsScorecard as Scorecard } from './ResultsPanel' import { InfoTooltip } from './InfoTooltip' +import { CellCard, ReplicateRow } from './CellPresentation' import { ReceivedPromptPanel, RunStepTrace, UnresolvedReferencesNote } from './NodeRunOutputPanel' function formatNumber(value: number | null, maximumFractionDigits = 0): string { @@ -195,19 +196,6 @@ function nodeStatusClass(status: string): string { return 'border-transparent bg-muted text-muted-foreground' } -function Scorecard({ label, help, value, note, icon: Icon }: { label: string; help: string; value: string; note?: string; icon: typeof Coins }) { - return ( -
    -
    - {label}{help} -
    -

    {value}

    - {note &&

    {note}

    } -
    - ) -} - function usageSummarizesMetric( key: string, result: Pick, @@ -620,7 +608,6 @@ export function ResultsTab({ }) { const [metricPreference, setMetricPreference] = useState(null) const [expandedResultCells, setExpandedResultCells] = useState>(() => new Set()) - const [downloading, setDownloading] = useState(false) const scope: ResultsScope = { protocol_id: protocolId, protocol_revision_id: versionId || undefined } const resultsQuery = useQuery({ queryKey: ['experiments', experimentId, 'run-results', scope], queryFn: () => experimentsApi.getRunResults(experimentId, scope), refetchInterval: 5000 }) if (resultsQuery.isLoading) return
    @@ -663,25 +650,9 @@ export function ResultsTab({ const bestCell = metricKey && rankable ? sortedCells.find((cell) => typeof cell.metric_means[metricKey] === 'number') ?? null : null const directionLabel = selectedDirection === 'minimize' ? 'Lowest' : 'Highest' - async function downloadResults() { - setDownloading(true) - try { - const blob = await experimentsApi.downloadRunResultsCsv(experimentId, scope) - const url = URL.createObjectURL(blob) - const link = document.createElement('a') - link.href = url - link.download = `${sanitizeFilename(experimentName, 'experiment')}-results.csv` - link.click() - URL.revokeObjectURL(url) - } finally { - setDownloading(false) - } - } - return ( -
    -
    - {metricKeys.length > 0 && ( + 0 && (
    Inspect metric
    )} - - {bestCell && metricKey ? ( + summary={bestCell && metricKey ? (
    Best current result

    {formatResultMetricValue(metricKey, bestCell.metric_means[metricKey], metricTypes)}

    @@ -705,25 +675,21 @@ export function ResultsTab({ ) : (

    Results are arriving

    Complete a run with reported metrics to rank conditions here.

    )} -
    - -
    - -
    - + scorecards={<> + -
    - {overview.obsolete_replicates > 0 &&
    {overview.obsolete_replicates} run{overview.obsolete_replicates === 1 ? '' : 's'} used an older canvas version and are excluded from current totals.
    } -
    -
    + } + notice={overview.obsolete_replicates > 0 &&
    {overview.obsolete_replicates} run{overview.obsolete_replicates === 1 ? '' : 's'} used an older canvas version and are excluded from current totals.
    } + comparison={ +

    {rankable ? 'Cell ranking' : 'Cell comparison'}{rankable ? 'Cells are ranked by the selected metric using current, non-obsolete replicates only.' : 'Metrics are shown for comparison without ordering cells or naming a winner.'}

    Current replicates only{rankable && isBinaryMetric(metricKey, metricTypes) ? ' · ranked by pass rate' : rankable && metricKey === primaryMetric ? ` · ${primaryMetricDirection === 'minimize' ? 'lower is better' : 'higher is better'}` : ''}

    {metricKeys.length === 0 ?

    No scalar metrics are declared or reported yet. Run details and provider usage are still available below.

    :
    {sortedCells.map((cell, index) => { const value = metricKey ? cell.metric_means[metricKey] : undefined; const measured = typeof value === 'number'; const ranked = rankable && measured; const note = measured && metricKey ? binaryMetricNote(metricKey, value, cell.metric_counts[metricKey], metricTypes) : null; const missingSummary = missingObservationSummary(replicatesByCell.get(cell.cell_label) ?? [], selectedMetricId, metricKey); return })}
    {rankable ? 'Rank' : '—'}{rankable && Rank among conditions with a measured value for the selected metric.}ConditionThe factor levels used for this group of replicates.{metricKey ? metricDisplayLabel(metricKey, metricTypes, metricAggregations) : ''}CostDurationRunsCompleted current replicates out of all generated replicates for this condition.
    {ranked ? index + 1 : '—'}

    {factorSummary(cell.factor_values, experiment.design_spec)}

    {cell.obsolete_count > 0 && {cell.obsolete_count} obsolete}    		{measured ? <>{formatResultMetricValue(metricKey!, value, metricTypes)}{note && {note}}{missingSummary && {missingSummary}} : {missingSummary ?? 'No measured observations'}}{formatCurrency(cell.cost_usd)}{formatDuration(cell.duration_seconds)}{cell.current_completed_count}/{cell.replicate_count}
    }
    -
    -

    Cell resultsEach card groups replicates that share the same experimental factor levels. Expand one to compare individual runs.

    Expand a condition to inspect its replicates, metrics, usage, outputs, and node activity.

    -
    + } + cells={ +
    {sortedCells.map((cell) => { const expanded = expandedResultCells.has(cell.cell_label) const cellReplicates = replicatesByCell.get(cell.cell_label) ?? [] @@ -749,35 +715,19 @@ export function ResultsTab({ return left.replicate_number - right.replicate_number }) return ( -
    -
    - - -
    - {expanded && ( -
    + setExpandedResultCells(current => { + const next = new Set(current) + if (next.has(cell.cell_label)) next.delete(cell.cell_label) + else next.add(cell.cell_label) + return next + })} + progress={<>{cell.current_completed_count}/{cell.replicate_count} current runs complete} + badges={cell.obsolete_count > 0 && {cell.obsolete_count} obsolete {cell.obsolete_count === 1 ? 'run' : 'runs'}} + metric={metric !== undefined && {metricDisplayLabel(metricKey!, metricTypes, metricAggregations)}{formatResultMetricValue(metricKey!, metric, metricTypes)}{binaryMetricNote(metricKey!, metric, cell.metric_counts[metricKey!], metricTypes) && {binaryMetricNote(metricKey!, metric, cell.metric_counts[metricKey!], metricTypes)}}} + actions={}> {cellReplicates.length === 0 ?

    No replicate results are available for this cell.

    : ( <> {metricKey && replicateMetricTotal !== null && ( @@ -795,23 +745,20 @@ export function ResultsTab({ const observation = observationForMetric(replicate, selectedMetricId, metricKey) const value = numericMetricValue(replicate, metricKey) return ( -
  • - -
    • + onSelectResult({ type: 'replicate', replicateLabel: replicate.replicate_label })} + detail={<>{replicate.cost_usd !== null ? formatCurrency(replicate.cost_usd) : 'Cost not reported'}{replicate.duration_seconds !== null ? ` · ${formatDuration(replicate.duration_seconds)}` : ''}} + status={<>{value !== null && {formatResultMetricValue(metricKey!, value, metricTypes, true)}}{observation ? {OBSERVATION_LABELS[observation.status]} : evaluationStatusLabel(replicate) && {evaluationStatusLabel(replicate)}}{statusLabel(replicate.status, replicate.obsolete)}} + /> )})} )} -
    - )} -
    + ) })}
    -
    -
    + } + /> ) } diff --git a/frontend/src/components/protocol/RunsTab.tsx b/frontend/src/components/protocol/RunsTab.tsx index fe3f2dc6..ab81dfb2 100644 --- a/frontend/src/components/protocol/RunsTab.tsx +++ b/frontend/src/components/protocol/RunsTab.tsx @@ -7,13 +7,14 @@ import { Button } from '@/components/ui/button' import { Checkbox } from '@/components/ui/checkbox' import { Skeleton } from '@/components/ui/skeleton' import { ApiError, experimentsApi, protocolsApi } from '@/api/client' -import { displayFactorLevel, factorValueKey, groupReplicatesIntoCells, type ExperimentalCell } from '@/lib/experiment' +import { displayFactorCondition, factorValueKey, groupReplicatesIntoCells, type ExperimentalCell } from '@/lib/experiment' import { factorBindingDiscrepancies, type FactorBindingDiscrepancy } from '@/lib/factorBindings' import { protocolForExperimentQueryKey, protocolGraphQueryKey, toPersistedGraph } from '@/lib/protocolGraph' import type { Experiment, ResultCell, ResultReplicate, Trial } from '@/types/experiments' import type { Protocol } from '@/types/protocols' import { DatasetRowRuns } from './DatasetRowRuns' import { RunConfirmDialog } from './RunConfirmDialog' +import { CellCard, ReplicateRow, ReplicateRunAction, CellRunAction, ReplicateStatus } from './CellPresentation' import { WarningBadge } from './nodes/WarningBadge' function factorEntries(cell: ExperimentalCell): [string, unknown][] { @@ -29,52 +30,6 @@ function cellSortKey(cell: ExperimentalCell): string { .join('|') } -// Factor names are stored as fully-qualified binding identifiers so two -// similarly named node fields never collide (e.g. Agent:Search:Enabled). -// That identity is useful to the engine, but it is not a readable treatment -// label. Keep the path as context and turn its final field/value pair into a -// sentence: "Agent · Search: Disabled" rather than "Agent:Search:Enabled: -// false". -function displayFactorCondition( - name: string, - value: unknown, - designSpec: Experiment['design_spec'] | undefined, - cellLabel?: string, -): string { - const parts = name.split(':').map((part) => part.trim()).filter(Boolean) - const field = parts.pop() ?? name - const displayedLevel = displayFactorLevel(designSpec, name, value, cellLabel) - - if (typeof value === 'boolean' && /enabled$/i.test(field)) { - if (displayedLevel !== String(value)) { - return [...parts, field].filter(Boolean).join(' · ') + `: ${displayedLevel}` - } - const subject = field.replace(/\s*enabled$/i, '').trim() - return [...parts, subject].filter(Boolean).join(' · ') + `: ${value ? 'Enabled' : 'Disabled'}` - } - - return [...parts, field].filter(Boolean).join(' · ') + `: ${displayedLevel}` -} - -function trialStatusBadge(status: Trial['status']) { - switch (status) { - case 'not_started': - return { label: 'Not started', className: 'border-transparent bg-muted text-muted-foreground' } - case 'queued': - return { label: 'Queued', className: 'border-transparent bg-[color:var(--primary)]/10 text-[color:var(--primary)]' } - case 'running': - return { label: 'Running', className: 'border-transparent bg-[color:var(--primary)]/10 text-[color:var(--primary)]' } - case 'finalizing': - return { label: 'Finalizing', className: 'border-transparent bg-[color:var(--primary)]/10 text-[color:var(--primary)]' } - case 'completed': - return { label: 'Completed', className: 'border-transparent bg-[color:var(--chart-3)]/10 text-[color:var(--chart-3)]' } - case 'failed': - return { label: 'Failed', className: 'border-transparent bg-destructive/10 text-destructive' } - case 'cancelled': - return { label: 'Cancelled', className: 'border-transparent bg-muted text-muted-foreground' } - } -} - function isActiveTrialStatus(status: Trial['status']): boolean { return status === 'queued' || status === 'running' || status === 'finalizing' } @@ -431,14 +386,7 @@ export function RunAllCellsButton({ return ( <> - + {compact ? {actionLabel} : } {dialogOpen && protocol && ( - + {dialogOpen && protocol && ( 0 || factorBindingDiscrepancies(designSpec, protocol.graph).length > 0} onInspect={onViewRowResult} /> + return 0 || factorBindingDiscrepancies(designSpec, protocol.graph).length > 0} onInspect={onViewRowResult} /> } if (protocol && resultsQuery.isLoading) return @@ -908,8 +854,6 @@ export function RunsTab({
    {cells.map((cell) => { const entries = factorEntries(cell) - const summary = entries.slice(0, 2) - const remaining = entries.slice(2) const expanded = expandedCells.has(cell.label) const replicateListId = `cell-${cell.label}-replicates` const obsoleteCount = cell.replicates.filter((replicate) => trialsByLabel.get(replicate.replicate_label)?.obsolete).length @@ -924,58 +868,16 @@ export function RunsTab({ hasFinishedRun(trialsByLabel.get(replicate.replicate_label)), ) return ( -
    -
    - -
    + displayFactorCondition(name, value, designSpec, cell.label))} + replicateCount={cell.replicates.length} expanded={expanded} listId={replicateListId} + onToggle={() => toggleCell(cell.label)} + progress={cellResult && <>{cellResult.current_completed_count}/{cellResult.replicate_count} current complete{cellUsage.map(value => {value})}} + badges={<> + {obsoleteCount > 0 && } + {truncatedCount > 0 && } + } + actions={<> -
    -
    - - {expanded && ( -
    + }> {trialsQuery.isLoading ? (
    @@ -1009,14 +907,11 @@ export function RunsTab({ ? OBSOLETE_TRIAL_BADGE : trial.truncated ? TRUNCATED_TRIAL_BADGE - : trialStatusBadge(trial.status) + : null : null return ( -
  • -
    -
    -
    -

    Replicate {replicate.replicate_number}

    + {trial?.obsolete && ( )} -
    - {replicateUsage.length > 0 &&

    {replicateUsage.map((value) => {value})}

    } -
    -
    - {badge ? {badge.label} : Status unavailable} - {trial?.run_id && ( - - )} - } + detail={replicateUsage.length > 0 && replicateUsage.map(value => {value})} + status={badge ? {badge.label} : trial ? : Status unavailable} + onView={trial?.run_id ? () => onViewResult(replicate.replicate_label) : undefined} + actions={ + -
    -
    - {trial?.error &&

    {trial.error}

    } -
    • + } + error={trial?.error} + /> ) })} )} -
    - )} -
    + ) })}
    diff --git a/frontend/src/components/protocol/StopRowRuns.tsx b/frontend/src/components/protocol/StopRowRuns.tsx new file mode 100644 index 00000000..27339ec9 --- /dev/null +++ b/frontend/src/components/protocol/StopRowRuns.tsx @@ -0,0 +1,18 @@ +import { useMutation, useQueryClient } from '@tanstack/react-query' +import { Square } from 'lucide-react' +import { protocolsApi } from '@/api/client' +import { Button } from '@/components/ui/button' +import type { Protocol } from '@/types/protocols' + +export function StopRowRuns({ protocol, runIds }: { protocol: Protocol; runIds: string[] }) { + const client = useQueryClient() + const stop = useMutation({ + mutationFn: async () => { + const outcomes = await Promise.allSettled([...new Set(runIds)].map(id => protocolsApi.cancelRun(protocol.id, id))) + if (outcomes.some(result => result.status === 'rejected')) throw new Error('Some runs could not be stopped. Try again.') + }, + onSettled: () => client.invalidateQueries({ queryKey: ['experiments', protocol.experiment_id] }), + }) + if (!runIds.length) return null + return
    {stop.error &&

    {stop.error.message}

    }
    +} diff --git a/frontend/src/components/protocol/edges/DatasetInputPanel.test.tsx b/frontend/src/components/protocol/edges/DatasetInputPanel.test.tsx index 499bcd18..6ed092dd 100644 --- a/frontend/src/components/protocol/edges/DatasetInputPanel.test.tsx +++ b/frontend/src/components/protocol/edges/DatasetInputPanel.test.tsx @@ -18,15 +18,18 @@ it('defaults to whole, requires explicit columns, preserves metadata and CSV ord mount() expect(screen.getByLabelText('Whole dataset')).toBeChecked() await screen.findByText(/3 original rows/) + expect(screen.getByRole('button', { name: 'Apply' })).toBeDisabled() fireEvent.click(screen.getByLabelText('Per row')) expect(screen.getByRole('button', { name: 'Apply' })).toBeDisabled() fireEvent.click(screen.getByRole('checkbox', { name: 'reference' })) fireEvent.click(screen.getByRole('checkbox', { name: 'question' })) fireEvent.click(screen.getByRole('button', { name: 'Apply' })) expect(state.edges[0].data).toEqual({ custom: 'keep', handoff: { mode: 'full' }, dataset_input: { mode: 'per_row', columns: ['question', 'reference'] } }) + expect(screen.getByRole('button', { name: 'Apply' })).toBeDisabled() fireEvent.click(screen.getByLabelText('Whole dataset')) fireEvent.click(screen.getByRole('button', { name: 'Apply' })) expect(state.edges[0].data).toEqual({ custom: 'keep', handoff: { mode: 'full' } }) + expect(screen.getByRole('button', { name: 'Apply' })).toBeDisabled() }) it('keeps unknown saved columns visible and blocks apply', async () => { state.edges[0].data.dataset_input = { mode: 'per_row', columns: ['missing'] } @@ -48,9 +51,25 @@ it('acknowledges Apply and explains that production needs publication in both mo fireEvent.click(screen.getByRole('checkbox', { name: 'question' })) fireEvent.click(screen.getByRole('button', { name: 'Apply' })) expect(screen.getByRole('status')).toHaveTextContent('Per row applied to the draft') - expect(screen.getByRole('status')).toHaveTextContent('Publish the canvas to update production runs and results') + expect(screen.getByRole('status')).toHaveTextContent('Publish the experiment to update production runs and results') fireEvent.click(screen.getByLabelText('Whole dataset')) expect(screen.queryByRole('status')).not.toBeInTheDocument() fireEvent.click(screen.getByRole('button', { name: 'Apply' })) expect(screen.getByRole('status')).toHaveTextContent('Whole dataset applied to the draft') }) + +it('disables Apply when column or mode edits are reverted', async () => { + state.edges[0].data.dataset_input = { mode: 'per_row', columns: ['question'] } + mount() + await screen.findByText(/3 original rows/) + const apply = screen.getByRole('button', { name: 'Apply' }) + expect(apply).toBeDisabled() + fireEvent.click(screen.getByRole('checkbox', { name: 'reference' })) + expect(apply).toBeEnabled() + fireEvent.click(screen.getByRole('checkbox', { name: 'reference' })) + expect(apply).toBeDisabled() + fireEvent.click(screen.getByLabelText('Whole dataset')) + expect(apply).toBeEnabled() + fireEvent.click(screen.getByLabelText('Per row')) + expect(apply).toBeDisabled() +}) diff --git a/frontend/src/components/protocol/edges/DatasetInputPanel.tsx b/frontend/src/components/protocol/edges/DatasetInputPanel.tsx index e9021050..44f6f414 100644 --- a/frontend/src/components/protocol/edges/DatasetInputPanel.tsx +++ b/frontend/src/components/protocol/edges/DatasetInputPanel.tsx @@ -16,13 +16,17 @@ export function DatasetInputPanel({ edgeId, source, disabled = false }: { edgeId const [mode, setMode] = useState(saved?.mode ?? 'whole_dataset') const [columns, setColumns] = useState(saved?.mode === 'per_row' ? saved.columns : []) const [applied, setApplied] = useState(null) + const [appliedKey, setAppliedKey] = useState(() => JSON.stringify(saved?.mode === 'per_row' + ? ['per_row', [...saved.columns].sort()] : ['whole_dataset'])) + const inputKey = JSON.stringify(mode === 'per_row' ? [mode, [...columns].sort()] : [mode]) + const hasChanges = inputKey !== appliedKey const id = (nodes.find(node => node.id === source)?.data.config as { dataset_id?: string })?.dataset_id const schema = useQuery({ queryKey: ['datasets', id, 'row-schema'], queryFn: () => datasetsApi.getRowSchema(id!), enabled: !!id }) const input: DatasetInput = mode === 'per_row' ? { mode, columns } : { mode } const graph = { nodes, edges: edges.map(item => item.id === edgeId ? { ...item, data: { ...item.data, dataset_input: input } } : item) } const error = !id ? 'Select a registered Dataset.' : schema.isError ? schema.error.message : schema.isLoading ? 'Loading original columns…' : !schema.data ? 'Source metadata unavailable.' : mode === 'per_row' && !schema.data.row_count ? 'The original Dataset has no rows.' : mode === 'per_row' && columns.some(column => !schema.data.columns.includes(column)) ? 'Saved columns are missing from the original CSV.' : datasetRowTopologyError(graph) function apply() { - if (error || disabled) return + if (error || disabled || !hasChanges) return setEdges(current => current.map(item => { if (item.id !== edgeId) return item const data = { ...item.data } @@ -31,6 +35,7 @@ export function DatasetInputPanel({ edgeId, source, disabled = false }: { edgeId return { ...item, data } })) setApplied(mode) + setAppliedKey(inputKey) } return

    Dataset input

    @@ -40,7 +45,7 @@ export function DatasetInputPanel({ edgeId, source, disabled = false }: { edgeId {schema.data &&

    {schema.data.row_count} original rows. {mode === 'per_row' ? 'Each row runs one complete protocol; Agents share one driver and select their own columns.' : 'Each replicate runs the protocol with the whole dataset.'}

    } {mode === 'per_row' &&
    {[...new Set([...(schema.data?.columns ?? []), ...columns])].map(column => )}
    } {error &&

    {error}

    } - - {applied ?

    {applied === 'per_row' ? 'Per row' : 'Whole dataset'} applied to the draft. Publish the canvas to update production runs and results.

    :

    Apply changes the draft connection. Publish the canvas to use it for production runs and results.

    } + + {applied ?

    {applied === 'per_row' ? 'Per row' : 'Whole dataset'} applied to the draft. Publish the experiment to update production runs and results.

    :

    Apply changes the draft connection. Publish the experiment to use it for production runs and results.

    }
    } diff --git a/frontend/src/lib/experiment.ts b/frontend/src/lib/experiment.ts index b593dd81..579d8336 100644 --- a/frontend/src/lib/experiment.ts +++ b/frontend/src/lib/experiment.ts @@ -91,6 +91,17 @@ export interface FactorSpec { level_labels?: string[] } +export function displayFactorCondition(name: string, value: unknown, designSpec: Experiment['design_spec'] | undefined, cellLabel?: string): string { + const parts = name.split(':').map(part => part.trim()).filter(Boolean) + const field = parts.pop() ?? name + const displayedLevel = displayFactorLevel(designSpec, name, value, cellLabel) + if (typeof value === 'boolean' && /enabled$/i.test(field) && displayedLevel === String(value)) { + const subject = field.replace(/\s*enabled$/i, '').trim() + return [...parts, subject].filter(Boolean).join(' · ') + `: ${value ? 'Enabled' : 'Disabled'}` + } + return [...parts, field].filter(Boolean).join(' · ') + `: ${displayedLevel}` +} + function getFactors(designSpec: Experiment['design_spec'] | undefined): FactorSpec[] | null { const factors = (designSpec as { factors?: unknown } | null)?.factors return Array.isArray(factors) ? (factors as FactorSpec[]) : null diff --git a/frontend/src/types/experiments.ts b/frontend/src/types/experiments.ts index d5a564b9..29365ed1 100644 --- a/frontend/src/types/experiments.ts +++ b/frontend/src/types/experiments.ts @@ -512,6 +512,7 @@ export interface RowSummary { } export interface ExperimentRunResults { + row_cells?: RowCell[] consumption_mode?: 'whole_dataset' | 'per_row' row_results?: RowResult[] row_summary?: RowSummary | null @@ -526,6 +527,14 @@ export interface ExperimentRunResults { replicates: ResultReplicate[] } +export interface RowCell { + cell_id: string + cell_label: string + factor_values: Record + replicate_count: number + replicates?: Array<{ replicate_result_id: string; replicate_label: string; replicate_number: number }> +} + export interface Replicate { // One independently runnable observation within the owning cell. id: string diff --git a/src/asaree/api/experiments.py b/src/asaree/api/experiments.py index f54d07aa..85553a4f 100644 --- a/src/asaree/api/experiments.py +++ b/src/asaree/api/experiments.py @@ -1080,6 +1080,7 @@ class RunResultsResponse(BaseModel): replicates: list[dict[str, Any]] consumption_mode: str row_results: list[dict[str, Any]] + row_cells: list[dict[str, Any]] = Field(default_factory=list) row_summary: dict[str, Any] | None diff --git a/src/asaree/services/experiment_run_results.py b/src/asaree/services/experiment_run_results.py index 1ab4faa6..79be7415 100644 --- a/src/asaree/services/experiment_run_results.py +++ b/src/asaree/services/experiment_run_results.py @@ -580,6 +580,19 @@ async def _summarize_row_results( return { "consumption_mode": "per_row", "row_results": row_results, + "row_cells": [ + { + "cell_id": str(cell.id), "cell_label": cell.cell_label, + "factor_values": cell.factor_values or {}, + "replicate_count": sum(parent.cell_id == cell.id for parent in replicates), + "replicates": [ + {"replicate_result_id": str(parent.id), "replicate_label": parent.replicate_label, + "replicate_number": parent.replicate_number} + for parent in replicates if parent.cell_id == cell.id + ], + } + for cell in sorted(cells, key=lambda cell: cell.cell_label) + ], "row_summary": row_summary, "overview": { "total_replicates": actual_parents, diff --git a/src/asaree/services/factorial_row_results.py b/src/asaree/services/factorial_row_results.py index d2a2e4c7..b345806e 100644 --- a/src/asaree/services/factorial_row_results.py +++ b/src/asaree/services/factorial_row_results.py @@ -264,6 +264,7 @@ async def claim_row_attempt( row_result_id: uuid.UUID, expected_run_id: uuid.UUID | None, create_kwargs: dict, + allow_completed: bool = False, ) -> ProtocolRun | None: """Atomically claim the latest-attempt slot for a row execution.""" slot = await db.scalar( @@ -278,7 +279,7 @@ async def claim_row_attempt( previous = await db.get(ProtocolRun, expected_run_id, populate_existing=True) if previous is None or previous.status not in TERMINAL_PROTOCOL_RUN_STATUSES: return None - if previous.status not in {"failed", "cancelled"}: + if previous.status not in {"failed", "cancelled"} and not allow_completed: return None from asaree.services.protocol_runs import create_protocol_run diff --git a/src/asaree/services/protocol_execution.py b/src/asaree/services/protocol_execution.py index a5eae9ed..82472965 100644 --- a/src/asaree/services/protocol_execution.py +++ b/src/asaree/services/protocol_execution.py @@ -4773,8 +4773,6 @@ async def plan_cell_runs( except DatasetRowInputError as exc: raise ProtocolValidationError(str(exc)) from exc if original_row_plan is not None: - if rerun_replicate_labels: - raise ProtocolValidationError("Row executions cannot be rerun through the cell batch; use row retry.") retry_targets = None if retry_row_result_ids is not None: if current_design is None: @@ -4908,6 +4906,9 @@ async def plan_cell_runs( unknown = requested_labels - labels if unknown: raise ProtocolValidationError(f"Unknown replicate label(s): {', '.join(sorted(unknown))}.") + requested_reruns = rerun_replicate_labels or set() + if requested_reruns - requested_labels: + raise ProtocolValidationError("Rerun replicates must belong to the selected batch.") selected_parents = [parent for parent in parents if parent.replicate_label in requested_labels] parent_values = [ { @@ -4957,12 +4958,14 @@ async def plan_cell_runs( raw_sha256=source.raw_sha256, row_index=candidate["dataset_row"]["row_index"], ) - claims.append((row_result.id, candidate, view)) - for row_result_id, candidate, view in sorted(claims, key=lambda item: item[0].int): + expected_run_id = row_result.run_id if candidate["replicate_label"] in requested_reruns else None + claims.append((row_result.id, candidate, view, expected_run_id)) + for row_result_id, candidate, view, expected_run_id in sorted(claims, key=lambda item: item[0].int): run = await claim_row_attempt( db, row_result_id=row_result_id, - expected_run_id=None, + expected_run_id=expected_run_id, + allow_completed=candidate["replicate_label"] in requested_reruns, create_kwargs={ "protocol_id": protocol_id, "owner_id": owner_id, @@ -4978,7 +4981,7 @@ async def plan_cell_runs( skipped += 1 else: created.append(run) - candidate_order = {row_id: index for index, (row_id, _, _) in enumerate(claims)} + candidate_order = {row_id: index for index, (row_id, _, _, _) in enumerate(claims)} created.sort(key=lambda run: candidate_order[run.row_result_id]) return created, skipped elif resolve_dataset_row_plan(revision.graph if revision is not None else graph, design_spec) is not None: diff --git a/tests/test_dataset_row_claims.py b/tests/test_dataset_row_claims.py index e7712a92..fbe1878c 100644 --- a/tests/test_dataset_row_claims.py +++ b/tests/test_dataset_row_claims.py @@ -55,6 +55,41 @@ async def test_claim_retry_forwards_row_result_id_once(monkeypatch, include_row_ assert create_kwargs == original_kwargs +@pytest.mark.parametrize( + ("status", "allow_completed", "stale", "claimed"), + [ + ("completed", False, False, False), + ("completed", True, False, True), + ("limit_reached", True, False, True), + ("limit_reached", False, False, False), + ("failed", False, False, True), + ("cancelled", False, False, True), + ("pending", True, False, False), + ("running", True, False, False), + ("finalizing", True, False, False), + ("completed", True, True, False), + ], +) +async def test_claim_requires_explicit_rerun_and_terminal_current_attempt( + monkeypatch, status, allow_completed, stale, claimed, +) -> None: + slot_id, previous_run_id = uuid.uuid4(), uuid.uuid4() + db = AsyncMock(spec=AsyncSession) + db.scalar.return_value = FactorialRowResult( + id=slot_id, run_id=uuid.uuid4() if stale else previous_run_id, + ) + db.get.return_value = ProtocolRun(id=previous_run_id, status=status) + create_run = AsyncMock(return_value=ProtocolRun(id=uuid.uuid4(), row_result_id=slot_id)) + monkeypatch.setattr("asaree.services.protocol_runs.create_protocol_run", create_run) + run = await claim_row_attempt( + db, row_result_id=slot_id, expected_run_id=previous_run_id, + allow_completed=allow_completed, + create_kwargs={"protocol_id": uuid.uuid4(), "owner_id": uuid.uuid4()}, + ) + assert (run is create_run.return_value) is claimed + assert create_run.await_count == int(claimed) + + @pytest_asyncio.fixture async def row_batch(tmp_path: Path) -> AsyncIterator[tuple[AsyncSession, dict]]: """Create row-batch setup committed for independently locking sessions.""" @@ -194,3 +229,49 @@ async def competing_claim(): ) assert count == 1 await engine.dispose() + + +@pytest.mark.asyncio +async def test_explicit_row_rerun_preserves_slots_and_history_and_skips_active(row_batch, monkeypatch) -> None: + db, ctx = row_batch + monkeypatch.setattr(protocol_api, "enqueue_protocol_run", AsyncMock()) + first = await create_cell_runs_endpoint(ctx["protocol"].id, ctx["user"], db, CellRunBatchRequest()) + assert len(first.protocol_run_ids) == 30 + for index, run_id in enumerate(first.protocol_run_ids): + run = await db.get(ProtocolRun, run_id) + assert run is not None + run.status = "running" if index == 0 else "completed" + run.attempt_result = {"output": f"original-{index}"} + slot = await db.get(FactorialRowResult, first.row_result_ids[index]) + assert slot is not None + slot.metric_values = {"score": index} + await db.commit() + + # Ordinary batches continue to skip previous attempts. + ordinary = await create_cell_runs_endpoint(ctx["protocol"].id, ctx["user"], db, CellRunBatchRequest()) + assert ordinary.protocol_run_ids == [] + rerun = await create_cell_runs_endpoint( + ctx["protocol"].id, ctx["user"], db, + CellRunBatchRequest(replicate_labels=["batch-parent"], rerun_replicate_labels=["batch-parent"]), + ) + assert len(rerun.protocol_run_ids) == 29 + assert set(rerun.row_result_ids) == set(first.row_result_ids[1:]) + for index, slot_id in enumerate(first.row_result_ids): + slot = await db.get(FactorialRowResult, slot_id, populate_existing=True) + assert slot is not None + attempts = list((await db.scalars(select(ProtocolRun).where(ProtocolRun.row_result_id == slot_id))).all()) + assert len(attempts) == (1 if index == 0 else 2) + original = next(attempt for attempt in attempts if attempt.id == first.protocol_run_ids[index]) + assert original.attempt_result == {"output": f"original-{index}"} + if index == 0: + assert slot.run_id == original.id + else: + assert slot.run_id != original.id + assert slot.metric_values is None + + # A second request sees active latest attempts and never creates duplicates. + duplicate = await create_cell_runs_endpoint( + ctx["protocol"].id, ctx["user"], db, + CellRunBatchRequest(replicate_labels=["batch-parent"], rerun_replicate_labels=["batch-parent"]), + ) + assert duplicate.protocol_run_ids == [] diff --git a/tests/test_dataset_row_results_projection.py b/tests/test_dataset_row_results_projection.py index f98b25ac..b94bd33f 100644 --- a/tests/test_dataset_row_results_projection.py +++ b/tests/test_dataset_row_results_projection.py @@ -196,6 +196,9 @@ async def no_enqueue(_run_id: uuid.UUID) -> None: summary = results["row_summary"] assert results["consumption_mode"] == "per_row" assert results["cells"] == [] and results["replicates"] == [] + assert len(results["row_cells"]) == 2 + assert {cell["cell_id"] for cell in results["row_cells"]} == {str(cell.id) for cell in ctx["cells"]} + assert all(cell["replicate_count"] == 1 for cell in results["row_cells"]) assert results["primary_metric"] is None and results["primary_metric_direction"] is None assert len(results["row_results"]) == 30 assert summary == { From 2a564b1d2dc2c50214c5dea94d7fe939ad8288c6 Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Tue, 6 Oct 2026 09:23:33 -0700 Subject: [PATCH 12/16] Hide publish experiment button when there are no changes --- frontend/src/pages/ProtocolCanvasPage.tsx | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/frontend/src/pages/ProtocolCanvasPage.tsx b/frontend/src/pages/ProtocolCanvasPage.tsx index 26b29881..0131ccfe 100644 --- a/frontend/src/pages/ProtocolCanvasPage.tsx +++ b/frontend/src/pages/ProtocolCanvasPage.tsx @@ -186,15 +186,17 @@ function ProtocolPublicationControl({ protocol, experimentId, draftBusy }: { pro {status} {error && {error}} - + {(protocol.has_unpublished_changes || draftBusy) && ( + + )}
    From 35658e03a47d5b7e639c42e31bc8d95d49832764 Mon Sep 17 00:00:00 2001 From: Jay Moran Date: Tue, 6 Oct 2026 09:31:07 -0700 Subject: [PATCH 13/16] Add searchable experiment version history with names and notes --- frontend/src/api/client.ts | 3 +- frontend/src/api/versionAnnotations.test.ts | 22 +++++ .../ExperimentVersionHistory.test.tsx | 51 ++++++++++ .../protocol/ExperimentVersionHistory.tsx | 90 ++++++++++++++++++ frontend/src/lib/experimentVersion.ts | 5 + .../src/pages/ProtocolCanvasPage.test.tsx | 55 ++++++++--- frontend/src/pages/ProtocolCanvasPage.tsx | 54 +++++------ frontend/src/types/protocols.ts | 4 + src/asaree/api/protocols.py | 93 ++++++++++++++++--- .../2c3d4e5f6071_version_annotations.py | 19 ++++ src/asaree/models/protocol_revision.py | 5 +- src/asaree/services/protocol_revisions.py | 5 +- tests/test_version_annotations.py | 63 +++++++++++++ 13 files changed, 415 insertions(+), 54 deletions(-) create mode 100644 frontend/src/api/versionAnnotations.test.ts create mode 100644 frontend/src/components/protocol/ExperimentVersionHistory.test.tsx create mode 100644 frontend/src/components/protocol/ExperimentVersionHistory.tsx create mode 100644 frontend/src/lib/experimentVersion.ts create mode 100644 src/asaree/migrations/versions/2c3d4e5f6071_version_annotations.py create mode 100644 tests/test_version_annotations.py diff --git a/frontend/src/api/client.ts b/frontend/src/api/client.ts index 0503eb29..ea31d26c 100644 --- a/frontend/src/api/client.ts +++ b/frontend/src/api/client.ts @@ -309,7 +309,7 @@ export const protocolsApi = { request(experimentId ? `/protocols?experiment_id=${experimentId}` : '/protocols'), update: (id: string, data: { name?: string; description?: string | null; graph?: ProtocolGraph }) => request(`/protocols/${id}`, { method: 'PATCH', body: data }), - publish: (id: string) => request(`/protocols/${id}/publish`, { method: 'POST' }), + publish: (id: string, annotations?: { name?: string | null; note?: string | null }) => request(`/protocols/${id}/publish`, { method: 'POST', body: annotations }), remove: (id: string) => request(`/protocols/${id}`, { method: 'DELETE' }), // 422 if the graph is empty or has a cycle -- returns immediately with // status "pending"; poll getRun for progress. cellLabel runs that one @@ -331,6 +331,7 @@ export const protocolsApi = { request(`/protocols/${id}/nodes/${nodeId}/prompt-preview`, { method: 'POST', body: { graph, ...rowSelection(options) } }), listRevisions: (id: string) => request(`/protocols/${id}/revisions`), getRevision: (id: string, revisionId: string) => request(`/protocols/${id}/revisions/${revisionId}`), + updateRevision: (id: string, revisionId: string, annotations: { name?: string | null; note?: string | null }) => request(`/protocols/${id}/revisions/${revisionId}`, { method: 'PATCH', body: annotations }), getRun: (id: string, runId: string) => request(`/protocols/${id}/runs/${runId}`), // Queued work is cancelled immediately. Active work raises // cancel_requested_at, which its executor honors at a safe interruption diff --git a/frontend/src/api/versionAnnotations.test.ts b/frontend/src/api/versionAnnotations.test.ts new file mode 100644 index 00000000..b64e2895 --- /dev/null +++ b/frontend/src/api/versionAnnotations.test.ts @@ -0,0 +1,22 @@ +import { afterEach, expect, it, vi } from 'vitest' +import { protocolsApi } from './client' + +afterEach(() => vi.unstubAllGlobals()) + +it.each(['publish', 'updateRevision'] as const)('sends %s annotations as a JSON object accepted by the API', async operation => { + const annotations = { name: 'Baseline', note: 'Keep for comparison' } + const fetch = vi.fn(async (_url: string, options: RequestInit) => { + const body: unknown = JSON.parse(options.body as string) + // FastAPI's VersionAnnotationsRequest rejects a JSON string with 422. + if (typeof body !== 'object' || body === null) { + return new Response(JSON.stringify({ detail: [{ msg: 'Input should be a valid dictionary' }] }), { status: 422 }) + } + return new Response(JSON.stringify({ id: 'p', ...body }), { status: 200 }) + }) + vi.stubGlobal('fetch', fetch) + const result = operation === 'publish' + ? protocolsApi.publish('p', annotations) + : protocolsApi.updateRevision('p', 'v1', annotations) + await expect(result).resolves.toMatchObject({ id: 'p', ...annotations }) + expect(JSON.parse(fetch.mock.calls[0][1].body as string)).toEqual(annotations) +}) diff --git a/frontend/src/components/protocol/ExperimentVersionHistory.test.tsx b/frontend/src/components/protocol/ExperimentVersionHistory.test.tsx new file mode 100644 index 00000000..df237312 --- /dev/null +++ b/frontend/src/components/protocol/ExperimentVersionHistory.test.tsx @@ -0,0 +1,51 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { fireEvent, render, screen, waitFor, within } from '@testing-library/react' +import { afterEach, expect, it, vi } from 'vitest' +import { protocolsApi } from '@/api/client' +import type { ProtocolRevision } from '@/types/protocols' +import { ExperimentVersionHistory } from './ExperimentVersionHistory' + +afterEach(() => vi.restoreAllMocks()) + +const versions: ProtocolRevision[] = [3, 2, 1].map(revision => ({ + id: `v${revision}`, protocol_id: 'p', revision, graph: { nodes: [], edges: [] }, + published_at: `2026-01-0${revision}`, name: revision === 1 ? 'Baseline' : null, + note: revision === 1 ? 'Lower critic threshold' : null, + run_count: revision === 2 ? 2 : 0, result_count: revision === 2 ? 1 : 0, +})) + +it('filters by provenance counts and searches notes while keeping production visible', () => { + const client = new QueryClient() + const onSelect = vi.fn() + const props = { versions, protocolId: 'p', publishedId: 'v3', selectedId: '', onSearch: vi.fn(), onSelect } + const { rerender } = render() + fireEvent.click(screen.getByRole('button', { name: 'With results' })) + expect(screen.getByRole('button', { name: 'View Version 3' })).toBeInTheDocument() + expect(screen.getByRole('button', { name: 'View Version 2' })).toBeInTheDocument() + expect(screen.queryByRole('button', { name: 'View Version 1 · Baseline' })).not.toBeInTheDocument() + fireEvent.click(screen.getByRole('button', { name: 'With runs' })) + expect(screen.getByText('1 of 3 versions · newest first')).toBeInTheDocument() + fireEvent.click(screen.getByRole('button', { name: 'All' })) + rerender() + expect(screen.getByRole('button', { name: 'View Version 1 · Baseline' })).toBeInTheDocument() + expect(screen.queryByRole('button', { name: 'View Version 2' })).not.toBeInTheDocument() + fireEvent.click(screen.getByRole('button', { name: 'View Version 1 · Baseline' })) + expect(onSelect).toHaveBeenCalledWith('v1') +}) + +it('edits annotations and updates the cached version without changing its definition', async () => { + const client = new QueryClient() + client.setQueryData(['protocols', 'p', 'revisions'], versions) + const updated = { ...versions[2], name: 'Reference', note: 'Keep for comparison' } + const save = vi.spyOn(protocolsApi, 'updateRevision').mockResolvedValue(updated) + render() + const entry = screen.getByRole('button', { name: 'View Version 1 · Baseline' }).closest('[data-slot="card"]') as HTMLElement + fireEvent.click(within(entry).getByRole('button', { name: 'Details' })) + expect(within(entry).getByText('Lower critic threshold')).toBeInTheDocument() + fireEvent.click(within(entry).getByRole('button', { name: 'Edit name and note' })) + fireEvent.change(within(entry).getByRole('textbox', { name: 'Version name (optional)' }), { target: { value: ' Reference ' } }) + fireEvent.change(within(entry).getByRole('textbox', { name: 'Version note (optional)' }), { target: { value: 'Keep for comparison' } }) + fireEvent.click(within(entry).getByRole('button', { name: 'Save details' })) + await waitFor(() => expect(save).toHaveBeenCalledWith('p', 'v1', { name: 'Reference', note: 'Keep for comparison' })) + await waitFor(() => expect(client.getQueryData(['protocols', 'p', 'revisions'])?.[2]).toEqual(updated)) +}) diff --git a/frontend/src/components/protocol/ExperimentVersionHistory.tsx b/frontend/src/components/protocol/ExperimentVersionHistory.tsx new file mode 100644 index 00000000..732c71a8 --- /dev/null +++ b/frontend/src/components/protocol/ExperimentVersionHistory.tsx @@ -0,0 +1,90 @@ +import { useEffect, useState } from 'react' +import { useMutation, useQueryClient } from '@tanstack/react-query' +import { protocolsApi } from '@/api/client' +import { Button } from '@/components/ui/button' +import { Card } from '@/components/ui/card' +import { Input } from '@/components/ui/input' +import { Textarea } from '@/components/ui/textarea' +import type { ProtocolRevision } from '@/types/protocols' +import { versionLabel } from '@/lib/experimentVersion' + +function changeSummary(version: ProtocolRevision, previous?: ProtocolRevision) { + if (!previous) return 'First published version' + const changes: string[] = [] + if (JSON.stringify(version.graph) !== JSON.stringify(previous.graph)) changes.push('Canvas') + for (const [key, label] of [['hypothesis', 'Hypothesis'], ['design_type', 'Design type'], ['design_spec', 'Design'], ['measurement_plan', 'Measurement plan'], ['task_brief', 'Task brief']] as const) { + if (JSON.stringify(version.experiment_snapshot?.[key]) !== JSON.stringify(previous.experiment_snapshot?.[key])) changes.push(label) + } + if (version.design_revision_id !== previous.design_revision_id) changes.push('Generated cells') + return changes.length ? `Changed: ${changes.join(', ')}` : 'Experiment definition unchanged' +} + +function VersionEntry({ version, previous, current, selected, protocolId, onSelect }: { + version: ProtocolRevision; previous?: ProtocolRevision; current: boolean; selected: boolean; + protocolId: string; onSelect: (id: string) => void +}) { + const [expanded, setExpanded] = useState(false) + const [editing, setEditing] = useState(false) + const [name, setName] = useState(version.name ?? '') + const [note, setNote] = useState(version.note ?? '') + const queryClient = useQueryClient() + const mutation = useMutation({ + mutationFn: () => protocolsApi.updateRevision(protocolId, version.id, { name: name.trim() || null, note: note.trim() || null }), + onSuccess: updated => { + queryClient.setQueryData(['protocols', protocolId, 'revisions'], versions => versions?.map(item => item.id === updated.id ? updated : item)) + setEditing(false) + }, + }) + return +
    + Version {version.revision}{version.name ? ` · ${version.name}` : ''} + {current && Production} + {new Date(version.published_at).toLocaleDateString()} +
    +

    {version.run_count ?? 0} runs · {version.result_count ?? 0} with results

    +
    + + +
    + {expanded &&
    +

    {changeSummary(version, previous)}

    + {editing ?
    { event.preventDefault(); mutation.mutate() }}> + +