Skip to content

Commit bba0d00

Browse files
wangtsiaowangtsiao
authored
Merge pull request #113 from 7df-lab/dev/research_tui_0623
Dev/research tui 0623
2 parents b18d327 + 9b01127 commit bba0d00

60 files changed

Lines changed: 2139 additions & 450 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎crates/cli/src/main.rs‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ fn format_with_separators(value: usize) -> String {
8383
}
8484

8585
fn format_token_usage_line(exit: &devo_tui::AppExit, color_enabled: bool) -> Option<String> {
86-
let total = exit.total_input_tokens + exit.total_output_tokens;
86+
let total = exit.total_tokens;
8787
let non_cached_input = exit
8888
.total_input_tokens
8989
.saturating_sub(exit.total_cache_read_tokens);
@@ -567,6 +567,7 @@ mod tests {
567567
turn_count: 1,
568568
total_input_tokens: 10,
569569
total_output_tokens: 2,
570+
total_tokens: 12,
570571
total_cache_read_tokens: 5,
571572
};
572573

@@ -590,6 +591,7 @@ mod tests {
590591
turn_count: 1,
591592
total_input_tokens: 10,
592593
total_output_tokens: 2,
594+
total_tokens: 12,
593595
total_cache_read_tokens: 5,
594596
};
595597

@@ -600,6 +602,24 @@ mod tests {
600602
assert!(lines[1].contains("\u{1b}["));
601603
}
602604

605+
#[test]
606+
fn exit_usage_uses_accumulated_display_total() {
607+
let exit = devo_tui::AppExit {
608+
session_id: Some(SessionId::new()),
609+
onboarding_completed: false,
610+
turn_count: 1,
611+
total_input_tokens: 10,
612+
total_output_tokens: 2,
613+
total_tokens: 25,
614+
total_cache_read_tokens: 0,
615+
};
616+
617+
assert_eq!(
618+
format_token_usage_line(&exit, /*color_enabled*/ false),
619+
Some("Token usage: total=25 input=10 output=2".to_string())
620+
);
621+
}
622+
603623
#[test]
604624
fn onboarding_exit_messages_include_next_step_after_success() {
605625
let session_id = SessionId::new();
@@ -609,6 +629,7 @@ mod tests {
609629
turn_count: 0,
610630
total_input_tokens: 0,
611631
total_output_tokens: 0,
632+
total_tokens: 0,
612633
total_cache_read_tokens: 0,
613634
};
614635

@@ -635,6 +656,7 @@ mod tests {
635656
turn_count: 0,
636657
total_input_tokens: 0,
637658
total_output_tokens: 0,
659+
total_tokens: 0,
638660
total_cache_read_tokens: 0,
639661
};
640662

@@ -652,6 +674,7 @@ mod tests {
652674
turn_count: 0,
653675
total_input_tokens: 0,
654676
total_output_tokens: 0,
677+
total_tokens: 0,
655678
total_cache_read_tokens: 0,
656679
};
657680

‎crates/cli/src/prompt_command.rs‎

Lines changed: 19 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -271,6 +271,7 @@ fn prompt_turn_config(
271271
struct PromptUsage {
272272
input_tokens: usize,
273273
output_tokens: usize,
274+
total_tokens: usize,
274275
cache_creation_input_tokens: usize,
275276
cache_read_input_tokens: usize,
276277
}
@@ -280,6 +281,7 @@ impl PromptUsage {
280281
Self {
281282
input_tokens: session.total_input_tokens,
282283
output_tokens: session.total_output_tokens,
284+
total_tokens: session.total_tokens,
283285
cache_creation_input_tokens: session.total_cache_creation_tokens,
284286
cache_read_input_tokens: session.total_cache_read_tokens,
285287
}
@@ -390,23 +392,24 @@ struct PromptUsageDelta {
390392
input_tokens: usize,
391393
output_tokens: usize,
392394
#[serde(skip_serializing_if = "Option::is_none")]
395+
reasoning_output_tokens: Option<usize>,
396+
#[serde(skip_serializing_if = "Option::is_none")]
397+
total_tokens: Option<usize>,
398+
#[serde(skip_serializing_if = "Option::is_none")]
393399
cache_creation_input_tokens: Option<usize>,
394400
#[serde(skip_serializing_if = "Option::is_none")]
395401
cache_read_input_tokens: Option<usize>,
396402
}
397403

398404
impl PromptUsageDelta {
399-
fn new(
400-
input_tokens: usize,
401-
output_tokens: usize,
402-
cache_creation_input_tokens: Option<usize>,
403-
cache_read_input_tokens: Option<usize>,
404-
) -> Self {
405+
fn new(usage: &devo_protocol::Usage) -> Self {
405406
Self {
406-
input_tokens,
407-
output_tokens,
408-
cache_creation_input_tokens,
409-
cache_read_input_tokens,
407+
input_tokens: usage.input_tokens,
408+
output_tokens: usage.output_tokens,
409+
reasoning_output_tokens: usage.reasoning_output_tokens,
410+
total_tokens: usage.total_tokens,
411+
cache_creation_input_tokens: usage.cache_creation_input_tokens,
412+
cache_read_input_tokens: usage.cache_read_input_tokens,
410413
}
411414
}
412415
}
@@ -442,19 +445,9 @@ fn write_query_event_jsonl(session_id: &str, event: &QueryEvent) -> Result<()> {
442445
session_id,
443446
item_type: "reasoning",
444447
}),
445-
QueryEvent::UsageDelta {
446-
input_tokens,
447-
output_tokens,
448-
cache_creation_input_tokens,
449-
cache_read_input_tokens,
450-
} => write_jsonl(&PromptJsonlEvent::UsageDelta {
448+
QueryEvent::UsageDelta { usage } => write_jsonl(&PromptJsonlEvent::UsageDelta {
451449
session_id,
452-
usage: PromptUsageDelta::new(
453-
*input_tokens,
454-
*output_tokens,
455-
*cache_creation_input_tokens,
456-
*cache_read_input_tokens,
457-
),
450+
usage: PromptUsageDelta::new(usage),
458451
}),
459452
QueryEvent::ToolUseStart { id, name, input } => {
460453
write_jsonl(&PromptJsonlEvent::ToolCallStarted {
@@ -512,19 +505,9 @@ fn write_query_event_jsonl(session_id: &str, event: &QueryEvent) -> Result<()> {
512505
session_id,
513506
stop_reason,
514507
}),
515-
QueryEvent::Usage {
516-
input_tokens,
517-
output_tokens,
518-
cache_creation_input_tokens,
519-
cache_read_input_tokens,
520-
} => write_jsonl(&PromptJsonlEvent::Usage {
508+
QueryEvent::Usage { usage } => write_jsonl(&PromptJsonlEvent::Usage {
521509
session_id,
522-
usage: PromptUsageDelta::new(
523-
*input_tokens,
524-
*output_tokens,
525-
*cache_creation_input_tokens,
526-
*cache_read_input_tokens,
527-
),
510+
usage: PromptUsageDelta::new(usage),
528511
}),
529512
}
530513
}
@@ -591,6 +574,7 @@ mod tests {
591574
usage: PromptUsage {
592575
input_tokens: 3,
593576
output_tokens: 5,
577+
total_tokens: 8,
594578
cache_creation_input_tokens: 0,
595579
cache_read_input_tokens: 2,
596580
},
@@ -608,6 +592,7 @@ mod tests {
608592
"usage": {
609593
"input_tokens": 3,
610594
"output_tokens": 5,
595+
"total_tokens": 8,
611596
"cache_creation_input_tokens": 0,
612597
"cache_read_input_tokens": 2
613598
}

‎crates/client/src/stdio.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1284,6 +1284,7 @@ fn acp_session_metadata_from_start_params(
12841284
reasoning_effort: None,
12851285
total_input_tokens: 0,
12861286
total_output_tokens: 0,
1287+
total_tokens: 0,
12871288
total_cache_creation_tokens: 0,
12881289
total_cache_read_tokens: 0,
12891290
prompt_token_estimate: 0,
@@ -1318,6 +1319,7 @@ fn acp_session_metadata_from_session_info(session_info: &AcpSessionInfo) -> Sess
13181319
reasoning_effort: None,
13191320
total_input_tokens: 0,
13201321
total_output_tokens: 0,
1322+
total_tokens: 0,
13211323
total_cache_creation_tokens: 0,
13221324
total_cache_read_tokens: 0,
13231325
prompt_token_estimate: 0,
@@ -1479,6 +1481,7 @@ mod tests {
14791481
reasoning_effort: None,
14801482
total_input_tokens: 0,
14811483
total_output_tokens: 0,
1484+
total_tokens: 0,
14821485
total_cache_creation_tokens: 0,
14831486
total_cache_read_tokens: 0,
14841487
prompt_token_estimate: 0,
@@ -1585,6 +1588,7 @@ mod tests {
15851588
reasoning_effort: None,
15861589
total_input_tokens: 0,
15871590
total_output_tokens: 0,
1591+
total_tokens: 0,
15881592
total_cache_creation_tokens: 0,
15891593
total_cache_read_tokens: 0,
15901594
prompt_token_estimate: 0,
@@ -1649,6 +1653,7 @@ mod tests {
16491653
reasoning_effort: None,
16501654
total_input_tokens: 0,
16511655
total_output_tokens: 0,
1656+
total_tokens: 0,
16521657
total_cache_creation_tokens: 0,
16531658
total_cache_read_tokens: 0,
16541659
prompt_token_estimate: 0,

‎crates/core/src/durable_record.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1454,6 +1454,8 @@ mod tests {
14541454
output_tokens: 50,
14551455
cache_creation_input_tokens: Some(0),
14561456
cache_read_input_tokens: Some(0),
1457+
reasoning_output_tokens: None,
1458+
total_tokens: None,
14571459
};
14581460
let record = DurableRecord::TurnCompleted(TurnCompletedRecord {
14591461
schema_version: 1,

‎crates/core/src/query.rs‎

Lines changed: 7 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -123,12 +123,7 @@ pub enum QueryEvent {
123123
ReasoningCompleted,
124124
/// Incremental token usage update from the provider stream.
125125
/// TODO: Review the mechanism from the OpenAI API / Anthropic API documentation.
126-
UsageDelta {
127-
input_tokens: usize,
128-
output_tokens: usize,
129-
cache_creation_input_tokens: Option<usize>,
130-
cache_read_input_tokens: Option<usize>,
131-
},
126+
UsageDelta { usage: devo_protocol::Usage },
132127
/// The assistant started a tool call.
133128
ToolUseStart {
134129
/// Stable provider-issued tool use identifier.
@@ -162,12 +157,7 @@ pub enum QueryEvent {
162157
/// A turn is complete (model stopped generating).
163158
TurnComplete { stop_reason: StopReason },
164159
/// Token usage update.
165-
Usage {
166-
input_tokens: usize,
167-
output_tokens: usize,
168-
cache_creation_input_tokens: Option<usize>,
169-
cache_read_input_tokens: Option<usize>,
170-
},
160+
Usage { usage: devo_protocol::Usage },
171161
}
172162

173163
/// Callback for streaming query events to the UI layer.
@@ -1050,30 +1040,20 @@ pub async fn query(
10501040
// Accumulate all usage counters at completion time.
10511041
session.total_input_tokens += response.usage.input_tokens;
10521042
session.total_output_tokens += response.usage.output_tokens;
1043+
session.total_tokens += response.usage.display_total_tokens();
10531044
session.total_cache_creation_tokens +=
10541045
response.usage.cache_creation_input_tokens.unwrap_or(0);
10551046
session.total_cache_read_tokens +=
10561047
response.usage.cache_read_input_tokens.unwrap_or(0);
10571048
session.last_input_tokens = response.usage.input_tokens;
1058-
session.last_turn_tokens = response
1059-
.usage
1060-
.input_tokens
1061-
.saturating_add(response.usage.output_tokens);
1049+
session.last_turn_tokens = response.usage.display_total_tokens();
10621050

10631051
emit(QueryEvent::Usage {
1064-
input_tokens: response.usage.input_tokens,
1065-
output_tokens: response.usage.output_tokens,
1066-
cache_creation_input_tokens: response.usage.cache_creation_input_tokens,
1067-
cache_read_input_tokens: response.usage.cache_read_input_tokens,
1052+
usage: response.usage.clone(),
10681053
});
10691054
}
10701055
Ok(StreamEvent::UsageDelta(usage)) => {
1071-
emit(QueryEvent::UsageDelta {
1072-
input_tokens: usage.input_tokens,
1073-
output_tokens: usage.output_tokens,
1074-
cache_creation_input_tokens: usage.cache_creation_input_tokens,
1075-
cache_read_input_tokens: usage.cache_read_input_tokens,
1076-
});
1056+
emit(QueryEvent::UsageDelta { usage });
10771057
}
10781058
Err(e) => {
10791059
warn!(
@@ -1597,6 +1577,7 @@ fn retry_backoff_duration(attempt: usize) -> Duration {
15971577

15981578
#[cfg(test)]
15991579
mod tests {
1580+
use devo_protocol::Usage;
16001581
use std::collections::HashMap;
16011582
use std::pin::Pin;
16021583
use std::sync::Arc;
@@ -1632,7 +1613,6 @@ mod tests {
16321613
use devo_protocol::StreamEvent;
16331614
use devo_protocol::ThreadGoal;
16341615
use devo_protocol::ThreadGoalStatus;
1635-
use devo_protocol::Usage;
16361616
use devo_provider::ModelProviderSDK;
16371617
use devo_safety::PermissionMode;
16381618
use futures::Stream;

‎crates/core/src/replay.rs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ pub struct PendingItemProjection {
7777
pub struct UsageTotals {
7878
pub total_input_tokens: i64,
7979
pub total_output_tokens: i64,
80+
pub total_tokens: i64,
8081
pub total_cache_read_tokens: i64,
8182
pub total_cache_creation_tokens: i64,
8283
pub total_reasoning_tokens: i64,
@@ -264,7 +265,9 @@ pub fn build_replay_projection(
264265
crate::durable_record::UsageMetricKind::ReasoningOutputTokens => {
265266
usage_totals.total_reasoning_tokens += m.value;
266267
}
267-
_ => {}
268+
crate::durable_record::UsageMetricKind::TotalTokens => {
269+
usage_totals.total_tokens += m.value;
270+
}
268271
}
269272
}
270273
}
@@ -431,8 +434,10 @@ fn accumulate_usage(totals: &mut UsageTotals, usage: &Option<TurnUsage>) {
431434
if let Some(u) = usage {
432435
totals.total_input_tokens += u.input_tokens as i64;
433436
totals.total_output_tokens += u.output_tokens as i64;
437+
totals.total_tokens += u.display_total_tokens() as i64;
434438
totals.total_cache_creation_tokens += u.cache_creation_input_tokens.unwrap_or(0) as i64;
435439
totals.total_cache_read_tokens += u.cache_read_input_tokens.unwrap_or(0) as i64;
440+
totals.total_reasoning_tokens += u.reasoning_output_tokens.unwrap_or(0) as i64;
436441
}
437442
}
438443

@@ -509,6 +514,8 @@ mod tests {
509514
output_tokens: 50,
510515
cache_creation_input_tokens: None,
511516
cache_read_input_tokens: None,
517+
reasoning_output_tokens: None,
518+
total_tokens: None,
512519
}),
513520
workspace_change_set_id: None,
514521
completed_at: now(),

‎crates/core/src/session.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,7 @@ pub struct SessionState {
242242
pub turn_count: usize,
243243
pub total_input_tokens: usize,
244244
pub total_output_tokens: usize,
245+
pub total_tokens: usize,
245246
pub total_cache_creation_tokens: usize, // TODO: from Anthropic Messages API, indicate how many tokens utlized to create cache.
246247
pub total_cache_read_tokens: usize, // TODO: same with `total_input_cached_tokens`.
247248
pub prompt_token_estimate: usize,
@@ -278,6 +279,7 @@ impl SessionState {
278279
turn_count: 0,
279280
total_input_tokens: 0,
280281
total_output_tokens: 0,
282+
total_tokens: 0,
281283
total_cache_creation_tokens: 0,
282284
total_cache_read_tokens: 0,
283285
prompt_token_estimate: 0,

‎crates/core/tests/real_llm_e2e.rs‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -106,11 +106,7 @@ fn collect_text_messages(session: &SessionState) -> Vec<String> {
106106

107107
fn last_usage(events: &[QueryEvent]) -> Option<(usize, usize)> {
108108
events.iter().rev().find_map(|event| match event {
109-
QueryEvent::Usage {
110-
input_tokens,
111-
output_tokens,
112-
..
113-
} => Some((*input_tokens, *output_tokens)),
109+
QueryEvent::Usage { usage } => Some((usage.input_tokens, usage.output_tokens)),
114110
_ => None,
115111
})
116112
}

‎crates/protocol/src/acp.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -995,9 +995,12 @@ mod tests {
995995
output_tokens: 4,
996996
cache_creation_input_tokens: None,
997997
cache_read_input_tokens: None,
998+
reasoning_output_tokens: None,
999+
total_tokens: None,
9981000
},
9991001
total_input_tokens: 30,
10001002
total_output_tokens: 12,
1003+
total_tokens: 42,
10011004
total_cache_read_tokens: 0,
10021005
last_query_input_tokens: 3,
10031006
context_window: Some(200_000),

0 commit comments

Comments
 (0)