ai-space/goose
3
0
Fork 0
mirror of https://github.com/aaif-goose/goose.git synced 2026-07-03 14:10:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(acp): propagate token usage from PromptResponse (#8271)

Browse Source 
Signed-off-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
This commit is contained in:
Matt Van Horn
2026-04-09 07:48:44 -07:00
committed by GitHub
parent 2cfdae1553
commit c9a48221a9
1 changed files with 19 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/acp/provider.rs
+19 -5
View File
@@ -1,3 +1,4 @@
use agent_client_protocol_schema::Usage as AcpUsage;
use agent_client_protocol_schema::AGENT_METHOD_NAMES;
use anyhow::{Context, Result};
use async_stream::try_stream;
@@ -30,7 +31,7 @@ use crate::conversation::message::{Message, MessageContent, TOOL_META_EXTERNAL_D
use crate::model::ModelConfig;
use crate::permission::permission_confirmation::PrincipalType;
use crate::permission::{Permission, PermissionConfirmation};
use crate::providers::base::{MessageStream, PermissionRouting, Provider};
use crate::providers::base::{MessageStream, PermissionRouting, Provider, ProviderUsage, Usage};
use crate::providers::errors::ProviderError;
use crate::subprocess::configure_subprocess;
@@ -102,7 +103,7 @@ enum AcpUpdate {
request: Box<RequestPermissionRequest>,
response_tx: oneshot::Sender<RequestPermissionResponse>,
},
Complete(StopReason),
Complete(StopReason, Option<AcpUsage>),
Error(String),
}
@@ -391,7 +392,7 @@ impl Provider for AcpProvider {
async fn stream(
&self,
_model_config: &ModelConfig,
model_config: &ModelConfig,
_session_id: &str,
_system: &str,
messages: &[Message],
@@ -417,6 +418,8 @@ impl Provider for AcpProvider {
.map_err(|_| ProviderError::RequestFailed("goose_mode lock poisoned".into()))?;
let reject_all_tools = goose_mode == GooseMode::Chat;
let model_name = model_config.model_name.clone();
Ok(Box::pin(try_stream! {
let mut suppress_text = false;
let mut rejected_tool_calls: HashSet<String> = HashSet::new();
@@ -529,7 +532,18 @@ impl Provider for AcpProvider {
}
let _ = response_tx.send(map_permission_response(&request, decision));
}
AcpUpdate::Complete(_reason) => {
AcpUpdate::Complete(_reason, usage) => {
if let Some(usage) = usage {
let provider_usage = ProviderUsage::new(
model_name.clone(),
Usage::new(
Some(usage.input_tokens as i32),
Some(usage.output_tokens as i32),
Some(usage.total_tokens as i32),
),
);
yield (None, Some(provider_usage));
}
break;
}
AcpUpdate::Error(e) => {
@@ -967,7 +981,7 @@ async fn handle_requests(
match response {
Ok(r) => {
log_undelivered(
response_tx.try_send(AcpUpdate::Complete(r.stop_reason)),
response_tx.try_send(AcpUpdate::Complete(r.stop_reason, r.usage)),
AGENT_METHOD_NAMES.session_prompt,
);
}