Skip to content

feat(rocky-cli): pre-merge budget projection on rocky preview cost#343

Merged
hugocorreia90 merged 1 commit intomainfrom
feat/preview-cost-budget-projection
May 2, 2026
Merged

feat(rocky-cli): pre-merge budget projection on rocky preview cost#343
hugocorreia90 merged 1 commit intomainfrom
feat/preview-cost-budget-projection

Conversation

@hugocorreia90
Copy link
Copy Markdown
Contributor

@hugocorreia90 hugocorreia90 commented May 2, 2026

Summary

Closes the Arc 2 residual on the active wave plan: surface run-level [budget] breaches at preview time, so a PR reviewer (and a CI gate) can see "this PR would breach max_usd / max_duration_ms / max_bytes_scanned if merged" before the merge happens.

Project-level [budget] already exists (#288). What's missing today is the projection. rocky preview cost runs against the branch run that rocky preview create produced — those branch totals are the projection. This PR feeds them through BudgetConfig::check_breaches (the same function rocky run uses at end-of-run) and surfaces the result on the JSON output and in the existing rocky preview Action's PR-comment markdown.

Surface

  • PreviewCostOutput.projected_budget_breaches: Vec<BudgetBreachOutput> — populated only when the project declares a [budget] block. Empty when no budget or projected totals stay within every limit. Mirrors the RunOutput.budget_breaches shape so the same downstream consumers (PR-comment templates, JSON listeners) can process both with one code path.
  • PreviewCostSummary.total_branch_duration_ms and total_branch_bytes_scanned — gives the projection the run-level totals it needs across all three budget dimensions. Bytes are Option<u64> because the non-BigQuery adapters still inherit the default stub on WarehouseAdapter::execute_statement_with_stats.
  • project_budget_breaches(budget, summary) -> Vec<BudgetBreachOutput> — pure read-only helper at rocky_cli::commands::preview. Reusable outside the preview flow if a future rocky budget check subcommand wants the same logic.

Markdown

  • New "Budget projection" section in cost.markdown — only renders when breaches are projected, so projects with a benign budget don't get noise on every PR.
  • Section text flips between advisory ("would breach … warnings") and blocking ("would fail the run") based on [budget].on_breach. Closes the open question of whether the projection text should imply CI-blocking — it now reflects exactly what would happen at run-time.

Action

  • .github/actions/rocky-preview/action.yml already prints cost.markdown verbatim, so it picks up the new section automatically. No Action changes needed.

Codegen

Schema change cascades regenerated — schemas/preview_cost.schema.json, integrations/dagster/src/dagster_rocky/types_generated/preview_cost_schema.py, editors/vscode/src/types/generated/preview_cost.ts. Both new fields are additive (Option/Vec-default) so direct JSON consumers stay forward-compatible.

Notes for reviewers

  • Per-model [budget] blocks are out of scope for this PR — that's its own UX call and a follow-up. This PR projects only the run-level totals (the BudgetConfig already has).
  • The cost-dimension projection short-circuits on None cost (e.g., when no adapter cost params resolve), matching BudgetConfig::check_breaches's Option-cost rule. So a "no cost data" preview never false-positives a max_usd breach.

Test plan

  • cargo test -p rocky-cli --lib preview:: — 32 preview tests pass, six new for budget projection
  • cargo check --workspace — clean
  • cargo fmt --check + cargo clippy -p rocky-cli --all-targets — clean
  • pytest -p integrations/dagster tests/test_types.py — 24 type tests pass on the regenerated Pydantic shape

@hugocorreia90
feat(rocky-cli): pre-merge budget projection on rocky preview cost
ac7e332 
Closes the Arc 2 residual on the active wave plan: surface run-level
[budget] breaches at preview time so a PR reviewer (and a CI gate) can
see "this PR would breach max_usd / max_duration_ms / max_bytes_scanned
if merged" before the merge happens.

Surface
- `PreviewCostOutput.projected_budget_breaches: Vec<BudgetBreachOutput>`
  — populated only when [budget] is configured. Mirrors the
  `RunOutput.budget_breaches` shape so PR-comment templates and JSON
  listeners can process both with one code path.
- `PreviewCostSummary.total_branch_duration_ms` and
  `total_branch_bytes_scanned` — gives the projection check the run-
  level totals it needs across all three budget dimensions.
- `project_budget_breaches(budget, summary) -> Vec<BudgetBreachOutput>`
  — pure read-only helper; reusable outside the preview flow.

Markdown
- New "Budget projection" section in the cost markdown — only renders
  when breaches are projected, so projects with a benign budget don't
  get noise in every PR comment.
- Section text flips between advisory ("would breach … warnings") and
  blocking ("would fail the run") based on `[budget].on_breach`.

Action
- The .github/actions/rocky-preview Action prints `cost.markdown`
  verbatim, so it picks up the new section automatically — no Action
  changes needed.

Codegen
- Schema regenerated: schemas/preview_cost.schema.json,
  integrations/dagster/.../preview_cost_schema.py,
  editors/vscode/.../preview_cost.ts. Dagster type tests pass; vscode
  no-op shim absorbs the additive fields.

Tests
- Six new unit tests on `project_budget_breaches` and the markdown
  renderer covering each limit dimension, the unset / within-budget
  paths, the cost-dimension `None`-skip, and both `on_breach` framings.
@hugocorreia90 hugocorreia90 merged commit b220420 into main May 2, 2026
12 of 15 checks passed
@hugocorreia90 hugocorreia90 deleted the feat/preview-cost-budget-projection branch May 2, 2026 00:08
This was referenced May 2, 2026
Merged
Merged
hugocorreia90 added a commit that referenced this pull request May 2, 2026
@hugocorreia90
feat(rocky-cli): per-model [budget] blocks + cost-projection coverage
a426967 
Closes the per-model side of the pre-merge budget surface. Project-level
[budget] (PR #288) and project-level pre-merge cost projection (PR #343)
shipped earlier; this lands the per-model parts and the matching
PreviewCostOutput / Markdown extensions.

Per-model [budget] lives in the model
sidecar, alongside materialization / intent / unique_key. Same fields
as project-level (max_usd, max_duration_ms, max_bytes_scanned,
on_breach), each Option so a partial sidecar block (e.g. only max_usd)
inherits missing fields from the project-level config.

Parsing surface. ModelConfig gains an optional `budget:
Option<ModelBudgetConfig>` field. ModelBudgetConfig sits next to
BudgetConfig in rocky-core::config with a `resolve(&BudgetConfig) ->
BudgetConfig` API that performs field-level inheritance into a fully
resolved BudgetConfig the existing check_breaches pipeline can consume
unchanged. `#[serde(deny_unknown_fields)]` matches BudgetConfig.

Precedence rule for on_breach. Per-model is the local authority — when
explicitly set on the sidecar it wins, even if project-level differs.
Per-model `on_breach = "warn"` overrides project-level `error` for that
one model; per-model `on_breach = "error"` overrides project-level
`warn`. When the per-model sidecar omits `on_breach`, the project-level
value applies. Implemented by typing per-model `on_breach` as
`Option<BudgetBreachAction>` so absence vs explicit value is
distinguishable.

JSON shape (additive, no breaking change). PreviewCostOutput grows
`projected_per_model_budget_breaches: Vec<PerModelBudgetBreachOutput>`,
serialized as a separate field. The existing `projected_budget_breaches`
field continues to surface only project-level breaches and is unchanged
for direct JSON consumers. PerModelBudgetBreachOutput carries
model_name + limit_type + the resolved limit/actual + the resolved
on_breach (so PR readers see the limit they actually crossed and
whether it's blocking).

Action Markdown shape. The existing "Budget projection" section in
`rocky preview cost --markdown` (consumed verbatim by
.github/actions/rocky-preview/) is extended to include a per-model
breach subtable when per-model breaches exist. The section header
flips from advisory to "would fail the run" when any per-model breach
has `on_breach = "error"`, even if the project-level breach is merely
advisory. Action YAML needs no code change — it already pulls
`.markdown` from the cost.json output.

`rocky preview cost` gains `--models` (default `models`) so it can
load sidecars; missing/malformed model directories silently degrade to
project-level-only projection (consistent with existing behaviour for
malformed rocky.toml).

Codegen cascade. preview_cost.schema.json + the dagster Pydantic + the
vscode TypeScript bindings now carry PerModelBudgetBreachOutput and the
new `projected_per_model_budget_breaches` field.

Tests added (10):
  - rocky-core::config: ModelBudgetConfig defaults / parse / unknown
    fields / resolve field inheritance / resolve on_breach precedence
    (5)
  - rocky-cli::commands::preview: ModelBudgetConfig::resolve field
    inheritance, on_breach per-model authority, project_per_model_
    budget_breaches walks per-model deltas with copied-skip,
    cost_markdown surfaces per-model section + flips header on error,
    sidecar parses [budget] block end-to-end (5)
hugocorreia90 added a commit that referenced this pull request May 2, 2026
@hugocorreia90
feat(rocky-cli): per-model [budget] blocks + cost-projection coverage
93768ad 
Closes the per-model side of the pre-merge budget surface. Project-level
[budget] (PR #288) and project-level pre-merge cost projection (PR #343)
shipped earlier; this lands the per-model parts and the matching
PreviewCostOutput / Markdown extensions.

Per-model [budget] lives in the model
sidecar, alongside materialization / intent / unique_key. Same fields
as project-level (max_usd, max_duration_ms, max_bytes_scanned,
on_breach), each Option so a partial sidecar block (e.g. only max_usd)
inherits missing fields from the project-level config.

Parsing surface. ModelConfig gains an optional `budget:
Option<ModelBudgetConfig>` field. ModelBudgetConfig sits next to
BudgetConfig in rocky-core::config with a `resolve(&BudgetConfig) ->
BudgetConfig` API that performs field-level inheritance into a fully
resolved BudgetConfig the existing check_breaches pipeline can consume
unchanged. `#[serde(deny_unknown_fields)]` matches BudgetConfig.

Precedence rule for on_breach. Per-model is the local authority — when
explicitly set on the sidecar it wins, even if project-level differs.
Per-model `on_breach = "warn"` overrides project-level `error` for that
one model; per-model `on_breach = "error"` overrides project-level
`warn`. When the per-model sidecar omits `on_breach`, the project-level
value applies. Implemented by typing per-model `on_breach` as
`Option<BudgetBreachAction>` so absence vs explicit value is
distinguishable.

JSON shape (additive, no breaking change). PreviewCostOutput grows
`projected_per_model_budget_breaches: Vec<PerModelBudgetBreachOutput>`,
serialized as a separate field. The existing `projected_budget_breaches`
field continues to surface only project-level breaches and is unchanged
for direct JSON consumers. PerModelBudgetBreachOutput carries
model_name + limit_type + the resolved limit/actual + the resolved
on_breach (so PR readers see the limit they actually crossed and
whether it's blocking).

Action Markdown shape. The existing "Budget projection" section in
`rocky preview cost --markdown` (consumed verbatim by
.github/actions/rocky-preview/) is extended to include a per-model
breach subtable when per-model breaches exist. The section header
flips from advisory to "would fail the run" when any per-model breach
has `on_breach = "error"`, even if the project-level breach is merely
advisory. Action YAML needs no code change — it already pulls
`.markdown` from the cost.json output.

`rocky preview cost` gains `--models` (default `models`) so it can
load sidecars; missing/malformed model directories silently degrade to
project-level-only projection (consistent with existing behaviour for
malformed rocky.toml).

Codegen cascade. preview_cost.schema.json + the dagster Pydantic + the
vscode TypeScript bindings now carry PerModelBudgetBreachOutput and the
new `projected_per_model_budget_breaches` field.

Tests added (10):
  - rocky-core::config: ModelBudgetConfig defaults / parse / unknown
    fields / resolve field inheritance / resolve on_breach precedence
    (5)
  - rocky-cli::commands::preview: ModelBudgetConfig::resolve field
    inheritance, on_breach per-model authority, project_per_model_
    budget_breaches walks per-model deltas with copied-skip,
    cost_markdown surfaces per-model section + flips header on error,
    sidecar parses [budget] block end-to-end (5)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@hugocorreia90