MythEclipse/zesdex
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: Enhance token usage tracking and improve chat UI with emojis

Browse Source
This commit is contained in:
asepharyana
2026-07-15 03:39:15 +07:00
parent 1a6a0c628e
commit 188e7cc9a9
4 changed files with 144 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/app/runtime/actions/mod.rs
+26 -2
View File
@@ -1125,9 +1125,22 @@ fn run_agent_turn(
Remember: Cycle 0 MUST be investigation-only (access: read). Cycle 1 MUST be planning-only (access: read). Only subsequent cycles can perform modifications (access: write/full)."
));
let planner_prompt_chars = system_msg.content.as_deref().map_or(0, str::len)
+ user_msg.content.as_deref().map_or(0, str::len);
let planner_result = tc.client.chat_with_tools_non_streaming(&[system_msg, user_msg], None);
let pipeline_result = match planner_result {
Ok((reply, _)) => {
Ok((reply, usage_opt)) => {
let (mut tok_in, mut tok_out) = usage_opt.unwrap_or((0, 0));
if tok_in == 0 {
tok_in = (planner_prompt_chars / 4).max(1) as u64;
}
if tok_out == 0 {
let response_chars = reply.content.as_deref().map_or(0, str::len);
tok_out = (response_chars / 4).max(1) as u64;
}
if let Ok(mut q) = events_q.lock() {
q.push_back(TurnEvent::Usage { tokens_in: tok_in, tokens_out: tok_out });
}
let reply_text = reply.content.as_deref().unwrap_or("").trim();
let clean_json = if reply_text.starts_with("```") {
let mut lines = reply_text.lines();
@@ -1350,7 +1363,18 @@ fn run_agent_turn(
}
};
let (tok_in, tok_out) = final_usage.unwrap_or((0, 0));
let (mut tok_in, mut tok_out) = final_usage.unwrap_or((0, 0));
if tok_in == 0 {
let total_chars: usize = wire_msgs.iter()
.filter_map(|m| m.content.as_deref())
.map(str::len)
.sum();
tok_in = (total_chars / 4).max(1) as u64;
}
if tok_out == 0 {
let response_chars = response.content.as_deref().map_or(0, str::len);
tok_out = (response_chars / 4).max(1) as u64;
}
if let Ok(mut q) = events_q.lock() {
q.push_back(TurnEvent::Usage { tokens_in: tok_in, tokens_out: tok_out });
}
src/app/runtime/stream/mod.rs
+79 -66
View File
@@ -107,6 +107,9 @@ impl SseParser {
return vec![];
}
};
let mut events = Vec::new();
if let Some(usage) = value.get("usage") {
if !usage.is_null() {
let prompt_tokens = usage.get("prompt_tokens").and_then(serde_json::Value::as_u64).unwrap_or_else(|| {
@@ -122,84 +125,73 @@ impl SseParser {
tracing::warn!("[stream] total_tokens missing in usage chunk");
prompt_tokens + completion_tokens
});
// Only emit Usage as a standalone event if this chunk
// contains nothing else (no choices, no delta). Some
// non-standard providers may bundle usage WITH content
// in the same chunk; emitting both prevents content loss.
let has_other_content = value.get("choices")
.and_then(|c| c.as_array())
.is_some_and(|arr| arr.iter().any(|ch| {
ch.get("delta").and_then(|d| d.get("content")).is_some()
|| ch.get("delta").and_then(|d| d.get("reasoning_content")).is_some()
|| ch.get("delta").and_then(|d| d.get("tool_calls")).is_some()
}));
if !has_other_content {
return vec![StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens }];
}
events.push(StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens });
}
}
match event_type.as_str() {
let mut other_events = match event_type.as_str() {
"message.stop" => vec![StreamEvent::Done],
"message.delta" | "" => {
let Some(delta) = value.get("delta").or_else(|| value.get("choices")) else { return vec![] };
if let Some(choices) = delta.as_array() {
let Some(choice) = choices.first() else { return vec![] };
let Some(d) = choice.get("delta") else { return vec![] };
let mut d_events = Vec::new();
if let Some(delta) = value.get("delta").or_else(|| value.get("choices")) {
if let Some(choices) = delta.as_array() {
if let Some(choice) = choices.first() {
if let Some(d) = choice.get("delta") {
// Content token
if let Some(content) = d.get("content").and_then(|c| c.as_str()) {
d_events.push(StreamEvent::Token(content.to_string()));
}
// Content token
if let Some(content) = d.get("content").and_then(|c| c.as_str()) {
return vec![StreamEvent::Token(content.to_string())];
}
// Reasoning token
if let Some(reasoning) = d.get("reasoning_content").and_then(|r| r.as_str()) {
d_events.push(StreamEvent::Reasoning(reasoning.to_string()));
}
// Reasoning token
if let Some(reasoning) = d.get("reasoning_content").and_then(|r| r.as_str()) {
return vec![StreamEvent::Reasoning(reasoning.to_string())];
}
// Tool calls — iterate ALL entries, not just first()
if let Some(tool_calls) = d.get("tool_calls").and_then(|tc| tc.as_array()) {
for tc in tool_calls {
let index = tc.get("index").and_then(serde_json::Value::as_u64).unwrap_or_else(|| {
tracing::warn!("[stream] tool call delta missing index, defaulting to 0");
0
}) as usize;
let id = tc.get("id").and_then(|i| i.as_str()).map(std::string::ToString::to_string);
let name = tc.get("function")
.and_then(|f| f.get("name"))
.and_then(|n| n.as_str())
.map(std::string::ToString::to_string);
let args_delta = tc.get("function")
.and_then(|f| f.get("arguments"))
.and_then(|a| a.as_str())
.unwrap_or("")
.to_string();
d_events.push(StreamEvent::ToolCallDelta {
index,
id,
name,
arguments_delta: args_delta,
});
}
}
// Tool calls — iterate ALL entries, not just first()
if let Some(tool_calls) = d.get("tool_calls").and_then(|tc| tc.as_array()) {
let mut events = Vec::with_capacity(tool_calls.len());
for tc in tool_calls {
let index = tc.get("index").and_then(serde_json::Value::as_u64).unwrap_or_else(|| {
tracing::warn!("[stream] tool call delta missing index, defaulting to 0");
0
}) as usize;
let id = tc.get("id").and_then(|i| i.as_str()).map(std::string::ToString::to_string);
let name = tc.get("function")
.and_then(|f| f.get("name"))
.and_then(|n| n.as_str())
.map(std::string::ToString::to_string);
let args_delta = tc.get("function")
.and_then(|f| f.get("arguments"))
.and_then(|a| a.as_str())
.unwrap_or("")
.to_string();
events.push(StreamEvent::ToolCallDelta {
index,
id,
name,
arguments_delta: args_delta,
});
}
if !events.is_empty() {
return events;
}
}
// Finish reason
if let Some(reason) = choice.get("finish_reason").and_then(|r| r.as_str()) {
if reason == "stop" || reason == "tool_calls" {
return vec![StreamEvent::Done];
// Finish reason
if let Some(reason) = choice.get("finish_reason").and_then(|r| r.as_str()) {
if reason == "stop" || reason == "tool_calls" {
d_events.push(StreamEvent::Done);
}
}
}
}
} else if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
d_events.push(StreamEvent::Token(content.to_string()));
}
}
if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
return vec![StreamEvent::Token(content.to_string())];
}
vec![]
d_events
}
_ => vec![],
}
};
events.append(&mut other_events);
events
}
/// Clears any partially-buffered SSE frame. Reserved for reconnect/retry flows that
@@ -350,6 +342,27 @@ mod tests {
}
}
#[test]
fn feed_parses_usage_and_content_bundled_chunk() {
let mut p = SseParser::new();
let events = p.feed(
"data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}\n\n",
);
assert_eq!(events.len(), 2);
match (&events[0], &events[1]) {
(
StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens },
StreamEvent::Token(t),
) => {
assert_eq!(*prompt_tokens, 10);
assert_eq!(*completion_tokens, 5);
assert_eq!(*total_tokens, 15);
assert_eq!(t, "hello");
}
other => panic!("expected [Usage, Token], got {other:?}"),
}
}
#[test]
fn feed_ignores_empty_data_lines() {
let mut p = SseParser::new();
src/app/subagent/engine.rs
+15 -5
View File
@@ -449,12 +449,22 @@ pub fn run_subagent(ctx: &SubagentContext, tx: &mpsc::Sender<SubagentEvent>) ->
// drain thread can accumulate it and update the Usage panel.
// Without this, the Usage panel always shows zeros because the
// subagent never tells the parent about the tokens consumed.
if let Some((tokens_in, tokens_out)) = returned_usage {
let _ = tx.blocking_send(SubagentEvent::Usage {
tokens_in,
tokens_out,
});
let (mut tok_in, mut tok_out) = returned_usage.unwrap_or((0, 0));
if tok_in == 0 {
let prompt_chars: usize = messages.iter()
.filter_map(|m| m.content.as_deref())
.map(str::len)
.sum();
tok_in = (prompt_chars / 4).max(1) as u64;
}
if tok_out == 0 {
let response_chars = response.content.as_deref().map_or(0, str::len);
tok_out = (response_chars / 4).max(1) as u64;
}
let _ = tx.blocking_send(SubagentEvent::Usage {
tokens_in: tok_in,
tokens_out: tok_out,
});
let has_tool_calls = response.tool_calls.is_some()
&& response.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty());