Chore/stream handler cleanup - #64
Conversation
|
@CodeRabbit review |
✅ Action performedReview finished.
|
|
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:
📝 WalkthroughWalkthroughChangesThe PR adds typed provider responses, object-safe loop memory, configurable token counting, revised stream retries and timeouts, bounded provider reads, terminal-state reporting, and non-consuming tool availability checks. ChangesTyped provider and streaming flow
Loop memory and context accounting
Loop terminal-state reporting
Tool health availability inspection
Supporting updates
Possibly related PRs
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ 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: 4
🤖 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 `@src/stream/handler.rs`:
- Around line 285-291: Update jittered_base_delay to incorporate per-stream or
per-request entropy when calculating the jitter factor, rather than deriving it
solely from attempt. Preserve the zero-jitter fast path and ensure concurrent
streams with the same attempt can produce different retry delays.
- Around line 1727-1729: Update the transport retry path around
retry_config.jittered_base_delay and sleep_cancellable to pass the jittered
delay through clamp_delay_to_deadline using total_deadline before sleeping.
Preserve the existing attempt increment and cancellation/error propagation while
ensuring the backoff cannot extend beyond the total stream deadline.
- Around line 1456-1460: Update the invalid-configuration branches in
src/stream/handler.rs:1456-1460 and src/stream/handler.rs:1484-1488 so they
assign StreamTimeoutConfig::default() and StreamRetryConfig::default(),
respectively, when validation fails. Preserve the existing warning logs and
valid-configuration assignments.
In `@src/tool/health.rs`:
- Around line 614-631: Update would_be_half_open so it reports whether the next
allow_request call will transition into or operate as a HalfOpen probe: return
true only for an Open breaker whose recovery duration has elapsed, and false
when already HalfOpen or Closed. Revise the method’s rustdoc to describe this
next-request behavior and clarify that an existing HalfOpen state has a probe in
flight.
🪄 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 Plus
Run ID: c40d2150-878c-4d35-a98f-a3a1accc591c
📒 Files selected for processing (14)
CHANGELOG.mdsrc/compact.rssrc/engine/bare.rssrc/engine/bare/emission.rssrc/engine/core/lifecycle.rssrc/managers.rssrc/provider.rssrc/provider/anthropic.rssrc/provider/gemini.rssrc/provider/openai.rssrc/reflection.rssrc/stream/handler.rssrc/tool/health.rssrc/tool/registry.rs
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/engine/bare.rs (1)
1510-1555: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMemory retrieval query is empty on every turn after a tool call.
turn_input(used as thememory.retrieve(&turn_input, 3)query) is derived by filtering the previous history message's parts forText— but on any turn following a tool call, that message is aToolResult-only message, soturn_inputis"". InInMemoryStore::retrieve, an empty query makesquery_wordsempty, so the word-overlap bonus is always0.0, and retrieval degenerates to "top base-relevance entries" instead of anything actually relevant to the current turn — defeating the documented purpose ("retrieves relevant entries as context before each turn").💡 Suggested fix: fall back to the run's original input when turn_input is empty
- if let Some(memory) = self.managers.memory() { - match memory.retrieve(&turn_input, 3).await { + if let Some(memory) = self.managers.memory() { + let query = if turn_input.is_empty() { + self.current_run().map_or(String::new(), |r| r.input.clone()) + } else { + turn_input.clone() + }; + match memory.retrieve(&query, 3).await {🤖 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 `@src/engine/bare.rs` around lines 1510 - 1555, Update turn_input construction in handle_call_llm so that when the history-based input after a tool call is empty, it falls back to the current run’s original input from current_run().input. Preserve the existing first-turn and non-empty history behavior, ensuring memory.retrieve receives a meaningful query on every turn.
🤖 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 `@src/engine/bare/dispatch.rs`:
- Around line 786-812: Update record_tool_memory to bound the persisted
trajectory content before constructing MemoryEntry, truncating both tc.input and
tool_result.output to an explicit maximum length while preserving the existing
tool name, category, and store behavior. Use a clear shared cap and ensure
truncation cannot produce unbounded prompt content when entries are later
reused.
---
Outside diff comments:
In `@src/engine/bare.rs`:
- Around line 1510-1555: Update turn_input construction in handle_call_llm so
that when the history-based input after a tool call is empty, it falls back to
the current run’s original input from current_run().input. Preserve the existing
first-turn and non-empty history behavior, ensuring memory.retrieve receives a
meaningful query on every 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: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: a9fd9531-20b7-414d-836f-0d5c3ad088b9
📒 Files selected for processing (10)
CHANGELOG.mdCargo.tomlsrc/capabilities.rssrc/engine/bare.rssrc/engine/bare/dispatch.rssrc/managers.rssrc/memory.rssrc/memory/builtin.rssrc/stream/handler.rssrc/tool/health.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- CHANGELOG.md
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
src/stream/handler.rs (3)
1531-1540: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winRemove the leftover doc fragment.
Line 1536 is a remnant of the removed
with_configdocumentation. It has an unmatched parenthesis and contradicts the new independent setters.📝 Proposed doc fix
/// Read-only access to the [`StreamRetryConfig`] stored on the handler. /// Mutate via /// [`with_retry_config`](Self::with_retry_config). - /// both retry and timeout together); there is no per-field setter. #[must_use]🤖 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 `@src/stream/handler.rs` around lines 1531 - 1540, Remove the stray “both retry and timeout together” documentation fragment, including its unmatched parenthesis, from the `retry_config` method docs. Keep the remaining description of read-only access and mutation via `with_retry_config` intact.
1765-1790: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winMove the
rate_limit_retrydoc block back to its function.Lines 1765-1777 document
rate_limit_retry(they describeRateLimitRetry::HardStop,Escalate, andRetry). They are now contiguous with thedecide_rate_limit_errordoc block that starts at line 1778, so both blocks attach todecide_rate_limit_error.rate_limit_retryat line 1850 is left undocumented, and the merged rustdoc describes two different decisions.Move lines 1765-1777 directly above
fn rate_limit_retry.🤖 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 `@src/stream/handler.rs` around lines 1765 - 1790, Move the rustdoc paragraph describing count increments and the RateLimitRetry::HardStop, Escalate, and Retry outcomes from above decide_rate_limit_error to directly above rate_limit_retry. Keep the decide_rate_limit_error documentation focused only on mapping rate-limit stream errors to ErrorAction.
2221-2233: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winPropagate fallback usage in
drive_turnso the test asserts the real contract.
drive_turndiscards the newusagefield and hardcodesusage: None. The test at lines 3503-3506 then assertsresult.usage.is_none(), and its comment states that the non-streaming path does not carry token counts. The mock returnsSome(Usage::default()), and the production consumerBareLoop::stream_turninsrc/engine/bare/stream.rsreturnsfallback_usage. The assertion now hides a regression in usage propagation.💚 Proposed test fix
HandlerEvent::Fallback { message, stop_reason: fallback_stop_reason, - .. + usage: fallback_usage, } => { from_fallback = true; return Ok(DriveResult { message, - usage: None, + usage: fallback_usage, stop_reason: fallback_stop_reason, from_fallback, }); }And at lines 3503-3506:
- // Usage is None on the fallback path (non-streaming JSON doesn't - // reliably carry token counts). - assert!(result.usage.is_none()); + // Usage comes from the typed non-streaming response. + assert_eq!(result.usage, Some(Usage::default()));Also applies to: 3503-3506
🤖 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 `@src/stream/handler.rs` around lines 2221 - 2233, Update the HandlerEvent::Fallback branch in drive_turn to propagate the fallback event’s usage value into DriveResult instead of hardcoding usage: None, preserving the existing message, stop_reason, and from_fallback fields. Update the related test assertion and comment to expect the mock’s Some(Usage::default()) result, validating real usage propagation.src/engine/bare.rs (1)
1596-1601: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winDo not inject tool-derived memory as system instructions.
Lines 1596-1601 add
MemoryEntry.memoryasRole::System.record_tool_memorystores successful tool input and output, so untrusted tool content can become high-priority instructions on a later model call.Encode retrieved memory as lower-trust data and explicitly delimit it as untrusted reference content.
🤖 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 `@src/engine/bare.rs` around lines 1596 - 1601, Update the contributor_messages construction around record_tool_memory retrieval so the memory is not emitted with Role::System. Encode the retrieved summary as lower-trust reference content using the appropriate non-instruction role, and explicitly delimit it as untrusted data while preserving the existing “Relevant memory” context.
🧹 Nitpick comments (1)
src/testing.rs (1)
736-748: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAlign the non-streaming mock with the streaming mock shape.
stream_messagesemits a text part and a tool-call part, and reportsUsage::new(50, 25).create_messagedrops the text whentool_callis set, and reportsUsage::default()(zeros). Tests that exercise the streaming-to-non-streaming fallback then observe a different message shape and different token counts depending on which path ran. That weakens coverage for the fallback behavior this PR changes.♻️ Proposed fix
let response = self.pop_response(); Box::pin(async move { let stop_reason = crate::stream::StreamStopReason::from_api_str(&response.stop_reason) .unwrap_or(crate::stream::StreamStopReason::EndTurn); - let parts = if let Some(tc) = response.tool_call { - vec![MessagePart::tool_call(tc.id, tc.name, tc.input)] - } else { - vec![MessagePart::text(response.text)] - }; + let mut parts = vec![MessagePart::text(response.text)]; + if let Some(tc) = response.tool_call { + parts.push(MessagePart::tool_call(tc.id, tc.name, tc.input)); + } Ok(crate::api::NonStreamingResponse { message: Message::new(Role::Assistant, parts), stop_reason, - usage: Some(crate::stream::Usage::default()), + usage: Some(Usage::new(50, 25)), }) })🤖 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 `@src/testing.rs` around lines 736 - 748, Update the non-streaming mock in create_message to match stream_messages: always include the response text part, append the tool-call part when response.tool_call is present, and report Usage::new(50, 25) instead of Usage::default(). Preserve the existing stop-reason conversion and response construction.
🤖 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 `@src/compact.rs`:
- Around line 113-132: Clean up the HeuristicTokenCounter documentation by
removing the duplicated introductory block and the misleading “Presets” section.
Keep one accurate description that refers to the implemented CHARS_PER_TOKEN
constant rather than the nonexistent bytes_per_token parameter.
In `@src/engine/bare.rs`:
- Around line 677-680: Update set_token_counter so the newly installed counter
is also propagated to the configured ContextManager, preserving synchronization
regardless of setter order; alternatively reject calls made after
context-manager configuration. Add tests covering set_token_counter before and
after set_context_manager.
- Around line 1679-1680: Update the context token calculation in the
model-response handling flow around model_response so it includes the current
model message before estimating context size. Ensure the message that
model_response appends to pending history is represented in the counted history,
allowing compaction to trigger before the context budget is exceeded.
In `@src/engine/bare/dispatch.rs`:
- Around line 7-15: Update truncate_to to determine whether truncation is needed
using character indexing rather than s.len() byte length. Check for a character
at max_len and return s unchanged when none exists; otherwise retain only the
first max_len characters and append the ellipsis.
In `@src/message.rs`:
- Around line 412-419: Add the serde default attribute to the name field in
MessagePart::ToolResult so deserialization of older serialized tool-result data
missing this field succeeds. Preserve the existing String type and current
behavior for newly serialized results.
In `@src/provider/gemini.rs`:
- Around line 169-181: Update the function-call id sentence in the
NonStreamingResponse documentation to state that the parsed functionCall.id is
preserved, matching the implementation and
build_response_parses_function_call_id assertion; remove the claim that Gemini
calls have no id or that it is left empty.
---
Outside diff comments:
In `@src/engine/bare.rs`:
- Around line 1596-1601: Update the contributor_messages construction around
record_tool_memory retrieval so the memory is not emitted with Role::System.
Encode the retrieved summary as lower-trust reference content using the
appropriate non-instruction role, and explicitly delimit it as untrusted data
while preserving the existing “Relevant memory” context.
In `@src/stream/handler.rs`:
- Around line 1531-1540: Remove the stray “both retry and timeout together”
documentation fragment, including its unmatched parenthesis, from the
`retry_config` method docs. Keep the remaining description of read-only access
and mutation via `with_retry_config` intact.
- Around line 1765-1790: Move the rustdoc paragraph describing count increments
and the RateLimitRetry::HardStop, Escalate, and Retry outcomes from above
decide_rate_limit_error to directly above rate_limit_retry. Keep the
decide_rate_limit_error documentation focused only on mapping rate-limit stream
errors to ErrorAction.
- Around line 2221-2233: Update the HandlerEvent::Fallback branch in drive_turn
to propagate the fallback event’s usage value into DriveResult instead of
hardcoding usage: None, preserving the existing message, stop_reason, and
from_fallback fields. Update the related test assertion and comment to expect
the mock’s Some(Usage::default()) result, validating real usage propagation.
---
Nitpick comments:
In `@src/testing.rs`:
- Around line 736-748: Update the non-streaming mock in create_message to match
stream_messages: always include the response text part, append the tool-call
part when response.tool_call is present, and report Usage::new(50, 25) instead
of Usage::default(). Preserve the existing stop-reason conversion and response
construction.
🪄 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 Plus
Run ID: 0f0b8d47-c94e-4506-a0a3-404a890e2e4b
📒 Files selected for processing (21)
CHANGELOG.mdsrc/api.rssrc/compact.rssrc/compact/truncating.rssrc/compact/types.rssrc/engine/bare.rssrc/engine/bare/dispatch.rssrc/engine/bare/message.rssrc/engine/bare/stream.rssrc/engine/core/lifecycle.rssrc/engine/core/machine.rssrc/message.rssrc/provider.rssrc/provider/anthropic.rssrc/provider/gemini.rssrc/provider/openai.rssrc/reflection/llm.rssrc/stream.rssrc/stream/handler.rssrc/structured.rssrc/testing.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- src/provider.rs
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/provider/gemini.rs (1)
715-740: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftValidate legacy
ToolResultnames before building Gemini requests.Old tool results deserialize with
name == "", and this serializer copies that value unconditionally intofunctionResponse.name. Recover the missing function name from a matchingToolCallbycall_id, or reject the request before sending it. Add a round-trip test that deserializes legacy tool-result JSON and passes it throughbuild_request_body; the existing test only checks deserialization.🤖 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 `@src/provider/gemini.rs` around lines 715 - 740, Update the Gemini request-building flow around MessagePart::ToolResult to validate empty result names before constructing functionResponse. Resolve a missing name by matching call_id against the corresponding MessagePart::ToolCall, or reject the request if no match exists; preserve existing names. Extend the relevant round-trip test to deserialize legacy tool-result JSON and pass it through build_request_body, asserting the recovered name or expected rejection.src/provider/openai.rs (2)
190-195: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftReject invalid OpenAI tool argument JSON.
build_responsecurrently turns anyfunction.argumentsparse failure intoToolCall { input: {} }, so malformed JSON such as"hi",{, or{"q"becomes the same no-argument call as"". In non-streaming responses the argument parsing takes place insidebuild_response, so return anApiResponseInvalidfor non-empty invalid JSON; only missing or empty arguments should default to{}`. Update the OpenAI argument tests to cover both malformed and omitted arguments.🤖 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 `@src/provider/openai.rs` around lines 190 - 195, Update build_response’s function.arguments handling to default to an empty object only when arguments are missing or empty, and return ApiResponseInvalid when non-empty JSON fails to parse. Ensure this behavior applies to non-streaming responses, and extend the OpenAI argument tests to cover malformed and omitted arguments.
208-211: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winNormalize zero OpenAI usage to
Nonein both response paths.
Map’susagefield is missing by OpenAI default, but theUsagestruct is zeroed by default when a provider sends empty or absent usage values. The streaming path still carries that zeroed value from the emptychoicesusage-only chunk. If a downstream model supplies zero input and output counts, returnNoneinstead ofSome(Usage::default())so it is not treated as a final token count.🤖 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 `@src/provider/openai.rs` around lines 208 - 211, Normalize zero-valued OpenAI usage to None in both response paths, including the raw response mapping around OpenAiUsage::deserialize and the streaming usage-only chunk handling. After constructing Usage, return Some only when input and output token counts are not both zero; otherwise return None, preserving existing behavior for nonzero usage.
🤖 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.
Outside diff comments:
In `@src/provider/gemini.rs`:
- Around line 715-740: Update the Gemini request-building flow around
MessagePart::ToolResult to validate empty result names before constructing
functionResponse. Resolve a missing name by matching call_id against the
corresponding MessagePart::ToolCall, or reject the request if no match exists;
preserve existing names. Extend the relevant round-trip test to deserialize
legacy tool-result JSON and pass it through build_request_body, asserting the
recovered name or expected rejection.
In `@src/provider/openai.rs`:
- Around line 190-195: Update build_response’s function.arguments handling to
default to an empty object only when arguments are missing or empty, and return
ApiResponseInvalid when non-empty JSON fails to parse. Ensure this behavior
applies to non-streaming responses, and extend the OpenAI argument tests to
cover malformed and omitted arguments.
- Around line 208-211: Normalize zero-valued OpenAI usage to None in both
response paths, including the raw response mapping around
OpenAiUsage::deserialize and the streaming usage-only chunk handling. After
constructing Usage, return Some only when input and output token counts are not
both zero; otherwise return None, preserving existing behavior for nonzero
usage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: a47bf00f-5768-4064-bc90-cafb7dd24af8
📒 Files selected for processing (12)
CHANGELOG.mdsrc/compact.rssrc/engine/bare.rssrc/engine/bare/dispatch.rssrc/engine/core/machine.rssrc/message.rssrc/provider.rssrc/provider/anthropic.rssrc/provider/gemini.rssrc/provider/openai.rssrc/provider/sse.rssrc/stream/handler.rs
🚧 Files skipped from review as they are similar to previous changes (6)
- src/provider.rs
- src/engine/bare/dispatch.rs
- src/compact.rs
- src/engine/core/machine.rs
- src/stream/handler.rs
- src/engine/bare.rs
1 participant
No description provided.