diff --git a/composer/cli/cache_autoprove.py b/composer/cli/cache_autoprove.py new file mode 100644 index 00000000..8a7c6cf4 --- /dev/null +++ b/composer/cli/cache_autoprove.py @@ -0,0 +1,604 @@ +"""Cache & Memory Explorer for the Auto-Prove pipeline. + +Browses the cache + memory namespaces produced by the ``cli_pipeline`` +drivers (``tui-autoprove`` / ``console-autoprove`` / the foundry +entries). Wired as ``cache-autoprove`` in ``pyproject.toml``. + +Usage:: + + # by run id (recommended — recovers the cache root, memory namespace, + # threat-model digest and plugin manifest from the run's ``cache_root`` + # tags; works even when the design doc was auto-discovered): + cache-autoprove run + + # by reconstructing the namespace from the original inputs (requires + # the design doc, so it does NOT work for auto-discovered runs): + cache-autoprove inputs \\ + --cache-ns [--memory-ns ] [--threat-model ] \\ + [--plugins ...] +""" + +import argparse +import asyncio +import json +import pathlib +import sys +from dataclasses import dataclass +from typing import AsyncGenerator + +from langgraph.store.base import BaseStore + +from composer.input.types import DEFAULT_RECURSION_LIMIT +from composer.input.files import file_digest +from composer.ui.cache_explorer import ( + CacheNode, StoreNode, CacheTreeNode, CacheExplorerApp, DummyServices, + node, section, node_for, leaf, memory, collect_tree, +) +from composer.spec.context import WorkflowContext, CVLGeneration, CVLJudge, CacheKey +from composer.spec.source.harness import ( + config_key, + system_setup_key, + harness_generation_key, + HARNESS_ANALYSIS_KEY, + ContractSetup, + SystemDescriptionHarnessed, + AgentSystemDescription, + HarnessResult, +) +from composer.pipeline.cli import root_cache_key, user_ns +from composer.pipeline.core import ( + COMMON_SYSTEM_CACHE_KEY, PROPERTIES_KEY, + _component_cache_key, _batch_cache_key, _pre_property_cache_key, +) +from composer.pipeline.plugins import installed_plugin_manifest, manifest_digest +from composer.pipeline.run_tags import AutoProveCacheTags, CACHE_ROOT_RECORD +from composer.core.user import get_uid +from composer.workflow.services import store_context +from composer.io.run_index import get_run_data +from composer.spec.source.summarizer import _summary_key, _SummaryCache +from composer.spec.source.struct_invariant import STRUCTURAL_INV_KEY, Invariants +from composer.spec.source.pipeline import INV_CVL_KEY +from composer.spec.prop_inference import ( + _BugAnalysisCache, _AgentResult, _AgentRoundWithHistory, + bug_analysis_key_from_digest, agent_round_key, AGENT_RESULT_KEY, +) +from composer.spec.cvl_generation import GeneratedCVL, _LastAttemptCache, LAST_ATTEMPT_KEY, CVL_JUDGE_KEY +from composer.spec.system_model import ( + SourceApplication, SourceExplicitContract, SourceExternalActor, + HarnessedApplication, HarnessedExplicitContract, HarnessDefinition, + ContractInstance, ContractComponentInstance, +) + + +# The driver writes the analyzed SourceApplication under CacheKey(COMMON_SYSTEM_CACHE_KEY) +# (pipeline.core.run_pipeline); mirror that key here to read it back. +SYSTEM_ANALYSIS_KEY = CacheKey[None, SourceApplication](COMMON_SYSTEM_CACHE_KEY) + + +# --------------------------------------------------------------------------- +# Cache value type +# --------------------------------------------------------------------------- + +@dataclass +class PluginCacheRaw: + """Wrapper for raw values pulled from a plugin's pre-inference cache + subtree. Plugins cache under their own keys with their own value + shapes, so these can't be rehydrated into typed models — display the + stored payload as-is.""" + payload: dict + + +type AutoProveCachedValue = ( + SourceApplication + | ContractSetup + | SystemDescriptionHarnessed + | AgentSystemDescription + | HarnessResult + | _SummaryCache + | Invariants + | GeneratedCVL + | _LastAttemptCache + | _BugAnalysisCache + | _AgentResult + | _AgentRoundWithHistory + | CVLJudge + | PluginCacheRaw +) + + +# --------------------------------------------------------------------------- +# Tree construction +# --------------------------------------------------------------------------- + +def _build_harnessed_app( + sa: SourceApplication, + config_val: ContractSetup | None, +) -> HarnessedApplication: + """Reconstruct the HarnessedApplication the same way pipeline.py does.""" + contract_to_harness: dict[str, list[HarnessDefinition]] = {} + if config_val is not None: + for c in config_val.system_description.transitive_closure: + if not c.harness_definition: + continue + contract_to_harness.setdefault(c.harness_definition.harness_of, []).append( + HarnessDefinition(name=c.solidity_identifier, path=c.path) + ) + + comp: list[SourceExternalActor | HarnessedExplicitContract] = [] + for c in sa.components: + if not isinstance(c, SourceExplicitContract): + comp.append(c) + continue + comp.append(HarnessedExplicitContract( + sort=c.sort, + name=c.name, + solidity_identifier=c.solidity_identifier, + components=c.components, + description=c.description, + path=c.path, + harnesses=contract_to_harness.get(c.solidity_identifier, []), + )) + + return HarnessedApplication( + application_type=sa.application_type, + description=sa.description, + components=comp, + ) + + +async def _enumerate_raw_subtree( + store: BaseStore, + ctx: WorkflowContext, +) -> AsyncGenerator[CacheTreeNode[AutoProveCachedValue], None]: + """Surface every store slot under ``ctx``'s cache namespace as raw + ``StoreNode`` entries. Used for plugin pre-inference subtrees, whose + keys and value shapes are plugin-private.""" + ns = ctx.cache_namespace + assert ns is not None + items = await store.asearch(ns, limit=10_000) + if not items: + yield StoreNode[AutoProveCachedValue]( + label="(no cached entries)", slot=(ns, ""), value=None, + ) + return + for item in items: + rel = "/".join((*item.namespace[len(ns):], item.key)) + yield StoreNode[AutoProveCachedValue]( + label=rel, + slot=(tuple(item.namespace), item.key), + value=PluginCacheRaw(payload=item.value), + ) + + +async def _resolve_bug_key( + feat_ctx: WorkflowContext, + tags: AutoProveCacheTags, +) -> CacheKey: + """The bug-analysis key is parameterized on the threat-model digest and + the refinement (interactive) flag. Both come from the run tags; records + written before the ``interactive`` tag existed leave it ``None``, in + which case probe both variants and use whichever was written.""" + if tags.interactive is not None: + return bug_analysis_key_from_digest( + tags.threat_model_digest, with_refinement=tags.interactive, + ) + for refine in (False, True): + candidate = bug_analysis_key_from_digest(tags.threat_model_digest, refine) + if await feat_ctx.child(candidate).cache_get(_BugAnalysisCache) is not None: + return candidate + return bug_analysis_key_from_digest(tags.threat_model_digest, with_refinement=False) + + +async def _build_cvl_gen_nodes( + ctx: WorkflowContext[CVLGeneration], +) -> AsyncGenerator[CacheTreeNode[AutoProveCachedValue], None]: + yield memory(ctx, child=CVL_JUDGE_KEY, label="Feedback") + yield await leaf(ctx, LAST_ATTEMPT_KEY, "Last Attempt", _LastAttemptCache) + + +async def _build_component_nodes( + prop_ctx: WorkflowContext, + feat: ContractComponentInstance, + store: BaseStore, + tags: AutoProveCacheTags, +) -> AsyncGenerator[CacheTreeNode[AutoProveCachedValue], None]: + comp_key = _component_cache_key(feat, manifest_digest(tags.plugins)) + async with node_for(prop_ctx, comp_key, feat.component.name) as feat_ctx: + # Per-plugin pre-inference namespaces are siblings of the component + # key under PROPERTIES_KEY, but they're per-component work — surface + # them inside the component's subtree. + for plugin in tags.plugins: + pre_ctx = prop_ctx.child(_pre_property_cache_key(feat, plugin)) + with node(CacheNode(label=f"Plugin pre-inference: {plugin}", ctx=pre_ctx)): + async for n in _enumerate_raw_subtree(store, pre_ctx): + yield n + + # Bug analysis is layered: aggregate (_BugAnalysisCache) → agent result + # (_AgentResult) → per-round (_AgentRoundWithHistory). + bug_key = await _resolve_bug_key(feat_ctx, tags) + async with node_for(feat_ctx, bug_key, "Bug Analysis", _BugAnalysisCache) as bug_ctx: + async with node_for(bug_ctx, AGENT_RESULT_KEY, "Agent result", _AgentResult) as agent_ctx: + # Round indices are dense; probe 0..N until the first miss. + i = 0 + while True: + round_node = await leaf( + agent_ctx, agent_round_key(i), + f"Round {i + 1}", _AgentRoundWithHistory, + ) + if round_node.value is None: + break + yield round_node + i += 1 + + bug_cache = await feat_ctx.child(bug_key).cache_get(_BugAnalysisCache) + if bug_cache is None: + return + batch_key = _batch_cache_key(bug_cache.items) + async with node_for(feat_ctx, batch_key, "CVL Generation", GeneratedCVL) as cvl_ctx: + async for n in _build_cvl_gen_nodes(cvl_ctx.abstract(CVLGeneration)): + yield n + + +async def build_tree_inner( + root_ctx: WorkflowContext[None], + store: BaseStore, + tags: AutoProveCacheTags, +) -> AsyncGenerator[CacheTreeNode[AutoProveCachedValue], None]: + sa_leaf = await leaf(root_ctx, SYSTEM_ANALYSIS_KEY, "system-analysis", SourceApplication) + yield sa_leaf + + # Read config value upfront so we can derive the summary key outside the with block + config_val = await root_ctx.child(config_key).cache_get(ContractSetup) + + async with node_for(root_ctx, config_key, "config", ContractSetup) as config_ctx: + if sa_leaf.value is not None: + async with node_for(config_ctx, system_setup_key(sa_leaf.value), "setup", SystemDescriptionHarnessed) as setup_ctx: + ha_leaf = await leaf(setup_ctx, HARNESS_ANALYSIS_KEY, "harness-analysis", AgentSystemDescription) + yield ha_leaf + if ha_leaf.value is not None and ha_leaf.value.needs_harnessing(): + yield await leaf( + setup_ctx, + harness_generation_key(ha_leaf.value), + "harness-generation", + HarnessResult, + ) + + # Summary — key derivable only once ContractSetup is cached + if config_val is not None: + yield await leaf(root_ctx, _summary_key(config_val), "summary", _SummaryCache) + + yield await leaf(root_ctx, STRUCTURAL_INV_KEY, "structural-inv", Invariants) + async with node_for(root_ctx, INV_CVL_KEY, "invariant-cvl", GeneratedCVL) as inv_cvl_ctx: + async for n in _build_cvl_gen_nodes(inv_cvl_ctx.abstract(CVLGeneration)): + yield n + + # Properties — per-component plugin pre-inference + bug analysis + CVL generation + if sa_leaf.value is None: + with section("properties (no source analysis)"): + pass + return + + harnessed_app = _build_harnessed_app(sa_leaf.value, config_val) + + # Find the main contract. The pipeline matches the entry point by + # solidity_identifier (pipeline.core.main_instance), so the explorer does too. + contract_ind = -1 + for i, c in enumerate(harnessed_app.contract_components): + if c.solidity_identifier == tags.contract_name: + contract_ind = i + break + + if contract_ind == -1: + with section(f"properties (contract '{tags.contract_name}' not found)"): + pass + return + + contract_instance = ContractInstance(contract_ind, app=harnessed_app) + prop_ctx = root_ctx.child(PROPERTIES_KEY) + + with section("properties"): + for comp_idx in range(len(contract_instance.contract.components)): + feat = ContractComponentInstance(_contract=contract_instance, ind=comp_idx) + async for n in _build_component_nodes(prop_ctx, feat, store, tags): + yield n + + +async def build_tree( + root_ctx: WorkflowContext[None], + store: BaseStore, + tags: AutoProveCacheTags, +) -> CacheNode[AutoProveCachedValue]: + return await collect_tree("root", root_ctx, build_tree_inner(root_ctx, store, tags)) + + +# --------------------------------------------------------------------------- +# Value formatting +# --------------------------------------------------------------------------- + +def format_value(val: AutoProveCachedValue) -> list[str]: + lines: list[str] = [] + + match val: + case SourceApplication(application_type=app_type, description=desc, components=comps): + lines.append(f"Type: {app_type}") + lines.append(f"Description: {desc}") + lines.append("") + for c in comps: + match c: + case SourceExplicitContract(name=name, sort=sort, path=path, description=cdesc): + lines.append(f"[{sort}] {name} ({path})") + lines.append(f" {cdesc}") + case SourceExternalActor(name=name, path=path, description=cdesc): + loc = f" ({path})" if path else "" + lines.append(f"[external] {name}{loc}") + lines.append(f" {cdesc}") + + case ContractSetup(system_description=sys_desc, config=cfg): + lines.append("Pre-audit setup: OK") + lines.append(f"Summaries path: {cfg.summaries_path}") + lines.append(f"User types: {len(cfg.user_types)}") + lines.append(f"Closure contracts: {len(sys_desc.transitive_closure)}") + + case AgentSystemDescription( + non_trivial_state=nts, + erc20_contracts=erc20s, + external_interfaces=ext_ifaces, + transitive_closure=closure, + ): + lines.append(f"Non-trivial state: {nts}") + lines.append(f"ERC20 contracts: {', '.join(erc20s) if erc20s else 'none'}") + lines.append(f"Needs harnessing: {val.needs_harnessing()}") + lines.append("") + lines.append(f"Transitive closure ({len(closure)}):") + for c in closure: + instances = f" x{c.num_instances}" if c.num_instances else "" + lines.append(f" {c.solidity_identifier}{instances}") + for lf in c.link_fields: + lines.append(f" links → {', '.join(lf.target)}") + if ext_ifaces: + lines.append("") + lines.append(f"External interfaces ({len(ext_ifaces)}):") + for ei in ext_ifaces: + lines.append(f" {ei.name}: {ei.behavioral_spec}") + + case SystemDescriptionHarnessed( + non_trivial_state=nts, + erc20_contracts=erc20s, + external_interfaces=ext_ifaces, + transitive_closure=closure, + ): + lines.append(f"Non-trivial state: {nts}") + lines.append(f"ERC20 contracts: {', '.join(erc20s) if erc20s else 'none'}") + lines.append("") + lines.append(f"Transitive closure ({len(closure)}):") + for c in closure: + harnessed = " [harnessed]" if c.harness_definition else "" + lines.append(f" {c.solidity_identifier} ({c.path}){harnessed}") + if c.harness_definition: + lines.append(f" harness of: {c.harness_definition.harness_of}") + for lf in c.link_fields: + lines.append(f" links → {', '.join(lf.target)} via {", ".join(lf.link_paths)}") + if ext_ifaces: + lines.append("") + lines.append(f"External interfaces ({len(ext_ifaces)}):") + for ei in ext_ifaces: + lines.append(f" {ei.name}: {ei.behavioral_spec}") + + case HarnessResult(identifier_to_source=identifier_to_source): + for target, harnesses in identifier_to_source.items(): + lines.append(f"{target}:") + for h in harnesses: + lines.append(f" {h.harness_name} → {h.path}") + for harnesses in identifier_to_source.values(): + for h in harnesses: + lines.append("") + lines.append(f"--- {h.harness_name} ({h.path}) ---") + lines.extend(h.source.splitlines()) + + case _SummaryCache(content=content): + lines.extend(content.splitlines()) + + case Invariants(inv=invs): + lines.append(f"Invariants ({len(invs)}):") + for inv in invs: + lines.append(f" {inv.description}") + + case GeneratedCVL(commentary=commentary, cvl=cvl, skipped=skipped): + lines.append(f"Commentary: {commentary}") + if skipped: + lines.append(f"Skipped: {len(skipped)}") + lines.append("") + lines.extend(cvl.splitlines()) + + # Subclass-of-_BugAnalysisCache cases first — match order matters. + case _AgentRoundWithHistory(items=items, reasoning=reasoning, agent_conversation=history): + lines.append(f"Properties this round ({len(items)}):") + for p in items: + lines.append(f" - [{p.sort}] {p.description}") + lines.append("") + lines.append("--- Reasoning ---") + lines.append(reasoning) + lines.append("") + lines.append(f"Agent history: {len(history)} message(s)") + + case _AgentResult(items=items, final_history=history): + lines.append(f"Cumulative properties ({len(items)}):") + for p in items: + lines.append(f" - [{p.sort}] {p.description}") + lines.append("") + lines.append(f"Final-round history: {len(history)} message(s)") + + case _BugAnalysisCache(items=items): + lines.append(f"Properties ({len(items)}):") + for p in items: + lines.append(f" - [{p.sort}] {p.description}") + + case _LastAttemptCache(cvl=cvl): + lines.append("LAST ATTEMPT") + lines.append(cvl) + + case PluginCacheRaw(payload=payload): + lines.extend(json.dumps(payload, indent=2, default=str).splitlines()) + + case CVLJudge(): + ... + + return lines + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def _resolve_from_inputs(args: argparse.Namespace) -> AutoProveCacheTags | None: + """Reconstruct the tags from the original CLI inputs. Requires the design + doc (it feeds the byte-hash root key), so it does not work for auto-discovered + runs — use the ``run`` subcommand for those. Returns ``None`` if the doc is + unreadable.""" + project_root = pathlib.Path(args.project_root).resolve() + main_contract_path, contract_name = args.main_contract.split(":", 1) + full_contract_path = pathlib.Path(main_contract_path).resolve() + relative_path = str(full_contract_path.relative_to(project_root)) + + sys_path = pathlib.Path(args.system_doc) + if not sys_path.is_file(): + print(f"Error: cannot read {sys_path}", file=sys.stderr) + return None + + root_ns = user_ns( + args.cache_ns, + root_cache_key(str(project_root), sys_path, relative_path, contract_name), + ) + memory_ns = args.memory_ns + if memory_ns: + memory_ns = get_uid() + "/" + memory_ns + + plugins: list[str] = ( + sorted(args.plugins) if args.plugins is not None else installed_plugin_manifest() + ) + tm_digest = ( + file_digest(pathlib.Path(args.threat_model)) + if args.threat_model is not None else None + ) + + return AutoProveCacheTags( + cache_root=list(root_ns), + contract_name=contract_name, + memory_ns=memory_ns, + plugins=plugins, + threat_model_digest=tm_digest, + # Not recoverable from the inputs — the tree builder probes both + # refinement variants of the bug-analysis key. + interactive=None, + ) + + +async def _resolve_from_run( + store: BaseStore, run_id: str, uid: str | None, +) -> AutoProveCacheTags | None: + """Look up the tags the pipeline recorded for ``run_id`` — the + ``cache_root`` run-data record written from ``cli_pipeline`` once the + design doc (hence cache root) is resolved. Returns ``None`` if the run + isn't found.""" + rec = await get_run_data(store, run_id, CACHE_ROOT_RECORD, uid=uid) + if rec is None: + print( + f"Error: no cache metadata for run {run_id!r} (looked under uid={uid!r}).", + file=sys.stderr, + ) + return None + return AutoProveCacheTags.model_validate(rec) + + +async def _async_main(args: argparse.Namespace) -> int: + async with store_context() as store: + tags = ( + await _resolve_from_run(store, args.run_id, args.uid) + if args.mode == "run" + else _resolve_from_inputs(args) + ) + if tags is None: + return 1 + if tags.cache_root is None: + print( + "Error: the run ran without caching — nothing to explore.", + file=sys.stderr, + ) + return 1 + root_ns = tuple(tags.cache_root) + + print(f"Root namespace: {root_ns}", file=sys.stderr) + + root_ctx: WorkflowContext[None] = WorkflowContext.create( + services=DummyServices(), # type: ignore[arg-type] + thread_id="explorer", + store=store, + recursion_limit=DEFAULT_RECURSION_LIMIT, + memory_namespace=tags.memory_ns, + cache_namespace=root_ns, + ) + + status = f"Cache NS: {root_ns}" + if tags.memory_ns: + status += f" | Memory NS: {tags.memory_ns}" + if tags.plugins: + status += f" | Plugins: {', '.join(tags.plugins)}" + if tags.threat_model_digest: + status += f" | TM digest: {tags.threat_model_digest}" + + app = CacheExplorerApp( + build_tree=lambda: build_tree(root_ctx, store, tags), + format_value=format_value, + store=store, + status=status, + ) + await app.run_async() + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Cache & Memory Explorer for the Auto-Prove pipeline" + ) + sub = parser.add_subparsers(dest="mode", required=True) + + p_run = sub.add_parser( + "run", + help="Explore a run by id (recommended). Reads the tags the run recorded " + "in its metadata (cache root, memory ns, plugins, threat-model digest) " + "— works even when the design doc was auto-discovered.", + ) + p_run.add_argument("run_id", help="Run id (from the autoprove logs / ap-trail).") + p_run.add_argument( + "--uid", default=None, + help="User-id namespace the run was logged under (default: the run's default namespace).", + ) + + p_inputs = sub.add_parser( + "inputs", + help="Reconstruct the cache namespace from the original CLI inputs. Requires " + "the supplied design doc, so it does NOT work for auto-discovered runs.", + ) + p_inputs.add_argument("project_root", help="Root directory of the Solidity project") + p_inputs.add_argument("main_contract", help="Main contract as path:ContractName") + p_inputs.add_argument("system_doc", help="Path to the design document (text or PDF)") + p_inputs.add_argument("--cache-ns", required=True, dest="cache_ns", + help="Cache namespace (same as passed to autoprove)") + p_inputs.add_argument("--memory-ns", dest="memory_ns", default=None, + help="Memory namespace (enables memory browsing)") + p_inputs.add_argument("--threat-model", dest="threat_model", default=None, + help="Path to the threat model used for the original run — its " + "digest parameterizes the bug-analysis cache key. Omit for " + "runs without one.") + p_inputs.add_argument("--plugins", dest="plugins", nargs="*", default=None, + help="Plugin names active for the original run — the manifest " + "digest is suffixed onto per-component cache keys. Defaults " + "to the plugins installed in this environment; pass with no " + "names to force plugin-free keys.") + + args = parser.parse_args() + return asyncio.run(_async_main(args)) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/composer/input/files.py b/composer/input/files.py index 87aa5179..a5b354ec 100644 --- a/composer/input/files.py +++ b/composer/input/files.py @@ -85,6 +85,14 @@ def _bytes_digest(b: bytes) -> str: return hashlib.sha256(b).hexdigest()[:16] +def file_digest(path: pathlib.Path) -> str: + """The digest ``Document.to_digest()`` reports for a document loaded + from ``path``, computed straight from the file bytes — every concrete + document shape digests its raw bytes, so this matches without going + through an uploader.""" + return _bytes_digest(path.read_bytes()) + + # Suffixes treated as binary regardless of byte content; short-circuits # the byte-scan heuristic for common cases. _KNOWN_BINARY_SUFFIXES = {".pdf"} diff --git a/composer/pipeline/PLUGIN.md b/composer/pipeline/PLUGIN.md new file mode 100644 index 00000000..c169974c --- /dev/null +++ b/composer/pipeline/PLUGIN.md @@ -0,0 +1,245 @@ +# AutoProve Pipeline Plugins + +Plugins let an out-of-tree package participate in the property-inference +phase of an autoprove pipeline run: contribute extra prompt material to the +per-component property-inference agent, post-process the property list it +produces, run their own (cached, task-tracked) agents to do either, and ship +their own Jinja templates. Plugins +are discovered from the installed environment — the host needs no +configuration to pick one up. + +## Shipping a plugin + +Register a **loader class** under the entry-point group +`certora.autoprove.plugins`: + +```toml +[project.entry-points."certora.autoprove.plugins"] +my-plugin = "my_pkg.plugin:MyPluginLoader" +``` + +The entry-point **name** (`my-plugin`) is the plugin's identity: it appears in +cache keys, task ids, and the run's recorded plugin manifest. It must be +unique across the installed environment — a duplicate name fails the whole +pipeline at startup. Renaming it invalidates cached work (see +[Cache-key effects](#cache-key-effects)). + +The loader: + +```python +class MyPluginLoader(PipelinePluginLoader): + @asynccontextmanager + async def initialize(self) -> AsyncIterator[PipelinePlugin]: + async with open_my_resources() as res: + yield MyPlugin(res) +``` + +Contract: + +* The loader class is instantiated with **no arguments** during entry-point + scanning, before any event-loop work. Construction must be trivial; all + real setup belongs in `initialize()`. +* `initialize()` returns an async context manager yielding the + `PipelinePlugin` instance. Its scope **is** the pipeline run: it is entered + before the pipeline body starts and + exited when the pipeline returns *or raises*. Any resource the plugin's + hooks or tools touch (DB pools, subprocesses, caches) must be acquired here + and remain valid exactly while the scope is open. +* An exception from `initialize()` (either side of the `yield`) fails the run. + +The plugin class: + +```python +class MyPlugin(PipelinePlugin): + NAME = "My Plugin" # display name (task labels); the entry-point + # name, not NAME, is the identity +``` + +## Hooks + +### `property_inference_input_hook(comp, run) -> AnyPropertyGenerationInput | None` + +Called once per **functional component** of the main contract, per plugin, +*before* that component's property-inference agent runs. Return extra input +for the inference prompt, or `None` to contribute nothing for this component. + +Invocation model: + +* All plugins × all components run **concurrently**. Hook + implementations and any state shared on the plugin instance must tolerate + concurrent calls. Size shared resources accordingly (e.g. a DB pool with + `max_size=1` serializes every component's probe). +* An exception from any hook **aborts the entire pipeline run**. Treat + recoverable conditions ("nothing relevant found", "corpus unavailable") as + `None`, not as raises. +* The host does **not** cache hook results. Long-running work must implement + its own caching under `run.ctx` (see below) or it re-executes every run. + +### `post_process_property_inference(comp, run, props) -> list[PropertyFormulation]` + +Called once per **functional component**, per plugin, after that component's +property inference delivers a non-empty property list. Receives the current +list and returns its replacement — filter, rewrite, or extend. The base +implementation returns `props` unchanged; only override it to actually +transform the list. + +Invocation model: + +* Within a component, plugins run **sequentially, chained, in lexicographic + entry-point-name order**: each plugin receives the previous plugin's + output. If ordering relative to another plugin matters, that ordering is + controlled by the entry-point names. Different components' chains still + run concurrently with each other. +* Returning an **empty list drops the component** from the rest of the + pipeline (no CVL generation for it), exactly as if inference had produced + nothing. +* The returned list is what downstream CVL generation keys its cache on. + Output must be **deterministic for a given input list** (stable order, + stable text) — otherwise every run lands in a fresh CVL-generation cache + entry even when nothing changed. +* The exception and caching rules from the pre hook apply unchanged: a raise + aborts the run, and the host does not cache the hook's result. Note that + `run.ctx`'s namespace incorporates a digest of the *incoming* property + list (see below), so plugin-side caching is input-addressed — a change in + what inference (or an earlier plugin in the chain) produced automatically + lands you in a fresh namespace. + +## `PluginContext` — what a hook receives + +Both hooks receive `run: PluginContext[C]`, where `C` is the hook's marker +type (`PrePropertyInference` / `PostPropertyInference`). It carries four +things: + +* **`run.ctx: WorkflowContext[C]`** — a cache/memory context rooted at a + namespace private to this *(component, plugin)* pair; for the post hook + the namespace additionally incorporates a digest of the incoming property + list, so it moves whenever the input changes. Use it for plugin-side + caching and agent memory. + * `ctx.get_memory_tool()` — memory tool bound to this namespace, for + handing to an agent. + * `ctx.thread_id`, `ctx.recursion_limit` — pass through to + `run_to_completion` when running your own graph. + * `PrePropertyInference`/`PostPropertyInferece` is a generic "marker" context type, you should + implement a `CacheKey[PrePropertyInference, MyModelType]` and navigate + (via `child_ctx = ctx.child(...)`) to the `MyModelType` context + * `await child_ctx.cache_get(MyModelType)` / `await child_ctx.cache_put(value)` - + typed slot at the context's child key (values are pydantic models). + * It is strongly recommended to use the caching layer. If you use the caching + layer, it is recommended to use the same context object for the + memory tool/thread id as for caching. + That is, instead of `child_ctx.cache_get(MyModelType); ... ctx.get_memory_tool()` + do `child_ctx.cache_get(MyModelType); ... child_ctx.get_memory_tool()` + + Cached values persist across runs (same `--cache-ns`). The namespace is + keyed by the component digest, the plugin's entry-point name, and (post + hook only) the incoming property list — **not** by your plugin's version + or logic. If your output format changes in a way that stale cache entries + would poison, version your own child keys. + +* **`run.env: ServiceHost`** — model builders (`env.models.builder_heavy()`, + `builder_lite()`), analysis tools, sort. The template loader inside + `env.models` is already the plugin's own (see [Templates](#templates)), so + `with_sys_prompt_template(...)` / `TypedTemplate` resolve plugin templates. + +* **`run.source: SourceCode`** — the target project (root, main contract, + relative path). + +* **`await run.runner(label, job)`** — run `job` (a nullary async callable) + as a pipeline task. **All substantial work — anything that runs an agent, + hits an LLM, or takes real wall-clock time — must go through this.** It + registers the work with the pipeline's task display and concurrency + semaphore; work done outside it is invisible and unthrottled. The label is + shown as `() Plugin {NAME}: {label}` (sub-phase: + `Property Pre-Inference` / `Property Post-Process`); task ids are + namespaced per phase/plugin/component, so labels only need to be + human-meaningful, not unique. + +## Property-generation inputs + +The hook's return value is spliced into the inference agent's prompt. Two +shapes, sharing `uid` / `sort` / `when`: + +```python +PropertyGenerationInput(uid=..., sort=..., when=..., input=payload) +CacheablePropertyGenerationInput(uid=..., sort=..., when=..., provide=fn) +``` + +* **`uid`** — stable identifier (reverse-DNS style, e.g. + `"com.certora.corpus"`). It is the sort key for deterministic placement: + it must not vary across runs or components. +* **`sort`** — + * `"generic"`: the payload is **byte-identical for every component** of the + application (e.g. a threat model). Generic inputs are placed before all + component-specific material so the prompt prefix is shared across + components. Do not mark component-dependent content generic; it silently + destroys that prefix sharing. + * `"specific"`: component-dependent content. +* **`when`** — + * `"initial"`: included only in the agent's first round. + * `"always"`: included in every round's prompt. +* **payload** (`MessagePayloadType`): a string, a content-block dict, or a + list of either. +* **`CacheablePropertyGenerationInput.provide(is_boundary: bool)`** — called + at prompt-assembly time. When `is_boundary` is true, this input is the last + block before a prompt-cache breakpoint and the final content block it + returns should carry `cache_control` (e.g. + `document.to_dict(with_cache=True)`). When false, return the same content + without the marker. + +Placement order in the assembled prompt: + +``` +[system prompt] +generic/always · generic/initial · specific/always · specific/initial +[host initial prompt (component, prior rounds)] +``` + +Within each group: non-cacheable inputs (sorted by `uid`), then cacheable +inputs (sorted by `uid`). Rounds after the first drop the `initial` groups. + +## Cache-key effects + +The host folds the **installed plugin manifest** (sorted entry-point names) +into every per-component cache key: + +* Installing, removing, or renaming *any* plugin changes the digest and + invalidates all per-component inference/CVL cache entries — for every + plugin, not just the changed one. +* The manifest is recorded in the run's `cache_root` tags. +* The per-plugin namespace handed to the hook (`run.ctx`) is keyed by + *(component digest, plugin name)* and is **not** part of the component key + — its contents are the plugin's own to version and invalidate. + +## Templates + +`self.load_jinja_template(template_name, **params)` on the plugin renders +Jinja templates with this resolution: + +* **Bare names** (`"my_prompt.j2"`) resolve from the plugin's own template + directory first: `PackageLoader()`, i.e. a + `templates/` directory in the package containing the plugin class. Override + `plugin_loader()` to supply any other `jinja2.BaseLoader`; return `None` + to use the host loader unchanged (this is also the silent fallback when + `PackageLoader` can't locate a template directory). +* **`autoprover/`-prefixed names** (`"autoprover/tools/edit_prompting.j2"`) + resolve from the host's template set. Host templates are resolved via + namespacing magic, so you needn't worry about your template names colliding + with host template names. + +The same loader is spliced into `run.env.models`, so agents built from the +plugin's `PluginContext` resolve `with_sys_prompt_template(...)` / +`TypedTemplate` names by the same rules. + +## Checklist + +1. `pyproject.toml`: entry point in `certora.autoprove.plugins`. +2. Loader: trivial `__init__`, resources in `initialize()`, yield the plugin. +3. Plugin: set `NAME`; implement `property_inference_input_hook` and/or + `post_process_property_inference`. +4. Cache expensive work under `run.ctx`; run it via `run.runner`. +5. Return `None` for "nothing to contribute"; never raise for recoverable + conditions. +6. Choose `uid`/`sort`/`when` per the placement rules; use the cacheable + variant for large payloads that deserve a prompt-cache breakpoint. +7. Templates in `templates/` next to the plugin class; reference host + templates with the `autoprover/` prefix. diff --git a/composer/pipeline/cli.py b/composer/pipeline/cli.py index c398bf9f..d51537b3 100644 --- a/composer/pipeline/cli.py +++ b/composer/pipeline/cli.py @@ -23,6 +23,8 @@ from composer.spec.service_host import ModelProvider from composer.spec.system_analysis import SolidityIdentifier from .core import PipelineBackend, run_pipeline +from .plugins import installed_plugin_manifest +from .run_tags import AutoProveCacheTags, CACHE_ROOT_RECORD from composer.io.multi_job import HandlerFactory, run_task, TaskInfo from composer.diagnostics.timing import RunSummary, install_run_summary from composer.llm.registry import get_provider_for @@ -235,11 +237,14 @@ async def cli_pipeline[P: enum.Enum, H]( await conns.uploader.get_document(pathlib.Path(threat_path)) if (threat_path := args.threat_model) is not None else None ) - await data_logger("cache_root", { - "cache_root": list(cache_root) if cache_root is not None else None, - "contract_name": str(contract_name), - "memory_ns": memory_ns, - }) + await data_logger(CACHE_ROOT_RECORD, AutoProveCacheTags( + cache_root=list(cache_root) if cache_root is not None else None, + contract_name=str(contract_name), + memory_ns=memory_ns, + plugins=installed_plugin_manifest(), + threat_model_digest=threat_model.to_digest() if threat_model is not None else None, + interactive=args.interactive, + ).model_dump()) full_source = SourceCode( content=system_doc_doc, contract_name=init_source.contract_name, diff --git a/composer/pipeline/core.py b/composer/pipeline/core.py index 97ce1b0b..7adb4ac0 100644 --- a/composer/pipeline/core.py +++ b/composer/pipeline/core.py @@ -17,7 +17,9 @@ import enum import logging from dataclasses import dataclass -from typing import Protocol +from dataclasses import dataclass, field +from pathlib import Path +from typing import Protocol, TypedDict, NotRequired from abc import ABC, abstractmethod from pydantic import BaseModel @@ -32,7 +34,7 @@ ) from composer.spec.types import PropertyFormulation, ArtifactIdentifier from composer.spec.system_analysis import run_component_analysis -from composer.spec.prop_inference import run_property_inference +from composer.spec.prop_inference import run_property_inference, AnyPropertyGenerationInput from composer.spec.util import string_hash from composer.input.files import Document from composer.spec.source.report.build import build_report @@ -43,6 +45,8 @@ from .ptypes import ( BackendJob, BackendResult, ComponentOutcome, CorePhases, CorePipelineResult, Delivered, GaveUp, PipelineRun, SystemAnalysisSpec ) +from .plugin_api import PrePropertyInference, PostPropertyInference +from .plugins import load_plugins, PluginManager, PluginPhaseManager, PluginPhaseRunner COMMON_SYSTEM_CACHE_KEY = "system-analysis" @@ -130,8 +134,14 @@ def main_instance(app: AnyApplication, source: SourceCode) -> ContractInstance: class _Batch(BackendJob): feat_ctx: WorkflowContext[ComponentGroup] -def _component_cache_key(c: ContractComponentInstance) -> CacheKey[Properties, ComponentGroup]: - return CacheKey(string_hash("|".join([c.app.model_dump_json(), str(c.ind), str(c._contract.ind)]))) +def _component_digest(c: ContractComponentInstance) -> str: + return string_hash("|".join([c.app.model_dump_json(), str(c.ind), str(c._contract.ind)])) + +def _component_cache_key(c: ContractComponentInstance, plugin_digest: str | None) -> CacheKey[Properties, ComponentGroup]: + raw_digest = _component_digest(c) + if plugin_digest is not None: + raw_digest += f"-{plugin_digest}" + return CacheKey(raw_digest) def _batch_cache_key[FormT: BaseModel](props: list[PropertyFormulation]) -> CacheKey[ComponentGroup, FormT]: @@ -145,7 +155,6 @@ def extract_task_id(idx: int) -> str: def formalize_task_id(idx: int) -> str: return f"formalize-{idx}" -# ---- the driver -------------------------------------------------------------- async def run_pipeline[P: enum.Enum, FormT: BackendResult, H, A: ArtifactIdentifier]( backend: PipelineBackend[P, FormT, H, A], run: PipelineRun[P, H], @@ -153,6 +162,21 @@ async def run_pipeline[P: enum.Enum, FormT: BackendResult, H, A: ArtifactIdentif interactive: bool = False, threat_model: Document | None = None, max_bug_rounds: int = 3, +) -> CorePipelineResult[FormT]: + async with load_plugins(run) as plugins: + return await run_pipeline_inner( + backend, run, plugins, interactive=interactive, threat_model=threat_model, max_bug_rounds=max_bug_rounds + ) + +# ---- the driver -------------------------------------------------------------- +async def run_pipeline_inner[P: enum.Enum, FormT: BackendResult, H, A: ArtifactIdentifier]( + backend: PipelineBackend[P, FormT, H, A], + run: PipelineRun[P, H], + plugin_manager: PluginManager, + *, + interactive: bool = False, + threat_model: Document | None = None, + max_bug_rounds: int = 3, ) -> CorePipelineResult[FormT]: spec, phases = backend.analysis_spec, backend.core_phases source = run.source @@ -172,7 +196,7 @@ async def run_pipeline[P: enum.Enum, FormT: BackendResult, H, A: ArtifactIdentif ) if analyzed is None: raise ValueError("System analysis produced no result.") - + # 2. Backend transform + main-contract location (prover: harness lift; foundry: identity). prepared = await backend.prepare_system(analyzed, run) @@ -180,8 +204,19 @@ async def run_pipeline[P: enum.Enum, FormT: BackendResult, H, A: ArtifactIdentif # this preserves the prover's autosetup ∥ bug-analysis overlap, generically. formalizer_task = asyncio.create_task(prepared.prepare_formalization(run)) - batches = await _extract_all(prepared.main, backend.backend_guidance, run, - phases["extraction"], interactive, threat_model, max_bug_rounds) + batches = await _extract_all( + prepared.main, + backend.backend_guidance, + run, + phases["extraction"], + interactive, + threat_model, + max_bug_rounds, + plugin_manager.bind_phase( + backend.core_phases.get("extraction_plugin") or backend.core_phases["extraction"], + "Property Extraction" + ) + ) formalizer = await formalizer_task if not batches: raise ValueError("No properties extracted from any component.") @@ -250,23 +285,80 @@ async def _run(batch: _Batch) -> ComponentOutcome[FormT]: return _tally(outcomes) +def _pre_property_cache_key(feat: ContractComponentInstance, plugin: str) -> CacheKey[Properties, PrePropertyInference]: + key = f"{_component_digest(feat)}-{string_hash(plugin)}-pre" + return CacheKey(key) + +def _post_property_cache_key(feat: ContractComponentInstance, plugin: str, curr_props: list[PropertyFormulation]) -> CacheKey[Properties, PostPropertyInference]: + props = string_hash("|".join( + p.model_dump_json() for p in curr_props + )) + key = f"{_component_digest(feat)}-{string_hash(plugin)}-{props}" + return CacheKey(key) + async def _extract_all[P: enum.Enum, H]( main: ContractInstance, backend_guidance: str, run: PipelineRun[P, H], phase: P, interactive: bool, threat_model: Document | None, max_rounds: int, + plugins: PluginPhaseManager[P] ) -> list[_Batch]: prop_ctx = run.ctx.child(PROPERTIES_KEY) async def _one(idx: int) -> _Batch | None: feat = ContractComponentInstance(_contract=main, ind=idx) - feat_ctx = await prop_ctx.child(_component_cache_key(feat), - {"component": feat.component.model_dump()}) + async def run_plugin_pre(runner: PluginPhaseRunner[P]) -> AnyPropertyGenerationInput | None: + p = runner.plugin_id + ctxt = await prop_ctx.child(_pre_property_cache_key(feat, p), { + "plugin-name": p + }) + run_ctxt = runner.bind(str(idx), ctxt) + return await runner.plugin.property_inference_input_hook( + feat, run_ctxt + ) + + pre_process = await asyncio.gather(*[ + run_plugin_pre(plug_runner) for plug_runner in plugins.runners( + sub_phase_id="pre-inference", sub_phase_label="Property Pre-Inference" + ) + ]) + + feat_ctx = await prop_ctx.child(_component_cache_key(feat, plugins.plugin_digest), + {"component": feat.component.model_dump(), "plugins": plugins.plugin_manifest}) + props = await run.runner( TaskInfo(extract_task_id(idx), feat.component.name, phase), lambda conv: run_property_inference( feat_ctx, run.env, feat, refinement=conv if interactive else None, - threat_model=threat_model, max_rounds=max_rounds, backend_guidance=backend_guidance), + threat_model=threat_model, max_rounds=max_rounds, backend_guidance=backend_guidance, + extra_input=[ t for t in pre_process if t ] + ), ) - return _Batch(feat, props, feat_ctx) if props else None + if not props: + return None + + async def run_plugin_post( + runner: PluginPhaseRunner, props: list[PropertyFormulation] + ) -> list[PropertyFormulation]: + p = runner.plugin_id + ctxt = await prop_ctx.child( + _post_property_cache_key(feat, p, props), + { + "plugin-name": p, + "props": [p.model_dump() for p in props] + } + ) + run_ctxt = runner.bind(str(idx), ctxt) + return await runner.plugin.post_process_property_inference( + feat, run_ctxt, props + ) + accum = props + for runner in plugins.runners( + sub_phase_id="post-inference", sub_phase_label="Property Post-Process", sorted_run=True + ): + accum = await run_plugin_post( + runner, accum + ) + + return _Batch(feat, accum, feat_ctx) if accum else None got = await asyncio.gather(*[_one(i) for i in range(len(main.contract.components))]) return [b for b in got if b is not None] diff --git a/composer/pipeline/plugin_api.py b/composer/pipeline/plugin_api.py new file mode 100644 index 00000000..82bda24c --- /dev/null +++ b/composer/pipeline/plugin_api.py @@ -0,0 +1,122 @@ +from typing import AsyncContextManager, Protocol, Callable, Awaitable, Any +from functools import cached_property +from abc import ABC, abstractmethod +from graphcore.graph import TemplateLoader +from jinja2.loaders import BaseLoader, PackageLoader, ChoiceLoader, PrefixLoader + +from composer.templates.loader import base_loader, load_jinja_template, _autoescape +from composer.spec.system_model import ContractComponentInstance +from composer.pipeline.ptypes import PipelineRun +from composer.spec.context import WorkflowContext, SourceCode +from composer.spec.service_host import ServiceHost +from composer.spec.types import PropertyFormulation +from composer.spec.prop_inference import AnyPropertyGenerationInput + +class PluginContext[C](Protocol): + @property + def ctx(self) -> WorkflowContext[C]: + ... + + @property + def env(self) -> ServiceHost: + ... + + @property + def source(self) -> SourceCode: + ... + + async def runner[T]( + self, + label: str, + job: Callable[[], Awaitable[T]] + ) -> T: + ... + +class PrePropertyInference: + pass + +class PostPropertyInference: + pass + +import jinja2 +from jinja2 import Environment, ChoiceLoader, PrefixLoader + +class _NonStrippingPrefixLoader(PrefixLoader): + """Like PrefixLoader, but the compiled template keeps its prefix in .name, + so join_path can see which namespace a template came from.""" + def load(self, environment, name, globals=None): + loader, local_name = self.get_loader(name) + if globals is None: + globals = {} + try: + source, filename, uptodate = loader.get_source(environment, local_name) + except jinja2.TemplateNotFound as e: + raise jinja2.TemplateNotFound(name) from e + code = environment.compile(source, name, filename) # full name, not local_name + return environment.template_class.from_code(environment, code, globals, uptodate) + +class _PluginEnvironment(Environment): + namespace_prefixes = ("autoprover/",) + def join_path(self, template, parent): + # already namespaced -> leave alone + if template.startswith(self.namespace_prefixes): + return template + # inherit the parent's namespace for bare references + if parent and parent.startswith(self.namespace_prefixes): + prefix = next(p for p in self.namespace_prefixes if parent.startswith(p)) + return prefix + template + return template + +class PipelinePlugin(ABC): + NAME: str + + def plugin_loader(self) -> BaseLoader | None: + t = type(self).__module__ + try: + return PackageLoader(t) + except ValueError: + return None + + @cached_property + def load_jinja_template(self) -> TemplateLoader: + loader = self.plugin_loader() + if loader is None: + return load_jinja_template + + + new_jinja_loader = ChoiceLoader([ + loader, + _NonStrippingPrefixLoader({ + "autoprover": base_loader + }) + ]) + compilation_env = _PluginEnvironment(loader=new_jinja_loader, autoescape=_autoescape) + def _load_jinja_template(template_name: str, **kwargs: Any) -> str: + """Load and render a Jinja template from the script directory""" + template = compilation_env.get_template(template_name) + return template.render(**kwargs) + return _load_jinja_template + + + async def property_inference_input_hook( + self, + comp: ContractComponentInstance, + run: PluginContext[PrePropertyInference] + ) -> AnyPropertyGenerationInput | None: + return None + + async def post_process_property_inference( + self, + comp: ContractComponentInstance, + run: PluginContext[PostPropertyInference], + props: list[PropertyFormulation] + ) -> list[PropertyFormulation]: + return props + + +class PipelinePluginLoader(ABC): + @abstractmethod + def initialize( + self + ) -> AsyncContextManager[PipelinePlugin]: + ... diff --git a/composer/pipeline/plugins.py b/composer/pipeline/plugins.py new file mode 100644 index 00000000..6933d2ab --- /dev/null +++ b/composer/pipeline/plugins.py @@ -0,0 +1,146 @@ + +import enum +from dataclasses import dataclass, replace +from typing import Protocol, Callable, Awaitable, Any, Iterable, AsyncIterator, Never +import importlib.metadata +from contextlib import AsyncExitStack, asynccontextmanager +from functools import cached_property + + +from composer.io.multi_job import TaskInfo +from composer.spec.context import ( + WorkflowContext, SourceCode, +) +from composer.spec.service_host import ServiceHost +from composer.spec.util import string_hash +from .ptypes import PipelineRun +from .plugin_api import PipelinePluginLoader, PipelinePlugin, PluginContext + +class _RunnerFun(Protocol): + async def __call__[T](self, label: str, job: Callable[[], Awaitable[T]]) -> T: ... + +@dataclass +class PluginRunner[C]: + ctx: WorkflowContext[C] + env: ServiceHost + source: SourceCode + + runner: _RunnerFun + +class DisplayStrings(tuple[str, str]): + __slots__ = () + + @property + def display_str(self) -> str: + return self[1] + + @property + def id_str(self) -> str: + return self[0] + + def __new__( + cls, id: str, display: str + ): + return super().__new__(cls, (id, display)) + +@dataclass +class PluginPhaseRunner[P: enum.Enum]: + plugin: PipelinePlugin + _run: PipelineRun[P, Any] + _phase: tuple[P, str] + _sub_phase: DisplayStrings + plugin_id: str + + def bind[C]( + self, + uid: str, + ctxt: WorkflowContext[C] + ) -> PluginContext[C]: + new_loader = self.plugin.load_jinja_template + env = self._run.env + env = replace(env, models=replace(env.models, _loader=new_loader)) + async def run[T]( + label: str, + job: Callable[[], Awaitable[T]] + ) -> T: + label = f"(Plugin {self._sub_phase.display_str}) {self.plugin.NAME}: {label}" + return await self._run.runner( + TaskInfo(f"{self._phase[0].name}-{self._sub_phase.id_str}-{self.plugin_id}-{uid}", label, self._phase[0]), + job, + ) + return PluginRunner( + ctxt, + env, + self._run.source, + run + ) + + +PLUGIN_ENTRY_POINT_GROUP = "certora.autoprove.plugins" + + +def installed_plugin_manifest() -> list[str]: + """Sorted names of the autoprove plugins installed in this environment — + the manifest ``load_plugins`` will end up loading, computable without + initializing anything.""" + return sorted( + ep.name for ep in importlib.metadata.entry_points(group=PLUGIN_ENTRY_POINT_GROUP) + ) + + +def manifest_digest(manifest: list[str]) -> str | None: + """Digest of a sorted plugin manifest as suffixed onto per-component + cache keys; ``None`` when no plugins are active.""" + if not manifest: + return None + return string_hash("|".join(manifest)) + + +@dataclass +class PluginManager[P: enum.Enum]: + _plugins: dict[str, PipelinePlugin] + _run: PipelineRun[P, Any] + + @cached_property + def plugin_digest(self) -> None | str: + return manifest_digest(self.plugin_manifest) + + @cached_property + def plugin_manifest(self) -> list[str]: + return sorted(self._plugins.keys()) + + def bind_phase( + self, phase: P, label: str + ) -> "PluginPhaseManager[P]": + return PluginPhaseManager( + self._plugins, self._run, (phase, label) + ) + +@dataclass +class PluginPhaseManager[P: enum.Enum](PluginManager): + _phase: tuple[P, str] + + def runners(self, *, sub_phase_id: str, sub_phase_label: str, sorted_run: bool = False) -> Iterable[PluginPhaseRunner[P]]: + to_iter : Iterable[tuple[str, PipelinePlugin]] = sorted( + self._plugins.items(), key=lambda r: r[0] + ) if sorted_run else self._plugins.items() + for (k,v) in to_iter: + yield PluginPhaseRunner(v, self._run, self._phase, DisplayStrings(sub_phase_id, sub_phase_label), k) + +@asynccontextmanager +async def load_plugins[P: enum.Enum](run: PipelineRun[P, Never]) -> AsyncIterator[PluginManager[P]]: + plugins : dict[str, PipelinePluginLoader] = {} + for ep in importlib.metadata.entry_points(group=PLUGIN_ENTRY_POINT_GROUP): + if ep.name in plugins: + raise RuntimeError(f"Multiple plugins with name: {ep.name}, failing") + loader = ep.load() + if not isinstance(loader, type) or not issubclass(loader, PipelinePluginLoader): + raise RuntimeError(f"Bad plugin declaration: {ep.name}: {ep.module}.{ep.attr} is not a PipelinePluginLoader") + plugins[ep.name] = loader() + async with AsyncExitStack() as stack: + loaded_plugins = { + k: await stack.enter_async_context(v.initialize()) for (k, v) in plugins.items() + } + yield PluginManager( + loaded_plugins, run + ) diff --git a/composer/pipeline/run_tags.py b/composer/pipeline/run_tags.py new file mode 100644 index 00000000..9eb10106 --- /dev/null +++ b/composer/pipeline/run_tags.py @@ -0,0 +1,41 @@ +"""Tag schema for auto-prove pipeline runs. + +Records the inputs that drive the auto-prove cache-key layout into the +run's ``cache_root`` run-data record — written once the design doc +(hence the cache root) is resolved — so downstream tooling +(``cache-autoprove``) can rehydrate every namespace, including the +plugin-suffixed per-component keys, from a run id alone. + +This is the *canonical* shape — both the writer (``cli_pipeline``) and +the reader (``cache-autoprove``) reference this model; round-trip via +``model_dump()`` into the record and ``model_validate(...)`` back. +""" + +from pydantic import BaseModel, Field + + +CACHE_ROOT_RECORD = "cache_root" +"""Run-data key the tags are stored under (``get_run_data(store, run_id, CACHE_ROOT_RECORD)``).""" + + +class AutoProveCacheTags(BaseModel): + cache_root: list[str] | None + """Resolved cache namespace tuple; ``None`` when the run had no ``--cache-ns``.""" + + contract_name: str + + memory_ns: str | None + """Fully-qualified (uid-prefixed) memory namespace, if memory was enabled.""" + + plugins: list[str] = Field(default_factory=list) + """Sorted plugin manifest active for the run — the input to + ``manifest_digest``, which is suffixed onto per-component cache keys.""" + + threat_model_digest: str | None = None + """``Document.to_digest()`` of the threat model, which parameterizes + ``bug_analysis_key``; ``None`` for runs without one.""" + + interactive: bool | None = None + """Whether the run used interactive refinement (selects the ``|refine`` + bug-analysis key variant). ``None`` on records written before this + field existed — readers must probe both variants.""" diff --git a/composer/spec/natspec/pipeline.py b/composer/spec/natspec/pipeline.py index b69e9329..45bd5863 100644 --- a/composer/spec/natspec/pipeline.py +++ b/composer/spec/natspec/pipeline.py @@ -39,7 +39,7 @@ Contract ) from composer.spec.util import string_hash -from composer.spec.prop_inference import run_property_inference +from composer.spec.prop_inference import run_property_inference, CacheablePropertyGenerationInput from composer.spec.types import PropertyFormulation from composer.spec.natspec.interface_gen import generate_interface, DESCRIPTION as INTERFACE_GEN_DESC from composer.spec.natspec.stub_gen import generate_stub @@ -235,8 +235,13 @@ async def _analyze_component(component_idx: int) -> _ComponentBatch | None: feat_ctx, services.env, feat, refinement=conv if interactive else None, extra_input=[ - "For reference, the system document describing the entire application is as follows.", - system_doc.content.to_dict(), + CacheablePropertyGenerationInput( + "certora:system-doc", "generic", "always", lambda cache: [ + "For reference, the system document describing the entire application is as follows.", + system_doc.content.to_dict(cache) + ] + ) + ], max_rounds=services.max_bug_rounds, ), diff --git a/composer/spec/prop_inference.py b/composer/spec/prop_inference.py index 3edc92c0..7f6dd69f 100644 --- a/composer/spec/prop_inference.py +++ b/composer/spec/prop_inference.py @@ -4,18 +4,14 @@ Parameterized by source availability via AnalysisInput tuple. """ -from typing import Any, Callable, NotRequired, override, Literal, Sequence -import re -from difflib import SequenceMatcher +from typing import Any, Callable, NotRequired, Sequence, Literal from pydantic import BaseModel, Field +from dataclasses import dataclass -from langchain_core.tools import BaseTool -from langgraph.types import interrupt, Command -from langchain_core.messages import AnyMessage, SystemMessage, AIMessage, ToolMessage, HumanMessage +from langchain_core.messages import AnyMessage, HumanMessage -from graphcore.graph import MessagesState, FlowInput -from graphcore.tools.schemas import WithImplementation +from graphcore.graph import MessagesState, FlowInput, MessagePayloadType, RawMessageType from composer.input.files import Document from composer.spec.context import WorkflowContext, CacheKey, ComponentGroup @@ -25,19 +21,33 @@ from composer.tools.thinking import RoughDraftState, get_rough_draft_tools from composer.spec.service_host import Sort, ServiceHost from composer.io.conversation import ConversationContextProvider -from composer.spec.refinement import refinement_loop, EndConversation, SyncStateUpdateTool from composer.templates.loader import load_jinja_template -from composer.ui.tool_display import tool_display -from composer.spec.util import string_hash - -from rich.markdown import Markdown -from rich.console import Group -from rich.text import Text +from composer.spec.prop_refinement import user_property_refinement class _BugAnalysisCache(BaseModel): items: list[PropertyFormulation] = Field(description="The security properties you have extracted about the component. Do NOT include any properties " \ "mentioned in (if any were provided to you). If you have not extracted any novel properties, return an empty list") +@dataclass +class PropertyGenerationInputBase: + uid: str + sort: Literal["specific", "generic"] + when: Literal["initial", "always"] + + +@dataclass +class CacheablePropertyGenerationInput(PropertyGenerationInputBase): + provide: Callable[[bool], MessagePayloadType] + +@dataclass +class PropertyGenerationInput(PropertyGenerationInputBase): + input: MessagePayloadType + + @property + def provide(self) -> Callable[[bool], MessagePayloadType]: + return lambda _x: self.input + +type AnyPropertyGenerationInput = PropertyGenerationInput | CacheablePropertyGenerationInput class _AgentRoundResult(_BugAnalysisCache): """ @@ -53,16 +63,25 @@ class _AgentRoundResult(_BugAnalysisCache): class _AgentRoundWithHistory(_AgentRoundResult): agent_conversation: list[AnyMessage] -def bug_analysis_key( - threat_model: Document | None, +def bug_analysis_key_from_digest( + threat_model_digest: str | None, with_refinement: bool ) -> CacheKey[ComponentGroup, _BugAnalysisCache]: base_key = "bug_analysis" if with_refinement: base_key += "|refine" - if threat_model is None: + if threat_model_digest is None: return CacheKey[ComponentGroup, _BugAnalysisCache](base_key) - return CacheKey[ComponentGroup, _BugAnalysisCache](base_key + "-tm-" + threat_model.to_digest()) + return CacheKey[ComponentGroup, _BugAnalysisCache](base_key + "-tm-" + threat_model_digest) + +def bug_analysis_key( + threat_model: Document | None, + with_refinement: bool +) -> CacheKey[ComponentGroup, _BugAnalysisCache]: + return bug_analysis_key_from_digest( + threat_model.to_digest() if threat_model is not None else None, + with_refinement, + ) class _AgentResult(_BugAnalysisCache): final_history: list[AnyMessage] @@ -76,9 +95,6 @@ def agent_round_key( DESCRIPTION = "Property extraction" -class RefinementState(MessagesState): - properties: list[PropertyFormulation] - CERTORA_BACKEND_GUIDANCE: str = """\ You *must* limit your invariants/vectors/properties to those that can be plausibly formally stated or reasoned about in a symbolic reasoning tool @@ -99,149 +115,6 @@ class RefinementState(MessagesState): field not exceeding 2^128 - 1, etc.) """ - -def _get_initial_prompt( - context: ContractComponentInstance, - sort: Sort, - prev_results: list[_AgentRoundResult], - backend_guidance: str, -) -> str: - return load_jinja_template( - "property_analysis_prompt.j2", - context=context, - backend_guidance=backend_guidance, - sort=sort, - prior_properties=prev_results - ) - -@tool_display("Ending conversation...", None) -class Exit(WithImplementation[str]): - """ - Call this when the user has indicated they are happy with the properties you have generated - """ - @override - def run(self) -> str: - return interrupt(EndConversation()) - -@tool_display("Updating requirements", None) -class SetRequirements(SyncStateUpdateTool[list[PropertyFormulation]]): - """ - Called with the new properties as requested by the user - """ - - new_requirements: list[PropertyFormulation] = Field(description="The new requirements after taking into account user feedback.") - - @override - def run(self) -> Command: - return self._update(self.new_requirements) - - -# GitHub-ish dark theme -LINE_DEL = "red on #3a1d1d" -LINE_ADD = "green on #1d3a1d" -WORD_DEL = "bold white on #802020" -WORD_ADD = "bold white on #206020" -DIM = "grey50" - - -def _word_diff(a: str, b: str) -> tuple[Text, Text]: - """Return (minus, plus) Text with word-level highlights.""" - a_toks = re.findall(r"\S+|\s+", a) - b_toks = re.findall(r"\S+|\s+", b) - sm = SequenceMatcher(None, a_toks, b_toks) - - minus, plus = Text(), Text() - for tag, i1, i2, j1, j2 in sm.get_opcodes(): - a_chunk = "".join(a_toks[i1:i2]) - b_chunk = "".join(b_toks[j1:j2]) - if tag == "equal": - minus.append(a_chunk, style=LINE_DEL) - plus.append(b_chunk, style=LINE_ADD) - elif tag == "delete": - minus.append(a_chunk, style=WORD_DEL) - elif tag == "insert": - plus.append(b_chunk, style=WORD_ADD) - elif tag == "replace": - minus.append(a_chunk, style=WORD_DEL) - plus.append(b_chunk, style=WORD_ADD) - return minus, plus - - -def _diff_replace_block(a_block: list[str], b_block: list[str]) -> list[Text]: - """Nested line-level diff inside an outer 'replace' block. - - Only inner 'replace' pairs get word-diffed; inner insert/delete become - whole-line adds/removes. This avoids noisy word-diffs when lines shift. - """ - out: list[Text] = [] - sm = SequenceMatcher(None, a_block, b_block) - - for tag, i1, i2, j1, j2 in sm.get_opcodes(): - if tag == "equal": - for line in a_block[i1:i2]: - out.append(Text(f" {line}", style=DIM)) - elif tag == "delete": - for line in a_block[i1:i2]: - out.append(Text(f"- {line}", style=LINE_DEL)) - elif tag == "insert": - for line in b_block[j1:j2]: - out.append(Text(f"+ {line}", style=LINE_ADD)) - elif tag == "replace": - a_lines, b_lines = a_block[i1:i2], b_block[j1:j2] - n = min(len(a_lines), len(b_lines)) - for k in range(n): - m, p = _word_diff(a_lines[k], b_lines[k]) - out.append(Text("- ", style=LINE_DEL) + m) - out.append(Text("+ ", style=LINE_ADD) + p) - for line in a_lines[n:]: - out.append(Text(f"- {line}", style=LINE_DEL)) - for line in b_lines[n:]: - out.append(Text(f"+ {line}", style=LINE_ADD)) - return out - - -def diff_states(state_a: list[str], state_b: list[str]) -> Group: - sm = SequenceMatcher(None, state_a, state_b) - out: list[Text] = [] - - for tag, i1, i2, j1, j2 in sm.get_opcodes(): - if tag == "equal": - for line in state_a[i1:i2]: - out.append(Text(f" {line}", style=DIM)) - elif tag == "delete": - for line in state_a[i1:i2]: - out.append(Text(f"- {line}", style=LINE_DEL)) - elif tag == "insert": - for line in state_b[j1:j2]: - out.append(Text(f"+ {line}", style=LINE_ADD)) - elif tag == "replace": - out.extend(_diff_replace_block(state_a[i1:i2], state_b[j1:j2])) - - return Group(*out) - -def _front_matter_message(items: Sequence[str | dict]) -> HumanMessage | None: - """Pack reference material (system doc, threat model) into a single - ``HumanMessage`` that lives in ``FlowInput.front_matter`` — placed - between the system prompt and the initial prompt — with the last - content block marked as a prompt-cache breakpoint. - - All blocks before and including the breakpoint are cached and reused - across rounds, so the per-round delta is just the small initial prompt - + the agent's per-round work, not the (large) system doc / threat - model on every call. - """ - if not items: - return None - blocks: list[dict] = [] - for item in items: - if isinstance(item, str): - blocks.append({"type": "text", "text": item}) - else: - blocks.append(dict(item)) - blocks[-1] = {**blocks[-1], "cache_control": {"type": "ephemeral"}} - return HumanMessage(content=[ *blocks ]) - - def _unique_titles_validator( prev: list[_AgentRoundResult], ) -> Callable[[Any, _AgentRoundResult], str | None]: @@ -276,15 +149,78 @@ def validate(_state: Any, result: _AgentRoundResult) -> str | None: return validate +def _partition[S](s: Sequence[S], pred: Callable[[S], bool]) -> tuple[list[S], list[S]]: + a = [] + b = [] + for t in s: + if pred(t): + a.append(t) + else: + b.append(t) + return a, b + +def get_initial_prompt_builder( + extra_inputs: Sequence[AnyPropertyGenerationInput], + component: ContractComponentInstance +) -> Callable[[list[_AgentRoundResult]], MessagePayloadType]: + # Order priority (to facilitate caching) + # Generic-always -> generic-first -> component-always -> component-first -> initial-prompt + # within each group we sort cacheable things last, and then within THOSE groups we sort by UID + generic_input, component_input = _partition(extra_inputs, lambda d: d.sort == "generic") + + generic_always, generic_first = _partition(generic_input, lambda d: d.when == "always") + + component_always, component_first = _partition(component_input, lambda d: d.when == "always") + + def cache_stable_sort(s: list[AnyPropertyGenerationInput]): + cacheable, uncacheable = _partition(s, lambda d: isinstance(d, CacheablePropertyGenerationInput)) + cacheable_sorted = sorted(cacheable, key=lambda d: d.uid) + uncacheable_sorted = sorted(uncacheable, key=lambda d: d.uid) + return [*uncacheable_sorted, *cacheable_sorted] + + def extend(s: list[RawMessageType], to_extend: list[AnyPropertyGenerationInput], cache_last: bool): + for (i, t) in enumerate(to_extend, 1): + r = t.provide(i == len(to_extend) and cache_last) + if isinstance(r, list): + s.extend(r) + else: + s.append(r) + + first_round_prefix : list[RawMessageType] = [] + + extend(first_round_prefix, cache_stable_sort(generic_always), cache_last=True) + + later_round_prefix = first_round_prefix.copy() + + stable_component_always = cache_stable_sort(component_always) + + extend(first_round_prefix, cache_stable_sort(generic_first), cache_last=True) # want to match (S, GA, GF) prefix across components + extend(first_round_prefix, stable_component_always, cache_last=len(generic_first) == 0) # if there is no GF, then the first round is (S, GA, CA), so warm that prefix for later rounds + extend(first_round_prefix, cache_stable_sort(component_first), cache_last=False) + + extend(later_round_prefix, stable_component_always, cache_last=True) + + def renderer(prev_results: list[_AgentRoundResult]) -> MessagePayloadType: + rendered = load_jinja_template( + "property_analysis_prompt.j2", + prior_properties=prev_results, + context=component, + ) + if len(prev_results) == 0: + # first round + return [*first_round_prefix, rendered] + else: + return [*later_round_prefix, rendered] + + return renderer async def _run_bug_round( env: ServiceHost, - component: ContractComponentInstance, - front_matter_items: Sequence[str | dict], ctx: WorkflowContext[_AgentResult], round: int, + prompt_render: Callable[[list[_AgentRoundResult]], MessagePayloadType], prev: list[_AgentRoundResult], - backend_guidance: str, + system_prompt: str ) -> _AgentRoundWithHistory: round_ctx = ctx.child(agent_round_key(round)) if (cached := await round_ctx.cache_get(_AgentRoundWithHistory)) is not None: @@ -305,21 +241,18 @@ class ST(MessagesState, RoughDraftState): ).with_input( BugAnalysisInput ).with_initial_prompt( - _get_initial_prompt(component, env.sort, prev, backend_guidance) + prompt_render(prev) ).with_tools( get_rough_draft_tools(ST) ).with_tools( env.analysis_tools - ).with_sys_prompt_template( - "property_analysis_system_prompt.j2", sort=env.sort + ).with_sys_prompt( + system_prompt ).compile_async() flow_input: BugAnalysisInput = BugAnalysisInput( input=[], memory=None, did_read=False, ) - front = _front_matter_message(front_matter_items) - if front is not None: - flow_input["front_matter"] = [front] r = await run_to_completion( d, @@ -342,32 +275,27 @@ async def _run_bug_analysis_inner( agent_component_analysis: WorkflowContext[_AgentResult], env: ServiceHost, component: ContractComponentInstance, - extra_input: Sequence[str | dict], - threat_model: Document | None, + extra_input: Sequence[AnyPropertyGenerationInput], max_rounds: int, backend_guidance: str, ) -> _AgentResult: if (cached := await agent_component_analysis.cache_get(_AgentResult)) is not None: return cached - - front_matter_items: list[str | dict] = [ *extra_input ] - if threat_model is not None: - front_matter_items.extend([ - "In addition, a coworker has already written a 'threat model' for this application, which may include vulnerabilities/issues that" - "are common in this type of application. This threat model is written for the entire application (not just the component you are analyzing) " - "so some of the issues/vulnerabilities/attacks may not be relevant to your analysis. Do *NOT* overfit to this threat model; carefully " - "analyze what content of the provided threat model is worth considering vs out of scope. Further, this threat model is just a starting point, " - "you should ALSO look for threats *not* mentioned in this document.", - threat_model.to_dict(with_cache=True) - ]) + + initial_prompt_builder = get_initial_prompt_builder( + extra_inputs=extra_input, component=component + ) prev_rounds : list[_AgentRoundResult] = [] last_round_convo : list[AnyMessage] | None = None + system_prompt = load_jinja_template( + "property_analysis_system_prompt.j2", sort=env.sort, backend_guidance=backend_guidance + ) + for i in range(0, max_rounds): next_result = await _run_bug_round( - env, component, front_matter_items, agent_component_analysis, i, prev_rounds, - backend_guidance, + env, agent_component_analysis, i, initial_prompt_builder, prev_rounds, system_prompt ) if len(next_result.items) == 0: assert last_round_convo is not None @@ -390,7 +318,7 @@ async def run_property_inference( ctx: WorkflowContext[ComponentGroup], env: ServiceHost, component: ContractComponentInstance, - extra_input : Sequence[str | dict] = tuple(), + extra_input : Sequence[AnyPropertyGenerationInput] = tuple(), threat_model: Document | None = None, refinement: ConversationContextProvider | None = None, max_rounds: int = 3, @@ -411,12 +339,28 @@ async def run_property_inference( if (cached := await component_analysis.cache_get(_BugAnalysisCache)) is not None: return cached.items + actual_extra_input = [ + *extra_input + ] + if threat_model is not None: + actual_extra_input.append(CacheablePropertyGenerationInput( + "certora:thread_model", "generic", "always", + provide=lambda cache: [ + "In addition, a coworker has already written a 'threat model' for this application, which may include vulnerabilities/issues that" + "are common in this type of application. This threat model is written for the entire application (not just the component you are analyzing) " + "so some of the issues/vulnerabilities/attacks may not be relevant to your analysis. Do *NOT* overfit to this threat model; carefully " + "analyze what content of the provided threat model is worth considering vs out of scope. Further, this threat model is just a starting point, " + "you should ALSO look for threats *not* mentioned in this document.", + threat_model.to_dict(with_cache=cache) + ] + )) + + agent_attempt = await _run_bug_analysis_inner( component_analysis.child(AGENT_RESULT_KEY), env, component, - extra_input, - threat_model, + actual_extra_input, max_rounds=max_rounds, backend_guidance=backend_guidance, ) @@ -425,67 +369,8 @@ async def run_property_inference( await component_analysis.cache_put(_BugAnalysisCache(items=to_ret)) return to_ret - msg_history = agent_attempt.final_history - assert isinstance(msg_history[0], SystemMessage) and isinstance(msg_history[-1], ToolMessage) - import uuid - edited_history = [ - SystemMessage(load_jinja_template("bug_refinement_chat_system_prompt.j2", sort=env.sort)), - *msg_history[1:], - AIMessage("", id=uuid.uuid4().hex) - ] - - def sort_to_string( - s: Literal["attack_vector", "invariant", "safety_property"] - ) -> str: - match s: - case "attack_vector": - return "Attack Vector" - case "invariant": - return "Invariant" - case "safety_property": - return "Safety Property" - - def property_as_text( - prop: PropertyFormulation - ) -> str: - return f"* [{sort_to_string(prop.sort)}] {prop.description}" - - def property_as_md( - prop: PropertyFormulation - ) -> str: - sort_str = sort_to_string(prop.sort) - return f"* \\[{sort_str}\\] {prop.description}" - - def properties_as_text( - l: list[PropertyFormulation] - ) -> list[str]: - return [ property_as_text(p) for p in l ] - - def properties_as_md( - l: list[PropertyFormulation] - ) -> list[str]: - return [ property_as_md(p) for p in l ] - - def render_properties_as_md( - l: list[PropertyFormulation] - ) -> Markdown: - md = "## Current Properties\n" - return Markdown(md + "\n".join(properties_as_md(l))) - - async with refinement(render_properties_as_md(agent_attempt.items)) as client: - res = await refinement_loop( - llm=env.llm_heavy(), - client=client, - init_messages=edited_history, - init_data=agent_attempt.items, - tools=[*env.analysis_tools, Exit.as_tool("finalize_properties"), SetRequirements.as_tool("update_requirements")], - state_renderer=render_properties_as_md, - diff_renderer=lambda a, b: \ - Group( - Text("Properties changed"), - diff_states(properties_as_text(a), properties_as_text(b)) - ) - ) - to_ret = res["extra_data"] - await component_analysis.cache_put(_BugAnalysisCache(items = to_ret)) - return to_ret + refined_props = await user_property_refinement( + env, agent_attempt, refinement + ) + await component_analysis.cache_put(_BugAnalysisCache(items = refined_props)) + return refined_props diff --git a/composer/spec/prop_refinement.py b/composer/spec/prop_refinement.py new file mode 100644 index 00000000..f14d2d11 --- /dev/null +++ b/composer/spec/prop_refinement.py @@ -0,0 +1,207 @@ + +from typing import override, Literal, Sequence, Protocol +import re +from difflib import SequenceMatcher +from pydantic import Field + +from langgraph.types import interrupt, Command + +from langchain_core.messages import AnyMessage, SystemMessage, AIMessage, ToolMessage + +from graphcore.tools.schemas import WithImplementation + +from composer.spec.types import PropertyFormulation +from composer.spec.service_host import ServiceHost +from composer.io.conversation import ConversationContextProvider +from composer.spec.refinement import refinement_loop, EndConversation, SyncStateUpdateTool +from composer.templates.loader import load_jinja_template +from composer.ui.tool_display import tool_display + +from rich.markdown import Markdown +from rich.console import Group +from rich.text import Text + +class AgenticAttempt(Protocol): + @property + def final_history(self) -> Sequence[AnyMessage]: + ... + + @property + def items(self) -> list[PropertyFormulation]: + ... + + +@tool_display("Ending conversation...", None) +class Exit(WithImplementation[str]): + """ + Call this when the user has indicated they are happy with the properties you have generated + """ + @override + def run(self) -> str: + return interrupt(EndConversation()) + +@tool_display("Updating requirements", None) +class SetRequirements(SyncStateUpdateTool[list[PropertyFormulation]]): + """ + Called with the new properties as requested by the user + """ + + new_requirements: list[PropertyFormulation] = Field(description="The new requirements after taking into account user feedback.") + + @override + def run(self) -> Command: + return self._update(self.new_requirements) + + +# GitHub-ish dark theme +LINE_DEL = "red on #3a1d1d" +LINE_ADD = "green on #1d3a1d" +WORD_DEL = "bold white on #802020" +WORD_ADD = "bold white on #206020" +DIM = "grey50" + + +def _word_diff(a: str, b: str) -> tuple[Text, Text]: + """Return (minus, plus) Text with word-level highlights.""" + a_toks = re.findall(r"\S+|\s+", a) + b_toks = re.findall(r"\S+|\s+", b) + sm = SequenceMatcher(None, a_toks, b_toks) + + minus, plus = Text(), Text() + for tag, i1, i2, j1, j2 in sm.get_opcodes(): + a_chunk = "".join(a_toks[i1:i2]) + b_chunk = "".join(b_toks[j1:j2]) + if tag == "equal": + minus.append(a_chunk, style=LINE_DEL) + plus.append(b_chunk, style=LINE_ADD) + elif tag == "delete": + minus.append(a_chunk, style=WORD_DEL) + elif tag == "insert": + plus.append(b_chunk, style=WORD_ADD) + elif tag == "replace": + minus.append(a_chunk, style=WORD_DEL) + plus.append(b_chunk, style=WORD_ADD) + return minus, plus + + +def _diff_replace_block(a_block: list[str], b_block: list[str]) -> list[Text]: + """Nested line-level diff inside an outer 'replace' block. + + Only inner 'replace' pairs get word-diffed; inner insert/delete become + whole-line adds/removes. This avoids noisy word-diffs when lines shift. + """ + out: list[Text] = [] + sm = SequenceMatcher(None, a_block, b_block) + + for tag, i1, i2, j1, j2 in sm.get_opcodes(): + if tag == "equal": + for line in a_block[i1:i2]: + out.append(Text(f" {line}", style=DIM)) + elif tag == "delete": + for line in a_block[i1:i2]: + out.append(Text(f"- {line}", style=LINE_DEL)) + elif tag == "insert": + for line in b_block[j1:j2]: + out.append(Text(f"+ {line}", style=LINE_ADD)) + elif tag == "replace": + a_lines, b_lines = a_block[i1:i2], b_block[j1:j2] + n = min(len(a_lines), len(b_lines)) + for k in range(n): + m, p = _word_diff(a_lines[k], b_lines[k]) + out.append(Text("- ", style=LINE_DEL) + m) + out.append(Text("+ ", style=LINE_ADD) + p) + for line in a_lines[n:]: + out.append(Text(f"- {line}", style=LINE_DEL)) + for line in b_lines[n:]: + out.append(Text(f"+ {line}", style=LINE_ADD)) + return out + + +def diff_states(state_a: list[str], state_b: list[str]) -> Group: + sm = SequenceMatcher(None, state_a, state_b) + out: list[Text] = [] + + for tag, i1, i2, j1, j2 in sm.get_opcodes(): + if tag == "equal": + for line in state_a[i1:i2]: + out.append(Text(f" {line}", style=DIM)) + elif tag == "delete": + for line in state_a[i1:i2]: + out.append(Text(f"- {line}", style=LINE_DEL)) + elif tag == "insert": + for line in state_b[j1:j2]: + out.append(Text(f"+ {line}", style=LINE_ADD)) + elif tag == "replace": + out.extend(_diff_replace_block(state_a[i1:i2], state_b[j1:j2])) + + return Group(*out) + +def sort_to_string( + s: Literal["attack_vector", "invariant", "safety_property"] +) -> str: + match s: + case "attack_vector": + return "Attack Vector" + case "invariant": + return "Invariant" + case "safety_property": + return "Safety Property" + +def property_as_text( + prop: PropertyFormulation +) -> str: + return f"* [{sort_to_string(prop.sort)}] {prop.description}" + +def property_as_md( + prop: PropertyFormulation +) -> str: + sort_str = sort_to_string(prop.sort) + return f"* \\[{sort_str}\\] {prop.description}" + +def properties_as_text( + l: list[PropertyFormulation] +) -> list[str]: + return [ property_as_text(p) for p in l ] + +def properties_as_md( + l: list[PropertyFormulation] +) -> list[str]: + return [ property_as_md(p) for p in l ] + +def render_properties_as_md( + l: list[PropertyFormulation] +) -> Markdown: + md = "## Current Properties\n" + return Markdown(md + "\n".join(properties_as_md(l))) + + +async def user_property_refinement( + env: ServiceHost, + agent_attempt: AgenticAttempt, + refinement: ConversationContextProvider +) -> list[PropertyFormulation]: + msg_history = agent_attempt.final_history + assert isinstance(msg_history[0], SystemMessage) and isinstance(msg_history[-1], ToolMessage) + import uuid + edited_history = [ + SystemMessage(load_jinja_template("bug_refinement_chat_system_prompt.j2", sort=env.sort)), + *msg_history[1:], + AIMessage("", id=uuid.uuid4().hex) + ] + + async with refinement(render_properties_as_md(agent_attempt.items)) as client: + res = await refinement_loop( + llm=env.llm_heavy(), + client=client, + init_messages=edited_history, + init_data=agent_attempt.items, + tools=[*env.analysis_tools, Exit.as_tool("finalize_properties"), SetRequirements.as_tool("update_requirements")], + state_renderer=render_properties_as_md, + diff_renderer=lambda a, b: \ + Group( + Text("Properties changed"), + diff_states(properties_as_text(a), properties_as_text(b)) + ) + ) + to_ret = res["extra_data"] + return to_ret diff --git a/composer/spec/service_host.py b/composer/spec/service_host.py index d899c098..3bfb4df5 100644 --- a/composer/spec/service_host.py +++ b/composer/spec/service_host.py @@ -24,10 +24,10 @@ """ from typing import Literal, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from langgraph.types import Checkpointer -from graphcore.graph import Builder +from graphcore.graph import Builder, TemplateLoader from langchain_core.language_models.chat_models import BaseChatModel as LLM from langchain_core.tools import BaseTool @@ -54,6 +54,7 @@ class ModelProvider: heavy_model: CoreModelProvider lite_model: CoreModelProvider checkpointer: Checkpointer + _loader: TemplateLoader | None = field(default=None) def llm_heavy( self, *, cache_level: CacheLevel = CacheLevel.SHORT, disable_thinking: bool = False @@ -78,10 +79,10 @@ def builder_lite( def _builder(self, llm: LLM) -> Builder[None, None, None]: return Builder[None, None, None]().with_llm( llm - ).with_loader(load_jinja_template).with_checkpointer(self.checkpointer) + ).with_loader(self._loader or load_jinja_template).with_checkpointer(self.checkpointer) -@dataclass +@dataclass(frozen=True) class PureServiceHost: """``ServiceHost`` without source tools — the pre-stub natspec phase uses this directly. Call :meth:`bind_source_tools` once a usable @@ -123,7 +124,7 @@ def bind_source_tools(self, tools: Sequence[BaseTool]) -> "ServiceHost": ) -@dataclass +@dataclass(frozen=True) class ServiceHost(PureServiceHost): """Concrete per-agent environment carrying both tool tuples and the workflow ``sort``. The single env type for all per-agent callsites — diff --git a/composer/spec/source/autosetup.py b/composer/spec/source/autosetup.py index 1c501b16..338c064b 100644 --- a/composer/spec/source/autosetup.py +++ b/composer/spec/source/autosetup.py @@ -167,7 +167,7 @@ def log_complete(self, returncode: int): _drain(proc.stdout, cb.log_stdout, logging.INFO), _drain(proc.stderr, stderr_lines.append, logging.ERROR), ) - except Exception: + except BaseException: _logger.exception("Error while draining AutoSetup subprocess output") # A sink raising (or cancellation) would leave the child running with # its pipes half-drained; kill it so wait() can reap. diff --git a/composer/spec/types.py b/composer/spec/types.py index 368bc6b0..01641693 100644 --- a/composer/spec/types.py +++ b/composer/spec/types.py @@ -56,10 +56,24 @@ def artifact_text(self) -> str: ... invariant. Shared so every layer (inference, report, grouping) addresses the same vocabulary instead of redeclaring the literal.""" -class PropertyFormulation(BaseModel): +class UntitledPropertyFormulation(BaseModel): + sort: PropertyType = Field(description="The type of property you are describing.") + description: str = Field(description="The description of the property") + + @property + def sort_description(self) -> str: + match self.sort: + case "attack_vector": + return "Attack Vector" + case "invariant": + return "Invariant" + case "safety_property": + return "Safety Property" + +class PropertyFormulation(UntitledPropertyFormulation): """ A property or invariant that must hold for the component """ title: str = Field(description="A short, descriptive snake_case identifier for the property (e.g. 'total_supply_preserved'). Must be unique within the batch of properties.") - sort: PropertyType = Field(description="The type of property you are describing.") - description: str = Field(description="The description of the property") + + diff --git a/composer/templates/loader.py b/composer/templates/loader.py index d9708d7e..d38f3b43 100644 --- a/composer/templates/loader.py +++ b/composer/templates/loader.py @@ -1,9 +1,13 @@ from typing import Any from jinja2 import Environment, FileSystemLoader +from jinja2.loaders import BaseLoader import pathlib +from graphcore.graph import TemplateLoader script_dir = pathlib.Path(__file__).parent +base_loader = FileSystemLoader(script_dir) + def _autoescape(template_name: str | None) -> bool: # HTML templates (``*.html.j2``) must autoescape interpolated values; prompt templates @@ -11,7 +15,15 @@ def _autoescape(template_name: str | None) -> bool: return template_name is not None and template_name.endswith(".html.j2") -env = Environment(loader=FileSystemLoader(script_dir), autoescape=_autoescape) +env = Environment(loader=base_loader, autoescape=_autoescape) + +def make_loader(jinja_loader: BaseLoader) -> TemplateLoader: + my_env = Environment(loader=jinja_loader, autoescape=_autoescape) + def load(template_name: str, **kwargs: Any) -> str: + template = my_env.get_template(template_name) + return template.render(**kwargs) + return load + def load_jinja_template(template_name: str, **kwargs: Any) -> str: """Load and render a Jinja template from the script directory""" diff --git a/composer/templates/property_analysis_prompt.j2 b/composer/templates/property_analysis_prompt.j2 index 88f264b7..22ea8898 100644 --- a/composer/templates/property_analysis_prompt.j2 +++ b/composer/templates/property_analysis_prompt.j2 @@ -38,66 +38,6 @@ answer. {% endif %} - -{% if sort != "greenfield" %} -Input: The implementation of the application, and a description of a specific component of that application -{% else %} -Input: A design document describing the application, and a description of a specific component of that application -{% endif %} -Output: A list of "security properties" - -Follow these steps exactly: - -## Step 1 - -{% if sort != "greenfield" %} -Analyze the provided implementation and component description and formulate a set of security properties that are PLAUSIBLE -for this component. Each security property description must be precise and concise. -{% else %} -Analyze the provided design document and component description and formulate a set of security properties that are PLAUSIBLE -for this component. Each security property description must be precise and concise. -{% endif %} - -Security properties fall into one of three categories: -1. Invariants are representational invariants that should always hold from the point of construction of the smart contract. - Good example: "The sum of all user balances should equal the total supply" - Good example: "The pool's balance in backing asset should be greater than or equal to the total supply" - Bad example: "The state fields should be correct" (overly broad) - Bad example: "It should be possible to withdraw funds" (not an invariant, also overly broad) -2. Safety properties: A safety property is a concrete statement about what should not be possible in a correctly implemented protocol - Good example: "A user without the admin role should not be able to approve other admins" - Good example: "A user should not be able to change anyone's balance but their own" - Bad example: "A user should not be able to hack the protocol" (overly broad) -3. Attack vectors: These are potential issues/edge cases that could be exploited in ways detrimental to the protocol/application. - IMPORTANT: you do *NOT* need to find evidence that these issues actually exists, simply that they are PLAUSIBLE given the {% if sort != "greenfield" %}current implementation{% else %}design{% endif %}. - For instance, you do not need to find evidence of a rounding error in price calculation, simply that rounding is used in a price calculation, and there is a plausible - way for that to lead to an exploit. - Good example: "A malicious actor could manipulate a price oracle to unbalance the pool, allowing free minting" - Good example: "There could be a rounding discrepancy between the swap function and the previewSwap function which could provide arbitrage opportunities" - Bad example: "There could be a rounding error, and that would be bad somehow" (does not lead to a concrete issue/exploit) - Bad example: "The protocol could be hacked" (overly broad) - Bad example: "The underlying EVM consensus layer could be compromised allowing attackers to set arbitrary storage states" (implausible) - -{{ backend_guidance }} - -Note: Safety properties can (and should) include properties describing the intended, normal behavior of the smart contract. In other words, the properties -described do not necessarily need to have some *immediate* security implication to be interesting. For example, a safety property which states a user depositing liquidity -receives an appropriate amount of liquidity tokens is both valid and interesting. - -NB: you do **NOT** need to be able to formalize the properties yourself, that is out of scope for your task. Simply limit your analysis to those properties which can be -reasonably formalized in the downstream verification tool. - -## Step 2 -Write a rough draft list of your properties using the provided "write_rough_draft" tool. - -## Step 3 -Review your rough draft. For each property/invariant in your rough draft, make sure it meets the guidelines and definitions described in step 1. Pay particular -attention for properties/invariants explicitly described as being uninteresting or impossible to verify in the downstream verification tool. - -## Step 4 -Output the results of your analysis using the result tool. - - {% if sort != "greenfield" %} diff --git a/composer/templates/property_analysis_system_prompt.j2 b/composer/templates/property_analysis_system_prompt.j2 index f0e5dcc2..b1d5b6ad 100644 --- a/composer/templates/property_analysis_system_prompt.j2 +++ b/composer/templates/property_analysis_system_prompt.j2 @@ -9,6 +9,66 @@ component. The work happens in a multi-round loop: each round produces only the NEW properties not surfaced by prior rounds, plus a written reasoning record that later rounds and a human reviewer will read. +# Task + +{% if sort != "greenfield" %} +You will be provided the implementation of the application, and a description of a specific component of that application. +{% else %} +You will be provided a design document describing the application, and a description of a specific component of that application. +{% endif %} +Your output is a list of "security properties". + +Follow these steps exactly: + +## Step 1 + +{% if sort != "greenfield" %} +Analyze the provided implementation and component description and formulate a set of security properties that are PLAUSIBLE +for this component. Each security property description must be precise and concise. +{% else %} +Analyze the provided design document and component description and formulate a set of security properties that are PLAUSIBLE +for this component. Each security property description must be precise and concise. +{% endif %} + +Security properties fall into one of three categories: +1. Invariants are representational invariants that should always hold from the point of construction of the smart contract. + Good example: "The sum of all user balances should equal the total supply" + Good example: "The pool's balance in backing asset should be greater than or equal to the total supply" + Bad example: "The state fields should be correct" (overly broad) + Bad example: "It should be possible to withdraw funds" (not an invariant, also overly broad) +2. Safety properties: A safety property is a concrete statement about what should not be possible in a correctly implemented protocol + Good example: "A user without the admin role should not be able to approve other admins" + Good example: "A user should not be able to change anyone's balance but their own" + Bad example: "A user should not be able to hack the protocol" (overly broad) +3. Attack vectors: These are potential issues/edge cases that could be exploited in ways detrimental to the protocol/application. + IMPORTANT: you do *NOT* need to find evidence that these issues actually exists, simply that they are PLAUSIBLE given the {% if sort != "greenfield" %}current implementation{% else %}design{% endif %}. + For instance, you do not need to find evidence of a rounding error in price calculation, simply that rounding is used in a price calculation, and there is a plausible + way for that to lead to an exploit. + Good example: "A malicious actor could manipulate a price oracle to unbalance the pool, allowing free minting" + Good example: "There could be a rounding discrepancy between the swap function and the previewSwap function which could provide arbitrage opportunities" + Bad example: "There could be a rounding error, and that would be bad somehow" (does not lead to a concrete issue/exploit) + Bad example: "The protocol could be hacked" (overly broad) + Bad example: "The underlying EVM consensus layer could be compromised allowing attackers to set arbitrary storage states" (implausible) + +{{ backend_guidance }} + +Note: Safety properties can (and should) include properties describing the intended, normal behavior of the smart contract. In other words, the properties +described do not necessarily need to have some *immediate* security implication to be interesting. For example, a safety property which states a user depositing liquidity +receives an appropriate amount of liquidity tokens is both valid and interesting. + +NB: you do **NOT** need to be able to formalize the properties yourself, that is out of scope for your task. Simply limit your analysis to those properties which can be +reasonably formalized in the downstream verification tool. + +## Step 2 +Write a rough draft list of your properties using the provided "write_rough_draft" tool. + +## Step 3 +Review your rough draft. For each property/invariant in your rough draft, make sure it meets the guidelines and definitions described in step 1. Pay particular +attention for properties/invariants explicitly described as being uninteresting or impossible to verify in the downstream verification tool. + +## Step 4 +Output the results of your analysis using the result tool. + # Behavior ## Quality over quantity — and "nothing new" is the right answer when it's true @@ -64,8 +124,8 @@ resolving ambiguity in the inputs, favor self-consistency. {% with draft_subject = "your property list" %} {% include "rough_draft_protocol.j2" %} {% endwith %} -When reviewing your draft, verify each property against the criteria in your -task prompt — formal-verifiability, non-triviality, clear violation +When reviewing your draft, verify each property against the criteria in the +Task section above — formal-verifiability, non-triviality, clear violation consequence, defensibility in your `reasoning`. {% if sort != "greenfield" %} diff --git a/graphcore b/graphcore index 8b9d3203..ba8ddccd 160000 --- a/graphcore +++ b/graphcore @@ -1 +1 @@ -Subproject commit 8b9d320316a2cfbfa775da0a961501539a4124f0 +Subproject commit ba8ddccd2eaccd0aa765aa22179310492c40d6b2 diff --git a/pyproject.toml b/pyproject.toml index 05832d4c..bba48bcd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -134,6 +134,7 @@ tui-foundry = "composer.cli.tui_foundry:main" autoprove-report-render = "composer.spec.source.report.render:main" tui-natspec = "composer.cli.tui_pipeline:main" cache-natspec = "composer.cli.cache_natspec:main" +cache-autoprove = "composer.cli.cache_autoprove:main" console-codegen = "composer.cli.console_codegen:main" tui-codegen = "composer.cli.tui_codegen:main" # certora_autosetup entry points diff --git a/tests/test_autoprove_report.py b/tests/test_autoprove_report.py index 9849699a..edae052a 100644 --- a/tests/test_autoprove_report.py +++ b/tests/test_autoprove_report.py @@ -88,17 +88,10 @@ def _gen(mapping: dict[str, list[str]] | None = None, ) -def _input(name, unit_file, props, result: GeneratedCVL | None, link : str | None="L1") -> ReportComponentInput[GeneratedCVL]: - """``link`` is the result's prover run link (``GeneratedCVL.final_link``); the prover fetcher - keys its verdicts off it. ``None`` (or a ``None`` result) means no run link, so no verdicts.""" - return ReportComponentInput( - name=name, - props=props, - formalized=Delivered( - deliverable=pathlib.Path(unit_file), - result=result.model_copy(update={"final_link": link}) - ) if result is not None else None - ) +def _input(name, unit_file, props, result: GeneratedCVL | None, + link: str | None = "L1") -> ReportComponentInput[GeneratedCVL]: + return ReportComponentInput(name=name, props=props, + formalized=Delivered(result, unit_file) if result is not None else None) def _fp(component, title, refs, desc="d", sort: PropertyType = "safety_property") -> FormalizedProperty: