Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

agent2: Token count (#36496)

Browse source 
Release Notes:

- N/A

---------

Co-authored-by: Agus Zubiaga <agus@zed.dev>
This commit is contained in:
Bennet Bo Fenner 2025-08-19 22:40:31 +02:00 • committed by GitHub
parent 6825715503
commit 5fb68cb8be
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 321 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

44
crates/agent2/src/agent.rs
View file

@ -1,8 +1,8 @@
use crate::HistoryStore;
use crate::{
ContextServerRegistry, Thread, ThreadEvent, ToolCallAuthorization, UserMessageContent,
templates::Templates,
ContextServerRegistry, Thread, ThreadEvent, ThreadsDatabase, ToolCallAuthorization,
UserMessageContent, templates::Templates,
};
use crate::{HistoryStore, ThreadsDatabase};
use acp_thread::{AcpThread, AgentModelSelector};
use action_log::ActionLog;
use agent_client_protocol as acp;
@ -673,6 +673,11 @@ impl NativeAgentConnection {
thread.update_tool_call(update, cx)
})??;
}
ThreadEvent::TokenUsageUpdate(usage) => {
acp_thread.update(cx, |thread, cx| {
thread.update_token_usage(Some(usage), cx)
})?;
}
ThreadEvent::TitleUpdate(title) => {
acp_thread
.update(cx, |thread, cx| thread.update_title(title, cx))??;
@ -895,10 +900,12 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
cx: &mut App,
) -> Option<Rc<dyn acp_thread::AgentSessionEditor>> {
self.0.update(cx, |agent, _cx| {
agent
.sessions
.get(session_id)
.map(|session| Rc::new(NativeAgentSessionEditor(session.thread.clone())) as _)
agent.sessions.get(session_id).map(|session| {
Rc::new(NativeAgentSessionEditor {
thread: session.thread.clone(),
acp_thread: session.acp_thread.clone(),
}) as _
})
})
}
@ -907,14 +914,27 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
}
}
struct NativeAgentSessionEditor(Entity<Thread>);
struct NativeAgentSessionEditor {
thread: Entity<Thread>,
acp_thread: WeakEntity<AcpThread>,
}
impl acp_thread::AgentSessionEditor for NativeAgentSessionEditor {
fn truncate(&self, message_id: acp_thread::UserMessageId, cx: &mut App) -> Task<Result<()>> {
Task::ready(
self.0
.update(cx, |thread, cx| thread.truncate(message_id, cx)),
)
match self.thread.update(cx, |thread, cx| {
thread.truncate(message_id.clone(), cx)?;
Ok(thread.latest_token_usage())
}) {
Ok(usage) => {
self.acp_thread
.update(cx, |thread, cx| {
thread.update_token_usage(usage, cx);
})
.ok();
Task::ready(Ok(()))
}
Err(error) => Task::ready(Err(error)),
}
}
}