strukto-ai--mirage/python/mirage/workspace/provision/command.py

# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

import dataclasses

from mirage.cache.file.mixin import FileCacheMixin
from mirage.commands.builtin.generic.crossmount import is_cross_mount
from mirage.commands.resolve import get_extension
from mirage.commands.spec import parse_command, parse_to_kwargs
from mirage.provision import Precision, ProvisionResult, combine_sum
from mirage.types import PathSpec
from mirage.utils.key_prefix import mount_key
from mirage.workspace.mount import MountRegistry
from mirage.workspace.mount.namespace import Namespace
from mirage.workspace.session import Session


async def _check_cache_hits(
    cache: FileCacheMixin | None,
    parts: list[str | PathSpec],
) -> int:
    """Count how many path args are already cached."""
    if cache is None:
        return 0
    hits = 0
    for p in parts[1:]:
        if isinstance(p, PathSpec) and await cache.exists(p.virtual):
            hits += 1
    return hits


def _mount_groups(registry: MountRegistry,
                  parts: list[str | PathSpec]) -> list[list[PathSpec]]:
    """Group path args by their own mount, in first-appearance order.

    Args that resolve to no mount (glob patterns, expression operands)
    ride with the first group: they scoped to the primary mount before
    and must not fabricate a cross-mount split.
    """
    groups: list[list[PathSpec]] = []
    seen: dict[str, int] = {}
    unresolved: list[PathSpec] = []
    for p in parts[1:]:
        if not isinstance(p, PathSpec):
            continue
        if p.pattern:
            # Globs are not expanded during planning; a pattern operand
            # (find -name, ls *.txt) must not fabricate a mount group.
            unresolved.append(p)
            continue
        try:
            prefix = registry.mount_for(p.virtual).prefix
        except ValueError:
            unresolved.append(p)
            continue
        idx = seen.get(prefix)
        if idx is None:
            seen[prefix] = len(groups)
            groups.append([p])
        else:
            groups[idx].append(p)
    if unresolved:
        if groups:
            groups[0].extend(unresolved)
        else:
            groups.append(unresolved)
    return groups


async def handle_command_provision(
    registry: MountRegistry,
    parts: list[str | PathSpec],
    session: Session,
    namespace: Namespace | None = None,
) -> ProvisionResult:
    """Estimate cost of a simple command.

    Paths are namespace-followed first (a symlinked read costs its
    target, and the cache-hit check sees the entry the executor would
    actually serve), then grouped by mount: a command spanning mounts
    is estimated per mount against each mount's own backend and the
    results summed, instead of statting foreign paths against the
    first path's backend.
    """
    if not parts:
        return ProvisionResult(precision=Precision.EXACT)

    if namespace is not None:
        for i, p in enumerate(parts):
            if isinstance(p, PathSpec):
                followed = namespace.follow(p.virtual)
                if followed != p.virtual:
                    parts[i] = PathSpec.from_str_path(followed)

    cmd_name = str(parts[0])
    cmd_str = " ".join(p.virtual if isinstance(p, PathSpec) else p
                       for p in parts)

    groups = _mount_groups(registry, parts)
    if len(groups) > 1:
        path_parts = [p for p in parts[1:] if isinstance(p, PathSpec)]
        if not is_cross_mount(cmd_name, path_parts, registry):
            # The executor rejects this command across mounts, so an
            # aggregated byte estimate would cost a run that errors.
            return ProvisionResult(command=cmd_str,
                                   precision=Precision.UNKNOWN)
        texts = [p for p in parts[1:] if not isinstance(p, PathSpec)]
        children = []
        for group in groups:
            sub: list[str | PathSpec] = [cmd_name, *texts, *group]
            children.append(await
                            handle_command_provision(registry, sub, session))
        combined = combine_sum(";", children)
        combined.command = cmd_str
        return combined

    first_scope = None
    for p in parts[1:]:
        if isinstance(p, PathSpec):
            first_scope = p
            break
    mount_path = first_scope.virtual if first_scope else session.cwd

    try:
        mount = registry.mount_for(mount_path)
    except ValueError:
        # Pathless commands (seq, date, ...) still need a mount to
        # resolve their registration; any mount carries the general
        # commands, so fall back to the first one.
        mounts = registry.mounts()
        if not mounts:
            return ProvisionResult(command=cmd_str,
                                   precision=Precision.UNKNOWN)
        mount = mounts[0]

    extension = get_extension(first_scope.virtual) if first_scope else None
    cmd = mount.resolve_command(cmd_name, extension)
    if cmd is None or cmd.provision_fn is None:
        return ProvisionResult(command=cmd_str, precision=Precision.UNKNOWN)

    mount_prefix = mount.prefix.rstrip("/")
    resource_scopes = []
    for i, p in enumerate(parts[1:], start=1):
        if isinstance(p, PathSpec):
            scoped = dataclasses.replace(p,
                                         resource_path=mount_key(
                                             p.virtual, mount_prefix))
            parts[i] = scoped
            resource_scopes.append(scoped)

    # Parse flags so plan functions receive them as kwargs (e.g. r=True)
    argv = [p.virtual if isinstance(p, PathSpec) else p for p in parts[1:]]
    spec = mount.spec_for(cmd_name)
    if spec is not None:
        parsed = parse_command(spec, argv, cwd=session.cwd)
        flag_kwargs = parse_to_kwargs(parsed)
        text_args = parsed.texts()
    else:
        flag_kwargs = {}
        text_args = [p for p in parts[1:] if not isinstance(p, PathSpec)]

    result = await cmd.provision_fn(mount.resource.accessor,
                                    resource_scopes,
                                    *text_args,
                                    command=cmd_str,
                                    prefix=mount.prefix.rstrip("/"),
                                    index=mount.resource.index,
                                    **flag_kwargs)
    if not result.command:
        result.command = cmd_str

    hits = await _check_cache_hits(registry.file_cache, parts)
    if hits > 0:
        result.cache_hits = hits
        result.cache_read_low = result.network_read_low
        result.cache_read_high = result.network_read_high
        result.network_read_low = 0
        result.network_read_high = 0

    return result