diff --git a/rust/crates/api/src/client.rs b/rust/crates/api/src/client.rs index a4ac1c0..2adaca4 100644 --- a/rust/crates/api/src/client.rs +++ b/rust/crates/api/src/client.rs @@ -1,4 +1,5 @@ use crate::error::ApiError; +use crate::prompt_cache::{PromptCache, PromptCacheRecord, PromptCacheStats}; use crate::providers::anthropic::{self, AnthropicClient, AuthSource}; use crate::providers::openai_compat::{self, OpenAiCompatClient, OpenAiCompatConfig}; use crate::providers::{self, Provider, ProviderKind}; @@ -58,6 +59,30 @@ impl ProviderClient { } } + #[must_use] + pub fn with_prompt_cache(self, prompt_cache: PromptCache) -> Self { + match self { + Self::Anthropic(client) => Self::Anthropic(client.with_prompt_cache(prompt_cache)), + other => other, + } + } + + #[must_use] + pub fn prompt_cache_stats(&self) -> Option { + match self { + Self::Anthropic(client) => client.prompt_cache_stats(), + Self::Xai(_) | Self::OpenAi(_) => None, + } + } + + #[must_use] + pub fn take_last_prompt_cache_record(&self) -> Option { + match self { + Self::Anthropic(client) => client.take_last_prompt_cache_record(), + Self::Xai(_) | Self::OpenAi(_) => None, + } + } + pub async fn send_message( &self, request: &MessageRequest, diff --git a/rust/crates/runtime/src/lib.rs b/rust/crates/runtime/src/lib.rs index 63e6a3f..47e48af 100644 --- a/rust/crates/runtime/src/lib.rs +++ b/rust/crates/runtime/src/lib.rs @@ -32,8 +32,9 @@ pub use config::{ CLAW_SETTINGS_SCHEMA_NAME, }; pub use conversation::{ - auto_compaction_threshold_from_env, ApiClient, ApiRequest, AssistantEvent, AutoCompactionEvent, - ConversationRuntime, RuntimeError, StaticToolExecutor, ToolError, ToolExecutor, TurnSummary, + auto_compaction_threshold_from_env, ApiClient, ApiRequest, AssistantEvent, + AutoCompactionEvent, ConversationRuntime, PromptCacheEvent, RuntimeError, + StaticToolExecutor, ToolError, ToolExecutor, TurnSummary, }; pub use file_ops::{ edit_file, glob_search, grep_search, read_file, write_file, EditFileOutput, GlobSearchOutput, diff --git a/rust/crates/rusty-claude-cli/src/main.rs b/rust/crates/rusty-claude-cli/src/main.rs index 35ad552..1633f1c 100644 --- a/rust/crates/rusty-claude-cli/src/main.rs +++ b/rust/crates/rusty-claude-cli/src/main.rs @@ -17,7 +17,7 @@ use std::time::{Duration, Instant, UNIX_EPOCH}; use api::{ resolve_startup_auth_source, AnthropicClient, AuthSource, ContentBlockDelta, InputContentBlock, - InputMessage, MessageRequest, MessageResponse, OutputContentBlock, + InputMessage, MessageRequest, MessageResponse, OutputContentBlock, PromptCache, SessionTracer, StreamEvent as ApiStreamEvent, ToolChoice, ToolDefinition, ToolResultContentBlock, }; @@ -34,7 +34,7 @@ use runtime::{ clear_oauth_credentials, generate_pkce_pair, generate_state, load_system_prompt, parse_oauth_callback_request_target, resolve_sandbox_status, save_oauth_credentials, ApiClient, ApiRequest, AssistantEvent, CompactionConfig, ConfigLoader, ConfigSource, - ContentBlock, ConversationMessage, ConversationRuntime, MessageRole, + ContentBlock, ConversationMessage, ConversationRuntime, MessageRole, PromptCacheEvent, OAuthAuthorizationRequest, OAuthConfig, OAuthTokenExchangeRequest, PermissionMode, PermissionPolicy, ProjectContext, RuntimeError, Session, TokenUsage, ToolError, ToolExecutor, UsageTracker, @@ -3154,6 +3154,7 @@ fn build_runtime( let mut runtime = ConversationRuntime::new_with_features( session, AnthropicRuntimeClient::new( + session_id, model, enable_tools, emit_output, @@ -3267,6 +3268,7 @@ struct AnthropicRuntimeClient { impl AnthropicRuntimeClient { fn new( + session_id: &str, model: String, enable_tools: bool, emit_output: bool, @@ -3277,7 +3279,8 @@ impl AnthropicRuntimeClient { Ok(Self { runtime: tokio::runtime::Runtime::new()?, client: AnthropicClient::from_auth(resolve_cli_auth_source()?) - .with_base_url(api::read_base_url()), + .with_base_url(api::read_base_url()) + .with_prompt_cache(PromptCache::new(session_id)), model, enable_tools, emit_output, @@ -4036,7 +4039,24 @@ fn response_to_events( Ok(events) } -fn push_prompt_cache_record(_client: &AnthropicClient, _events: &mut Vec) {} +fn push_prompt_cache_record(client: &AnthropicClient, events: &mut Vec) { + if let Some(record) = client.take_last_prompt_cache_record() { + if let Some(event) = prompt_cache_record_to_runtime_event(record) { + events.push(AssistantEvent::PromptCache(event)); + } + } +} + +fn prompt_cache_record_to_runtime_event(record: api::PromptCacheRecord) -> Option { + let cache_break = record.cache_break?; + Some(PromptCacheEvent { + unexpected: cache_break.unexpected, + reason: cache_break.reason, + previous_cache_read_input_tokens: cache_break.previous_cache_read_input_tokens, + current_cache_read_input_tokens: cache_break.current_cache_read_input_tokens, + token_drop: cache_break.token_drop, + }) +} struct CliToolExecutor { renderer: TerminalRenderer, diff --git a/rust/crates/tools/src/lib.rs b/rust/crates/tools/src/lib.rs index d097fbe..f37c636 100644 --- a/rust/crates/tools/src/lib.rs +++ b/rust/crates/tools/src/lib.rs @@ -14,7 +14,7 @@ use runtime::{ edit_file, execute_bash, glob_search, grep_search, load_system_prompt, read_file, write_file, ApiClient, ApiRequest, AssistantEvent, BashCommandInput, ContentBlock, ConversationMessage, ConversationRuntime, GrepSearchInput, MessageRole, PermissionMode, PermissionPolicy, - RuntimeError, Session, ToolError, ToolExecutor, + PromptCacheEvent, RuntimeError, Session, ToolError, ToolExecutor, }; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; @@ -2049,7 +2049,24 @@ fn response_to_events(response: MessageResponse) -> Vec { events } -fn push_prompt_cache_record(_client: &ProviderClient, _events: &mut Vec) {} +fn push_prompt_cache_record(client: &ProviderClient, events: &mut Vec) { + if let Some(record) = client.take_last_prompt_cache_record() { + if let Some(event) = prompt_cache_record_to_runtime_event(record) { + events.push(AssistantEvent::PromptCache(event)); + } + } +} + +fn prompt_cache_record_to_runtime_event(record: api::PromptCacheRecord) -> Option { + let cache_break = record.cache_break?; + Some(PromptCacheEvent { + unexpected: cache_break.unexpected, + reason: cache_break.reason, + previous_cache_read_input_tokens: cache_break.previous_cache_read_input_tokens, + current_cache_read_input_tokens: cache_break.current_cache_read_input_tokens, + token_drop: cache_break.token_drop, + }) +} fn final_assistant_text(summary: &runtime::TurnSummary) -> String { summary