API Reference¶

def main():
    """Deploy the Marimo notebook app to Flyte."""
    flyte.init_from_config(root_dir=Path(__file__).parent)
    app = flyte.serve(note_env)
    print(f"App URL: {app.url}")

    def _shutdown(signum, frame):
        """Handle SIGINT/SIGTERM by killing the entire process group."""
        pid = app._process.pid
        try:
            # Kill the whole process group (marimo + its children)
            os.killpg(os.getpgid(pid), signal.SIGTERM)
        except (ProcessLookupError, PermissionError):
            pass
        app.deactivate(wait=True)
        sys.exit(0)

    signal.signal(signal.SIGINT, _shutdown)
    signal.signal(signal.SIGTERM, _shutdown)
    app._process.wait()

def main():
    """Build and push images for every Flyte task environment."""
    flyte.init_from_config(root_dir=Path(PROJECT_ROOT))

    for env in (scrna_env, gatk_env):
        logger.info(f"Building image for env: {env.name}")
        cache = flyte.build_images(env)
        logger.info(f"{env.name}: {cache!r}")

def log_execution() -> str:
    """Start a per-execution log sink and return the execution ID.

    Derives the workflow name from the calling function, fetches the current
    git commit hash, and creates a dedicated logfile for this execution.
    Warns if the git tree has uncommitted changes.
    """
    workflow = inspect.currentframe().f_back.f_code.co_name
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")

    result = subprocess.run(
        ["git", "rev-parse", "--short", "HEAD"],
        capture_output=True,
        text=True,
    )
    commit = result.stdout.strip() or "unknown"

    status = subprocess.run(
        ["git", "status", "--porcelain"],
        capture_output=True,
        text=True,
    )
    if status.stdout.strip():
        commit += "-dirty"
        logger.warning("Git tree is dirty — uncommitted changes present")

    execution_id = f"{workflow}-{commit}-{timestamp}"
    logger.add(_log_dir / f"{execution_id}.log")
    logger.info(f"Execution started: {execution_id}")
    return execution_id

def marshal_output(value: Any) -> Any:
    """Convert a typed Python object to a JSON-friendly structure for MCP transport."""
    if value is None:
        return None

    if hasattr(value, "to_dict"):
        return value.to_dict()

    if isinstance(value, Path):
        return str(value)

    if isinstance(value, tuple):
        return {f"o{i}": marshal_output(item) for i, item in enumerate(value)}

    if isinstance(value, list):
        return [marshal_output(item) for item in value]

    if isinstance(value, dict):
        return {k: marshal_output(v) for k, v in value.items()}

    return value

@mcp.tool()
async def delete_file(cid: str) -> str:
    """Delete a file by CID."""
    comp = Asset(cid=cid)
    await default_client.delete(comp)
    return f"Deleted file {cid}"

@mcp.tool()
async def download_file(cid: str) -> str:
    """Download a file by CID to local cache. Returns the local path."""
    comp = Asset(cid=cid)
    await default_client.download(comp)
    return str(comp.path)

@mcp.tool()
async def fetch_resource_bundle(bundle_name: str) -> list[dict]:
    """Download a predefined resource bundle into local storage.

    Bundles are curated sets of files (e.g. reference genomes, demo datasets)
    defined in the codebase. Each file is identified by CID and downloaded
    via the standard storage path (signed URL with JWT, or public IPFS gateway).

    When PINATA_JWT is set, remote metadata is authoritative and overwrites
    local records. Without a JWT, the bundle manifest provides the metadata.

    Args:
        bundle_name: Name of the bundle (from list_bundles).

    Returns:
        List of fetched files with cid, keyvalues, and local path.
    """
    from stargazer.bundles import fetch_bundle

    return await fetch_bundle(bundle_name)

@mcp.tool()
def list_bundles() -> list[dict]:
    """List available resource bundles.

    Returns:
        List of bundles with name, description, and file_count.
    """
    from stargazer.bundles import list_bundles as _list_bundles

    return _list_bundles()

@mcp.tool()
def list_tasks(category: str | None = None) -> list[dict]:
    """List available tasks and workflows with their parameter signatures.

    Args:
        category: Filter by "task" or "workflow". Omit for all.

    Returns:
        Catalog of tasks with name, category, description, params, and outputs.
    """
    return _registry.to_catalog(category=category)

def main():
    """Run the Stargazer MCP server."""
    import sys

    transport = "stdio"
    if "--http" in sys.argv:
        transport = "streamable-http"
    mcp.run(transport=transport)

@mcp.tool()
async def query_files(keyvalues: dict[str, str]) -> list[dict]:
    """Query files by metadata key-value pairs. Returns matching files."""
    return await default_client.query(keyvalues)

@mcp.tool()
async def run_task(task_name: str, filters: dict, inputs: dict | None = None) -> dict:
    """Run a single task by name for ad-hoc experimentation.

    Use this for testing individual tools in isolation. Asset parameters
    are assembled from storage using the provided filters — one call to
    assemble() resolves all required assets. Scalar and Path parameters
    are passed separately via inputs.

    For reproducible pipeline runs, use run_workflow instead.

    Args:
        task_name: Name of the task (from list_tasks with category="task").
        filters: Keyvalue filters for assemble() to resolve asset parameters
                 (e.g. {"build": "GRCh38", "sample_id": "NA12878"}).
        inputs: Optional scalar/Path keyword arguments (str, int, bool, list[str]).

    Returns:
        Serialized task output. Single outputs returned directly,
        multi-outputs as {"o0": ..., "o1": ...}.
    """
    info = _registry.get(task_name)
    if info is None:
        available = [t.name for t in _registry.list_tasks(category="task")]
        raise ValueError(f"Unknown task: {task_name!r}. Available: {available}")
    if info.category != "task":
        raise ValueError(f"{task_name!r} is a workflow — use run_workflow instead.")

    inputs = inputs or {}

    # Assemble all assets from storage in one query
    assets = await assemble(**filters) if filters else []

    # Build kwargs: match Asset params from the assembled list, scalars from inputs
    kwargs = {}
    for p in info.params:
        asset_key = _asset_key_for_hint(p.type_hint)
        if asset_key:
            matched = [a for a in assets if a._asset_key == asset_key]
            if not matched and p.required:
                raise ValueError(
                    f"Task {task_name!r} requires {p.name} ({asset_key}) "
                    f"but no matching asset found for filters: {filters}"
                )
            if matched:
                kwargs[p.name] = (
                    matched if _is_list_asset_hint(p.type_hint) else matched[-1]
                )
        elif p.name in inputs:
            value = inputs[p.name]
            if p.type_hint is Path and isinstance(value, str):
                value = Path(value)
            kwargs[p.name] = value

    return await _execute(info, kwargs)

@mcp.tool()
async def run_workflow(workflow_name: str, inputs: dict) -> dict:
    """Run a workflow by name for reproducible pipeline execution.

    Workflows accept scalar parameters (str, int, bool, list[str]) and
    handle their own asset assembly internally. Pass inputs exactly as
    the workflow signature defines them — no automatic resolution is
    performed.

    For ad-hoc experimentation with individual tools, use run_task instead.

    Args:
        workflow_name: Name of the workflow (from list_tasks with category="workflow").
        inputs: Keyword arguments as a JSON dict (scalars only).

    Returns:
        Serialized workflow output. Single outputs returned directly,
        multi-outputs as {"o0": ..., "o1": ...}.
    """
    info = _registry.get(workflow_name)
    if info is None:
        available = [t.name for t in _registry.list_tasks(category="workflow")]
        raise ValueError(f"Unknown workflow: {workflow_name!r}. Available: {available}")
    if info.category != "workflow":
        raise ValueError(f"{workflow_name!r} is a task — use run_task instead.")

    return await _execute(info, dict(inputs))

@mcp.resource("stargazer://config")
async def show_config() -> str:
    """Show current Stargazer configuration and available task counts."""
    tasks = _registry.list_tasks(category="task")
    workflows = _registry.list_tasks(category="workflow")
    config = {
        "pinata_jwt": "set" if os.environ.get("PINATA_JWT") else "unset",
        "pinata_visibility": os.environ["PINATA_VISIBILITY"],
        "local_dir": str(default_client.local_dir),
        "tasks": len(tasks),
        "workflows": len(workflows),
    }
    return json.dumps(config, indent=2)

@mcp.tool()
async def upload_file(path: str, keyvalues: dict[str, str]) -> dict:
    """Upload a file with metadata key-value pairs.

    keyvalues must include "asset". Valid asset keys are derived
    from the Asset registry (e.g. asset=reference component=fasta).

    When displaying results, always show a table with the CID and all keyvalues.
    """
    asset_key = keyvalues.get("asset")
    if asset_key not in ASSET_REGISTRY:
        valid = sorted(ASSET_REGISTRY.keys())
        raise ValueError(f"Invalid asset key {asset_key!r}. Valid keys: {valid}")
    cls = ASSET_REGISTRY[asset_key]
    import dataclasses
    from stargazer.assets.asset import _BASE_FIELDS

    declared = {f.name for f in dataclasses.fields(cls)} - _BASE_FIELDS
    unknown = set(keyvalues) - declared - {"asset"}
    if unknown:
        raise ValueError(
            f"Unknown keys for {asset_key!r}: {unknown}. Allowed: {sorted(declared)}"
        )
    comp = cls.from_keyvalues(keyvalues, path=Path(path))
    await default_client.upload(comp)
    return comp.to_dict()

@gatk_env.task
async def apply_bqsr(
    alignment: Alignment,
    ref: Reference,
    bqsr_report: BQSRReport,
) -> Alignment:
    """
    Apply Base Quality Score Recalibration to a BAM file.

    Args:
        alignment: Input BAM asset
        ref: Reference FASTA asset
        bqsr_report: Recalibration table from base_recalibrator

    Returns:
        Alignment asset with recalibrated BAM file

    Reference:
        https://gatk.broadinstitute.org/hc/en-us/articles/360037055712-ApplyBQSR
    """
    logger.info(alignment.to_dict())
    logger.info(ref.to_dict())
    logger.info(bqsr_report.to_dict())
    # fetch() auto-downloads companions (.fai, .dict, .bai, etc.)
    await alignment.fetch()
    await ref.fetch()
    await bqsr_report.fetch()

    ref_path = ref.path
    bam_path = alignment.path
    recal_path = bqsr_report.path

    if not recal_path or not recal_path.exists():
        raise FileNotFoundError("BQSR recalibration report not found in cache.")

    output_dir = _storage.default_client.local_dir
    output_bam = output_dir / f"{alignment.sample_id}_recalibrated.bam"

    cmd = [
        "gatk",
        "ApplyBQSR",
        "-R",
        str(ref_path),
        "-I",
        str(bam_path),
        "--bqsr-recal-file",
        str(recal_path),
        "-O",
        str(output_bam),
    ]

    await _run(cmd, cwd=str(output_dir))

    if not output_bam.exists():
        raise FileNotFoundError(f"ApplyBQSR did not create output BAM at {output_bam}")

    recal_bam = Alignment()
    await recal_bam.update(
        output_bam,
        sample_id=alignment.sample_id,
        format="bam",
        sorted="coordinate",
        duplicates_marked=alignment.duplicates_marked,
        bqsr_applied=True,
        tool="gatk_apply_bqsr",
        reference_cid=ref.cid,
    )

    logger.info(recal_bam.to_dict())
    return recal_bam

@gatk_env.task
async def apply_vqsr(
    vcf: Variants,
    ref: Reference,
    vqsr_model: VQSRModel,
    truth_sensitivity_filter_level: float | None = None,
) -> Variants:
    """
    Apply VQSR recalibration to a VCF using GATK ApplyVQSR.

    The recalibration mode (SNP or INDEL) is read from vqsr_model.keyvalues["mode"].
    If truth_sensitivity_filter_level is not provided, defaults to 99.5 for SNP
    and 99.0 for INDEL.

    Args:
        vcf: Raw (or SNP-filtered) VCF Variants asset
        ref: Reference FASTA asset
        vqsr_model: Recalibration model from variant_recalibrator
        truth_sensitivity_filter_level: VQSLOD filter threshold (optional)

    Returns:
        Variants asset with VQSR-filtered VCF

    Reference:
        https://gatk.broadinstitute.org/hc/en-us/articles/360035531612-Variant-Quality-Score-Recalibration-VQSR
    """
    logger.info(vcf.to_dict())
    logger.info(ref.to_dict())
    logger.info(vqsr_model.to_dict())
    mode = vqsr_model.mode or "SNP"
    if mode not in ("SNP", "INDEL"):
        raise ValueError(f"VQSRModel mode must be 'SNP' or 'INDEL', got {mode!r}")

    if not vqsr_model.tranches_path:
        raise ValueError("VQSRModel is missing tranches_path")
    tranches_path = Path(vqsr_model.tranches_path)

    filter_level = truth_sensitivity_filter_level or _DEFAULT_FILTER_LEVEL[mode]

    await vcf.fetch()
    await ref.fetch()
    await vqsr_model.fetch()

    output_dir = _storage.default_client.local_dir
    sample_id = vcf.sample_id or "cohort"
    output_vcf = output_dir / f"{sample_id}_vqsr_{mode.lower()}.vcf"

    cmd = [
        "gatk",
        "ApplyVQSR",
        "-R",
        str(ref.path),
        "-V",
        str(vcf.path),
        "--recal-file",
        str(vqsr_model.path),
        "--tranches-file",
        str(tranches_path),
        "--truth-sensitivity-filter-level",
        str(filter_level),
        "--create-output-variant-index",
        "true",
        "-mode",
        mode,
        "-O",
        str(output_vcf),
    ]

    _, stderr = await _run(cmd, cwd=str(output_dir))

    if not output_vcf.exists():
        raise FileNotFoundError(
            f"ApplyVQSR did not create output VCF at {output_vcf}. stderr: {stderr}"
        )

    source_samples = vcf.source_samples or [sample_id]
    filtered_vcf = Variants()
    await filtered_vcf.update(
        output_vcf,
        sample_id=sample_id,
        caller="apply_vqsr",
        variant_type="vcf",
        vqsr_mode=mode,
        build=ref.build,
        sample_count=len(source_samples),
        source_samples=source_samples,
    )

    idx_path = output_dir / f"{output_vcf.name}.idx"
    if idx_path.exists():
        vidx = VariantsIndex()
        await vidx.update(idx_path, sample_id=sample_id, variants_cid=filtered_vcf.cid)

    logger.info(filtered_vcf.to_dict())
    return filtered_vcf

@gatk_env.task
async def base_recalibrator(
    alignment: Alignment,
    ref: Reference,
    known_sites: list[KnownSites],
) -> BQSRReport:
    """
    Generate a Base Quality Score Recalibration report.

    Uses GATK BaseRecalibrator to analyze patterns of covariation in the
    sequence dataset and produce a recalibration table.

    Args:
        alignment: Input BAM asset (should be sorted and have duplicates marked)
        ref: Reference FASTA asset
        known_sites: List of KnownSites VCF assets (dbSNP, known indels, etc.)

    Returns:
        BQSRReport asset containing the recalibration table

    Reference:
        https://gatk.broadinstitute.org/hc/en-us/articles/360036898312-BaseRecalibrator
    """
    logger.info(alignment.to_dict())
    logger.info(ref.to_dict())
    logger.info([x.to_dict() for x in known_sites])
    if not known_sites:
        raise ValueError("known_sites list cannot be empty for BQSR")

    # fetch() auto-downloads companions (.fai, .dict, .bai, etc.)
    await alignment.fetch()
    await ref.fetch()
    for site in known_sites:
        await site.fetch()

    ref_path = ref.path
    bam_path = alignment.path
    output_dir = _storage.default_client.local_dir
    output_recal = output_dir / f"{alignment.sample_id}_bqsr.table"

    cmd = [
        "gatk",
        "BaseRecalibrator",
        "-R",
        str(ref_path),
        "-I",
        str(bam_path),
        "-O",
        str(output_recal),
    ]
    for site in known_sites:
        cmd.extend(["--known-sites", str(site.path)])

    await _run(cmd, cwd=str(output_dir))

    if not output_recal.exists():
        raise FileNotFoundError(
            f"BaseRecalibrator did not create recalibration report at {output_recal}"
        )

    report = BQSRReport()
    await report.update(
        output_recal,
        sample_id=alignment.sample_id,
        tool="gatk_base_recalibrator",
        alignment_cid=alignment.cid,
    )

    logger.info(report.to_dict())
    return report

@gatk_env.task
async def combine_gvcfs(
    gvcfs: list[Variants],
    ref: Reference,
    cohort_id: str = "cohort",
) -> Variants:
    """
    Combine multiple per-sample GVCFs into a single multi-sample GVCF.

    Args:
        gvcfs: List of Variants assets, each containing a GVCF from a single sample
        ref: Reference FASTA asset
        cohort_id: Identifier for the combined cohort (default: "cohort")

    Returns:
        Variants asset with combined multi-sample GVCF

    Reference:
        https://gatk.broadinstitute.org/hc/en-us/articles/360037053272-CombineGVCFs
    """
    logger.info([x.to_dict() for x in gvcfs])
    logger.info(ref.to_dict())
    if not gvcfs:
        raise ValueError("gvcfs list cannot be empty")

    for i, gvcf in enumerate(gvcfs):
        if gvcf.variant_type != "gvcf":
            raise ValueError(
                f"combine_gvcfs requires GVCF files, but gvcfs[{i}] is not a GVCF. "
                f"sample_id={gvcf.sample_id}"
            )

    # fetch() auto-downloads companions (.fai, .dict for ref; .idx for each gvcf)
    await ref.fetch()

    gvcf_paths: list[Path] = []
    sample_ids: list[str] = []

    for gvcf in gvcfs:
        await gvcf.fetch()
        gvcf_paths.append(gvcf.path)
        sample_ids.append(gvcf.sample_id)

    ref_path = ref.path
    output_dir = _storage.default_client.local_dir
    output_gvcf = output_dir / f"{cohort_id}_combined.g.vcf"

    cmd = [
        "gatk",
        "CombineGVCFs",
        "-R",
        str(ref_path),
        "-O",
        str(output_gvcf),
    ]

    for gvcf_path in gvcf_paths:
        cmd.extend(["-V", str(gvcf_path)])

    stdout, stderr = await _run(cmd, cwd=str(output_dir))

    if not output_gvcf.exists():
        raise FileNotFoundError(
            f"CombineGVCFs did not create output GVCF at {output_gvcf}. "
            f"stderr: {stderr}"
        )

    combined_vcf = Variants()
    await combined_vcf.update(
        output_gvcf,
        sample_id=cohort_id,
        caller="combine_gvcfs",
        variant_type="gvcf",
        build=ref.build,
        sample_count=len(sample_ids),
        source_samples=sample_ids,
    )

    # Upload implicit index file produced by GATK
    idx_path = output_dir / f"{output_gvcf.name}.idx"
    if idx_path.exists():
        vidx = VariantsIndex()
        await vidx.update(idx_path, sample_id=cohort_id, variants_cid=combined_vcf.cid)

    logger.info(combined_vcf.to_dict())
    return combined_vcf

@gatk_env.task
async def create_sequence_dictionary(ref: Reference) -> SequenceDict:
    """
    Create a sequence dictionary (.dict file) using GATK CreateSequenceDictionary.

    Args:
        ref: Reference FASTA asset

    Returns:
        SequenceDict asset containing the .dict file
    """
    logger.info(ref.to_dict())
    await ref.fetch()
    ref_path = ref.path

    if not ref_path or not ref_path.exists():
        raise FileNotFoundError(f"Reference file not found at {ref_path}")

    output_dir = _storage.default_client.local_dir
    dict_path = output_dir / f"{ref_path.stem}.dict"
    dict_path.unlink(missing_ok=True)
    cmd = [
        "gatk",
        "CreateSequenceDictionary",
        "-R",
        str(ref_path),
        "-O",
        str(dict_path),
    ]
    await _run(cmd, cwd=str(output_dir))

    if not dict_path.exists():
        raise FileNotFoundError(
            f"Sequence dictionary file {dict_path.name} was not created"
        )

    seq_dict = SequenceDict()
    await seq_dict.update(
        dict_path,
        build=ref.build,
        tool="gatk_CreateSequenceDictionary",
        reference_cid=ref.cid,
    )

    logger.info(seq_dict.to_dict())
    return seq_dict