From b4d01a73fc11b9d62c4ed36325fa5651faf525cf Mon Sep 17 00:00:00 2001 From: John Toman Date: Tue, 30 Jun 2026 11:06:50 -0700 Subject: [PATCH 01/10] Wires in monitoring Adds nudges to the agent to give up/try something new; refuse to run the prover on "let's just try a timeout again" case --- composer/prover/core.py | 20 +++- composer/spec/source/author.py | 5 + composer/spec/source/monitor.py | 15 +++ composer/spec/source/prover.py | 172 ++++++++++++++++++++++++++------ graphcore | 2 +- 5 files changed, 180 insertions(+), 34 deletions(-) create mode 100644 composer/spec/source/monitor.py diff --git a/composer/prover/core.py b/composer/prover/core.py index cf1a492b..980e3c66 100644 --- a/composer/prover/core.py +++ b/composer/prover/core.py @@ -45,7 +45,7 @@ from composer.prover.analysis import analyze_cex_raw from composer.prover.cloud import CloudJobError, cloud_results -from composer.prover.ptypes import RuleResult +from composer.prover.ptypes import RuleResult, RulePath, StatusCodes from composer.prover.results import read_and_format_run_result from composer.templates.loader import load_jinja_template from composer.prover.prover_protocol import ProverResult @@ -118,10 +118,20 @@ class ProverReport: ``link`` is the prover run's URL (cloud) or local results directory. """ - rule_status: dict[str, bool] + raw_rule_status: dict[RulePath, StatusCodes] + result_str: str link: str + @property + def rule_status(self) -> dict[str, bool]: + to_ret = {} + for (k, v) in self.raw_rule_status.items(): + if k.rule in to_ret and not to_ret[k.rule]: + continue + to_ret[k.rule] = v == "VERIFIED" + return to_ret + @property def all_verified(self) -> bool: return all(self.rule_status.values()) @@ -528,8 +538,12 @@ async def run_prover( continue prover_report[rule_name] = i.status == "VERIFIED" + raw_rule_results : dict[RulePath, StatusCodes] = { + k.path: k.status for k in parsed.values() + } + return ProverReport( - rule_status=prover_report, + raw_rule_status=raw_rule_results, result_str=result_str, link=run_result["link"], ) diff --git a/composer/spec/source/author.py b/composer/spec/source/author.py index d0bf3cd3..269854f1 100644 --- a/composer/spec/source/author.py +++ b/composer/spec/source/author.py @@ -32,6 +32,7 @@ from langgraph.types import Command from composer.spec.feedback import property_feedback_judge, FeedbackTemplate from composer.ui.tool_display import tool_display +from .monitor import monitor from graphcore.graph import FlowInput @@ -376,6 +377,8 @@ async def batch_cvl_generation( [prover_tool, ExpectRulePassage.as_tool("expect_rule_passage"), ExpectRuleFailure.as_tool("expect_rule_failure"), GiveUpTool.as_tool("give_up"), PublishResultTool.as_tool("result"), ctx.get_memory_tool()] ).with_state( SourceCVLGenerationState + ).with_monitor( + monitor ).with_output_key( "result" ).with_input( @@ -411,6 +414,8 @@ async def batch_cvl_generation( property_rules=[], validations={}, failed=None, + prover_history=[], + reminders_channel=[] ) ) diff --git a/composer/spec/source/monitor.py b/composer/spec/source/monitor.py new file mode 100644 index 00000000..3ed56d74 --- /dev/null +++ b/composer/spec/source/monitor.py @@ -0,0 +1,15 @@ +from langgraph.config import get_config +from langchain_core.messages import HumanMessage +from .author import SourceCVLGenerationState +from graphcore.graph import MonitorReturn +from composer.prover.ptypes import StatusCodes + +def monitor( + curr_state: SourceCVLGenerationState +) -> MonitorReturn: + if not curr_state["reminders_channel"]: + return None, None + + return [HumanMessage(f"{'\n'.join(curr_state["reminders_channel"])}")], { + "reminders_channel": None + } diff --git a/composer/spec/source/prover.py b/composer/spec/source/prover.py index 63b96422..a5424c7b 100644 --- a/composer/spec/source/prover.py +++ b/composer/spec/source/prover.py @@ -14,22 +14,24 @@ import time from contextlib import contextmanager, nullcontext from pathlib import Path -from typing import Annotated, Callable, Iterator, override, AsyncContextManager +from typing import Annotated, Callable, Iterator, override, AsyncContextManager, Sequence, Literal from typing_extensions import TypedDict, NotRequired from langchain_core.tools import InjectedToolCallId, tool, BaseTool +from langchain_core.messages import AIMessage from langgraph.prebuilt import InjectedState -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, Discriminator from langgraph.config import get_stream_writer from langgraph.types import Command -from composer.prover.ptypes import RuleResult +from composer.prover.ptypes import RuleResult, RulePath from graphcore.graph import LLM from composer.prover.core import ( - ProverOptions, ProverCallbacks, run_prover, DefaultCexHandler + ProverOptions, run_prover, DefaultCexHandler ) from composer.prover.callbacks import ProverEventCallbacks +from composer.prover.ptypes import StatusCodes from composer.ui.tool_display import tool_display from composer.diagnostics.stream import ( ProverOutputEvent, CloudPollingEvent, RuleAnalysisResult, @@ -40,6 +42,7 @@ from graphcore.graph import tool_state_update from composer.spec.util import temp_certora_file from composer.spec.gen_types import CERTORA_DIR, SPECS_DIR +from composer.spec.util import string_hash _logger = logging.getLogger("composer.prover") @@ -76,6 +79,23 @@ def _merge_rule_skips(left: dict[str, str], right: dict[str, str]) -> dict[str, to_ret[k] = v return to_ret +class ProverRunLog(TypedDict): + tool_call_id: str + prover_results: list[tuple[RulePath, StatusCodes]] + spec_digest: str + rules: list[str] | None + sort: Literal["run"] + +class NagMarker(TypedDict): + nagged_rules: list[RulePath] + sort: Literal["nag"] + +type ProverHistoryItem = Annotated[ProverRunLog | NagMarker, Discriminator("sort")] + +def _merge_prover_history(left: list[ProverHistoryItem], right: list[ProverHistoryItem]) -> list[ProverHistoryItem]: + to_ret = left.copy() + to_ret.extend(right) + return to_ret class ProverStateExtra(TypedDict): rule_skips: Annotated[dict[str, str], _merge_rule_skips] @@ -86,6 +106,9 @@ class ProverStateExtra(TypedDict): # Basename the spec is materialized/persisted under (e.g. "autospec_"). # NotRequired so other ProverStateExtra injectors (e.g. config_edit) needn't set it. spec_stem: NotRequired[str] + prover_history: Annotated[list[ProverHistoryItem], _merge_prover_history] + reminders_channel: list[str] + type ProverEvents = CEXAnalysisStart | CloudPollingEvent | ProverOutputEvent | RuleAnalysisResult | ProverRun | ProverLink | ProverResult @@ -201,15 +224,8 @@ def tmp_spec( def _prover_sem(cloud: bool) -> AsyncContextManager[None]: if not cloud: return asyncio.Semaphore(1) - - class ToRet(): - async def __aenter__(self): - return - - async def __aexit__(self, exc_type, exc, tb): - return - - return ToRet() + else: + return nullcontext() def get_prover_tool( llm: LLM, @@ -235,8 +251,25 @@ async def verify_spec( state: Annotated[StateWithSkips, InjectedState], rules: list[str] | None = None ) -> str | Command: + last_msg = state["messages"][-1] + if isinstance(last_msg, AIMessage) and any( + i["id"] != tool_call_id for i in last_msg.tool_calls + ): + return "Cannot call the verify_spec tool in parallel with other tool calls. verify_spec must be the only tool you call in a turn" + if state["curr_spec"] is None: return "Specification not yet put on VFS" + + spec_hash = string_hash( + state["curr_spec"] + ) + + if len(state['prover_history']) > 0: + last_item = state["prover_history"][-1] + if last_item["sort"] == "run" and any(i == "TIMEOUT" for (_,i) in last_item) and last_item["spec_digest"] == spec_hash: + return "Refusing to re-run prover on identical spec with a known TIMEOUT result; timeouts are not transient " \ + "errors and will not go away by re-running the tool." + conf = state["config"] # With a seeded stem, name the spec/conf after it (so on-disk names match the # dump) under a lock; else fall back to unique uid names (no lock needed). @@ -255,7 +288,7 @@ async def verify_spec( summary = get_run_summary() component = (spec_stem or main_contract).removeprefix("autospec_") - iteration = summary.prover_total_calls + 1 + iteration = len(state["prover_history"]) + 1 config["msg"] = f"{component} #{iteration}" with temp_certora_file( @@ -276,22 +309,101 @@ async def verify_spec( DefaultCexHandler(llm, state, summarization_threshold=10) ) - if isinstance(result, str): - return result - all_verified = True - for (r, stat) in result.rule_status.items(): - if r in state["rule_skips"]: - continue - if not stat: - all_verified = False - break - if rules is None and all_verified: - return tool_state_update( - tool_call_id=tool_call_id, content=result.result_str, - prover_link=result.link, validations=stamper(state), - ) - return tool_state_update( - tool_call_id=tool_call_id, content=result.result_str, prover_link=result.link + if isinstance(result, str): + return result + + all_verified = True + for (r, stat) in result.rule_status.items(): + if r in state["rule_skips"]: + continue + if not stat: + all_verified = False + break + + spec_hash = string_hash( + state["curr_spec"] + ) + + stuck_rules = { + k: v for (k,v) in result.raw_rule_status.items() if v in ("TIMEOUT", "ERROR", "SANITY_FAILED") + } + + stuck_count = { + k: 1 for k in stuck_rules.keys() + } + + to_warn : set[RulePath] = set() + + known_tc_ids = { + l["id"] + for msg in state["messages"] if isinstance(msg, AIMessage) + for l in msg.tool_calls if l["name"] == "verify_spec" + } + + seen_post_compaction_history = False + + history_ind = len(state["prover_history"]) - 1 + while history_ind >= 0 and len(stuck_count) > 0: + it = state["prover_history"][history_ind] + if it["sort"] == "nag": + for r in it["nagged_rules"]: + del stuck_count[r] + history_ind -= 1 + continue + assert it["sort"] == "run" + if it["tool_call_id"] not in known_tc_ids: + seen_post_compaction_history = True + target_rules = set(it["rules"]) if it["rules"] else None + for k in stuck_count.keys(): + if target_rules is not None and k.rule not in target_rules: + continue + if not any( + rp == k and stuck_rules[k] == stat for (rp, stat) in it["prover_results"] + ): + del stuck_count[k] + else: + stuck_count[k] += 1 + if stuck_count[k] == 3: + to_warn.add(k) + del stuck_count[k] + + prover_update : list[ProverHistoryItem] = [ + ProverRunLog( + tool_call_id=tool_call_id, + prover_results=[(k, v) for (k,v) in result.raw_rule_status.items()], + rules=rules, + spec_digest=spec_hash, + sort="run" + ) + ] + nag_channel = { + + } + if len(to_warn) > 0: + prover_update.append(NagMarker( + sort="nag", + nagged_rules=list(to_warn) + )) + nag_channel["reminders_channel"] = [ + "The following rule(s) have had identical failures on the last 3 runs of the prover:", + *(f"- {it.pprint()}" for it in to_warn), + "You may need to significantly change your approach, or skip the property if this is a persistent issue (you may need to use rebuttals to communicate" + " these failures to the feedback judge)." + ] + if seen_post_compaction_history: + nag_channel["reminders_channel"].append( + "(NB: Some of these prover calls happened before your most recent task history summarization)" ) + + if rules is None and all_verified: + return tool_state_update( + tool_call_id=tool_call_id, content=result.result_str, + prover_link=result.link, validations=stamper(state), + prover_history=prover_update + ) + return tool_state_update( + tool_call_id=tool_call_id, content=result.result_str, prover_link=result.link, + prover_history=prover_update, **nag_channel + ) return verify_spec diff --git a/graphcore b/graphcore index 8b9d3203..967c99af 160000 --- a/graphcore +++ b/graphcore @@ -1 +1 @@ -Subproject commit 8b9d320316a2cfbfa775da0a961501539a4124f0 +Subproject commit 967c99afc126ee83ec43aab048fb3b2e9f11e0a7 From 7f60a8e98c556165b50da99500a51bb4394c19bc Mon Sep 17 00:00:00 2001 From: John Toman Date: Thu, 16 Jul 2026 19:45:40 -0700 Subject: [PATCH 02/10] pricing info moved --- composer/llm/pricing.py | 110 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 composer/llm/pricing.py diff --git a/composer/llm/pricing.py b/composer/llm/pricing.py new file mode 100644 index 00000000..60c16f75 --- /dev/null +++ b/composer/llm/pricing.py @@ -0,0 +1,110 @@ +from dataclasses import dataclass + +@dataclass(frozen=True) +class PriceTier: + """Per-million-token prices in USD for one model at one + context tier. + + ``input`` is the price for *fresh* input tokens (the bucket left + after subtracting ``cache_read`` and ``cache_write`` from the + total). ``output`` is the price for output tokens, which on the + OpenAI side already includes reasoning tokens (billed at the + output rate by both providers). ``cache_read`` / ``cache_write`` + are the cache-bucket rates. + + Anthropic-specific note: the ``cache_write`` rate here is the + 5-minute ephemeral rate. Anthropic also has a 1-hour ephemeral + rate that's higher (~2× base input); our normalized usage rolls + both into a single ``cache_write_tokens`` bucket, so workloads + that heavily use 1h caching will be slightly under-billed by + this banner. Approximation is fine for a banner.""" + input: float + output: float + cache_read: float + cache_write: float + + +@dataclass(frozen=True) +class ModelPricing: + """Pricing entry for one model family. ``long`` is the + long-context tier (used when input token count exceeds the + threshold) and applies only to OpenAI models that publish a + separate >272K-input rate; Anthropic models keep ``long = None`` + and bill everything at ``short`` rates.""" + short: PriceTier + long: PriceTier | None = None + + +# OpenAI's published >272K input-token threshold for long-context +# pricing. Once an individual call's input crosses this, the long +# tier applies *for the full session* per OpenAI's terms; we +# approximate that by switching on a per-message basis (a session +# that drifts above 272K will mostly stay there). +_OPENAI_LONG_CONTEXT_THRESHOLD = 272_000 + + +# Pricing tables transcribed from Anthropic + OpenAI rate cards. +# Sources should be re-checked when new model families ship. +_PRICING: list[tuple[str, ModelPricing]] = [ + # ---- Anthropic ---- + # claude-opus-4.5 / 4.6 / 4.7 share a rate card; older 4 / 4.1 + # are pricier. Matching by prefix-of-prefix so "claude-opus-4-7" + # and "claude-opus-4-7-20260301" both hit the right entry. + ("claude-opus-4-7", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25))), + ("claude-opus-4-6", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25))), + ("claude-opus-4-5", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25))), + ("claude-opus-4-1", ModelPricing(short=PriceTier(15.00, 75.00, 1.50, 18.75))), + ("claude-opus-4", ModelPricing(short=PriceTier(15.00, 75.00, 1.50, 18.75))), + + ("claude-sonnet-4-6", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75))), + ("claude-sonnet-4-5", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75))), + ("claude-sonnet-4", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75))), + + ("claude-haiku-4-5", ModelPricing(short=PriceTier(1.00, 5.00, 0.10, 1.25))), + + # ---- OpenAI ---- + # gpt-5.5 / 5.4 publish short (≤272K input) and long (>272K) tiers. + # Pro variants don't publish a cached-in discount (cache_read = + # base input). Mini/nano don't publish a long tier; we use short + # for everything on those. + ("gpt-5.5-pro", ModelPricing( + short=PriceTier(30.00, 180.00, 30.00, 30.00), + long=PriceTier(60.00, 270.00, 60.00, 60.00), + )), + ("gpt-5.5", ModelPricing( + short=PriceTier(5.00, 30.00, 0.50, 5.00), + long=PriceTier(10.00, 45.00, 1.00, 10.00), + )), + ("gpt-5.4-pro", ModelPricing( + short=PriceTier(30.00, 180.00, 30.00, 30.00), + long=PriceTier(60.00, 270.00, 60.00, 60.00), + )), + ("gpt-5.4-mini", ModelPricing(short=PriceTier(0.75, 4.50, 0.075, 0.75))), + ("gpt-5.4-nano", ModelPricing(short=PriceTier(0.20, 1.25, 0.02, 0.20))), + ("gpt-5.4", ModelPricing( + short=PriceTier(2.50, 15.00, 0.25, 2.50), + long=PriceTier(5.00, 22.50, 0.50, 5.00), + )), +] + + +def price_per_mtok(model: str | None, input_tokens: int) -> PriceTier | None: + """Look up per-MTok pricing by model name and call size. Returns + ``None`` for models with no table entry (cost contribution becomes + zero — better than guessing). + + Matched by prefix on the lowercased model name so dated revisions + (``claude-opus-4-7-20260301``, ``gpt-5.5-2026-...``) collapse into + the same family entry. Table is searched in order, so list more + specific prefixes (``gpt-5.5-pro``) before less specific + (``gpt-5.5``). For OpenAI models with a long tier, ``input_tokens`` + chooses short vs. long; Anthropic always uses the short tier.""" + if model is None: + return None + m = model.lower() + for prefix, pricing in _PRICING: + if m.startswith(prefix): + if pricing.long is not None and input_tokens > _OPENAI_LONG_CONTEXT_THRESHOLD: + return pricing.long + return pricing.short + return None From 12017e6298e4ed8acaa4656b0e1d100643a4398a Mon Sep 17 00:00:00 2001 From: John Toman Date: Thu, 16 Jul 2026 21:51:10 -0700 Subject: [PATCH 03/10] checkpoint --- composer/diagnostics/budget.py | 132 ++++++++++++++++++++++++++ composer/diagnostics/cost_callback.py | 84 ++++++++++++++++ composer/llm/anthropic.py | 16 +++- composer/llm/openai.py | 11 ++- composer/llm/pricing.py | 75 +++++++++------ composer/pipeline/core.py | 27 +++++- composer/pipeline/ptypes.py | 7 ++ composer/spec/system_analysis.py | 3 +- composer/ui/message_renderer.py | 112 +--------------------- 9 files changed, 321 insertions(+), 146 deletions(-) create mode 100644 composer/diagnostics/budget.py create mode 100644 composer/diagnostics/cost_callback.py diff --git a/composer/diagnostics/budget.py b/composer/diagnostics/budget.py new file mode 100644 index 00000000..d659fb67 --- /dev/null +++ b/composer/diagnostics/budget.py @@ -0,0 +1,132 @@ +from typing import Iterator, Callable, Any, Mapping +from typing_extensions import TypeVar +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass + +from langgraph.graph import MessagesState +from graphcore.graph import StateMonitor, MonitorReturn +from langchain.messages import HumanMessage + +StateVar = TypeVar("StateVar", default=MessagesState, bound=MessagesState) + +@dataclass +class BudgetCounter: + total_budget: float + curr_cost: float + + def overbudget(self) -> bool: + return self.curr_cost > self.total_budget + +_budget_accumulator = ContextVar[None | BudgetCounter]("_budget_accumulator", default=None) + +_cost_centers = ContextVar[None | dict[str, BudgetCounter]]("_cost_centers", default=None) + + +DEFAULT_BUDGET_PRESSURE_MESSAGE = """ + +You have almost exceeded the token cost budget allotted for this task. + +Finish your task in as orderly a fashion as possible; partial/incomplete results are better +than going over budget. + +""" + +@contextmanager +def total_budget( + costs: Mapping[str, float] +) -> Iterator[None]: + curr = _cost_centers.get() + if curr is not None: + raise RuntimeError("Not good") + prev = _cost_centers.set({ + k: BudgetCounter(total_budget=v, curr_cost=0.0) for (k, v) in costs.items() + }) + try: + yield None + finally: + _cost_centers.reset(prev) + +@contextmanager +def named_budget( + nm: str +) -> Iterator[None]: + if (res := _cost_centers.get()) is None: + raise RuntimeError("No costs installed") + if nm not in res: + raise RuntimeError(f"Named budget item not known: {nm}") + prev = _budget_accumulator.set(res[nm]) + try: + yield + finally: + _budget_accumulator.reset(prev) + +@contextmanager +def named_budget_or_nop( + nm: str +) -> Iterator[None]: + if (_cost_centers.get()) is None: + return + with named_budget(nm): + yield + +@contextmanager +def token_cost_budget( + total_cost: float, +) -> Iterator[None]: + if _budget_accumulator.get() is not None: + raise RuntimeError("Nested budgets not supported") + accum = BudgetCounter(total_budget=total_cost, curr_cost=0.0) + prev = _budget_accumulator.set(accum) + try: + yield None + finally: + _budget_accumulator.reset(prev) + +def accumulate_cost( + cost: float +): + accum = _budget_accumulator.get() + if accum is None: + return + accum.curr_cost += cost + +def budget_monitor( + *, + warn_threshold: float = 0.8, + warning_message: str | Callable[[StateVar], str] | None = None, + state_transformer: Callable[[StateVar], dict[str, Any]] | None = None, + on_overbudget: Callable[[], None] | None = None +) -> StateMonitor[StateVar]: + accum = _budget_accumulator.get() + if accum is None: + return lambda _ign: (None, None) + warned = False + def monitor( + curr_state: StateVar + ) -> MonitorReturn: + nonlocal warned + if accum.overbudget() and on_overbudget is not None: + on_overbudget() + if warned or not (accum.curr_cost >= accum.total_budget * warn_threshold): + return (None, None) + warned = True + msg : str + if warning_message is None: + msg = DEFAULT_BUDGET_PRESSURE_MESSAGE + elif isinstance(warning_message, str): + msg = warning_message + else: + msg = warning_message(curr_state) + + state_upd = None + if state_transformer is not None: + state_upd = state_transformer(curr_state) + return ([HumanMessage(msg)], state_upd) + return monitor + +def overbudget() -> bool: + res = _budget_accumulator.get() + if res is None: + return False + return res.overbudget() diff --git a/composer/diagnostics/cost_callback.py b/composer/diagnostics/cost_callback.py new file mode 100644 index 00000000..bb8217eb --- /dev/null +++ b/composer/diagnostics/cost_callback.py @@ -0,0 +1,84 @@ +"""LangChain callback that accumulates the running USD cost of an LLM's calls. + +Sibling to :class:`composer.diagnostics.usage_callback.UsageCallback`: attached at +model construction so it fires for *every* ``invoke`` / ``ainvoke`` through the +model. Where ``UsageCallback`` records raw token counts, this one prices each +response through a :data:`~composer.llm.pricing.PriceProvider` (the model's pricing +curve, curried on model name) and adds the result to a running total. + +Implemented as an :class:`~langchain_core.callbacks.AsyncCallbackHandler` with +``run_inline = True``. The async ``on_llm_end`` is awaited on the event-loop thread +either way, but ``run_inline`` also decides *which* context it runs in: without it, +``ahandle_event`` dispatches the handler through ``asyncio.gather`` — each coroutine +wrapped in a ``Task`` against a ``copy_context()`` snapshot, so any ``ContextVar`` +the handler *sets* lands in that throwaway copy and never reaches the caller. With +``run_inline`` the handler is instead awaited directly in the caller's task and +context (``manager.ahandle_event`` line ~437), so its contextvar reads and writes +are visible to the surrounding LLM call. This matters because the accumulator +participates in contextvar state, not just its own counter. On the sync ``invoke`` +path LangChain still drives the coroutine to completion.""" + +from typing import Any + +from langchain_core.callbacks import AsyncCallbackHandler +from langchain_core.messages import AIMessage +from langchain_core.outputs import ChatGeneration, LLMResult + +from graphcore.utils import get_normalized_token_usage +from composer.llm.pricing import PriceProvider +from .budget import accumulate_cost + + +class CostAccumulator(AsyncCallbackHandler): + """Prices each LLM response and accumulates the total into :attr:`total_cost`. + + ``price_provider`` maps a call's input-token count to its per-MTok + :class:`~composer.llm.pricing.PriceTier`. ``long_cache`` selects the 1-hour + cache-write rate over the 5-minute one; the whole conversation is assumed to + share a single cache TTL (see ``builder_for``).""" + + # Route through the direct-await dispatch branch so the handler runs in the + # caller's task/context (not a gather-spawned Task with a copied context): + # required for the handler's contextvar reads/writes to reach the caller. + run_inline = True + + def __init__(self, price_provider: PriceProvider, *, long_cache: bool = False) -> None: + self._price_provider = price_provider + self._long_cache = long_cache + self.total_cost: float = 0.0 + + async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: + try: + generation = response.generations[0][0] + except IndexError: + return + if not isinstance(generation, ChatGeneration): + return + msg = generation.message + if isinstance(msg, AIMessage): + accumulate_cost(self._cost_of(msg)) + + def _cost_of(self, msg: AIMessage) -> float: + """USD cost of a single response, in dollars. Zero for models with no + pricing-table entry.""" + usage = get_normalized_token_usage(msg) + tier = self._price_provider(usage["total_input_tokens"]) + if tier is None: + return 0.0 + + cache_read = usage["cache_read_tokens"] + cache_write = usage["cache_write_tokens"] + # Fresh input is the total minus the two cache buckets, which the tier + # prices separately. Clamp against provider rounding wobble. + fresh_input = max(0, usage["total_input_tokens"] - cache_read - cache_write) + cache_write_rate = tier.cache_write_1h if self._long_cache else tier.cache_write + + # thinking_tokens are a subset of total_output_tokens and bill at the + # output rate, so they need no separate term here. + per_mtok = ( + fresh_input * tier.input + + cache_read * tier.cache_read + + cache_write * cache_write_rate + + usage["total_output_tokens"] * tier.output + ) + return per_mtok / 1_000_000 diff --git a/composer/llm/anthropic.py b/composer/llm/anthropic.py index 61b86bc6..8b82dc8d 100644 --- a/composer/llm/anthropic.py +++ b/composer/llm/anthropic.py @@ -13,6 +13,7 @@ from composer.llm.provider import ( ProviderKind, CacheLevel, _ListIter, NoSuchElementError, ) +from composer.llm.pricing import PriceProvider, price_provider_for if TYPE_CHECKING: from langchain_core.language_models.chat_models import BaseChatModel @@ -126,17 +127,21 @@ class AnthropicModelProvider: model_name: str options: ModelConfiguration features: ModelFeatures - provider: ProviderKind = "anthropic" + price_provider: PriceProvider + provider: Literal["anthropic"] = "anthropic" @staticmethod def create(model_name: str, options: ModelConfiguration) -> "AnthropicModelProvider": - return AnthropicModelProvider(model_name, options, _model_parser(model_name)) + return AnthropicModelProvider( + model_name, options, _model_parser(model_name), price_provider_for(model_name) + ) def builder_for( self, *, cache_level: CacheLevel | None = None, disable_thinking: bool = False ) -> "BaseChatModel": from langchain_anthropic import ChatAnthropic from composer.diagnostics.usage_callback import UsageCallback + from composer.diagnostics.cost_callback import CostAccumulator opts = self.options thinking: dict[str, Any] | None @@ -174,5 +179,10 @@ def builder_for( betas=betas, thinking=thinking, model_kwargs=model_kwargs, - callbacks=[UsageCallback()], + callbacks=[ + UsageCallback(), + CostAccumulator( + self.price_provider, long_cache=cache_level == CacheLevel.LONG + ), + ], ) diff --git a/composer/llm/openai.py b/composer/llm/openai.py index ccca867c..0bbdc192 100644 --- a/composer/llm/openai.py +++ b/composer/llm/openai.py @@ -23,6 +23,7 @@ from composer.llm.provider import ( ProviderKind, CacheLevel, _ListIter, NoSuchElementError, ) +from composer.llm.pricing import PriceProvider, price_provider_for if TYPE_CHECKING: from langchain_core.language_models.chat_models import BaseChatModel @@ -175,16 +176,21 @@ class OpenAIModelProvider: model_name: str options: ModelConfiguration features: OpenAIModelFeatures + price_provider: PriceProvider provider: ProviderKind = "openai" @staticmethod def create(model_name: str, options: ModelConfiguration) -> "OpenAIModelProvider": - return OpenAIModelProvider(model_name, options, _model_parser(model_name)) + return OpenAIModelProvider( + model_name, options, _model_parser(model_name), price_provider_for(model_name) + ) def builder_for( self, *, cache_level: CacheLevel | None = None, disable_thinking: bool = False ) -> "BaseChatModel": from langchain_openai import ChatOpenAI + from composer.diagnostics.usage_callback import UsageCallback + from composer.diagnostics.cost_callback import CostAccumulator opts = self.options kwargs: dict[str, Any] = { @@ -205,5 +211,8 @@ def builder_for( temperature=1, timeout=None, max_retries=2, + # OpenAI has no cache-TTL knob, so long_cache stays False; cache_write_1h + # mirrors cache_write in the table anyway. + callbacks=[UsageCallback(), CostAccumulator(self.price_provider)], **kwargs, ) diff --git a/composer/llm/pricing.py b/composer/llm/pricing.py index 60c16f75..dc0c5ace 100644 --- a/composer/llm/pricing.py +++ b/composer/llm/pricing.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from typing import Callable @dataclass(frozen=True) class PriceTier: @@ -9,19 +10,19 @@ class PriceTier: after subtracting ``cache_read`` and ``cache_write`` from the total). ``output`` is the price for output tokens, which on the OpenAI side already includes reasoning tokens (billed at the - output rate by both providers). ``cache_read`` / ``cache_write`` - are the cache-bucket rates. - - Anthropic-specific note: the ``cache_write`` rate here is the - 5-minute ephemeral rate. Anthropic also has a 1-hour ephemeral - rate that's higher (~2× base input); our normalized usage rolls - both into a single ``cache_write_tokens`` bucket, so workloads - that heavily use 1h caching will be slightly under-billed by - this banner. Approximation is fine for a banner.""" + output rate by both providers). ``cache_read`` is the + cache-hit rate. + + Cache writes come in two rates: ``cache_write`` is the 5-minute + ephemeral rate and ``cache_write_1h`` is the 1-hour ephemeral rate + (~2× base input on Anthropic). Which one applies depends on the + conversation's cache TTL. OpenAI has no separate cache-write rate, + so both fields carry the same value there.""" input: float output: float cache_read: float cache_write: float + cache_write_1h: float @dataclass(frozen=True) @@ -50,40 +51,41 @@ class ModelPricing: # claude-opus-4.5 / 4.6 / 4.7 share a rate card; older 4 / 4.1 # are pricier. Matching by prefix-of-prefix so "claude-opus-4-7" # and "claude-opus-4-7-20260301" both hit the right entry. - ("claude-opus-4-7", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25))), - ("claude-opus-4-6", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25))), - ("claude-opus-4-5", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25))), - ("claude-opus-4-1", ModelPricing(short=PriceTier(15.00, 75.00, 1.50, 18.75))), - ("claude-opus-4", ModelPricing(short=PriceTier(15.00, 75.00, 1.50, 18.75))), + ("claude-opus-4-7", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25, 10.00))), + ("claude-opus-4-6", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25, 10.00))), + ("claude-opus-4-5", ModelPricing(short=PriceTier(5.00, 25.00, 0.50, 6.25, 10.00))), + ("claude-opus-4-1", ModelPricing(short=PriceTier(15.00, 75.00, 1.50, 18.75, 30.00))), + ("claude-opus-4", ModelPricing(short=PriceTier(15.00, 75.00, 1.50, 18.75, 30.00))), - ("claude-sonnet-4-6", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75))), - ("claude-sonnet-4-5", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75))), - ("claude-sonnet-4", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75))), + ("claude-sonnet-4-6", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75, 6.00))), + ("claude-sonnet-4-5", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75, 6.00))), + ("claude-sonnet-4", ModelPricing(short=PriceTier(3.00, 15.00, 0.30, 3.75, 6.00))), - ("claude-haiku-4-5", ModelPricing(short=PriceTier(1.00, 5.00, 0.10, 1.25))), + ("claude-haiku-4-5", ModelPricing(short=PriceTier(1.00, 5.00, 0.10, 1.25, 2.00))), # ---- OpenAI ---- # gpt-5.5 / 5.4 publish short (≤272K input) and long (>272K) tiers. # Pro variants don't publish a cached-in discount (cache_read = # base input). Mini/nano don't publish a long tier; we use short - # for everything on those. + # for everything on those. OpenAI has no separate cache-write rate, + # so cache_write_1h mirrors cache_write on every OpenAI entry. ("gpt-5.5-pro", ModelPricing( - short=PriceTier(30.00, 180.00, 30.00, 30.00), - long=PriceTier(60.00, 270.00, 60.00, 60.00), + short=PriceTier(30.00, 180.00, 30.00, 30.00, 30.00), + long=PriceTier(60.00, 270.00, 60.00, 60.00, 60.00), )), ("gpt-5.5", ModelPricing( - short=PriceTier(5.00, 30.00, 0.50, 5.00), - long=PriceTier(10.00, 45.00, 1.00, 10.00), + short=PriceTier(5.00, 30.00, 0.50, 5.00, 5.00), + long=PriceTier(10.00, 45.00, 1.00, 10.00, 10.00), )), ("gpt-5.4-pro", ModelPricing( - short=PriceTier(30.00, 180.00, 30.00, 30.00), - long=PriceTier(60.00, 270.00, 60.00, 60.00), + short=PriceTier(30.00, 180.00, 30.00, 30.00, 30.00), + long=PriceTier(60.00, 270.00, 60.00, 60.00, 60.00), )), - ("gpt-5.4-mini", ModelPricing(short=PriceTier(0.75, 4.50, 0.075, 0.75))), - ("gpt-5.4-nano", ModelPricing(short=PriceTier(0.20, 1.25, 0.02, 0.20))), + ("gpt-5.4-mini", ModelPricing(short=PriceTier(0.75, 4.50, 0.075, 0.75, 0.75))), + ("gpt-5.4-nano", ModelPricing(short=PriceTier(0.20, 1.25, 0.02, 0.20, 0.20))), ("gpt-5.4", ModelPricing( - short=PriceTier(2.50, 15.00, 0.25, 2.50), - long=PriceTier(5.00, 22.50, 0.50, 5.00), + short=PriceTier(2.50, 15.00, 0.25, 2.50, 2.50), + long=PriceTier(5.00, 22.50, 0.50, 5.00, 5.00), )), ] @@ -108,3 +110,18 @@ def price_per_mtok(model: str | None, input_tokens: int) -> PriceTier | None: return pricing.long return pricing.short return None + + +# A model's pricing curve as a function of a single call's input-token count: +# ``price_per_mtok`` with the model name curried away. The remaining argument +# picks the short/long context tier (OpenAI); Anthropic ignores it. +type PriceProvider = Callable[[int], PriceTier | None] + + +def price_provider_for(model: str | None) -> PriceProvider: + """Curry :func:`price_per_mtok` on ``model``: returns a callable mapping a + call's input-token count to its :class:`PriceTier` (or ``None`` when the + model has no table entry, so its cost contribution is zero).""" + def provider(input_tokens: int) -> PriceTier | None: + return price_per_mtok(model, input_tokens) + return provider diff --git a/composer/pipeline/core.py b/composer/pipeline/core.py index 97ce1b0b..8f614990 100644 --- a/composer/pipeline/core.py +++ b/composer/pipeline/core.py @@ -17,7 +17,7 @@ import enum import logging from dataclasses import dataclass -from typing import Protocol +from typing import Protocol, cast from abc import ABC, abstractmethod from pydantic import BaseModel @@ -41,8 +41,10 @@ from composer.spec.source.report import build as report_build from composer.spec.source.task_ids import SYSTEM_ANALYSIS_TASK_ID, REPORT_TASK_ID from .ptypes import ( - BackendJob, BackendResult, ComponentOutcome, CorePhases, CorePipelineResult, Delivered, GaveUp, PipelineRun, SystemAnalysisSpec + BackendJob, BackendResult, ComponentOutcome, CorePhases, CorePipelineResult, + Delivered, GaveUp, PipelineRun, SystemAnalysisSpec, PhaseBudget ) +from composer.diagnostics.budget import total_budget, named_budget_or_nop COMMON_SYSTEM_CACHE_KEY = "system-analysis" @@ -153,6 +155,27 @@ 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, + budget: PhaseBudget | None = None +) -> CorePipelineResult[FormT]: + if budget is not None: + with total_budget(cast(dict[str, float], budget)): + return await _run_pipeline_inner( + backend, run, interactive=interactive, + max_bug_rounds=max_bug_rounds, threat_model=threat_model + ) + else: + return await _run_pipeline_inner( + backend, run, interactive=interactive, + max_bug_rounds=max_bug_rounds, threat_model=threat_model + ) + +async def _run_pipeline_inner[P: enum.Enum, FormT: BackendResult, H, A: ArtifactIdentifier]( + backend: PipelineBackend[P, FormT, H, A], + run: PipelineRun[P, H], + *, + interactive: bool, + threat_model: Document | None, + max_bug_rounds: int ) -> CorePipelineResult[FormT]: spec, phases = backend.analysis_spec, backend.core_phases source = run.source diff --git a/composer/pipeline/ptypes.py b/composer/pipeline/ptypes.py index e665c790..dbf6683f 100644 --- a/composer/pipeline/ptypes.py +++ b/composer/pipeline/ptypes.py @@ -117,6 +117,13 @@ def all_failed(self) -> bool: driver raises before returning in the no-outcomes case anyway).""" return bool(self.outcomes) and self.n_delivered == 0 +class PhaseBudget(TypedDict): + formalization_preparation: float + system_analysis: float + system_preparation: float + property_extraction: float + formalization: float + __all__ = [ "CorePipelineResult", "ComponentOutcome", diff --git a/composer/spec/system_analysis.py b/composer/spec/system_analysis.py index 7666c90a..be55feb5 100644 --- a/composer/spec/system_analysis.py +++ b/composer/spec/system_analysis.py @@ -11,6 +11,7 @@ from composer.spec.service_host import ServiceHost from composer.spec.util import slugify_filename from composer.tools.thinking import RoughDraftState, get_rough_draft_tools +from composer.diagnostics.budget import budget_monitor DESCRIPTION = "Component analysis" @@ -128,7 +129,7 @@ def _validation_wrapper( ).with_initial_prompt_template( "application_analysis_prompt.j2", sort=env.sort, - ) + ).with_monitor(budget_monitor()) graph = b.compile_async() inputs : list[str | dict] = [ diff --git a/composer/ui/message_renderer.py b/composer/ui/message_renderer.py index f4d4d0f0..f1e0c0ea 100644 --- a/composer/ui/message_renderer.py +++ b/composer/ui/message_renderer.py @@ -24,6 +24,7 @@ from graphcore.graph import INITIAL_NODE, TOOL_RESULT_NODE, TOOLS_NODE from graphcore.utils import NormalizedTokenUsage, get_normalized_token_usage +from composer.llm.pricing import price_per_mtok KNOWN_NODES: set[str] = {INITIAL_NODE, TOOL_RESULT_NODE, TOOLS_NODE} @@ -42,115 +43,6 @@ def dot(style: str, text: Text | str) -> Text: return result -@dataclass(frozen=True) -class _PriceTier: - """Per-million-token prices in USD for one model at one - context tier. - - ``input`` is the price for *fresh* input tokens (the bucket left - after subtracting ``cache_read`` and ``cache_write`` from the - total). ``output`` is the price for output tokens, which on the - OpenAI side already includes reasoning tokens (billed at the - output rate by both providers). ``cache_read`` / ``cache_write`` - are the cache-bucket rates. - - Anthropic-specific note: the ``cache_write`` rate here is the - 5-minute ephemeral rate. Anthropic also has a 1-hour ephemeral - rate that's higher (~2× base input); our normalized usage rolls - both into a single ``cache_write_tokens`` bucket, so workloads - that heavily use 1h caching will be slightly under-billed by - this banner. Approximation is fine for a banner.""" - input: float - output: float - cache_read: float - cache_write: float - - -@dataclass(frozen=True) -class _ModelPricing: - """Pricing entry for one model family. ``long`` is the - long-context tier (used when input token count exceeds the - threshold) and applies only to OpenAI models that publish a - separate >272K-input rate; Anthropic models keep ``long = None`` - and bill everything at ``short`` rates.""" - short: _PriceTier - long: _PriceTier | None = None - - -# OpenAI's published >272K input-token threshold for long-context -# pricing. Once an individual call's input crosses this, the long -# tier applies *for the full session* per OpenAI's terms; we -# approximate that by switching on a per-message basis (a session -# that drifts above 272K will mostly stay there). -_OPENAI_LONG_CONTEXT_THRESHOLD = 272_000 - - -# Pricing tables transcribed from Anthropic + OpenAI rate cards. -# Sources should be re-checked when new model families ship. -_PRICING: list[tuple[str, _ModelPricing]] = [ - # ---- Anthropic ---- - # claude-opus-4.5 / 4.6 / 4.7 share a rate card; older 4 / 4.1 - # are pricier. Matching by prefix-of-prefix so "claude-opus-4-7" - # and "claude-opus-4-7-20260301" both hit the right entry. - ("claude-opus-4-7", _ModelPricing(short=_PriceTier(5.00, 25.00, 0.50, 6.25))), - ("claude-opus-4-6", _ModelPricing(short=_PriceTier(5.00, 25.00, 0.50, 6.25))), - ("claude-opus-4-5", _ModelPricing(short=_PriceTier(5.00, 25.00, 0.50, 6.25))), - ("claude-opus-4-1", _ModelPricing(short=_PriceTier(15.00, 75.00, 1.50, 18.75))), - ("claude-opus-4", _ModelPricing(short=_PriceTier(15.00, 75.00, 1.50, 18.75))), - - ("claude-sonnet-4-6", _ModelPricing(short=_PriceTier(3.00, 15.00, 0.30, 3.75))), - ("claude-sonnet-4-5", _ModelPricing(short=_PriceTier(3.00, 15.00, 0.30, 3.75))), - ("claude-sonnet-4", _ModelPricing(short=_PriceTier(3.00, 15.00, 0.30, 3.75))), - - ("claude-haiku-4-5", _ModelPricing(short=_PriceTier(1.00, 5.00, 0.10, 1.25))), - - # ---- OpenAI ---- - # gpt-5.5 / 5.4 publish short (≤272K input) and long (>272K) tiers. - # Pro variants don't publish a cached-in discount (cache_read = - # base input). Mini/nano don't publish a long tier; we use short - # for everything on those. - ("gpt-5.5-pro", _ModelPricing( - short=_PriceTier(30.00, 180.00, 30.00, 30.00), - long=_PriceTier(60.00, 270.00, 60.00, 60.00), - )), - ("gpt-5.5", _ModelPricing( - short=_PriceTier(5.00, 30.00, 0.50, 5.00), - long=_PriceTier(10.00, 45.00, 1.00, 10.00), - )), - ("gpt-5.4-pro", _ModelPricing( - short=_PriceTier(30.00, 180.00, 30.00, 30.00), - long=_PriceTier(60.00, 270.00, 60.00, 60.00), - )), - ("gpt-5.4-mini", _ModelPricing(short=_PriceTier(0.75, 4.50, 0.075, 0.75))), - ("gpt-5.4-nano", _ModelPricing(short=_PriceTier(0.20, 1.25, 0.02, 0.20))), - ("gpt-5.4", _ModelPricing( - short=_PriceTier(2.50, 15.00, 0.25, 2.50), - long=_PriceTier(5.00, 22.50, 0.50, 5.00), - )), -] - - -def _price_per_mtok(model: str | None, input_tokens: int) -> _PriceTier | None: - """Look up per-MTok pricing by model name and call size. Returns - ``None`` for models with no table entry (cost contribution becomes - zero — better than guessing). - - Matched by prefix on the lowercased model name so dated revisions - (``claude-opus-4-7-20260301``, ``gpt-5.5-2026-...``) collapse into - the same family entry. Table is searched in order, so list more - specific prefixes (``gpt-5.5-pro``) before less specific - (``gpt-5.5``). For OpenAI models with a long tier, ``input_tokens`` - chooses short vs. long; Anthropic always uses the short tier.""" - if model is None: - return None - m = model.lower() - for prefix, pricing in _PRICING: - if m.startswith(prefix): - if pricing.long is not None and input_tokens > _OPENAI_LONG_CONTEXT_THRESHOLD: - return pricing.long - return pricing.short - return None - class TokenStats: """Accumulates token usage across AI messages and updates a display widget. @@ -177,7 +69,7 @@ def __init__(self, display: Static): @staticmethod def _cost_of(usage: NormalizedTokenUsage) -> float: - price = _price_per_mtok(usage["model_name"], usage["total_input_tokens"]) + price = price_per_mtok(usage["model_name"], usage["total_input_tokens"]) if price is None: return 0.0 # Fresh (non-cached) input is the total minus the cached From c2a4e8a99d1cdc50f635513f7fd8e42255494eea Mon Sep 17 00:00:00 2001 From: John Toman Date: Thu, 16 Jul 2026 22:21:39 -0700 Subject: [PATCH 04/10] checkpoint --- composer/diagnostics/budget.py | 58 ++++++++++++++++- composer/foundry/author.py | 106 +++++++++++++++++++++++--------- composer/pipeline/core.py | 65 +++++++++++--------- composer/spec/cvl_generation.py | 11 ++++ composer/spec/feedback.py | 37 ++++++++--- composer/spec/prop_inference.py | 8 +++ composer/spec/source/author.py | 85 ++++++++++++++++++------- 7 files changed, 279 insertions(+), 91 deletions(-) diff --git a/composer/diagnostics/budget.py b/composer/diagnostics/budget.py index d659fb67..e9c974be 100644 --- a/composer/diagnostics/budget.py +++ b/composer/diagnostics/budget.py @@ -6,10 +6,29 @@ from langgraph.graph import MessagesState from graphcore.graph import StateMonitor, MonitorReturn -from langchain.messages import HumanMessage +from langchain_core.messages import HumanMessage StateVar = TypeVar("StateVar", default=MessagesState, bound=MessagesState) +# The fraction of a budget at which "budget pressure" begins: budget_monitor's +# warning fires and budget_pressure() flips true. One global so every accessor +# agrees on where the wrap-up window starts. +BUDGET_PRESSURE_THRESHOLD = 0.8 + + +class BudgetExceeded(Exception): + """Hard budget stop, raised cooperatively (from a monitor, between agent + turns) once the active budget is blown. Only monitors given an + ``on_overbudget`` callback raise; the workflow that launched the agent + catches this and converts it into its give-up result.""" + + +class BudgetPressureAbort(Exception): + """Raised by ``pressure_abort_monitor`` to terminate an auxiliary agent + (e.g. a feedback judge) whose output is worthless once the main agent is + in its wrap-up window. Caught by the tool that launched the agent.""" + + @dataclass class BudgetCounter: total_budget: float @@ -66,6 +85,9 @@ def named_budget_or_nop( nm: str ) -> Iterator[None]: if (_cost_centers.get()) is None: + # A @contextmanager generator must yield exactly once even on the + # nop path — a bare return raises "generator didn't yield". + yield return with named_budget(nm): yield @@ -93,7 +115,7 @@ def accumulate_cost( def budget_monitor( *, - warn_threshold: float = 0.8, + warn_threshold: float = BUDGET_PRESSURE_THRESHOLD, warning_message: str | Callable[[StateVar], str] | None = None, state_transformer: Callable[[StateVar], dict[str, Any]] | None = None, on_overbudget: Callable[[], None] | None = None @@ -130,3 +152,35 @@ def overbudget() -> bool: if res is None: return False return res.overbudget() + + +def raise_budget_exceeded() -> None: + """``on_overbudget`` callback for agents that opt into the hard stop.""" + raise BudgetExceeded( + "Token cost budget exhausted; the agent was cooperatively terminated." + ) + + +def budget_pressure() -> bool: + """Whether the active budget is inside its wrap-up window (accrued cost at + or past ``BUDGET_PRESSURE_THRESHOLD`` of the allotment). False when no + budget is installed. Use this to skip launching work that would only be + told to immediately pack it in (e.g. further property-extraction rounds).""" + res = _budget_accumulator.get() + if res is None: + return False + return res.curr_cost >= res.total_budget * BUDGET_PRESSURE_THRESHOLD + + +def pressure_abort_monitor() -> StateMonitor[MessagesState]: + """Monitor for auxiliary agents (feedback judges) that should not outlive + the main agent's wrap-up window: raises ``BudgetPressureAbort`` between + turns once budget pressure sets in. The tool that launched the agent + catches the exception and returns a canned "terminated for budget" + result. Reads the budget at call time, so it can be attached to a graph + compiled outside any budget scope.""" + def monitor(_curr_state: StateVar) -> MonitorReturn: + if budget_pressure(): + raise BudgetPressureAbort() + return (None, None) + return monitor diff --git a/composer/foundry/author.py b/composer/foundry/author.py index f7596455..4c03440d 100644 --- a/composer/foundry/author.py +++ b/composer/foundry/author.py @@ -43,7 +43,11 @@ from composer.spec.cvl_generation import ( PropertyFeedbackProtocol, RebuttalBase, SkippedProperty, ) -from composer.spec.feedback import PropertyFeedback +from composer.spec.feedback import BUDGET_ABORT_FEEDBACK, PropertyFeedback +from composer.diagnostics.budget import ( + BudgetExceeded, BudgetPressureAbort, budget_monitor, budget_pressure, + pressure_abort_monitor, raise_budget_exceeded, +) from composer.spec.gen_types import TypedTemplate from composer.spec.graph_builder import bind_standard, run_to_completion from composer.spec.types import PropertyFormulation @@ -372,6 +376,15 @@ class FeedbackTool( async def run(self) -> Command | str: if self.state["curr_test"] is None: return "No test written" + if budget_pressure(): + # Don't launch a judge that would be terminated on its first + # monitor tick; the author's budget warning already tells it + # feedback approval is no longer required. + return ( + "Good? False\nFeedback:\nThe feedback judge was not run due to " + "budget constraints. See the system alert: feedback approval is " + "no longer required for this task." + ) with self.tool_deps() as deps: res = await deps.thunk( self.state["curr_test"], @@ -427,6 +440,8 @@ def did_rough_draft_read(s: ST, _) -> str | None: "foundry_property_judge_system_prompt.j2" ).with_tools( [*get_rough_draft_tools(ST), judge_ctx.get_memory_tool(), GetTestTool.as_tool("get_test")] + ).with_monitor( + pressure_abort_monitor() ).compile_async() async def thunk( @@ -458,17 +473,20 @@ async def thunk( f" Addressing: {r.prior_feedback_reference}\n" f" Evidence: {r.evidence}" ) - res = await run_to_completion( - workflow, - TestJudgeInput( - input=input_parts, curr_test=test_source, - memory=None, did_read=False, - ), - thread_id=uniq_thread_id("foundry-feedback"), - recursion_limit=judge_ctx.recursion_limit, - description="Foundry test feedback judge", - within_tool=within_tool, - ) + try: + res = await run_to_completion( + workflow, + TestJudgeInput( + input=input_parts, curr_test=test_source, + memory=None, did_read=False, + ), + thread_id=uniq_thread_id("foundry-feedback"), + recursion_limit=judge_ctx.recursion_limit, + description="Foundry test feedback judge", + within_tool=within_tool, + ) + except BudgetPressureAbort: + return BUDGET_ABORT_FEEDBACK assert "result" in res return res["result"] @@ -514,17 +532,22 @@ class PublishResultTool( async def run(self) -> Command | str: if (err := check_foundry_completion(self.state)) is not None: return err - ran = self.state["last_test_names"] - if ran is None: - # Unreachable in practice — the forge_test stamp required above - # implies a run recorded its test names — but defend anyway. - return "Completion REJECTED: no forge_test run has been recorded." - with self.tool_deps() as titles: - err = validate_property_tests( - self.property_tests, self.state["skipped"], titles, ran, - ) - if err is not None: - return err + if not budget_pressure(): + # The forge-ground-truth cross-check requires a recorded run; in + # the budget wrap-up window (where the agent is told to delete + # failing tests and publish without re-running forge) the declared + # mapping is accepted as-is. + ran = self.state["last_test_names"] + if ran is None: + # Unreachable in practice — the forge_test stamp required above + # implies a run recorded its test names — but defend anyway. + return "Completion REJECTED: no forge_test run has been recorded." + with self.tool_deps() as titles: + err = validate_property_tests( + self.property_tests, self.state["skipped"], titles, ran, + ) + if err is not None: + return err return tool_state_update( self.tool_call_id, "Accepted", @@ -625,6 +648,21 @@ class FoundryPropertyGenParams(TypedDict): "foundry_property_generation_prompt.j2" ) +_BUDGET_WRAPUP_MESSAGE = """ + +You have almost exceeded the token cost budget for this task. Wrap up IMMEDIATELY; +a partial test file is better than going over budget. Concretely: + +- The forge-test and feedback validation requirements on publishing have been lifted. You + no longer need approval from the feedback judge — ignore any pending or future feedback, + including a judge response saying it was terminated. +- Do NOT start new forge runs or research. +- Delete any tests that do not currently compile or pass. +- Skip (`record_skip`) every property you have not gotten to work, citing budget exhaustion. +- Then publish what remains via the `result` tool. + +""" + async def batch_foundry_test_generation( ctx: WorkflowContext[FoundryGeneration], @@ -700,6 +738,11 @@ async def batch_foundry_test_generation( .with_sys_prompt_template("foundry_property_generation_system_prompt.j2") .inject(lambda b: bound_template.render_to(b.with_initial_prompt_template)) .with_summary_config(FoundryGenerationSummaryConfig()) + .with_monitor(budget_monitor( + warning_message=_BUDGET_WRAPUP_MESSAGE, + state_transformer=lambda _s: {"required_validations": []}, + on_overbudget=raise_budget_exceeded, + )) ) graph = builder.compile_async() @@ -716,13 +759,16 @@ async def batch_foundry_test_generation( ) tid, mnem = await ctx.thread_and_mnemonic() - res_state = await run_to_completion( - graph, - init_state, - thread_id=tid, - description=f"{description} ({mnem})", - recursion_limit=ctx.recursion_limit, - ) + try: + res_state = await run_to_completion( + graph, + init_state, + thread_id=tid, + description=f"{description} ({mnem})", + recursion_limit=ctx.recursion_limit, + ) + except BudgetExceeded as e: + return GaveUp(reason=f"Foundry test generation terminated: {e}") assert "result" in res_state assert res_state["failed"] is not None diff --git a/composer/pipeline/core.py b/composer/pipeline/core.py index 8f614990..3914f255 100644 --- a/composer/pipeline/core.py +++ b/composer/pipeline/core.py @@ -181,27 +181,34 @@ async def _run_pipeline_inner[P: enum.Enum, FormT: BackendResult, H, A: Artifact source = run.source # 1. System analysis (shared primitive, backend-parameterized; always yields SourceApplication). - analyzed = await run.runner( - TaskInfo(SYSTEM_ANALYSIS_TASK_ID, "System Analysis", phases["analysis"]), - lambda: run_component_analysis( - ty=SourceApplication, child_ctxt=run.ctx.child(CacheKey(spec.analysis_key)), - input=source, env=run.env, extra_input=[ - f"The main entry point of this application has been explicitly identified as {source.contract_name} at relative path {source.relative_path}. " - "Your output MUST contain an explicit contract instance with this solidity identifier.", - *spec.extra_input - ], - expected_main_id=source.contract_name, - ), - ) + with named_budget_or_nop("system_analysis"): + analyzed = await run.runner( + TaskInfo(SYSTEM_ANALYSIS_TASK_ID, "System Analysis", phases["analysis"]), + lambda: run_component_analysis( + ty=SourceApplication, child_ctxt=run.ctx.child(CacheKey(spec.analysis_key)), + input=source, env=run.env, extra_input=[ + f"The main entry point of this application has been explicitly identified as {source.contract_name} at relative path {source.relative_path}. " + "Your output MUST contain an explicit contract instance with this solidity identifier.", + *spec.extra_input + ], + expected_main_id=source.contract_name, + ), + ) 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) + with named_budget_or_nop("system_preparation"): + prepared = await backend.prepare_system(analyzed, run) # 3. Pre-formalization setup runs CONCURRENTLY with extraction (neither needs the other) — # this preserves the prover's autosetup ∥ bug-analysis overlap, generically. - formalizer_task = asyncio.create_task(prepared.prepare_formalization(run)) + # The budget scope is entered inside the coroutine (not around create_task) so the + # cost-center binding lives in the spawned task's own context. + async def _prepare_formalization() -> Formalizer[FormT]: + with named_budget_or_nop("formalization_preparation"): + return await prepared.prepare_formalization(run) + formalizer_task = asyncio.create_task(_prepare_formalization()) batches = await _extract_all(prepared.main, backend.backend_guidance, run, phases["extraction"], interactive, threat_model, max_bug_rounds) @@ -220,14 +227,15 @@ async def _run(batch: _Batch) -> ComponentOutcome[FormT]: result : FormT | GaveUp if cached_result is None: label = f"{batch.feat.component.name} ({len(batch.props)} properties)" - result : FormT | GaveUp = await run.runner( - TaskInfo( - formalize_task_id(batch.feat.ind), - f"{batch.feat.component.name} ({len(batch.props)} properties)", - phases["formalization"] - ), - lambda: formalizer.formalize(label, batch.feat, batch.props, child, run), - ) + with named_budget_or_nop("formalization"): + result : FormT | GaveUp = await run.runner( + TaskInfo( + formalize_task_id(batch.feat.ind), + f"{batch.feat.component.name} ({len(batch.props)} properties)", + phases["formalization"] + ), + lambda: formalizer.formalize(label, batch.feat, batch.props, child, run), + ) if not isinstance(result, GaveUp): await child.cache_put(result) else: @@ -283,12 +291,13 @@ 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()}) - 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), - ) + with named_budget_or_nop("property_extraction"): + 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), + ) return _Batch(feat, props, feat_ctx) if props else None got = await asyncio.gather(*[_one(i) for i in range(len(main.contract.components))]) diff --git a/composer/spec/cvl_generation.py b/composer/spec/cvl_generation.py index 3dd1a3d8..e2b6cd9e 100644 --- a/composer/spec/cvl_generation.py +++ b/composer/spec/cvl_generation.py @@ -31,6 +31,7 @@ from composer.spec.graph_builder import run_to_completion from composer.cvl.tools import put_cvl_raw, put_cvl, get_cvl, edit_cvl from composer.ui.tool_display import tool_display, suppress_ack +from composer.diagnostics.budget import budget_pressure class PropertyFeedbackProtocol(Protocol): @property @@ -310,6 +311,16 @@ async def run(self) -> Command: spec = st["curr_spec"] if spec is None: return tool_return(self.tool_call_id, "No spec put yet") + if budget_pressure(): + # Don't launch a judge that would be terminated on its first + # monitor tick; the author's budget warning already tells it + # feedback approval is no longer required. + return tool_return( + self.tool_call_id, + "Good? False\nFeedback The feedback judge was not run due to budget " + "constraints. See the system alert: feedback approval is no longer " + "required for this task.", + ) skipped = st["skipped"] t = await feedback(spec, skipped, self.rebuttals, self.tool_call_id) msg = f"Good? {t.good}\nFeedback {t.feedback}" diff --git a/composer/spec/feedback.py b/composer/spec/feedback.py index 1aceda7a..e6bada56 100644 --- a/composer/spec/feedback.py +++ b/composer/spec/feedback.py @@ -21,6 +21,7 @@ from composer.spec.cvl_generation import FeedbackToolContext, Rebuttal, SkippedProperty from composer.spec.system_model import ContractComponentInstance from composer.spec.util import uniq_thread_id +from composer.diagnostics.budget import BudgetPressureAbort, pressure_abort_monitor class PropertyFeedback(BaseModel): """ @@ -29,6 +30,19 @@ class PropertyFeedback(BaseModel): good: bool = Field(description="Whether the properties are good as is, or if there is room for improvement") feedback: str = Field(description="The feedback on the rule if work is needed. Can be empty if there is no feedback") + +# Canned judge verdict when the judge was not run (or was killed mid-run) +# because the author is in its budget wrap-up window. The author's own budget +# warning tells it feedback approval is no longer required, so a dead judge +# is expected there, not an error. +BUDGET_ABORT_FEEDBACK = PropertyFeedback( + good=False, + feedback=( + "The feedback judge was terminated due to budget constraints. " + "See the system alert: feedback approval is no longer required for this task." + ), +) + class Properties(TypedDict): properties: list[PropertyFormulation] @@ -98,7 +112,9 @@ def did_rough_draft_read(s: ST, _) -> str | None: lambda b: final_prompt.render_to(b.with_initial_prompt_template) ).inject( lambda g: system_prompt.render_to(g.with_sys_prompt_template) - ).with_tools([*rough_draft_tools, mem, get_cvl(ST), ]).compile_async() + ).with_tools([*rough_draft_tools, mem, get_cvl(ST), ]).with_monitor( + pressure_abort_monitor() + ).compile_async() async def the_tool( cvl: str, @@ -134,14 +150,17 @@ async def the_tool( f" Addressing: {r.prior_feedback_reference}\n" f" Evidence: {r.evidence}" ) - res = await run_to_completion( - workflow, - SpecJudgeInput(input=input_parts, curr_spec=cvl, memory=None, did_read=False), - thread_id=uniq_thread_id("feedback"), - recursion_limit=ctx.recursion_limit, - description="Property feedback judge", - within_tool=within_tool, - ) + try: + res = await run_to_completion( + workflow, + SpecJudgeInput(input=input_parts, curr_spec=cvl, memory=None, did_read=False), + thread_id=uniq_thread_id("feedback"), + recursion_limit=ctx.recursion_limit, + description="Property feedback judge", + within_tool=within_tool, + ) + except BudgetPressureAbort: + return BUDGET_ABORT_FEEDBACK assert "result" in res return res["result"] diff --git a/composer/spec/prop_inference.py b/composer/spec/prop_inference.py index 3edc92c0..b13fc6a8 100644 --- a/composer/spec/prop_inference.py +++ b/composer/spec/prop_inference.py @@ -25,6 +25,7 @@ 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.diagnostics.budget import budget_monitor, budget_pressure 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 @@ -312,6 +313,8 @@ class ST(MessagesState, RoughDraftState): env.analysis_tools ).with_sys_prompt_template( "property_analysis_system_prompt.j2", sort=env.sort + ).with_monitor( + budget_monitor() ).compile_async() flow_input: BugAnalysisInput = BugAnalysisInput( @@ -365,6 +368,11 @@ async def _run_bug_analysis_inner( last_round_convo : list[AnyMessage] | None = None for i in range(0, max_rounds): + # Under budget pressure a fresh round would be told to pack it in on + # its first monitor tick — don't bother launching it. Round 0 always + # runs (the loop's invariants require at least one round's history). + if i > 0 and budget_pressure(): + break next_result = await _run_bug_round( env, component, front_matter_items, agent_component_analysis, i, prev_rounds, backend_guidance, diff --git a/composer/spec/source/author.py b/composer/spec/source/author.py index 269854f1..f092db49 100644 --- a/composer/spec/source/author.py +++ b/composer/spec/source/author.py @@ -1,4 +1,4 @@ -from typing import NotRequired, override, Literal, Annotated +from typing import Callable, NotRequired, override, Literal, Annotated from typing_extensions import TypedDict import json @@ -32,9 +32,12 @@ from langgraph.types import Command from composer.spec.feedback import property_feedback_judge, FeedbackTemplate from composer.ui.tool_display import tool_display +from composer.diagnostics.budget import ( + BudgetExceeded, budget_monitor, raise_budget_exceeded, +) from .monitor import monitor -from graphcore.graph import FlowInput +from graphcore.graph import FlowInput, MonitorReturn class SourceAuthorExtra(TypedDict): failed: bool | None @@ -332,6 +335,41 @@ async def run(self) -> Command | str: _PropertyGenTemplate = TypedTemplate[PropertyGenParams]("property_generation_prompt.j2") +_BUDGET_WRAPUP_MESSAGE = """ + +You have almost exceeded the token cost budget for this task. Wrap up IMMEDIATELY; +a partial spec is better than going over budget. Concretely: + +- The feedback and prover validation requirements on publishing have been lifted. You no + longer need approval from the feedback judge — ignore any pending or future feedback, + including a judge response saying it was terminated. +- Do NOT start new prover runs or research. +- Delete any rules/invariants that do not currently work. +- Skip (`record_skip`) every property you have not gotten to work, citing budget exhaustion. +- Then publish what remains via the `result` tool. + +""" + + +def _author_monitor() -> "Callable[[SourceCVLGenerationState], MonitorReturn]": + """The author's monitor: budget wrap-up takes precedence; otherwise the + usual reminders-channel drain. On the (single) turn the budget warning + fires any pending reminders are dropped — moot, since the warning tells + the agent to ignore prover/feedback outcomes anyway.""" + b_monitor = budget_monitor( + warning_message=_BUDGET_WRAPUP_MESSAGE, + state_transformer=lambda _s: {"required_validations": []}, + on_overbudget=raise_budget_exceeded, + ) + + def combined(curr_state: SourceCVLGenerationState) -> MonitorReturn: + msgs, upd = b_monitor(curr_state) + if msgs is None and upd is None: + return monitor(curr_state) + return msgs, upd + + return combined + async def batch_cvl_generation( ctx: WorkflowContext[CVLGeneration], init_config: dict, @@ -378,7 +416,7 @@ async def batch_cvl_generation( ).with_state( SourceCVLGenerationState ).with_monitor( - monitor + _author_monitor() ).with_output_key( "result" ).with_input( @@ -398,26 +436,29 @@ async def batch_cvl_generation( }), props ) - res_state = await run_cvl_generator( - ctx = ctx, - d = task_graph, - description=description, - ctxt=feedback_env, - in_state=SourceCVLGenerationInput( - curr_spec=None, - config=init_config, - spec_stem=spec_stem, - input=[], - required_validations=[FEEDBACK_VALIDATION_KEY, PROVER_VALIDATION_KEY], - rule_skips={}, - skipped=[], - property_rules=[], - validations={}, - failed=None, - prover_history=[], - reminders_channel=[] + try: + res_state = await run_cvl_generator( + ctx = ctx, + d = task_graph, + description=description, + ctxt=feedback_env, + in_state=SourceCVLGenerationInput( + curr_spec=None, + config=init_config, + spec_stem=spec_stem, + input=[], + required_validations=[FEEDBACK_VALIDATION_KEY, PROVER_VALIDATION_KEY], + rule_skips={}, + skipped=[], + property_rules=[], + validations={}, + failed=None, + prover_history=[], + reminders_channel=[] + ) ) - ) + except BudgetExceeded as e: + return GaveUp(reason=f"CVL generation terminated: {e}") assert "result" in res_state assert res_state["failed"] is not None From 48815ab57644a5bd84e7e186715010d16c81242e Mon Sep 17 00:00:00 2001 From: John Toman Date: Thu, 16 Jul 2026 22:41:45 -0700 Subject: [PATCH 05/10] Pool based accounting --- composer/diagnostics/budget.py | 52 ++++++++++++++++++++++++++-------- composer/pipeline/core.py | 6 ++-- composer/pipeline/ptypes.py | 11 +++++++ 3 files changed, 54 insertions(+), 15 deletions(-) diff --git a/composer/diagnostics/budget.py b/composer/diagnostics/budget.py index e9c974be..b6df628b 100644 --- a/composer/diagnostics/budget.py +++ b/composer/diagnostics/budget.py @@ -31,11 +31,26 @@ class BudgetPressureAbort(Exception): @dataclass class BudgetCounter: + """One node of the caps-over-pool scheme. Leaf counters are per-phase + *caps* whose ``parent`` is the run's shared *pool* (the real budget); + ``token_cost_budget`` creates parentless one-off counters. Cost accrues + up the chain, and both the hard stop and the pressure window trip on + whichever level is tighter — a phase can only starve later phases up to + its cap, while unspent phase money never leaves the pool (rollover is + automatic, not an explicit transfer).""" total_budget: float curr_cost: float + parent: "BudgetCounter | None" = None def overbudget(self) -> bool: - return self.curr_cost > self.total_budget + if self.curr_cost > self.total_budget: + return True + return self.parent.overbudget() if self.parent is not None else False + + def pressured(self, threshold: float = BUDGET_PRESSURE_THRESHOLD) -> bool: + if self.curr_cost >= self.total_budget * threshold: + return True + return self.parent.pressured(threshold) if self.parent is not None else False _budget_accumulator = ContextVar[None | BudgetCounter]("_budget_accumulator", default=None) @@ -53,17 +68,28 @@ def overbudget(self) -> bool: @contextmanager def total_budget( - costs: Mapping[str, float] + total: float, + caps: Mapping[str, float] ) -> Iterator[None]: + """Install the run's budget: ``total`` is the pool (the real bound on + spend) and ``caps`` are per-phase ceilings. Caps need not sum to the + pool — they only bound how much a single phase may hog, so each can be + generous; whatever a phase doesn't spend simply remains in the pool for + later phases.""" curr = _cost_centers.get() if curr is not None: raise RuntimeError("Not good") + pool = BudgetCounter(total_budget=total, curr_cost=0.0) prev = _cost_centers.set({ - k: BudgetCounter(total_budget=v, curr_cost=0.0) for (k, v) in costs.items() + k: BudgetCounter(total_budget=v, curr_cost=0.0, parent=pool) for (k, v) in caps.items() }) + # Work running outside any named center (e.g. the report phase) accrues + # to — and feels pressure from — the pool directly. + prev_accum = _budget_accumulator.set(pool) try: yield None finally: + _budget_accumulator.reset(prev_accum) _cost_centers.reset(prev) @contextmanager @@ -108,10 +134,11 @@ def token_cost_budget( def accumulate_cost( cost: float ): + # Accrue up the chain: the active center and (through parent) the pool. accum = _budget_accumulator.get() - if accum is None: - return - accum.curr_cost += cost + while accum is not None: + accum.curr_cost += cost + accum = accum.parent def budget_monitor( *, @@ -130,7 +157,7 @@ def monitor( nonlocal warned if accum.overbudget() and on_overbudget is not None: on_overbudget() - if warned or not (accum.curr_cost >= accum.total_budget * warn_threshold): + if warned or not accum.pressured(warn_threshold): return (None, None) warned = True msg : str @@ -162,14 +189,15 @@ def raise_budget_exceeded() -> None: def budget_pressure() -> bool: - """Whether the active budget is inside its wrap-up window (accrued cost at - or past ``BUDGET_PRESSURE_THRESHOLD`` of the allotment). False when no - budget is installed. Use this to skip launching work that would only be - told to immediately pack it in (e.g. further property-extraction rounds).""" + """Whether the active budget is inside its wrap-up window: accrued cost at + or past ``BUDGET_PRESSURE_THRESHOLD`` of the phase cap *or* of the run + pool, whichever trips first. False when no budget is installed. Use this + to skip launching work that would only be told to immediately pack it in + (e.g. further property-extraction rounds).""" res = _budget_accumulator.get() if res is None: return False - return res.curr_cost >= res.total_budget * BUDGET_PRESSURE_THRESHOLD + return res.pressured() def pressure_abort_monitor() -> StateMonitor[MessagesState]: diff --git a/composer/pipeline/core.py b/composer/pipeline/core.py index 3914f255..5f2d5a41 100644 --- a/composer/pipeline/core.py +++ b/composer/pipeline/core.py @@ -42,7 +42,7 @@ from composer.spec.source.task_ids import SYSTEM_ANALYSIS_TASK_ID, REPORT_TASK_ID from .ptypes import ( BackendJob, BackendResult, ComponentOutcome, CorePhases, CorePipelineResult, - Delivered, GaveUp, PipelineRun, SystemAnalysisSpec, PhaseBudget + Delivered, GaveUp, PipelineRun, SystemAnalysisSpec, RunBudget ) from composer.diagnostics.budget import total_budget, named_budget_or_nop @@ -155,10 +155,10 @@ 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, - budget: PhaseBudget | None = None + budget: RunBudget | None = None ) -> CorePipelineResult[FormT]: if budget is not None: - with total_budget(cast(dict[str, float], budget)): + with total_budget(budget.total, cast(dict[str, float], budget.caps)): return await _run_pipeline_inner( backend, run, interactive=interactive, max_bug_rounds=max_bug_rounds, threat_model=threat_model diff --git a/composer/pipeline/ptypes.py b/composer/pipeline/ptypes.py index dbf6683f..959e8253 100644 --- a/composer/pipeline/ptypes.py +++ b/composer/pipeline/ptypes.py @@ -118,12 +118,23 @@ def all_failed(self) -> bool: return bool(self.outcomes) and self.n_delivered == 0 class PhaseBudget(TypedDict): + """Per-phase spending *caps* (USD). Ceilings, not allotments: they bound + how much a single phase may hog, and need not sum to the run total — + unspent phase money stays in the run pool for later phases.""" formalization_preparation: float system_analysis: float system_preparation: float property_extraction: float formalization: float + +@dataclass(frozen=True) +class RunBudget: + """The run's token-cost budget: ``total`` is the pool (the real bound on + overall spend); ``caps`` are the per-phase ceilings drawn against it.""" + total: float + caps: PhaseBudget + __all__ = [ "CorePipelineResult", "ComponentOutcome", From 123622cd1b0d1c793757f7984096847c56298de2 Mon Sep 17 00:00:00 2001 From: John Toman Date: Thu, 16 Jul 2026 22:53:10 -0700 Subject: [PATCH 06/10] Circular import --- composer/spec/source/monitor.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/composer/spec/source/monitor.py b/composer/spec/source/monitor.py index 3ed56d74..2fcded3b 100644 --- a/composer/spec/source/monitor.py +++ b/composer/spec/source/monitor.py @@ -1,11 +1,16 @@ +from typing import TYPE_CHECKING + from langgraph.config import get_config from langchain_core.messages import HumanMessage -from .author import SourceCVLGenerationState from graphcore.graph import MonitorReturn from composer.prover.ptypes import StatusCodes +if TYPE_CHECKING: + # Runtime import would be circular: author.py imports this module. + from .author import SourceCVLGenerationState + def monitor( - curr_state: SourceCVLGenerationState + curr_state: "SourceCVLGenerationState" ) -> MonitorReturn: if not curr_state["reminders_channel"]: return None, None From 4769789452f1f85d864e87e222b2a5c2d33b7a50 Mon Sep 17 00:00:00 2001 From: John Toman Date: Fri, 17 Jul 2026 14:51:25 -0700 Subject: [PATCH 07/10] Budget aware caching --- composer/spec/context.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/composer/spec/context.py b/composer/spec/context.py index 623c3f9c..73439014 100644 --- a/composer/spec/context.py +++ b/composer/spec/context.py @@ -21,7 +21,7 @@ from composer.io.mnemonic_store import assign_mnemonic from composer.core.user import user_data_ns from composer.spec.system_model import SolidityIdentifier - +from composer.diagnostics.budget import budget_pressure # --------------------------------------------------------------------------- # Workflow input types @@ -240,6 +240,11 @@ async def cache_get(self, ty: type[K]) -> K | None: async def cache_put(self, value: K) -> None: """Put a typed value in the cache. No-op if caching disabled.""" + if budget_pressure(): + # refuse to cache anything produced under budget pressure + # It's likely incomplete or "rushed", and a future run with more money + # should try again + return None if not isinstance(value, BaseModel): raise ValueError("Caching not allowed for non-basemodel keys") if self.cache_namespace is None: @@ -257,8 +262,3 @@ async def thread_and_mnemonic(self) -> tuple[str, str]: tid = self.thread_id mnem = await assign_mnemonic(tid, self._store, user_data_ns() + MNEMONIC_KEYS) return (tid, mnem) - -# --------------------------------------------------------------------------- -# Utility -# --------------------------------------------------------------------------- - From 5a1dce79bd97857d024d2710a0938efb7471a97f Mon Sep 17 00:00:00 2001 From: John Toman Date: Tue, 21 Jul 2026 13:35:58 -0700 Subject: [PATCH 08/10] take skips into account when determining "stuck" status --- composer/spec/source/prover.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/composer/spec/source/prover.py b/composer/spec/source/prover.py index a5424c7b..de0f5dd2 100644 --- a/composer/spec/source/prover.py +++ b/composer/spec/source/prover.py @@ -325,7 +325,7 @@ async def verify_spec( ) stuck_rules = { - k: v for (k,v) in result.raw_rule_status.items() if v in ("TIMEOUT", "ERROR", "SANITY_FAILED") + k: v for (k,v) in result.raw_rule_status.items() if v in ("TIMEOUT", "ERROR", "SANITY_FAILED") and k.rule not in state["rule_skips"] } stuck_count = { From 442b1a8bd564a22ffd38b5a293cbdb505b3769f2 Mon Sep 17 00:00:00 2001 From: John Toman Date: Wed, 22 Jul 2026 11:19:48 -0700 Subject: [PATCH 09/10] sync graphcore pin --- graphcore | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/graphcore b/graphcore index 967c99af..f4584f66 160000 --- a/graphcore +++ b/graphcore @@ -1 +1 @@ -Subproject commit 967c99afc126ee83ec43aab048fb3b2e9f11e0a7 +Subproject commit f4584f667644f856191bc04dd46e0cd184036f55 diff --git a/pyproject.toml b/pyproject.toml index f999a88f..89d96261 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ authors = [ ] dependencies = [ - "graphcore @ git+ssh://git@github.com/Certora/graphcore.git@8b9d320316a2cfbfa775da0a961501539a4124f0", + "graphcore @ git+ssh://git@github.com/Certora/graphcore.git@f4584f667644f856191bc04dd46e0cd184036f55", "aiohttp>=3.13", "attrs>=26.1", "Jinja2>=3.1", From 28651f78c9f6dc709ec426dede49bf82ac59647c Mon Sep 17 00:00:00 2001 From: John Toman Date: Wed, 22 Jul 2026 11:20:36 -0700 Subject: [PATCH 10/10] okay fine, we'll use the old types --- composer/diagnostics/budget.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/composer/diagnostics/budget.py b/composer/diagnostics/budget.py index e633d765..b6df628b 100644 --- a/composer/diagnostics/budget.py +++ b/composer/diagnostics/budget.py @@ -1,4 +1,4 @@ -from typing import Iterator, Callable, Any, Mapping, Generator +from typing import Iterator, Callable, Any, Mapping from typing_extensions import TypeVar from contextlib import contextmanager from contextvars import ContextVar @@ -70,7 +70,7 @@ def pressured(self, threshold: float = BUDGET_PRESSURE_THRESHOLD) -> bool: def total_budget( total: float, caps: Mapping[str, float] -) -> Generator[None]: +) -> Iterator[None]: """Install the run's budget: ``total`` is the pool (the real bound on spend) and ``caps`` are per-phase ceilings. Caps need not sum to the pool — they only bound how much a single phase may hog, so each can be @@ -95,7 +95,7 @@ def total_budget( @contextmanager def named_budget( nm: str -) -> Generator[None]: +) -> Iterator[None]: if (res := _cost_centers.get()) is None: raise RuntimeError("No costs installed") if nm not in res: @@ -109,7 +109,7 @@ def named_budget( @contextmanager def named_budget_or_nop( nm: str -) -> Generator[None]: +) -> Iterator[None]: if (_cost_centers.get()) is None: # A @contextmanager generator must yield exactly once even on the # nop path — a bare return raises "generator didn't yield". @@ -121,7 +121,7 @@ def named_budget_or_nop( @contextmanager def token_cost_budget( total_cost: float, -) -> Generator[None]: +) -> Iterator[None]: if _budget_accumulator.get() is not None: raise RuntimeError("Nested budgets not supported") accum = BudgetCounter(total_budget=total_cost, curr_cost=0.0)