feat: load probe templates at runtime via --templates-dir - #174
praetorian-farida wants to merge 8 commits into
Conversation
Add a runtime template loader so new probes — including new multi-turn attack strategies — can be defined in YAML and run without recompiling. - pkg/templates: type-aware schema (static | multiturn) + validation - internal/runtimetemplates: TemplateStrategy (multiturn.Strategy built from YAML prompts, dry-run validated at load) + MultiTurnTemplateProbe + RegisterFromPath - cmd/augustus: --templates-dir flag, registered before glob expansion - examples/runtime-templates: static + multi-turn examples + README Static templates become single-turn prompt probes; multi-turn templates plug a YAML-defined strategy into the existing unified multi-turn engine (internal/multiturn). Detector is chosen via info.detector; engine parameters are overridable at scan time. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…onfig - multiturn: add a configurable Detector field to BaseMultiTurnProbe (defaults to DetectorJudge); existing multi-turn probes are unchanged. - runtimetemplates: embed BaseMultiTurnProbe instead of duplicating its Prober methods; MultiTurnTemplateProbe now implements ProbeDetectorConfig so multi-turn templates honor info.detector_config (parity with static). - lock the buildEngineConfigMap <-> multiturn config-key contract with a round-trip test so a key rename in multiturn fails here, not silently. - render(): on the (dry-run-guarded) runtime error path, log and return empty instead of a half-rendered prompt. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…orce RegisterFromPath now detects ID collisions with already-registered probes and fails closed (atomically — nothing is registered) unless override is allowed. A new scan --force flag opts into shadowing a built-in. Prevents a stray template file from silently replacing a built-in probe; mirrors Nuclei's custom-vs-official template distinction. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
CLAUDE.md Drift DetectionCode changes in this PR may have made documentation stale: CLAUDE.md
Automated drift check — please review and update if needed. |
There was a problem hiding this comment.
Claude Review
Critical issues
- Multi-turn strategy templates are NOT dry-run validated at load — contrary to the README/PR claim ("validated on load … a typo like `{{.Tpyo}}` fails immediately"). `RegisterFromPath` → `templates.LoadFromPath` only runs schema-level `Validate()` (engine/strategy blocks present, generator types set). The field-reference dry-run lives in `newTemplateStrategy`, which is only reached via the factory `newMultiTurnProbe` at `probes.Create` time. Net effect: a bad field ref in `attacker_system:`/`turn:` passes load (prints "Loaded N template(s)…") and only errors when the scan materializes the probe — fail-fast-at-load doesn't hold for multi-turn. Consider parsing the strategy once during `RegisterFromPath` (or in a load-time validation hook) so the typo aborts before the scan starts.
- `secondary_detectors` is silently ignored on multi-turn templates. `validateMultiTurn` doesn't run the secondary-detector validation that `validateStatic` does, and `MultiTurnTemplateProbe` doesn't implement `GetSecondaryDetectors`. A multi-turn YAML declaring `secondary_detectors` validates clean and silently has no effect — surprising given `info.detector_config` is threaded through. Either wire it up or reject it at validation with a clear error.
Security
No security concerns flagged. `--templates-dir` loads user-pointed local YAML and renders `text/template` (data interpolation, not code execution); no new external network calls or secret handling. `--force` override correctly fail-closes by default (no panic — registry replaces).
Test coverage
Thorough — 35+ new tests across schema validation, strategy rendering, factory construction, and atomic/fail-closed registration. Gap: no test asserts when multi-turn strategy template errors surface (load vs. create), which is exactly the behavior the first finding concerns.
There was a problem hiding this comment.
Codex Review
Critical Issues
info.detectoris not the sole post-hoc verdict for runtime multi-turn probes. The docs sayinfo.detectoris the reported verdict detector, andnewMultiTurnProbepasses that detector into the probe, butUnifiedEngine.buildUnifiedResultalways pre-populatesDetectorResults["judge.Judge"];Attempt.GetEffectiveScores()then takes the max across all detector results. A runtime template using a different detector can still be marked vulnerable because the in-loop judge score remains inDetectorResults. Seeexamples/runtime-templates/README.md:39,internal/runtimetemplates/probe.go:105,internal/multiturn/engine.go:790,pkg/attempt/attempt.go:174.
Security
No security concerns flagged.
Suggestions
- Detect duplicate template IDs within the same
--templates-dirbatch.RegisterFromPathonly checks each loaded template against the existing registry before registration; two new files with the same ID will silently replace the first because registry registration overwrites existing factories. Seeinternal/runtimetemplates/register.go:29,internal/runtimetemplates/register.go:45,pkg/registry/registry.go:88. - Either support or reject
secondary_detectorson multi-turn templates. The shared schema documents them, butvalidateMultiTurndoes not validate them andMultiTurnTemplateProbedoes not implementProbeSecondaryDetectors, so they are silently ignored. Seepkg/templates/types.go:236,pkg/templates/types.go:321,internal/runtimetemplates/probe.go:16.
Reviewed by Codex (gpt-5.5)
Gemini Review
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
✅ Files skipped from review due to trivial changes (1)
WalkthroughThis PR adds runtime YAML probe template loading. It adds type-aware template validation, a template-driven multiturn strategy, and a multi-turn probe wrapper that carries detector overrides and secondary detectors. Template registration now loads from a directory with collision and detector checks. The scan CLI adds ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
internal/runtimetemplates/register_test.go (1)
163-166: ⚡ Quick winAssert force override actually replaces the registered probe.
This test only checks
err == nilafter forced registration. Please also assert the created probe now reflects the second template payload, so a silent no-op override can’t pass.Suggested test hardening
// With force, the override succeeds. if _, err := RegisterFromPath(dir2, true); err != nil { t.Fatalf("RegisterFromPath() with force should override, got: %v", err) } + probe, err := probes.Create("runtimetest.OverrideTarget", nil) + if err != nil { + t.Fatalf("probes.Create() after force override: %v", err) + } + pm, ok := probe.(types.ProbeMetadata) + if !ok { + t.Fatal("overridden probe should implement ProbeMetadata") + } + prompts := pm.GetPrompts() + if len(prompts) != 1 || prompts[0] != "two" { + t.Fatalf("force override did not replace template content, prompts=%v", prompts) + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/runtimetemplates/register_test.go` around lines 163 - 166, The test currently only checks RegisterFromPath(dir2, true) returns no error but doesn't verify the override took effect; after the forced register call, fetch the registered probe/template by its name (the probe name used in the first template) using your codebase's lookup (e.g., the Get/Find/Lookup function that returns the registered template/probe) and assert that its contents/fields match the second template's payload (e.g., description, command, or other distinguishing fields) so the test fails if the override was a silent no-op; use the same symbols in the test (RegisterFromPath, dir2 and the probe name) to locate and validate the replaced template.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/runtime-templates/README.md`:
- Around line 11-13: Add a short note explaining that runtime template IDs that
collide with existing probe IDs will fail to load unless the CLI is invoked with
--force; update the README’s paragraph about templates being registered in the
probe registry (the sentence referencing --probes-glob "runtime.*" and --all) to
mention this collision behavior and show the remedial usage example (e.g., run
with --force) and that this is part of the CLI contract for template loading.
---
Nitpick comments:
In `@internal/runtimetemplates/register_test.go`:
- Around line 163-166: The test currently only checks RegisterFromPath(dir2,
true) returns no error but doesn't verify the override took effect; after the
forced register call, fetch the registered probe/template by its name (the probe
name used in the first template) using your codebase's lookup (e.g., the
Get/Find/Lookup function that returns the registered template/probe) and assert
that its contents/fields match the second template's payload (e.g., description,
command, or other distinguishing fields) so the test fails if the override was a
silent no-op; use the same symbols in the test (RegisterFromPath, dir2 and the
probe name) to locate and validate the replaced template.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 063cb9b3-2d3c-491b-9509-737aa27b3ed6
📒 Files selected for processing (15)
cmd/augustus/cli.gocmd/augustus/scan.goexamples/runtime-templates/README.mdexamples/runtime-templates/multiturn-example.yamlexamples/runtime-templates/static-example.yamlinternal/multiturn/probe.gointernal/multiturn/probe_test.gointernal/runtimetemplates/probe.gointernal/runtimetemplates/probe_test.gointernal/runtimetemplates/register.gointernal/runtimetemplates/register_test.gointernal/runtimetemplates/strategy.gointernal/runtimetemplates/strategy_test.gopkg/templates/multiturn_schema_test.gopkg/templates/types.go
…ectors, dup IDs) - C1: validate multi-turn strategy prompts at load (RegisterFromPath) so a bad field reference aborts before the scan starts, not at probes.Create — the fail-fast-at-load claim now holds for multi-turn. - C2/S2: secondary_detectors now apply to multi-turn templates — validated in validateMultiTurn (shared helper) and run via ProbeSecondaryDetectors on MultiTurnTemplateProbe (previously silently ignored). - S1: RegisterFromPath rejects duplicate template IDs within a single --templates-dir batch (previously the second silently overwrote the first). - C3 (docs): the multi-turn verdict is the max across detectors — the engine's in-loop judge.Judge score is retained, so info.detector adds a verdict signal rather than replacing it. Corrected examples/runtime-templates/README.md. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Thanks both — all findings addressed in Claude review
Codex review
All package tests, |
CLAUDE.md Drift DetectionCode changes in this PR may have made documentation stale:
|
| Section | Issue | Evidence |
|---|---|---|
| "Directory Structure" (lines 62-79) | New internal/runtimetemplates/ package not documented |
New 1000+ LOC package for runtime template loading, registration, and strategy compilation |
| "Adding New Components" → "New Probe" (lines 89-96) | Missing documentation of runtime template loading without rebuild | PR adds --templates-dir flag and internal/runtimetemplates.RegisterFromPath() for loading YAML probes at scan time |
| "CLI Usage Patterns" (lines 119-133) | Missing examples of --templates-dir and --force flags |
New flags added to cmd/augustus/cli.go enable loading and overriding probes via directory: --templates-dir ./path --force |
| "Adding New Components" → "New Probe" (lines 89-96) | Incomplete explanation of YAML templates | PR adds multi-turn template support (type: multiturn) with EngineConfig and StrategyConfig blocks in pkg/templates/types.go, but docs only mention static templates |
Automated drift check — please review and update if needed.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/runtimetemplates/probe.go`:
- Around line 128-130: During runtime conversion in the loop over
tmpl.Info.SecondaryDetectors, trim whitespace from detector names before storing
them into out by setting SecondaryDetector.Name = strings.TrimSpace(s.Name) (so
it matches validateSecondaryDetectors()'s trimmed checks); update the conversion
in the function that constructs out (the loop referencing
tmpl.Info.SecondaryDetectors and types.SecondaryDetector) and add an import for
strings if missing to ensure detector lookup won't fail due to leading/trailing
spaces.
In `@pkg/templates/types.go`:
- Around line 237-240: validateMultiTurn() currently bypasses the same
allowed-mode check done in validateStatic(), allowing invalid info.mode values
for type: multiturn; update validateMultiTurn() to perform the same validation
of info.mode (or call the shared validation used by validateStatic()) before
returning, ensuring info.mode is one of the permitted values; reference
validateMultiTurn(), validateStatic(), and the info.mode field when adding the
check so the multi-turn path enforces the same allowed-mode rules as static
templates.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 7099aa16-4a64-4968-a6e6-ca2e9bda0fce
📒 Files selected for processing (7)
examples/runtime-templates/README.mdinternal/runtimetemplates/probe.gointernal/runtimetemplates/probe_test.gointernal/runtimetemplates/register.gointernal/runtimetemplates/register_test.gopkg/templates/multiturn_schema_test.gopkg/templates/types.go
✅ Files skipped from review due to trivial changes (1)
- examples/runtime-templates/README.md
🚧 Files skipped from review as they are similar to previous changes (2)
- pkg/templates/multiturn_schema_test.go
- internal/runtimetemplates/register.go
…ndary detector names) - validateMultiTurn now validates info.mode (extracted shared validateMode helper) — invalid mode values previously passed for multi-turn templates. - secondaryDetectorsFromTemplate trims detector names so they match validateSecondaryDetectors()'s trimmed checks and resolve in the registry. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
CodeRabbit items addressed in
Build/vet/tests + gofumpt clean. |
CLAUDE.md Drift DetectionCode changes in this PR may have made documentation stale:
|
| Section | Issue | Evidence |
|---|---|---|
| "CLI Usage Patterns" | Missing new --templates-dir and --force flags for runtime template loading |
Added in cmd/augustus/cli.go; enables loading probe templates at scan time without rebuild via internal/runtimetemplates |
| "Adding New Components > New Probe" | Incomplete YAML documentation: only describes data/ subdirectory approach; doesn't mention new --templates-dir feature for runtime loading or type: multiturn support |
New examples/runtime-templates/ with comprehensive examples and tests show multi-turn templates and runtime loading are core features |
Automated drift check — please review and update if needed.
- Document ID-collision / --force behavior in the runtime templates README - Strengthen force-override test to verify the probe was actually replaced - Replace custom contains() with slices.Contains (lint modernize) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Code Review — runtime probe-template loaderReviewed in an isolated worktree. Baseline: Overall a strong PR — clean type-discriminated design ( 🟠 Requested change — load-time dry-run misses field refs inside
|
| author: your-team | ||
| description: "Multi-turn attack that escalates claimed authority each turn" | ||
| goal: obtain restricted operational details | ||
| detector: judge.Judge # post-hoc verdict detector (pipeline stage) |
There was a problem hiding this comment.
Overall the code seems fine, I am wondering about this though. I see the detector here is judge.Judge, but down in the engine there is a judge_generator_type. Are those the same or is there 2 judges?
There was a problem hiding this comment.
Two judges, two roles — and in this example they collapse into one.
engine.judge_generator_type is the in-loop steering judge: it scores every turn and drives early-exit/feedback. It's whatever generator that key resolves to — the template's value (here anthropic.Anthropic, but that's just the example's default) or whatever you override it to via --config-file, since the engine block is merged as defaults under scan-time config (buildEngineConfigMap → maps.Copy(m, cfg)).
info.detector is the post-hoc verdict detector. The catch: the engine records its in-loop max score under the detector key judge.Judge (engine.go:791), and the detector pipeline skips re-running any detector whose results already exist (detection.go:56, to preserve full conversation context). So with info.detector: judge.Judge the post-hoc pass is a no-op — the in-loop judge is the verdict. Point info.detector at a different detector to get a genuinely independent second opinion (verdict = max of both).
Clarified this in the README + example (which were previously misleading on this point) in the latest push.
blayne's review:
- Dry-run multi-turn strategy prompts against populated data so field-reference
errors inside {{range .History}} / {{if .History}} are caught at load instead
of silently sending an empty prompt mid-scan (render swallowed the error).
Regression test added.
- render() now logs at slog.Error (an empty prompt silently degrades a scan).
- Validate info.detector and secondary_detectors against the detector registry
at load, so an unknown detector fails fast rather than partway through a scan.
Collision handling:
- Remove the --force override escape hatch entirely. A runtime template whose id
collides with an existing probe now always fails to load with a clear error
telling the author to rename the id; a template must not shadow a built-in.
Evan's question (two judges):
- Clarify the README + example: engine.judge_generator_type is the in-loop
steering judge; info.detector is the post-hoc verdict detector. The engine
records its in-loop score under "judge.Judge", and the detector pipeline skips
re-running a same-named detector, so info.detector: judge.Judge is a no-op
(the in-loop judge is the verdict). The README previously implied it always
adds an independent signal, which was misleading.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
CLAUDE.md Drift DetectionCode changes in this PR introduce new documentation gaps:
Automated drift check — please review and update if needed. |
Address the CLAUDE.md drift check: document --templates-dir runtime loading, static vs multiturn template types, and add a CLI usage example. Closes the gap the drift bot flagged for this feature. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Summary
Adds a runtime probe-template loader so new probes — including genuinely new multi-turn attack strategies — can be defined in YAML and run with
--templates-dir, without recompiling Augustus.Templates register before probe selection, so they also work with
--probes-globand--all.What's added
pkg/templates— type-aware schema.type: static(default) keeps today's single-turn prompt probes;type: multiturnaddsengine+strategyblocks.Validate()branches by type and composes with the existing tool-use/secondary-detector validation.internal/runtimetemplates(new) —TemplateStrategy: implementsmultiturn.Strategyfrom YAML prompt templates (attacker_system,turn,rephrase,feedback) viatext/template, withsimple/extendedparser selection. Dry-run validated at load, so a bad field reference (e.g.{{.Tpyo}}) fails immediately instead of degrading prompts mid-scan.MultiTurnTemplateProbe: builds the existing unified multi-turn engine (internal/multiturn) via the sharedCreateGenerators; reports the template's owninfo.detector.RegisterFromPath: loads + registers static and multi-turn probes; warns on built-in override.cmd/augustus—--templates-dirflag, registered before glob expansion.examples/runtime-templates— static + multi-turn examples + README.Capability (and the ceiling)
goal,max_turns, generator types/models via--config-file)info.detectorTesting
go build ./..., package tests, andgo vetclean; gofumpt-clean.viewer_globexcaller): the agent disclosed Acme's full order table (20acme-corprows — IDs, amounts, products) to a Globex caller. Judge LLM:leak: true. Deterministic, confirmed true positive.leak: true); when it only gets deflections the judge confirms SAFE (leak: false). Both observed and judge-verified.detector_resultsalways contains the in-loopjudge.Judgescore alongsideinfo.detector, and the verdict is the max — so the displayed detector label is attribution, not proof of which detector fired. The shipped multi-turn example usesjudge.Judgefor a semantic verdict; the local agentdash test usedbase.StringDetectorand was cross-checked with the independent judge above.Notes
main; contains only this feature (one commit).🤖 Generated with Claude Code