fix(embed): one completeness contract for both backends, and a FASTA coverage check

apps/protspace/CLAUDE.md

@@ -46,7 +46,7 @@ Single entry point: `protspace = protspace.cli.app:app`
 | Command | Purpose |
 |---------|---------|
 | `protspace prepare` | Full pipeline: embed → reduce → annotate → bundle |
-| `protspace embed` | FASTA → HDF5 embeddings (Biocentral API or local GPU/CPU via `--backend local`) |
+| `protspace embed` | FASTA → HDF5 embeddings (Biocentral API or local GPU/CPU via `--backend local`). Exits non-zero on an incomplete embedding; capability-limited sequences (`--max-length`, GPU OOM) are skipped and named instead |
 | `protspace project` | HDF5 → dimensionality reduction |
 | `protspace annotate` | Fetch protein annotations |
 | `protspace bundle` | Combine projections + annotations → .parquetbundle |
@@ -156,6 +156,7 @@ src/protspace/
 │   │   ├── settings_converter.py # Settings table conversion
 │   │   └── writers.py          # Annotation output writers
 │   ├── embedding/
+│   │   ├── store.py            # Shared HDF5 layer + completeness contract (owned by neither backend)
 │   │   ├── biocentral.py       # Biocentral API client, model shortcut mappings
 │   │   └── local.py            # Local GPU/CPU backend (HF transformers, [local] extra)
 │   ├── parsers/
@@ -280,9 +281,10 @@ For a live count run `uv run pytest tests/ --collect-only -q`.
 | `test_stats_bundle.py` | Optional 5th (statistics) bundle part round-trip |
 | `test_annotation_select.py` | Annotation selection: suitability filter (cardinality/numeric/id-like exclusion), `auto` vs explicit-list label building (explicit names bypass the heuristic), missing-value dropping |
 | `test_annotation_validity.py` | `AnnotationValidityStatistic`: silhouette/DBI/CH scored per annotation on `ctx.coords`, embedding vs. projection `space_kind`, missing-value exclusion, single-category no-op, id-canonical subsample determinism |
-| `test_biocentral_embedder.py` | Biocentral API client, embedding flow |
+| `test_biocentral_embedder.py` | Biocentral API client, embedding flow, completeness gate (reads the .h5, not a counter), `/`-in-header rejection |
+| `test_embed_completeness.py` | Shared embed contract (`data/embedding/store.py`): `expected = requested - skipped`, skip-vs-fail, skip reporting, resume-covered runs, FASTA coverage direction + identifier normalisation |
 | `test_backend_switch.py` | Embedding backend switch: `resolve_default_backend` (Colab+GPU→local), `embed_fasta` local/biocentral dispatch (short key vs resolved name), `protspace embed --backend` CLI wiring + enum validation + non-positive batch_size rejection |
-| `test_local_embedder.py` | Local embedding backend: checkpoint resolution (12 short keys, Synthyra ESM-C), the notebook-gating sets pinned to the registry each constrains (`COLAB_OVERSIZED`→`LOCAL_CHECKPOINTS`, `BIOCENTRAL_INVALID`→`ALL_SHORT_KEYS`), per-family preprocessing/residue pooling, `/`-in-header guard, LocalEmbedConfig validation, empty-output guard, esm2_8m end-to-end + resume (slow) |
+| `test_local_embedder.py` | Local embedding backend: checkpoint resolution (12 short keys, Synthyra ESM-C), the notebook-gating sets pinned to the registry each constrains (`COLAB_OVERSIZED`→`LOCAL_CHECKPOINTS`, `BIOCENTRAL_INVALID`→`ALL_SHORT_KEYS`), per-family preprocessing/residue pooling, `/`-in-header guard, LocalEmbedConfig validation, over-length + OOM skips reported not failed, non-skip shortfall fails, esm2_8m end-to-end + resume (slow) |
 | `test_fasta.py` | FASTA parsing, edge cases, CSV annotation loading |
 | `test_biocentral_retriever.py` | Biocentral prediction retriever (TMbed parsing, per-sequence) |
 | `test_taxonomy_annotation_retriever.py` | Taxonomy via UniProt Taxonomy API (mocked + integration) |

apps/protspace/src/protspace/cli/common_options.py

@@ -202,13 +202,57 @@ def require_similarity_extra() -> None:
     ),
 ]
 
+Opt_MaxLength = Annotated[
+    int | None,
+    typer.Option(
+        min=1,
+        help=(
+            "Skip sequences longer than this (local backend only; default 2000). "
+            "Skipped sequences are named in the run summary."
+        ),
+        rich_help_panel="Embedding",
+    ),
+]
+
 # Input options (shared by prepare and project)
 Opt_Fasta = Annotated[
     Path | None,
     typer.Option(
         "-f",
         "--fasta",
         help="FASTA for -s/--similarity when input is HDF5.",
+        exists=True,
+        dir_okay=False,
         rich_help_panel="Input",
     ),
 ]
+
+
+def build_embed_config(
+    backend: "Backend",
+    batch_size: int | None = None,
+    max_length: int | None = None,
+):
+    """Build the embedding config matching *backend*, applying only what was given.
+
+    Both ``embed`` and ``prepare`` need this, and each needs it per backend, so
+    without it the same four-line conditional appears four times.
+    """
+    if backend == Backend.local:
+        from protspace.data.embedding.local import LocalEmbedConfig
+
+        opts = {}
+        if batch_size is not None:
+            opts["batch_size"] = batch_size
+        if max_length is not None:
+            opts["max_length"] = max_length
+        return LocalEmbedConfig(**opts)
+
+    from protspace.data.embedding.biocentral import EmbedConfig
+
+    if max_length is not None:
+        raise typer.BadParameter(
+            "--max-length applies to --backend local; the Biocentral backend "
+            "has no length cap."
+        )
+    return EmbedConfig(**({"batch_size": batch_size} if batch_size is not None else {}))

apps/protspace/src/protspace/cli/embed.py

@@ -13,7 +13,9 @@
     Backend,
     Opt_Backend,
     Opt_BatchSize,
+    Opt_MaxLength,
     Opt_Verbose,
+    build_embed_config,
 )
 
 logger = logging.getLogger(__name__)
@@ -47,6 +49,7 @@ def embed(
     ],
     backend: Opt_Backend = Backend.biocentral,
     batch_size: Opt_BatchSize = None,
+    max_length: Opt_MaxLength = None,
     verbose: Opt_Verbose = 0,
 ) -> None:
     """FASTA → per-model HDF5 embeddings.
@@ -68,29 +71,19 @@ def embed(
 
     output.mkdir(parents=True, exist_ok=True)
 
-    if backend == Backend.local:
-        from protspace.data.embedding.local import LocalEmbedConfig, embed_sequences
+    embed_config = build_embed_config(backend, batch_size, max_length)
 
-        embed_config = (
-            LocalEmbedConfig(batch_size=batch_size)
-            if batch_size is not None
-            else LocalEmbedConfig()
-        )
+    if backend == Backend.local:
+        from protspace.data.embedding.local import embed_sequences
 
         def resolve(name: str) -> str:
             return name  # local backend takes the short key directly
     else:
         from protspace.data.embedding.biocentral import (
-            EmbedConfig,
             embed_sequences,
             resolve_embedder,
         )
 
-        embed_config = (
-            EmbedConfig(batch_size=batch_size)
-            if batch_size is not None
-            else EmbedConfig()
-        )
         resolve = resolve_embedder
 
     failed_models: list[str] = []

apps/protspace/src/protspace/cli/prepare.py

@@ -32,6 +32,7 @@
     Opt_FpRatio,
     Opt_LearningRate,
     Opt_MaxIter,
+    Opt_MaxLength,
     Opt_Methods,
     Opt_Metric,
     Opt_MinDist,
@@ -42,6 +43,7 @@
     Opt_RandomState,
     Opt_Similarity,
     Opt_Verbose,
+    build_embed_config,
     require_similarity_extra,
 )
 
@@ -295,6 +297,7 @@ def prepare(
     embedder: Opt_Embedder = None,
     backend: Opt_Backend = Backend.biocentral,
     batch_size: Opt_BatchSize = None,
+    max_length: Opt_MaxLength = None,
     # Projection
     methods: Opt_Methods = None,
     similarity: Opt_Similarity = False,
@@ -422,22 +425,7 @@ def prepare(
         query_uniprot,
     )
 
-    if backend == Backend.local:
-        from protspace.data.embedding.local import LocalEmbedConfig
-
-        embed_config = (
-            LocalEmbedConfig(batch_size=batch_size)
-            if batch_size is not None
-            else LocalEmbedConfig()
-        )
-    else:
-        from protspace.data.embedding.biocentral import EmbedConfig
-
-        embed_config = (
-            EmbedConfig(batch_size=batch_size)
-            if batch_size is not None
-            else EmbedConfig()
-        )
+    embed_config = build_embed_config(backend, batch_size, max_length)
     embedding_sets: list[EmbeddingSet] = []
     fasta_for_similarity: Path | None = fasta
 
@@ -486,13 +474,9 @@ def prepare(
                     if not h5s:
                         logger.warning(f"No embedding files in: {path}")
                         continue
-                    embedding_sets.append(load_h5(h5s, name_override=name_override))
+                    emb_set = load_h5(h5s, name_override=name_override)
                 elif path.suffix.lower() in EMBEDDING_EXTENSIONS:
                     emb_set = load_h5([path], name_override=name_override)
-                    # Attach FASTA path from -f flag if provided (for sequence reuse)
-                    if fasta_for_similarity:
-                        emb_set.fasta_path = fasta_for_similarity
-                    embedding_sets.append(emb_set)
                 elif path.suffix.lower() in {".fasta", ".fa", ".faa"}:
                     _embed_all(
                         embedders,
@@ -504,12 +488,31 @@ def prepare(
                         force_reembed="embed" in refetch_stages,
                     )
                     fasta_for_similarity = path
+                    continue
                 else:
                     raise typer.BadParameter(f"Unsupported file: {path}")
 
+                # -f carries the sequences into the bundle, and it applies to
+                # every HDF5 input -- a directory of them as much as one file.
+                if fasta_for_similarity:
+                    emb_set.fasta_path = fasta_for_similarity
+                embedding_sets.append(emb_set)
+
         if not embedding_sets:
             raise typer.BadParameter("No valid input data found.")
 
+        # --- FASTA coverage ---
+        # Before similarity, not after: an uncovered protein inverts the whole
+        # MDS projection rather than degrading its own row.
+        if fasta_for_similarity is not None and embedding_sets:
+            from protspace.data.loaders.fasta import check_fasta_coverage
+
+            check_fasta_coverage(
+                fasta_for_similarity,
+                embedding_sets[0].headers,
+                required=bool(similarity),
+            )
+
         # --- Similarity ---
         # Both preconditions were checked before any input was read, so
         # `fasta_for_similarity` is set here whenever `similarity` is on.

apps/protspace/src/protspace/cli/project.py

@@ -57,6 +57,8 @@ def project(
             "-f",
             "--fasta",
             help="FASTA for -s/--similarity when input is HDF5.",
+            exists=True,
+            dir_okay=False,
             rich_help_panel="Input / Output",
         ),
     ] = None,
@@ -115,6 +117,9 @@ def project(
         raise typer.BadParameter("No valid HDF5 files found.")
 
     if similarity:
+        from protspace.data.loaders.fasta import check_fasta_coverage
+
+        check_fasta_coverage(fasta, embedding_sets[0].headers, required=True)
         sim_set = compute_similarity(fasta, embedding_sets[0].headers)
         embedding_sets.append(sim_set)
 

apps/protspace/src/protspace/data/embedding/biocentral.py

@@ -8,11 +8,19 @@
 from difflib import get_close_matches
 from pathlib import Path
 
-import h5py
 import numpy as np
 from biocentral_api import BiocentralAPI, CommonEmbedder, batched
 from tqdm import tqdm
 
+# Re-exported: the HDF5 layer moved to `store` so neither backend owns it, but
+# local.py, cli/annotate.py and existing importers still reach it from here.
+from protspace.data.embedding.store import (  # noqa: F401
+    finish_run,
+    load_existing_ids,
+    save_embeddings,
+    validate_headers,
+)
+
 logger = logging.getLogger(__name__)
 
 # Short aliases → CommonEmbedder enum member names.
@@ -114,27 +122,6 @@ def derive_h5_cache_path(fasta_path: Path, embedder: str) -> Path:
     return fasta_path.with_name(f"{fasta_path.stem}_{short}.h5")
 
 
-# ---------------------------------------------------------------------------
-# HDF5 helpers
-# ---------------------------------------------------------------------------
-
-
-def load_existing_ids(h5_path: Path) -> set[str]:
-    """Return the set of dataset keys already present in *h5_path*."""
-    if not h5_path.exists():
-        return set()
-    with h5py.File(h5_path, "r") as f:
-        return set(f.keys())
-
-
-def save_embeddings(h5_path: Path, embeddings: dict[str, np.ndarray]) -> None:
-    """Append embeddings to an HDF5 file (one dataset per protein)."""
-    with h5py.File(h5_path, "a") as f:
-        for protein_id, emb in embeddings.items():
-            if protein_id not in f:
-                f.create_dataset(protein_id, data=emb.astype(np.float32))
-
-
 # ---------------------------------------------------------------------------
 # Main orchestrator
 # ---------------------------------------------------------------------------
@@ -158,6 +145,9 @@ def embed_sequences(
     """
     cfg = embed_config or EmbedConfig()
 
+    # Reject HDF5-hostile identifiers before spending a single API call on them.
+    validate_headers(sequences)
+
     # Resume: skip already-embedded sequences
     existing_ids = load_existing_ids(h5_path)
     if existing_ids:
@@ -286,33 +276,12 @@ def embed_sequences(
 
     pbar.close()
 
-    # Gate on what landed in the file, not on a running total: save_embeddings skips
-    # IDs already present, and h5py turns an ID containing "/" into a group, so a
-    # counter can claim sequences the file does not hold.
-    missing = set(remaining) - load_existing_ids(h5_path)
-    if missing:
-        # ValueError, not RuntimeError: cli/embed.py and cli/prepare.py both catch
-        # (FileNotFoundError, ValueError) and render it as "ERROR: <msg>" + exit 1,
-        # whereas a RuntimeError escapes those handlers as a raw traceback.
-        embedded = len(remaining) - len(missing)
-        detail = (
-            f"{embedded:,} of {len(remaining):,} outstanding sequence(s) embedded, "
-            f"{len(missing):,} still missing "
-            f"({failed_batches} of {len(api_batches)} batch(es) failed)"
-        )
-        if embedded == 0:
-            raise ValueError(
-                f"No new embeddings were produced for {h5_path}: {detail}. "
-                f"Check the Biocentral server status and rerun."
-            )
-        raise ValueError(
-            f"Embedding incomplete for {h5_path}: {detail}. "
-            f"Partial results were kept — rerun to embed only what is missing."
-        )
-
-    print(f"\nDone. Embedded {len(remaining):,} sequence(s).")
-    print(f"Output: {h5_path}")
-    return h5_path
+    return finish_run(
+        h5_path,
+        remaining,
+        context=f"{failed_batches} of {len(api_batches)} batch(es) failed",
+        retry_hint="Check the Biocentral server status and rerun.",
+    )
 
 
 def probe_embedder(