Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 73 additions & 2 deletions src/apps/desktop/src/api/event_coalescer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
//! WebView costs one Tauri IPC message (JSON serialization, WebView2 boundary
//! crossing, JS parse + dispatch) and, when peer devices are attached, one
//! end-to-end encrypted relay message. This module merges chunks of the same
//! stream (session / turn / round / attempt / contentType) within a short
//! stream (session / turn / round / attempt / contentType / reasoningKind) within a short
//! window so the frontend still receives content-equivalent events at a
//! fraction of the message rate.
//!
Expand All @@ -24,6 +24,7 @@
//! emits thinking chunks before text chunks, and `flush` therefore emits the
//! merged thinking event before the merged text event.

use bitfun_core_types::ReasoningContentKind;
use bitfun_events::AgenticEvent;
use std::collections::HashMap;
use std::time::Duration;
Expand Down Expand Up @@ -118,7 +119,14 @@ pub fn update_rate_ema(previous: f64, flushed_chars: usize, elapsed: Duration) -
pub const INITIAL_RATE_EMA_CPS: f64 = WINDOW_REF_CPS;

/// Stable merge key for one streaming text/thinking stream.
type ChunkStreamKey = (String, String, String, String, bool);
type ChunkStreamKey = (
String,
String,
String,
String,
bool,
Option<ReasoningContentKind>,
);

fn resolve_attempt_token(attempt_id: &Option<String>, attempt_index: Option<u32>) -> String {
if let Some(id) = attempt_id {
Expand Down Expand Up @@ -148,6 +156,7 @@ enum PendingChunk {
attempt_id: Option<String>,
attempt_index: Option<u32>,
content: String,
reasoning_kind: Option<ReasoningContentKind>,
is_end: bool,
},
}
Expand Down Expand Up @@ -177,6 +186,7 @@ impl PendingChunk {
attempt_id,
attempt_index,
content,
reasoning_kind,
is_end,
} => AgenticEvent::ThinkingChunk {
session_id,
Expand All @@ -185,6 +195,7 @@ impl PendingChunk {
attempt_id,
attempt_index,
content,
reasoning_kind,
is_end,
},
}
Expand Down Expand Up @@ -251,6 +262,7 @@ impl TextChunkCoalescer {
round_id.clone(),
resolve_attempt_token(&attempt_id, attempt_index),
false,
None,
);
match self.pending.get_mut(&key) {
Some(PendingChunk::Text { text: pending, .. }) => {
Expand Down Expand Up @@ -284,6 +296,7 @@ impl TextChunkCoalescer {
attempt_id,
attempt_index,
content,
reasoning_kind,
is_end,
} => {
let key = (
Expand All @@ -292,6 +305,7 @@ impl TextChunkCoalescer {
round_id.clone(),
resolve_attempt_token(&attempt_id, attempt_index),
true,
reasoning_kind,
);
match self.pending.get_mut(&key) {
Some(PendingChunk::Thinking {
Expand All @@ -315,6 +329,7 @@ impl TextChunkCoalescer {
attempt_id,
attempt_index,
content,
reasoning_kind,
is_end,
},
);
Expand Down Expand Up @@ -389,6 +404,24 @@ mod tests {
attempt_id: attempt_id.map(str::to_string),
attempt_index,
content: content.to_string(),
reasoning_kind: None,
is_end,
}
}

fn typed_thinking_chunk(
content: &str,
reasoning_kind: ReasoningContentKind,
is_end: bool,
) -> AgenticEvent {
AgenticEvent::ThinkingChunk {
session_id: "s".to_string(),
turn_id: "t".to_string(),
round_id: "r".to_string(),
attempt_id: None,
attempt_index: None,
content: content.to_string(),
reasoning_kind: Some(reasoning_kind),
is_end,
}
}
Expand Down Expand Up @@ -474,6 +507,44 @@ mod tests {
}
}

#[test]
fn keeps_reasoning_text_and_summary_streams_separate() {
let mut coalescer = TextChunkCoalescer::new();
assert!(coalescer
.push(typed_thinking_chunk(
"private chain",
ReasoningContentKind::Reasoning,
false,
))
.is_empty());
assert!(coalescer
.push(typed_thinking_chunk(
"display summary",
ReasoningContentKind::Summary,
false,
))
.is_empty());

let events = coalescer.flush();
assert_eq!(events.len(), 2);
assert!(matches!(
&events[0],
AgenticEvent::ThinkingChunk {
content,
reasoning_kind: Some(ReasoningContentKind::Reasoning),
..
} if content == "private chain"
));
assert!(matches!(
&events[1],
AgenticEvent::ThinkingChunk {
content,
reasoning_kind: Some(ReasoningContentKind::Summary),
..
} if content == "display summary"
));
}

#[test]
fn flush_preserves_first_arrival_order_across_streams() {
let mut coalescer = TextChunkCoalescer::new();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ pub(crate) async fn aggregate_stream_response(
let UnifiedResponse {
text,
reasoning_content,
reasoning_content_kind: _,
thinking_signature: _,
tool_call,
usage: chunk_usage,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,18 @@ fn log_prompt_cache_diagnostics(request_body: &serde_json::Value) {
);
}

fn ensure_reasoning_summary_opt_in(request_body: &mut serde_json::Value) {
let Some(reasoning) = request_body
.get_mut("reasoning")
.and_then(serde_json::Value::as_object_mut)
else {
return;
};
reasoning
.entry("summary".to_string())
.or_insert_with(|| serde_json::Value::String("auto".to_string()));
}

fn try_build_request_body_with_context(
client: &AIClient,
instructions: Option<String>,
Expand Down Expand Up @@ -111,7 +123,7 @@ fn try_build_request_body_with_context(
if value.trim().is_empty() {
return Err(anyhow!("Responses reasoning effort must not be empty"));
}
body["reasoning"] = serde_json::json!({ "effort": value });
body["reasoning"] = serde_json::json!({ "effort": value, "summary": "auto" });
Ok(true)
}
ReasoningPresetAction::Toggle { .. } | ReasoningPresetAction::BudgetTokens { .. } => {
Expand Down Expand Up @@ -164,6 +176,7 @@ fn try_build_request_body_with_context(
);
shared::apply_reasoning_actions(preset, &mut request_body, protected_keys, &[], compile)?;
}
ensure_reasoning_summary_opt_in(&mut request_body);

shared::log_request_body(
TARGET,
Expand Down Expand Up @@ -301,6 +314,9 @@ mod tests {
use super::{build_request_body, build_request_body_with_context};
use crate::types::{ModelRequestContext, ToolDefinition};
use crate::{client::AIClient, types::AIConfig};
use bitfun_core_types::{
ReasoningPresetAction, ReasoningPresetDescriptor, ReasoningPresetSource,
};
use serde_json::json;

fn test_client() -> AIClient {
Expand Down Expand Up @@ -370,6 +386,51 @@ mod tests {
assert!(request_body.get("include").is_none());
}

#[test]
fn responses_reasoning_effort_requests_auto_summary() {
let client = test_client().with_reasoning_preset(&ReasoningPresetDescriptor {
id: "high".to_string(),
label: "High".to_string(),
order: 0,
actions: vec![ReasoningPresetAction::Effort {
value: "high".to_string(),
}],
source: ReasoningPresetSource::ModelConfig,
execution_provider: None,
execution_model: None,
});
let request_body = build_request_body(&client, None, Vec::new(), None, None);

assert_eq!(
request_body["reasoning"],
json!({ "effort": "high", "summary": "auto" })
);
}

#[test]
fn responses_reasoning_summary_preserves_explicit_override() {
let concise_request_body = build_request_body(
&test_client(),
None,
Vec::new(),
None,
Some(json!({ "reasoning": { "effort": "low", "summary": "concise" } })),
);
let disabled_request_body = build_request_body(
&test_client(),
None,
Vec::new(),
None,
Some(json!({ "reasoning": { "effort": "low", "summary": null } })),
);

assert_eq!(
concise_request_body["reasoning"]["summary"],
json!("concise")
);
assert_eq!(disabled_request_body["reasoning"]["summary"], json!(null));
}

#[test]
fn attaches_runtime_prompt_cache_key_after_custom_body_merge() {
let client = test_client();
Expand Down
Loading