# ========= 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. ========= from dataclasses import dataclass, field from functools import partial from typing import Any, Callable from mirage.provision import Precision, ProvisionResult from mirage.shell.node_kind import NodeKind, node_kind from mirage.shell.types import NodeType as NT from mirage.shell.types import RedirectKind from mirage.shell.types import ShellBuiltin as SB from mirage.types import PathSpec from mirage.workspace.expand import (classify_parts, expand_and_classify, expand_parts, expand_redirects) from mirage.workspace.mount import MountRegistry from mirage.workspace.mount.namespace import Namespace from mirage.workspace.provision.builtins import handle_builtin_provision from mirage.workspace.provision.command import handle_command_provision from mirage.workspace.provision.control import (handle_for_provision, handle_function_provision, handle_if_provision, handle_while_provision) from mirage.workspace.provision.pipes import (handle_connection_provision, handle_pipe_provision) from mirage.workspace.provision.redirect import handle_redirect_provision from mirage.workspace.provision.rollup import rollup_list, rollup_pipe from mirage.workspace.session import Session from mirage.shell.helpers import ( # isort: skip get_case_items, get_command_name, get_for_parts, get_function_body, get_function_name, get_if_branches, get_list_parts, get_negated_command, get_parts, get_pipeline_commands, get_redirects, get_subshell_body, get_text, get_while_parts, has_command_substitution, split_env_prefix) # eval / source execute their payload, so they are NOT free builtins: # leaving them out lets them fall through to command resolution, which # honestly reports UNKNOWN instead of a zero-cost EXACT. _BUILTIN_NAMES = frozenset({ SB.CD, SB.TRUE, SB.FALSE, SB.EXPORT, SB.UNSET, SB.LOCAL, SB.PRINTENV, SB.READ, SB.SET, SB.SHIFT, SB.TRAP, SB.TEST, SB.BRACKET, SB.DOUBLE_BRACKET, SB.WAIT, SB.FG, SB.KILL, SB.JOBS, SB.PS, SB.ECHO, SB.PRINTF, SB.SLEEP, "return", "break", "continue", }) @dataclass class PlanScope: """Walk-local planner state. Function definitions seen during this plan are recorded here (not on the session: planning must not mutate shell state), and `planning` guards recursive functions from looping the planner. """ functions: dict[str, list] = field(default_factory=dict) planning: set[str] = field(default_factory=set) async def _provision_redirected( recurse: Callable, registry: MountRegistry, namespace: Namespace | None, execute_fn: Callable, command: Any, redirects: list, session: Session, ) -> ProvisionResult: """Plan one redirected command: expand targets, cost, degrade. Args: recurse (Callable): the provision recursion. registry (MountRegistry): mount registry. namespace (Namespace | None): addressing authority. execute_fn (Callable): recursive execute (for expansions). command (Any): the redirected command node. redirects (list): parsed redirects. session (Session): shell session state. """ expanded, pipe_node = await expand_redirects(redirects, session, execute_fn, registry) # A cmdsub target expands empty under provision, so its # classification is garbage; the precision degrade below keeps # the plan honest without costing a phantom write (mirrors TS). targets = [ (r.kind, r.target) for r in expanded if r.kind in (RedirectKind.STDIN, RedirectKind.STDOUT) and isinstance( r.target, PathSpec) and not r.target.virtual.startswith("/dev/") and not (r.target_node is not None and has_command_substitution(r.target_node)) ] result = await handle_redirect_provision(recurse, registry, command, targets, session, namespace) if any(r.target_node is not None and has_command_substitution(r.target_node) for r in redirects): # A suppressed substitution hid the real redirect target. result.precision = Precision.UNKNOWN if pipe_node is not None: return rollup_pipe([result, await recurse(pipe_node, session)]) return result async def _provision_reassociated( recurse: Callable, registry: MountRegistry, namespace: Namespace | None, execute_fn: Callable, redirects: list, right: Any, node: Any, session: Session, ) -> ProvisionResult: """Provision recurse wrapper for a re-associated trailing redirect. Mirrors the executor: the list's last command carries the hoisted redirects; every other node provisions normally. Args: recurse (Callable): the provision recursion. registry (MountRegistry): mount registry. namespace (Namespace | None): addressing authority. execute_fn (Callable): recursive execute (for expansions). redirects (list): parsed redirects hoisted off the list. right (Any): the list's last command node. node (Any): node being provisioned by the connection handler. session (Session): shell session state. """ if node is not right: return await recurse(node, session) return await _provision_redirected(recurse, registry, namespace, execute_fn, right, redirects, session) async def provision_node( registry: MountRegistry, dispatch: Callable, execute_fn: Callable, namespace: Namespace | None, node: Any, session: Session, scope: PlanScope | None = None, ) -> ProvisionResult: """Walk tree-sitter AST and estimate execution cost. Dispatches on the same NodeKind classification as the executor (`mirage.shell.classify`), so every construct the executor runs has a planner branch; kinds neither walker supports fall through to an honest UNKNOWN. Args: registry (MountRegistry): mount registry for path resolution. dispatch (Callable): VFS op dispatcher (op, path, **kw). execute_fn (Callable): recursive execute (for expansions). node (Any): tree-sitter node to plan. session (Session): shell session state. scope (PlanScope | None): walk-local planner state; created at the root and threaded through recursion. """ plan_scope = scope if scope is not None else PlanScope() recurse = partial(provision_node, registry, dispatch, execute_fn, namespace, scope=plan_scope) kind = node_kind(node) if kind == NodeKind.COMMENT: return ProvisionResult(precision=Precision.EXACT) if kind in (NodeKind.PROGRAM, NodeKind.SUBSHELL, NodeKind.COMPOUND): if kind == NodeKind.SUBSHELL: body = get_subshell_body(node) else: body = [c for c in node.named_children if c.type != NT.COMMENT] children = [] for child in body: children.append(await recurse(child, session)) if not children: return ProvisionResult(precision=Precision.EXACT) return rollup_list(";", children) if kind == NodeKind.COMMAND: name = get_command_name(node) func_body = plan_scope.functions.get(name) if func_body is None: func_body = session.functions.get(name) if func_body is not None: return await handle_function_provision(recurse, name, func_body, plan_scope.planning, session) if name in _BUILTIN_NAMES: return await handle_builtin_provision() _, parts = split_env_prefix(get_parts(node)) if not parts: return ProvisionResult(precision=Precision.EXACT) expanded = await expand_parts(parts, session, execute_fn) classified = classify_parts(expanded, registry, session.cwd) result = await handle_command_provision(registry, classified, session, namespace) if any(has_command_substitution(p) for p in parts): # The plan walk suppressed the substitution, so the operand # list is incomplete: the totals are floors, not answers. result.precision = Precision.UNKNOWN return result if kind == NodeKind.PIPELINE: commands, _ = get_pipeline_commands(node) return await handle_pipe_provision(recurse, commands, session) if kind == NodeKind.LIST: left, op, right = get_list_parts(node) return await handle_connection_provision(recurse, left, op, right, session) if kind == NodeKind.REDIRECT: command, redirects = get_redirects(node) if command.type == NT.LIST: # Mirror the executor: a trailing redirect hoisted over an # &&/|| list binds to the last command. left, op, right = get_list_parts(command) wrapped = partial(_provision_reassociated, recurse, registry, namespace, execute_fn, redirects, right) return await handle_connection_provision(wrapped, left, op, right, session) return await _provision_redirected(recurse, registry, namespace, execute_fn, command, redirects, session) if kind == NodeKind.IF: branches, else_body = get_if_branches(node) return await handle_if_provision(recurse, branches, else_body, session) if kind == NodeKind.FOR: _, values, body = get_for_parts(node) if any(has_command_substitution(v) for v in values): # The iteration count comes from a suppressed substitution: # plan one pass as a floor and degrade. result = await handle_for_provision(recurse, body, 1, session) result.precision = Precision.UNKNOWN return result classified = await expand_and_classify(values, session, execute_fn, registry, session.cwd) n = len(classified) or 1 return await handle_for_provision(recurse, body, n, session) if kind == NodeKind.SELECT: # select re-prompts until break: unbounded like while. _, _, body = get_for_parts(node) return await handle_while_provision(recurse, body, session) if kind in (NodeKind.WHILE, NodeKind.UNTIL): _, body = get_while_parts(node) return await handle_while_provision(recurse, body, session) if kind == NodeKind.CASE: items = get_case_items(node) children = [] for _, body in items: if not body: continue stmts = [await recurse(stmt, session) for stmt in body] children.append(stmts[0] if len(stmts) == 1 else rollup_list(";", stmts)) if children: return rollup_list("||", children) return ProvisionResult(precision=Precision.EXACT) if kind == NodeKind.FUNCTION_DEF: name = get_function_name(node) body = get_function_body(node) if name and body is not None: plan_scope.functions[name] = body return await handle_builtin_provision() if kind in (NodeKind.DECLARATION, NodeKind.UNSET, NodeKind.TEST, NodeKind.VAR_ASSIGN): return await handle_builtin_provision() if kind == NodeKind.NEGATED: inner = get_negated_command(node) return await recurse(inner, session) return ProvisionResult(command=get_text(node), precision=Precision.UNKNOWN)