feat(aios-runtime): add loop detection middleware#5
Conversation
📝 WalkthroughWalkthroughThe PR introduces loop detection middleware to the AIOS runtime by adding a Changes
Sequence Diagram(s)sequenceDiagram
participant Model
participant Runtime as Runtime
participant Guard as Loop Detection Guard
participant Decision
participant ToolExec as Tool Execution
Model->>Runtime: Emits ToolCall
Runtime->>Runtime: Normalize & hash<br/>tool-call signature
Runtime->>Guard: on_tool_call(context, call)
Guard->>Guard: Check sliding window<br/>for repetitions
alt Repeated call exceeds threshold
Guard->>Decision: Block decision
else Repeated call in warning threshold
Guard->>Decision: Warn decision
else New/allowed call
Guard->>Decision: Allow decision
end
Decision-->>Runtime: ToolCallGuardDecision
Runtime->>Runtime: Persist custom event<br/>(warning/hard-stop)
alt Decision is Block
Runtime->>Runtime: Force text-only<br/>response (skip tool call)
else Decision is Allow/Warn
Runtime->>ToolExec: Proceed to policy<br/>evaluation & execution
end
ToolExec-->>Model: Tool result or<br/>guard message
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 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: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/aios-runtime/src/lib.rs (1)
61-95: 🛠️ Refactor suggestion | 🟠 MajorDecouple the generic guard hook from loop-detection-specific payloads.
ToolCallGuardDecisionbakes in loop-only fields (repetitions,signature), andpersist_loop_guard_event()hardcodesloop_detection.*for every guard decision. A second guard type cannot use this API without inventing bogus loop metadata or emitting mislabeled journal entries. As per coding guidelines, "Keep public APIs version-aware and backward compatible unless explicitly changing contracts".Also applies to: 1513-1561
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@crates/aios-runtime/src/lib.rs` around lines 61 - 95, ToolCallGuardDecision currently embeds loop-detection-specific fields (repetitions, signature) and the ToolCallGuard API plus journal emission (persist_loop_guard_event) assumes loop metadata for every guard; refactor by removing loop-only fields from the general ToolCallGuardDecision and replace them with a generic, extensible payload (e.g., an optional metadata map or a GuardMetadata enum) so non-loop guards can provide their own structured info, add a concrete LoopDetection variant or struct for loop-related data used by the loop-detection guard, and update ToolCallGuard::on_tool_call callers and the journal emission site (persist_loop_guard_event usage) to detect the LoopDetection variant and only persist loop_* entries for that case while other metadata is recorded under a generic guard metadata path; update types referenced in this change: ToolCallGuardDecision, ToolCallGuard, and the journal persistence call sites to handle the new metadata shape.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@crates/aios-runtime/src/lib.rs`:
- Around line 137-152: LoopDetectionConfig currently allows invalid combinations
via Default and public fields; add a constructor (e.g., LoopDetectionConfig::new
or a builder) that validates invariants and returns a Result/Err on invalid
input, and update Default to call that constructor or return a guaranteed-valid
instance: ensure warning_threshold > 0, warning_threshold <= hard_stop_limit,
and both warning_threshold and hard_stop_limit <= window_size + 1; reject or
panic on invalid combos per crate conventions and add unit tests covering
boundary cases (0, warning==hard_stop, warning>hard_stop, thresholds ==
window_size+1, and > window_size+1) to prevent silently widening safety
boundaries, referencing struct LoopDetectionConfig and its Default impl.
- Around line 154-158: The LoopDetectionMiddleware currently stores session
repetition state in memory via the history_by_session Arc<Mutex<HashMap<String,
VecDeque<String>>>> field, which makes loop decisions non-replayable; change the
implementation so the repetition window is persisted to durable storage (e.g.,
workspace files or the event journal) instead of kept only in RAM: replace or
wrap history_by_session with a storage-backed provider used by the
LoopDetectionMiddleware constructor and all places that push/pop history
(reference LoopDetectionMiddleware, history_by_session, LoopDetectionConfig, and
any methods that read/update the VecDeque) so the middleware loads history on
init and appends updates to the durable log/file whenever the window changes,
ensuring behavior is replayable after restarts and consistent across runtime
attachments.
---
Outside diff comments:
In `@crates/aios-runtime/src/lib.rs`:
- Around line 61-95: ToolCallGuardDecision currently embeds
loop-detection-specific fields (repetitions, signature) and the ToolCallGuard
API plus journal emission (persist_loop_guard_event) assumes loop metadata for
every guard; refactor by removing loop-only fields from the general
ToolCallGuardDecision and replace them with a generic, extensible payload (e.g.,
an optional metadata map or a GuardMetadata enum) so non-loop guards can provide
their own structured info, add a concrete LoopDetection variant or struct for
loop-related data used by the loop-detection guard, and update
ToolCallGuard::on_tool_call callers and the journal emission site
(persist_loop_guard_event usage) to detect the LoopDetection variant and only
persist loop_* entries for that case while other metadata is recorded under a
generic guard metadata path; update types referenced in this change:
ToolCallGuardDecision, ToolCallGuard, and the journal persistence call sites to
handle the new metadata shape.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: e051d650-2938-4417-9959-9e24d2eafffa
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (6)
crates/aios-kernel/src/lib.rscrates/aios-runtime/Cargo.tomlcrates/aios-runtime/README.mdcrates/aios-runtime/src/lib.rsdocs/ARCHITECTURE.mddocs/STATUS.md
| #[derive(Debug, Clone)] | ||
| pub struct LoopDetectionConfig { | ||
| pub warning_threshold: usize, | ||
| pub hard_stop_limit: usize, | ||
| pub window_size: usize, | ||
| } | ||
|
|
||
| impl Default for LoopDetectionConfig { | ||
| fn default() -> Self { | ||
| Self { | ||
| warning_threshold: 3, | ||
| hard_stop_limit: 5, | ||
| window_size: 20, | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Validate LoopDetectionConfig invariants at construction time.
warning_threshold == 0, warning_threshold > hard_stop_limit, or either threshold being greater than window_size + 1 all create broken safety behavior from a public config surface. Reject invalid combinations in new/a builder and cover the boundaries with a couple of tests. As per coding guidelines, "do not silently widen capabilities or safety boundaries".
Also applies to: 160-166
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@crates/aios-runtime/src/lib.rs` around lines 137 - 152, LoopDetectionConfig
currently allows invalid combinations via Default and public fields; add a
constructor (e.g., LoopDetectionConfig::new or a builder) that validates
invariants and returns a Result/Err on invalid input, and update Default to call
that constructor or return a guaranteed-valid instance: ensure warning_threshold
> 0, warning_threshold <= hard_stop_limit, and both warning_threshold and
hard_stop_limit <= window_size + 1; reject or panic on invalid combos per crate
conventions and add unit tests covering boundary cases (0, warning==hard_stop,
warning>hard_stop, thresholds == window_size+1, and > window_size+1) to prevent
silently widening safety boundaries, referencing struct LoopDetectionConfig and
its Default impl.
| #[derive(Clone)] | ||
| pub struct LoopDetectionMiddleware { | ||
| config: LoopDetectionConfig, | ||
| history_by_session: Arc<Mutex<HashMap<String, VecDeque<String>>>>, | ||
| } |
There was a problem hiding this comment.
Persist the repetition window outside process memory.
history_by_session only lives in RAM, so warnings/hard-stops reset after a restart or a fresh runtime attaching to the same session. That makes loop decisions non-replayable from the journal and weakens the safety boundary during recovery. As per coding guidelines, "Avoid hidden mutable state outside workspace files and event logs" and "do not silently widen capabilities or safety boundaries".
Also applies to: 168-191
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@crates/aios-runtime/src/lib.rs` around lines 154 - 158, The
LoopDetectionMiddleware currently stores session repetition state in memory via
the history_by_session Arc<Mutex<HashMap<String, VecDeque<String>>>> field,
which makes loop decisions non-replayable; change the implementation so the
repetition window is persisted to durable storage (e.g., workspace files or the
event journal) instead of kept only in RAM: replace or wrap history_by_session
with a storage-backed provider used by the LoopDetectionMiddleware constructor
and all places that push/pop history (reference LoopDetectionMiddleware,
history_by_session, LoopDetectionConfig, and any methods that read/update the
VecDeque) so the middleware loads history on init and appends updates to the
durable log/file whenever the window changes, ensuring behavior is replayable
after restarts and consistent across runtime attachments.
Summary
Validation
cargo fmt --allcargo clippy --workspace -- -D warningscargo test --workspaceLinear: BRO-423
Summary by CodeRabbit
Release Notes
New Features
Documentation
Tests