Flows

flows

FlowsHook — execution flow tracing with entry point detection.

Protocol hook that maps to the trace_flow and find_entry_points functionalities. Registered as ast:flows via axm.hooks entry point.

FlowsHook dataclass

Trace execution flows and detect entry points.

Reads working_dir from context, and entry, detail, max_depth, cross_module from params.

When entry contains multiple symbols (newline-separated), each is traced independently using a shared pre-computed callee index.

If entry is not provided, discovers entry points and traces them (excluding __all__ exports, capped at 20).

Source code in packages/axm-ast/src/axm_ast/hooks/flows.py

@dataclass
class FlowsHook:
    """Trace execution flows and detect entry points.

    Reads ``working_dir`` from *context*, and ``entry``, ``detail``,
    ``max_depth``, ``cross_module`` from *params*.

    When ``entry`` contains multiple symbols (newline-separated), each
    is traced independently using a shared pre-computed callee index.

    If ``entry`` is not provided, discovers entry points and traces
    them (excluding ``__all__`` exports, capped at 20).
    """

    # ``context``/``params`` are heterogeneous user payloads; ``object``
    # forces explicit narrowing at every read site (mirrors hooks/context.py).
    def execute(self, context: dict[str, object], **params: object) -> HookResult:
        """Execute the hook action.

        Args:
            context: Session context dictionary (must contain ``working_dir``).
            **params:
                Optional ``entry`` (symbol name, or newline-separated list).
                Optional ``detail`` ("source", "trace", or "compact").
                Optional ``max_depth`` (default 5).
                Optional ``cross_module`` (default False).

        Returns:
            HookResult with ``traces`` in metadata on success.
            When *detail* is ``"compact"``, ``traces`` is a single string
            (concatenated with entry-name headers for multi-symbol traces).
            Otherwise, ``traces`` is a dict of symbol → step-dicts.
        """
        raw_path = params.get("path") or context.get("working_dir", ".")
        if not isinstance(raw_path, (str, Path)):
            return HookResult.fail(
                f"path must be str or Path, got {type(raw_path).__name__}"
            )
        working_dir = Path(raw_path).resolve()

        if not working_dir.is_dir():
            return HookResult.fail(f"working_dir not a directory: {working_dir}")

        try:
            opts, is_compact = build_trace_opts(params)
            _ensure_flow_imports()
            assert get_package is not None
            pkg = get_package(working_dir)

            entry = params.get("entry")
            if entry is not None:
                if not isinstance(entry, str):
                    return HookResult.fail(
                        f"entry must be str, got {type(entry).__name__}"
                    )
                return self._trace_entries(pkg, entry, opts, compact=is_compact)
            return self._trace_all(pkg, opts, compact=is_compact)

        except Exception as exc:  # noqa: BLE001
            return HookResult.fail(f"Flow tracing failed: {exc}")

    @staticmethod
    def _deduplicate_entry_symbols(symbols: list[str]) -> list[str]:
        """Remove parent classes when qualified methods are present."""
        qualified = {s for s in symbols if "." in s}
        parents = {s.rsplit(".", 1)[0] for s in qualified}
        return [s for s in symbols if s not in parents]

    @staticmethod
    def _format_symbol_traces(
        steps: list[FlowStep],
        sym: str,
        compact: bool,
        format_fn: Callable[[list[FlowStep]], str],
    ) -> _TraceValue:
        """Format traced steps as compact string or dict list."""
        if compact:
            return format_fn(steps)
        # ``model_dump`` is upstream-typed ``dict[str, Any]`` (pydantic);
        # cast to ``dict[str, object]`` at this boundary.
        return [
            cast("dict[str, object]", s.model_dump(exclude_none=True)) for s in steps
        ]

    @staticmethod
    def _trace_entries(
        pkg: PackageInfo,
        entry: str,
        opts: _TraceOpts,
        *,
        compact: bool = False,
    ) -> HookResult:
        """Trace one or more explicitly-specified entry symbols.

        For multi-symbol traces, callees already seen in a previous
        symbol's trace are deduplicated (first-wins ordering).
        Deduplication runs before compact/dict formatting.
        """
        from axm_ast.core.flows import format_flow_compact

        assert trace_flow is not None

        symbols = _parse_entry_symbols(entry)
        symbols = FlowsHook._deduplicate_entry_symbols(symbols)
        kw = _trace_opts_kwargs(opts)

        if len(symbols) == 1:
            try:
                steps, _truncated = trace_flow(pkg, symbols[0], **kw)
            except ValueError as exc:
                return HookResult.fail(str(exc))
            return HookResult.ok(
                traces=FlowsHook._format_symbol_traces(
                    steps,
                    symbols[0],
                    compact,
                    format_flow_compact,
                ),
            )

        return FlowsHook._trace_multi_entries(
            pkg,
            symbols,
            kw,
            compact,
            format_flow_compact,
        )

    @staticmethod
    def _trace_multi_entries(
        pkg: PackageInfo,
        symbols: list[str],
        kw: _TraceKwargs,
        compact: bool,
        format_fn: Callable[[list[FlowStep]], str],
    ) -> HookResult:
        """Trace multiple entry symbols with cross-trace deduplication."""
        assert build_callee_index is not None
        assert trace_flow is not None

        index = build_callee_index(pkg)
        traces: dict[str, _TraceValue] = {}
        seen: set[str] = set()
        for sym in symbols:
            try:
                steps, _truncated = trace_flow(pkg, sym, callee_index=index, **kw)
            except ValueError:
                continue
            deduped = [s for s in steps if s.name == sym or s.name not in seen]
            seen.update(s.name for s in steps)
            traces[sym] = FlowsHook._format_symbol_traces(
                deduped,
                sym,
                compact,
                format_fn,
            )
        if compact:
            return HookResult.ok(
                traces=_concat_compact_traces(traces),
            )
        return HookResult.ok(traces=traces)

    @staticmethod
    def _trace_all(
        pkg: PackageInfo,
        opts: _TraceOpts,
        *,
        compact: bool = False,
    ) -> HookResult:
        """Discover entry points and trace them (with safety caps)."""
        from axm_ast.core.flows import format_flow_compact

        assert find_entry_points is not None
        assert build_callee_index is not None
        assert trace_flow is not None

        entries = _discover_entries(pkg)
        index = build_callee_index(pkg)
        kw = _trace_opts_kwargs(opts)
        traces = _trace_entries_to_values(
            pkg,
            entries,
            index,
            kw,
            _FormatOpts(compact=compact, format_fn=format_flow_compact),
        )

        if compact:
            return HookResult.ok(traces=_concat_compact_traces(traces))
        return HookResult.ok(traces=traces)

execute(context, **params)

Execute the hook action.

Parameters:

Name	Type	Description	Default
context	dict[str, object]	Session context dictionary (must contain working_dir).	required
**params	object	Optional entry (symbol name, or newline-separated list). Optional detail ("source", "trace", or "compact"). Optional max_depth (default 5). Optional cross_module (default False).	{}

Returns:

Type	Description
HookResult	HookResult with traces in metadata on success.
HookResult	When detail is "compact", traces is a single string
HookResult	(concatenated with entry-name headers for multi-symbol traces).
HookResult	Otherwise, traces is a dict of symbol → step-dicts.

Source code in packages/axm-ast/src/axm_ast/hooks/flows.py

def execute(self, context: dict[str, object], **params: object) -> HookResult:
    """Execute the hook action.

    Args:
        context: Session context dictionary (must contain ``working_dir``).
        **params:
            Optional ``entry`` (symbol name, or newline-separated list).
            Optional ``detail`` ("source", "trace", or "compact").
            Optional ``max_depth`` (default 5).
            Optional ``cross_module`` (default False).

    Returns:
        HookResult with ``traces`` in metadata on success.
        When *detail* is ``"compact"``, ``traces`` is a single string
        (concatenated with entry-name headers for multi-symbol traces).
        Otherwise, ``traces`` is a dict of symbol → step-dicts.
    """
    raw_path = params.get("path") or context.get("working_dir", ".")
    if not isinstance(raw_path, (str, Path)):
        return HookResult.fail(
            f"path must be str or Path, got {type(raw_path).__name__}"
        )
    working_dir = Path(raw_path).resolve()

    if not working_dir.is_dir():
        return HookResult.fail(f"working_dir not a directory: {working_dir}")

    try:
        opts, is_compact = build_trace_opts(params)
        _ensure_flow_imports()
        assert get_package is not None
        pkg = get_package(working_dir)

        entry = params.get("entry")
        if entry is not None:
            if not isinstance(entry, str):
                return HookResult.fail(
                    f"entry must be str, got {type(entry).__name__}"
                )
            return self._trace_entries(pkg, entry, opts, compact=is_compact)
        return self._trace_all(pkg, opts, compact=is_compact)

    except Exception as exc:  # noqa: BLE001
        return HookResult.fail(f"Flow tracing failed: {exc}")

build_trace_opts(params)

Build trace options and compact flag from hook parameters.

Source code in packages/axm-ast/src/axm_ast/hooks/flows.py

def build_trace_opts(params: dict[str, object]) -> tuple[_TraceOpts, bool]:
    """Build trace options and compact flag from hook parameters."""
    from axm_ast.core.flows import VALID_DETAILS

    detail = str(params.get("detail", "trace"))
    if detail not in VALID_DETAILS:
        msg = f"Invalid detail={detail!r}; must be one of {sorted(VALID_DETAILS)}"
        raise ValueError(msg)
    is_compact = detail == "compact"
    raw_max_depth = params.get("max_depth", 5)
    max_depth = int(raw_max_depth) if isinstance(raw_max_depth, (int, str)) else 5
    opts = _TraceOpts(
        max_depth=max_depth,
        cross_module=bool(params.get("cross_module", False)),
        detail=detail,
        exclude_stdlib=bool(params.get("exclude_stdlib", True)),
    )
    return opts, is_compact