Skip to content

ai-enhance: respect BigQuery cost guards in agent prompt#2057

Open
arsalann wants to merge 6 commits into
mainfrom
arsalann/ai-enhance-cost-guard
Open

ai-enhance: respect BigQuery cost guards in agent prompt#2057
arsalann wants to merge 6 commits into
mainfrom
arsalann/ai-enhance-cost-guard

Conversation

@arsalann
Copy link
Copy Markdown
Contributor

Summary

  • Adds a ## Query Cost Safety (BigQuery only) section to the bruin ai-enhance system prompt (pkg/enhance/prompt.go).
  • Instructs the agent to inspect .bruin.yml for max_query_cost / max_query_cost_soft / max_billable_bytes / max_billable_bytes_soft on the asset's google_cloud_platform connection before issuing any query.
  • If a limit is set: always bruin query --dry-run first, rewrite if the estimate exceeds the limit, and never pass --dangerously-bypass-soft-limits without explicit user confirmation.
  • If no limit is set: stop, warn the user that cost is uncapped, offer to add one, and require an explicit acknowledgement before proceeding.
  • Scope is explicitly BigQuery-only; other warehouses are called out as not currently supporting Bruin cost guards.

No changes to BigQuery enforcement code or the enhance command itself — prompt-only update.

Test plan

  • make format — clean (golangci-lint not installed locally; unrelated)
  • make test — full unit suite passes, including pkg/enhance (TestBuildEnhancePrompt)
  • Manual smoke: run bruin ai-enhance <bq-asset> against a pipeline with (a) no cost guard configured and (b) a max_query_cost_soft configured; confirm the agent surfaces the no-limit warning in (a) and dry-runs before querying in (b)

arsalann and others added 2 commits May 14, 2026 17:41
Add a Query Cost Safety section to the enhance system prompt instructing
the agent to check `.bruin.yml` for max_query_cost / max_billable_bytes
(and their _soft variants) before issuing queries against BigQuery
assets, dry-run when limits exist, and require explicit user
confirmation when they don't. Scoped to BigQuery only.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Mirror the new prompt-level Query Cost Safety section in the user
documentation, listing the four configurable keys and explaining what
the agent does when limits are configured versus absent. Scoped to
BigQuery.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
@greptile-apps
Copy link
Copy Markdown
Contributor

greptile-apps Bot commented May 14, 2026

Comments Outside Diff (1)

  1. pkg/enhance/prompt.go, line 1-38 (link)

    P2 No test coverage for the new cost-safety content

    prompt_test.go does not assert that the new section is present in the built prompt. Given that cost-guard enforcement is safety-critical, a test case verifying that key phrases (e.g., "max_query_cost", "--dry-run", "Query Cost Safety") appear in the output of BuildEnhancePrompt would prevent a future accidental deletion or truncation from going unnoticed.

    Prompt To Fix With AI
    This is a comment left during a code review.
    Path: pkg/enhance/prompt.go
    Line: 1-38
    
    Comment:
    **No test coverage for the new cost-safety content**
    
    `prompt_test.go` does not assert that the new section is present in the built prompt. Given that cost-guard enforcement is safety-critical, a test case verifying that key phrases (e.g., `"max_query_cost"`, `"--dry-run"`, `"Query Cost Safety"`) appear in the output of `BuildEnhancePrompt` would prevent a future accidental deletion or truncation from going unnoticed.
    
    How can I resolve this? If you propose a fix, please make it concise.
Prompt To Fix All With AI
Fix the following 3 code review issues. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 3
pkg/enhance/prompt.go:155-156
**Missing asset-to-connection resolution step**

The instructions jump straight to inspecting `.bruin.yml` without first telling the agent to read the asset file and extract its `connection` field. Without that step, the agent has no concrete connection name to look up — it may scan all `google_cloud_platform` connections or silently pick the first one, both of which can result in evaluating the wrong limits.

The instructions should begin with: read the asset file at `<asset-path>`, extract the value of its `connection` field (falling back to the pipeline default if unset), then look that name up in `.bruin.yml`.

### Issue 2 of 3
pkg/enhance/prompt.go:151-169
**Cost safety section placed after the modification and validation steps**

The "Context Discovery" section (step 1) already encourages the agent to read SQL queries and run tools to understand the asset. If the agent decides to issue a `bruin query` call during that discovery phase, it will have already planned that action before reaching this safety section — which appears as the very last substantive block in the prompt.

Placing this guard immediately after the "Context Discovery" section (or at least before "Asset Description Requirements") would make it a precondition the agent evaluates before any query is planned, not a post-hoc rule it reads after the fact.

### Issue 3 of 3
pkg/enhance/prompt.go:1-38
**No test coverage for the new cost-safety content**

`prompt_test.go` does not assert that the new section is present in the built prompt. Given that cost-guard enforcement is safety-critical, a test case verifying that key phrases (e.g., `"max_query_cost"`, `"--dry-run"`, `"Query Cost Safety"`) appear in the output of `BuildEnhancePrompt` would prevent a future accidental deletion or truncation from going unnoticed.

Reviews (1): Last reviewed commit: "docs(ai-enhance): document BigQuery cost..." | Re-trigger Greptile

Comment thread pkg/enhance/prompt.go Outdated
Comment thread pkg/enhance/prompt.go Outdated
arsalann and others added 4 commits May 14, 2026 17:50
Address review feedback:
- Add an explicit asset-to-connection resolution step so the agent looks
  up cost guards on the exact connection the asset uses, not the first
  google_cloud_platform connection it finds.
- Move the Query Cost Safety section above Context Discovery so it acts
  as a precondition the agent evaluates before planning any query.
- Add prompt_test cases asserting the cost guard phrases are present
  and that the section precedes Context Discovery.
- Mirror the resolution step in docs/commands/ai-enhance.md.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…proach

The earlier approach relied on instructing the AI agent (via prompt text) to
respect cost limits. That was unreliable — the agent could ignore or
misinterpret the precondition. Move the check into the command itself so it
runs deterministically before any AI work begins.

When the input scope includes any BigQuery asset (bq.sql/seed/source/sensor),
the command now resolves each asset's connection (asset.connection, falling
back to pipeline default_connections), reads .bruin.yml, and checks for any of
max_query_cost, max_query_cost_soft, max_billable_bytes, max_billable_bytes_soft.
If a connection has none, the user is prompted interactively for a USD limit;
pressing Enter applies a $5 default. The value is persisted to .bruin.yml so
subsequent runs don't re-prompt. Non-TTY / --output json modes apply the
default silently for CI safety.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
- gocritic: rewrite if-else chain in collectBigQueryConnections as switch.
- unparam: drop unused afero.Fs parameter from collectBigQueryConnections —
  the function delegates to DefaultPipelineBuilder, which uses its own fs.
- noctx: use exec.CommandContext in the gitInit test helper.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
golangci-lint's unparam flagged the error return as always nil — every code
path either returns no connections or appends to the slice. Drop the error
return and inline the check at the call site.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant