chore: import upstream snapshot with attribution

2026-07-13 12:30:44 +08:00
commit bcbd1bdb22
5748 changed files with 562488 additions and 0 deletions
@@ -0,0 +1,535 @@
+# ========= 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
+import inspect
+from typing import Any, Callable
+
+from mirage.cache.context import push_cache_manager
+from mirage.cache.manager import CacheManager
+from mirage.commands.builtin.utils.safeguard import (apply_op_safeguard,
+                                                     run_with_timeout)
+from mirage.commands.config import RegisteredCommand
+from mirage.commands.resolve import get_extension
+from mirage.commands.safeguard import CommandSafeguard, resolve_safeguard
+from mirage.commands.spec import CommandSpec
+from mirage.io.types import ByteSource, IOResult
+from mirage.observe.context import (push_mount_prefix, push_revisions,
+                                    reset_revisions, with_mount_prefix,
+                                    with_revisions)
+from mirage.ops.registry import RegisteredOp
+from mirage.resource.base import BaseResource
+from mirage.types import ConsistencyPolicy, MountMode, PathSpec
+from mirage.utils.key_prefix import mount_key
+
+
+def _wrap_cmd_streams(
+    result: tuple[ByteSource | None, IOResult],
+    mount_prefix: str,
+    revisions: dict[str, str] | None,
+) -> tuple[ByteSource | None, IOResult]:
+    """Wrap any async-iterator streams in ``result`` with the mount
+    prefix and active revisions, so ``record_stream`` and
+    ``revision_for`` calls inside the lazy backend body see the right
+    context when consumed after this frame exits.
+
+    Mirrors the ``exit_on_empty`` pattern: thin async-gen wrapper that
+    side-effects the recorder state as bytes flow through. Same object
+    appearing in both the primary stream and IOResult.reads/writes is
+    wrapped once (dedup by identity).
+
+    Args:
+        result: ``(stream, io)`` as returned by a command handler.
+        mount_prefix: prefix to push during stream consumption.
+        revisions: revisions map to push during stream consumption
+            (None when the mount has no pins installed).
+    """
+    stream, io = result
+    seen: dict[int, ByteSource] = {}
+
+    def _wrap(obj: ByteSource | None) -> ByteSource | None:
+        if obj is None or isinstance(obj, (bytes, bytearray)):
+            return obj
+        oid = id(obj)
+        if oid in seen:
+            return seen[oid]
+        wrapped = with_mount_prefix(mount_prefix, obj)
+        if revisions:
+            wrapped = with_revisions(revisions, wrapped)
+        seen[oid] = wrapped
+        return wrapped
+
+    stream = _wrap(stream)
+    for k, v in list(io.reads.items()):
+        io.reads[k] = _wrap(v)
+    for k, v in list(io.writes.items()):
+        io.writes[k] = _wrap(v)
+    return stream, io
+
+
+class MountEntry:
+    """A mounted resource with command and op dispatch.
+
+    Each mount has its own lookup tables for commands and ops.
+    Different mounts of the same resource type can have
+    different registered commands/ops.
+
+    Resolution hierarchy (same for commands and ops):
+    1. (name, extension) -- filetype-specific
+    2. (name, None) -- resource-specific
+    3. general[name] -- general fallback
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        resource: BaseResource,
+        mode: MountMode = MountMode.READ,
+        consistency: ConsistencyPolicy = ConsistencyPolicy.LAZY,
+    ) -> None:
+        if not prefix.startswith("/"):
+            raise ValueError(f"prefix must start with /: {prefix!r}")
+        if not prefix.endswith("/"):
+            raise ValueError(f"prefix must end with /: {prefix!r}")
+        if "//" in prefix:
+            raise ValueError(f"prefix must not contain //: {prefix!r}")
+        self.prefix = prefix
+        self.resource = resource
+        self.mode = mode
+        self.consistency = consistency
+        self.cache_manager: CacheManager | None = None
+        # Per-path revision pins installed at Workspace.load time. Read
+        # functions consult these via the ``revision_for`` contextvar
+        # lookup; on a hit, the backend GET pins to the recorded
+        # revision so replay serves the exact bytes the agent saw.
+        # Empty during normal runs; populated only by the snapshot
+        # loader.
+        self.revisions: dict[str, str] = {}
+        self._cmds: dict[tuple, RegisteredCommand] = {}
+        self._general_cmds: dict[str, RegisteredCommand] = {}
+        self._cmd_specs: dict[str, CommandSpec] = {}
+        self.command_safeguards: dict[str, CommandSafeguard] = {}
+        self._ops: dict[tuple, RegisteredOp] = {}
+        self._general_ops: dict[str, RegisteredOp] = {}
+        # key: (cmd_name, target_resource_type)
+        self._cross_cmds: dict[tuple, RegisteredCommand] = {}
+
+    # ── command registration ──────────────────────────
+
+    def register(self, cmd: RegisteredCommand) -> None:
+        """Register a resource-specific command."""
+        key = (cmd.name, cmd.filetype)
+        self._cmds[key] = cmd
+        if cmd.spec is not None:
+            self._cmd_specs[cmd.name] = cmd.spec
+
+    def register_general(
+        self,
+        cmd: RegisteredCommand,
+    ) -> None:
+        """Register a general command (resource=None).
+
+        General commands work on any resource (e.g. echo, pwd).
+        They are the last fallback in resolve_command().
+        """
+        self._general_cmds[cmd.name] = cmd
+        if cmd.spec is not None:
+            self._cmd_specs[cmd.name] = cmd.spec
+
+    def resolve_command(
+        self,
+        cmd_name: str,
+        extension: str | None = None,
+    ) -> RegisteredCommand | None:
+        """Resolve command with fallback hierarchy.
+
+        Lookup order:
+        1. (cmd_name, extension) -- filetype-specific
+        2. (cmd_name, None) -- resource-specific
+        3. general_cmds[cmd_name] -- general fallback
+        """
+        if extension:
+            cmd = self._cmds.get((cmd_name, extension))
+            if cmd is not None:
+                return cmd
+        cmd = self._cmds.get((cmd_name, None))
+        if cmd is not None:
+            return cmd
+        return self._general_cmds.get(cmd_name)
+
+    def spec_for(
+        self,
+        cmd_name: str,
+    ) -> CommandSpec | None:
+        """Get the spec for a command name."""
+        return self._cmd_specs.get(cmd_name)
+
+    def is_general_command(self, cmd_name: str) -> bool:
+        """Whether `cmd_name` is registered as a general command here."""
+        return cmd_name in self._general_cmds
+
+    def all_commands(self) -> list[RegisteredCommand]:
+        """All registered commands (per-mount + general), deduped by name."""
+        seen: set[str] = set()
+        out: list[RegisteredCommand] = []
+        for rc in self._cmds.values():
+            if rc.name in seen:
+                continue
+            seen.add(rc.name)
+            out.append(rc)
+        for rc in self._general_cmds.values():
+            if rc.name in seen:
+                continue
+            seen.add(rc.name)
+            out.append(rc)
+        return out
+
+    def filetype_handlers(
+        self,
+        cmd_name: str,
+    ) -> dict[str, Callable]:
+        """Get filetype-specific command handlers.
+
+        Example::
+
+            mount.register(generic_cat)   # ("cat", None)
+            mount.register(parquet_cat)   # ("cat", ".parquet")
+
+            mount.filetype_handlers("cat")
+            # -> {".parquet": parquet_cat_fn}
+
+        Args:
+            cmd_name (str): command name, e.g. "cat".
+        """
+        fns: dict[str, Callable] = {}
+        for (name, ft), rc in self._cmds.items():
+            if name == cmd_name and ft is not None:
+                if ft not in fns:
+                    fns[ft] = rc.fn
+        return fns
+
+    def register_fns(self, fns: list) -> None:
+        """Register commands and ops from decorated functions.
+
+        Args:
+            fns (list): Functions decorated with @command and/or @op.
+
+        Raises:
+            ValueError: If a command/op's resource doesn't match
+                this mount's resource.
+        """
+        pname = self.resource.name
+        for fn in fns:
+            if hasattr(fn, "_registered_commands"):
+                rcs = fn._registered_commands
+                matching = [
+                    rc for rc in rcs
+                    if rc.resource is None or rc.resource == pname
+                ]
+                if rcs and not matching:
+                    resources = sorted({rc.resource for rc in rcs})
+                    raise ValueError(
+                        f"command {rcs[0].name!r} is for resource(s) "
+                        f"{resources!r}, not {pname!r}")
+                for rc in matching:
+                    self.register(rc)
+            if hasattr(fn, "_registered_ops"):
+                ros = fn._registered_ops
+                matching_ops = [
+                    ro for ro in ros
+                    if ro.resource is None or ro.resource == pname
+                ]
+                if ros and not matching_ops:
+                    resources = sorted({ro.resource for ro in ros})
+                    raise ValueError(f"op {ros[0].name!r} is for resource(s) "
+                                     f"{resources!r}, not {pname!r}")
+                for ro in matching_ops:
+                    self.register_op(ro)
+
+    def unregister(self, names: list[str]) -> None:
+        """Remove all commands and ops with the given names.
+
+        Args:
+            names (list[str]): Command/op names to remove.
+        """
+        for name in names:
+            keys = [k for k in self._cmds if k[0] == name]
+            for k in keys:
+                del self._cmds[k]
+            self._general_cmds.pop(name, None)
+            self._cmd_specs.pop(name, None)
+            op_keys = [k for k in self._ops if k[0] == name]
+            for k in op_keys:
+                del self._ops[k]
+            self._general_ops.pop(name, None)
+
+    def commands(self) -> dict[str, list[str | None]]:
+        """List registered commands grouped by filetype variants.
+
+        Returns:
+            dict[str, list[str | None]]: Command name to filetype list.
+        """
+        result: dict[str, list[str | None]] = {}
+        for (name, filetype) in self._cmds:
+            result.setdefault(name, []).append(filetype)
+        for name in self._general_cmds:
+            result.setdefault(name, [])
+        for name in result:
+            result[name] = sorted(result[name],
+                                  key=lambda x: (x is not None, x or ""))
+        return dict(sorted(result.items()))
+
+    def registered_ops(self) -> dict[str, list[str | None]]:
+        """List registered ops grouped by filetype variants.
+
+        Returns:
+            dict[str, list[str | None]]: Op name to filetype list.
+        """
+        result: dict[str, list[str | None]] = {}
+        for (name, filetype) in self._ops:
+            result.setdefault(name, []).append(filetype)
+        for name in self._general_ops:
+            result.setdefault(name, [])
+        for name in result:
+            result[name] = sorted(result[name],
+                                  key=lambda x: (x is not None, x or ""))
+        return dict(sorted(result.items()))
+
+    # ── cross-mount registration ─────────────────────
+
+    def register_cross(
+        self,
+        cmd: RegisteredCommand,
+        target_resource_type: str,
+    ) -> None:
+        """Register a cross-mount command for a target.
+
+        Example::
+
+            mount.register_cross(cp_cmd, "ram")
+            # This mount can now cp to ram mounts
+
+        Args:
+            cmd: the cross-mount command.
+            target_resource_type: e.g. "ram", "s3".
+        """
+        key = (cmd.name, target_resource_type)
+        self._cross_cmds[key] = cmd
+
+    def resolve_cross(
+        self,
+        cmd_name: str,
+        target_resource_type: str,
+    ) -> RegisteredCommand | None:
+        """Find a cross-mount command for a target."""
+        return self._cross_cmds.get((cmd_name, target_resource_type))
+
+    # ── op registration ───────────────────────────────
+
+    def register_op(self, op: RegisteredOp) -> None:
+        """Register a resource-specific VFS op."""
+        key = (op.name, op.filetype)
+        self._ops[key] = op
+
+    def _resolve_cascade(
+        self,
+        name: str,
+        extension: str | None,
+        table: dict,
+        general: dict,
+    ) -> list:
+        """Resolve with cascade: try filetype, resource, general.
+
+        Returns list of matching entries to try in order.
+        First non-None result wins.
+        """
+        levels = []
+        if extension:
+            entry = table.get((name, extension))
+            if entry is not None:
+                levels.append(entry)
+        entry = table.get((name, None))
+        if entry is not None:
+            levels.append(entry)
+        entry = general.get(name)
+        if entry is not None:
+            levels.append(entry)
+        return levels
+
+    # ── execution ─────────────────────────────────────
+
+    async def execute_cmd(
+        self,
+        cmd_name: str,
+        paths: list[PathSpec],
+        texts: list[str],
+        flag_kwargs: dict,
+        *,
+        stdin: ByteSource | None = None,
+        cwd: str = "/",
+        dispatch: Callable | None = None,
+        session_id: str | None = None,
+        env: dict[str, str] | None = None,
+        exec_allowed: bool = True,
+    ) -> tuple[ByteSource | None, IOResult]:
+        """Execute a command on this mount's resource.
+
+        Pure dispatch — flag parsing is done upstream in
+        executor/command.py. This method just resolves the
+        command handler and calls it.
+
+        Args:
+            cmd_name (str): command name.
+            paths (list[PathSpec]): positional path args.
+            texts (list[str]): positional text args.
+            flag_kwargs (dict): parsed flags from upstream.
+            stdin (ByteSource | None): stdin data.
+            cwd (str): virtual cwd from session.
+        """
+        extension = get_extension(paths[0].virtual) if paths else None
+
+        handlers = self._resolve_cascade(cmd_name, extension, self._cmds,
+                                         self._general_cmds)
+        if not handlers:
+            return None, IOResult(
+                exit_code=127,
+                stderr=(f"{cmd_name}: command not found".encode()))
+
+        mount_prefix = self.prefix.rstrip("/")
+        filetype_fns = self.filetype_handlers(cmd_name)
+        is_filetype_cmd = extension is not None and (cmd_name,
+                                                     extension) in self._cmds
+
+        paths = [
+            dataclasses.replace(
+                p, resource_path=mount_key(p.virtual, mount_prefix))
+            if isinstance(p, PathSpec) else p for p in paths
+        ]
+
+        # Stamp this mount's backend key onto path-shaped flag values so
+        # backend reads can address them: a single PathSpec (e.g. awk -f,
+        # single grep -f) or a list of PathSpec (repeatable grep -f).
+        # Everything else (bools, strings, list[str] like repeated -e) is
+        # not a path and passes through unchanged.
+        kw = {}
+        for k, v in flag_kwargs.items():
+            if isinstance(v, PathSpec):
+                kw[k] = dataclasses.replace(v,
+                                            resource_path=mount_key(
+                                                v.virtual, mount_prefix))
+            elif isinstance(v, list) and v and all(
+                    isinstance(item, PathSpec) for item in v):
+                kw[k] = [
+                    dataclasses.replace(item,
+                                        resource_path=mount_key(
+                                            item.virtual, mount_prefix))
+                    for item in v
+                ]
+            else:
+                kw[k] = v
+        kw["index"] = self.resource.index
+        kw["cwd"] = PathSpec(
+            virtual=cwd,
+            directory=cwd,
+            resolved=False,
+            resource_path=mount_key(cwd, mount_prefix),
+        )
+        kw["filetype_fns"] = (filetype_fns if not is_filetype_cmd else None)
+        if stdin is not None:
+            kw["stdin"] = stdin
+        if dispatch is not None:
+            kw["dispatch"] = dispatch
+        if session_id is not None:
+            kw["session_id"] = session_id
+        if env is not None:
+            kw["env"] = env
+        kw["exec_allowed"] = exec_allowed
+
+        prev_prefix = push_mount_prefix(mount_prefix)
+        revs_token = push_revisions(self.revisions or None)
+        prev_manager = push_cache_manager(self.cache_manager)
+        try:
+            for cmd in handlers:
+                if cmd.write and self.mode == MountMode.READ:
+                    return None, IOResult(
+                        exit_code=1,
+                        stderr=(f"{cmd_name}: read-only mount "
+                                f"at {self.prefix}".encode()))
+                result = await cmd.fn(self.resource.accessor, paths, *texts,
+                                      **kw)
+                if result is not None:
+                    stream, io = _wrap_cmd_streams(result, mount_prefix,
+                                                   self.revisions or None)
+                    # TODO: hand back a finalization context separately
+                    # instead of stamping policy onto io.safeguard.
+                    io.safeguard = resolve_safeguard(
+                        cmd_name, cmd.safeguard,
+                        self.command_safeguards.get(cmd_name))
+                    return stream, io
+            return None, IOResult()
+        finally:
+            reset_revisions(revs_token)
+            push_mount_prefix(prev_prefix)
+            push_cache_manager(prev_manager)
+
+    async def execute_op(
+        self,
+        op_name: str,
+        path: str,
+        *args,
+        **kwargs,
+    ) -> Any:
+        """Execute a VFS op on this mount's resource.
+
+        Tries filetype-specific first, then resource-specific.
+        First non-None result wins.
+
+        Args:
+            op_name (str): operation name (e.g. "read", "stat").
+            path (str): virtual path.
+        """
+        filetype = get_extension(path)
+        levels = self._resolve_cascade(op_name, filetype, self._ops,
+                                       self._general_ops)
+        if not levels:
+            raise AttributeError(f"{self.resource.name}: "
+                                 f"no op {op_name!r}")
+
+        if self.mode == MountMode.READ and any(o.write for o in levels):
+            raise PermissionError(f"mount {self.prefix!r} is read-only")
+
+        mount_prefix = self.prefix.rstrip("/")
+        scope = PathSpec(
+            virtual=path,
+            directory=path.rsplit("/", 1)[0] or "/",
+            resource_path=mount_key(path, mount_prefix),
+        )
+        kwargs.setdefault("index", self.resource.index)
+        op_override = self.command_safeguards.get(op_name)
+        op_timeout = (op_override.timeout_seconds
+                      if op_override is not None else None)
+        prev_prefix = push_mount_prefix(mount_prefix)
+        revs_token = push_revisions(self.revisions or None)
+        try:
+            for op in levels:
+                result = op.fn(self.resource.accessor, scope, *args, **kwargs)
+                if inspect.isawaitable(result):
+                    result = await run_with_timeout(result, op_timeout,
+                                                    op_name)
+                if result is not None:
+                    return await apply_op_safeguard(result, op_override)
+            return None
+        finally:
+            reset_revisions(revs_token)
+            push_mount_prefix(prev_prefix)