diff --git a/composer/diagnostics/budget.py b/composer/diagnostics/budget.py
new file mode 100644
index 00000000..b6df628b
--- /dev/null
+++ b/composer/diagnostics/budget.py
@@ -0,0 +1,214 @@
+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_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:
+ """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:
+ 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)
+
+_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(
+ 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, 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
+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:
+ # 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
+
+@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
+):
+ # Accrue up the chain: the active center and (through parent) the pool.
+ accum = _budget_accumulator.get()
+ while accum is not None:
+ accum.curr_cost += cost
+ accum = accum.parent
+
+def budget_monitor(
+ *,
+ 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
+) -> 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.pressured(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()
+
+
+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 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.pressured()
+
+
+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/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/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/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
new file mode 100644
index 00000000..dc0c5ace
--- /dev/null
+++ b/composer/llm/pricing.py
@@ -0,0 +1,127 @@
+from dataclasses import dataclass
+from typing import Callable
+
+@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`` 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)
+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, 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, 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, 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. 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, 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, 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, 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, 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, 2.50),
+ long=PriceTier(5.00, 22.50, 0.50, 5.00, 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
+
+
+# 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 3c92a585..87142f66 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, RunBudget
)
+from composer.diagnostics.budget import total_budget, named_budget_or_nop
COMMON_SYSTEM_CACHE_KEY = "system-analysis"
@@ -154,32 +156,60 @@ 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: RunBudget | None = None
+) -> CorePipelineResult[FormT]:
+ if budget is not None:
+ 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
+ )
+ 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
# 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(
backend.analysis_spec.properties_key,
@@ -201,14 +231,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:
@@ -265,12 +296,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/pipeline/ptypes.py b/composer/pipeline/ptypes.py
index e665c790..959e8253 100644
--- a/composer/pipeline/ptypes.py
+++ b/composer/pipeline/ptypes.py
@@ -117,6 +117,24 @@ 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):
+ """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",
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/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
-# ---------------------------------------------------------------------------
-
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 d0bf3cd3..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,8 +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
@@ -331,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,
@@ -376,6 +415,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(
+ _author_monitor()
).with_output_key(
"result"
).with_input(
@@ -395,24 +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,
+ 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
diff --git a/composer/spec/source/monitor.py b/composer/spec/source/monitor.py
new file mode 100644
index 00000000..2fcded3b
--- /dev/null
+++ b/composer/spec/source/monitor.py
@@ -0,0 +1,20 @@
+from typing import TYPE_CHECKING
+
+from langgraph.config import get_config
+from langchain_core.messages import HumanMessage
+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"
+) -> 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 0b170fed..01c07a15 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 number {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") and k.rule not in state["rule_skips"]
+ }
+
+ 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/composer/spec/system_analysis.py b/composer/spec/system_analysis.py
index 57953825..3d713a1e 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"
@@ -132,7 +133,7 @@ def _validation_wrapper(
"application_analysis_prompt.j2",
sort=env.sort,
has_doc=input is not None
- )
+ ).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
diff --git a/graphcore b/graphcore
index 8b9d3203..f4584f66 160000
--- a/graphcore
+++ b/graphcore
@@ -1 +1 @@
-Subproject commit 8b9d320316a2cfbfa775da0a961501539a4124f0
+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",