Repository navigation
feat(aios-runtime): add turn middleware chain - #4
Merged
Merged
Conversation
📝 WalkthroughWalkthroughA middleware system is introduced to the kernel runtime. The Changes
Sequence DiagramsequenceDiagram
participant Client
participant KernelRuntime
participant TurnContext
participant Middleware1
participant MiddlewareN
participant TurnNext
participant ExecuteTurn
Client->>KernelRuntime: tick_on_branch(input)
activate KernelRuntime
KernelRuntime->>TurnContext: build(session, branch, input, state)
activate TurnContext
TurnContext->>TurnContext: set pending_approvals, estimated_mode
deactivate TurnContext
KernelRuntime->>TurnNext: run(middlewares[0], ctx)
activate TurnNext
TurnNext->>Middleware1: process(ctx, next)
activate Middleware1
note over Middleware1: inspect/mutate ctx
Middleware1->>TurnNext: next.run()
deactivate Middleware1
TurnNext->>MiddlewareN: process(ctx, next)
activate MiddlewareN
note over MiddlewareN: inspect/mutate ctx
MiddlewareN->>TurnNext: next.run()
deactivate MiddlewareN
TurnNext->>ExecuteTurn: execute_turn(ctx)
activate ExecuteTurn
note over ExecuteTurn: sense, estimate, gate, execute, commit, reflect
ExecuteTurn-->>TurnNext: TickOutput
deactivate ExecuteTurn
deactivate TurnNext
KernelRuntime-->>Client: TickOutput
deactivate KernelRuntime
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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 59-68: TurnContext's session/branch identity and workspace
metadata are publicly mutable, allowing middleware to change session_id,
branch_id or manifest.workspace_root; make these fields immutable by removing
pub from session_id, branch_id and manifest (or make TurnContext a
private/internal struct) and provide read-only accessors (e.g., session_id(),
branch_id(), manifest()) and a constructor/builder for creating TurnContext
instances; update execute_turn and any callers to use the new getters instead of
direct field access so all mutations must go through controlled runtime APIs.
- Around line 327-355: The code currently only checks the session and then runs
middlewares; instead, validate the branch lifecycle before constructing
TurnContext and invoking TurnNext::new(...).run(...) — ensure you look up the
branch by branch_id (using the same session lookup pattern around
self.sessions.lock() or the branch-store accessor), reject missing branches,
reject merged/read-only branches, enforce bounded fork sequence and disallow
using "main" as a merge source, and return an error early (before creating
TurnContext or calling turn_middlewares) so middleware never executes for
invalid/merged branches; modify the block that builds manifest/state and
immediately add a branch validation step referencing session_id and branch_id,
then only proceed to build TurnContext and call TurnNext::new(self,
&self.turn_middlewares).run(&mut ctx).await when the branch check passes.
- Around line 70-98: The current TurnMiddleware::process signature allows
middleware to return or mutate TickOutput and skip the terminal runtime path
(bypassing KernelRuntime::execute_turn), so change the middleware shape to
enforce the runtime path and explicit post-processing: replace the single
process method with explicit before/after hooks (e.g., before(&self, ctx: &mut
TurnContext) -> Result<()> and after(&self, ctx: &mut TurnContext, output: &mut
TickOutput) -> Result<()>) or otherwise require that process always call
TurnNext::run(ctx) and only perform side effects via an after hook; update
TurnNext/TurnNext::run to always delegate to KernelRuntime::execute_turn for the
terminal step and then invoke middleware after hooks in order, and adjust call
sites that implement TurnMiddleware to implement the new before/after methods
(references: TurnMiddleware, TurnNext, TurnNext::run, TurnContext, TickOutput,
KernelRuntime::execute_turn).
🪄 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: 922933a7-4253-4c3a-b931-6cd72728f725
⛔ 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
Comment on lines
+59
to
+68
| #[derive(Debug, Clone)] | ||
| pub struct TurnContext { | ||
| pub session_id: SessionId, | ||
| pub branch_id: BranchId, | ||
| pub manifest: SessionManifest, | ||
| pub input: TickInput, | ||
| pub state: AgentStateVector, | ||
| pub pending_approvals: Vec<ApprovalTicket>, | ||
| pub mode: OperatingMode, | ||
| } |
There was a problem hiding this comment.
Keep session/branch identity and workspace metadata immutable.
execute_turn() later trusts ctx.session_id, ctx.branch_id, and ctx.manifest.workspace_root for event appends, tool execution, checkpointing, and persisted state. With these fields publicly mutable, a buggy middleware can redirect a turn into another session/branch or arbitrary workspace path without going through any runtime guard.
🔒 Suggested shape
pub struct TurnContext {
- pub session_id: SessionId,
- pub branch_id: BranchId,
- pub manifest: SessionManifest,
+ session_id: SessionId,
+ branch_id: BranchId,
+ manifest: SessionManifest,
pub input: TickInput,
pub state: AgentStateVector,
pub pending_approvals: Vec<ApprovalTicket>,
pub mode: OperatingMode,
}
+
+impl TurnContext {
+ pub fn session_id(&self) -> &SessionId {
+ &self.session_id
+ }
+
+ pub fn branch_id(&self) -> &BranchId {
+ &self.branch_id
+ }
+
+ pub fn manifest(&self) -> &SessionManifest {
+ &self.manifest
+ }
+}As per coding guidelines, "Prefer deterministic file paths under session workspace" and "Route session state changes through aios-runtime."
🤖 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 59 - 68, TurnContext's
session/branch identity and workspace metadata are publicly mutable, allowing
middleware to change session_id, branch_id or manifest.workspace_root; make
these fields immutable by removing pub from session_id, branch_id and manifest
(or make TurnContext a private/internal struct) and provide read-only accessors
(e.g., session_id(), branch_id(), manifest()) and a constructor/builder for
creating TurnContext instances; update execute_turn and any callers to use the
new getters instead of direct field access so all mutations must go through
controlled runtime APIs.
Comment on lines
+70
to
+98
| #[async_trait] | ||
| pub trait TurnMiddleware: Send + Sync { | ||
| async fn process(&self, ctx: &mut TurnContext, next: TurnNext<'_>) -> Result<TickOutput>; | ||
| } | ||
|
|
||
| pub struct TurnNext<'a> { | ||
| runtime: &'a KernelRuntime, | ||
| middlewares: &'a [Arc<dyn TurnMiddleware>], | ||
| } | ||
|
|
||
| impl<'a> TurnNext<'a> { | ||
| fn new(runtime: &'a KernelRuntime, middlewares: &'a [Arc<dyn TurnMiddleware>]) -> Self { | ||
| Self { | ||
| runtime, | ||
| middlewares, | ||
| } | ||
| } | ||
|
|
||
| pub async fn run(self, ctx: &mut TurnContext) -> Result<TickOutput> { | ||
| match self.middlewares.split_first() { | ||
| Some((middleware, remaining)) => { | ||
| middleware | ||
| .process(ctx, TurnNext::new(self.runtime, remaining)) | ||
| .await | ||
| } | ||
| None => self.runtime.execute_turn(ctx).await, | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
This API still lets middleware bypass the terminal runtime path.
A middleware can return its own TickOutput or error without calling next.run(ctx), and it can also rewrite the TickOutput after next returns. That skips or desynchronizes event emission, policy/tool enforcement, sequence accounting, and persisted state, which is the opposite of the contract this feature is supposed to preserve.
🧭 Safer shape
#[async_trait]
pub trait TurnMiddleware: Send + Sync {
- async fn process(&self, ctx: &mut TurnContext, next: TurnNext<'_>) -> Result<TickOutput>;
+ async fn process(&self, ctx: &mut TurnContext) -> Result<()>;
}- TurnNext::new(self, &self.turn_middlewares)
- .run(&mut ctx)
- .await
+ for middleware in &self.turn_middlewares {
+ middleware.process(&mut ctx).await?;
+ }
+ self.execute_turn(&mut ctx).awaitIf you need post-processing, split it into explicit before/after hooks instead of letting middleware own TickOutput.
As per coding guidelines, "Emit events for meaningful transitions and preserve monotonic sequence per session and per branch" and "Route tool execution flows through aios-tools + aios-sandbox."
🤖 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 70 - 98, The current
TurnMiddleware::process signature allows middleware to return or mutate
TickOutput and skip the terminal runtime path (bypassing
KernelRuntime::execute_turn), so change the middleware shape to enforce the
runtime path and explicit post-processing: replace the single process method
with explicit before/after hooks (e.g., before(&self, ctx: &mut TurnContext) ->
Result<()> and after(&self, ctx: &mut TurnContext, output: &mut TickOutput) ->
Result<()>) or otherwise require that process always call TurnNext::run(ctx) and
only perform side effects via an after hook; update TurnNext/TurnNext::run to
always delegate to KernelRuntime::execute_turn for the terminal step and then
invoke middleware after hooks in order, and adjust call sites that implement
TurnMiddleware to implement the new before/after methods (references:
TurnMiddleware, TurnNext, TurnNext::run, TurnContext, TickOutput,
KernelRuntime::execute_turn).
Comment on lines
+327
to
+355
| let (manifest, state) = { | ||
| let sessions = self.sessions.lock(); | ||
| let session = sessions | ||
| .get(session_id.as_str()) | ||
| .with_context(|| format!("session not found: {session_id}"))?; | ||
| (session.manifest.clone(), session.state_vector.clone()) | ||
| }; | ||
|
|
||
| let pending_approvals = self | ||
| .approvals | ||
| .list_pending(session_id.clone()) | ||
| .await | ||
| .unwrap_or_default(); | ||
| let mode = self.estimate_mode(&state, pending_approvals.len()); | ||
|
|
||
| let mut ctx = TurnContext { | ||
| session_id: session_id.clone(), | ||
| branch_id: branch_id.clone(), | ||
| manifest, | ||
| input, | ||
| state, | ||
| pending_approvals, | ||
| mode, | ||
| }; | ||
|
|
||
| TurnNext::new(self, &self.turn_middlewares) | ||
| .run(&mut ctx) | ||
| .await | ||
| } |
There was a problem hiding this comment.
Reject invalid or merged branches before entering middleware.
This path only verifies that the session exists. A tick against a nonexistent or read-only branch still runs middleware and fails only when execute_turn() appends the first event, which widens the branch read-only boundary around user-supplied extension code.
🛑 Fail closed before middleware
- let (manifest, state) = {
+ let (manifest, state) = {
let sessions = self.sessions.lock();
let session = sessions
.get(session_id.as_str())
.with_context(|| format!("session not found: {session_id}"))?;
+ let branch = session
+ .branches
+ .get(branch_id)
+ .with_context(|| format!("branch not found: {}", branch_id.as_str()))?;
+ if let Some(target) = &branch.merged_into {
+ bail!(
+ "branch {} is merged into {} and is read-only",
+ branch_id.as_str(),
+ target.as_str()
+ );
+ }
(session.manifest.clone(), session.state_vector.clone())
};Based on learnings, "Enforce branch lifecycle invariants: bounded fork sequence, merged branches read-only, main not used as merge source."
🤖 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 327 - 355, The code currently
only checks the session and then runs middlewares; instead, validate the branch
lifecycle before constructing TurnContext and invoking
TurnNext::new(...).run(...) — ensure you look up the branch by branch_id (using
the same session lookup pattern around self.sessions.lock() or the branch-store
accessor), reject missing branches, reject merged/read-only branches, enforce
bounded fork sequence and disallow using "main" as a merge source, and return an
error early (before creating TurnContext or calling turn_middlewares) so
middleware never executes for invalid/merged branches; modify the block that
builds manifest/state and immediately add a branch validation step referencing
session_id and branch_id, then only proceed to build TurnContext and call
TurnNext::new(self, &self.turn_middlewares).run(&mut ctx).await when the branch
check passes.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
TurnMiddlewarechain to the canonicalKernelRuntime::tick_on_branchpathTurnContextplus a terminal executor so middleware can rewrite turn input/state without bypassing runtime behaviorValidation
cargo fmt --allcargo clippy --workspace -- -D warningscargo test --workspaceLinear: BRO-422
Summary by CodeRabbit
Release Notes
New Features
Documentation