1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
// Copyright 2025 Muvon Un Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// Message handling module - extracted from response.rs for better modularity
use crate::session::chat::session::ChatSession;
use crate::session::ProviderExchange;
use anyhow::Result;
pub struct MessageHandler;
impl MessageHandler {
/// Extract original tool calls from provider exchange with type safety
pub fn extract_original_tool_calls(exchange: &ProviderExchange) -> Option<serde_json::Value> {
// Check for unified format first (new clean approach)
if let Some(tool_calls) = exchange.response.get("tool_calls") {
return Some(tool_calls.clone());
}
// Use octolib's conversion method for provider-specific formats
match octolib::ProviderToolCalls::extract_from_exchange(exchange) {
Ok(Some(provider_calls)) => {
// Convert to unified GenericToolCall format using octolib
let generic_calls = provider_calls.to_generic_tool_calls();
Some(serde_json::to_value(&generic_calls).unwrap_or_default())
}
Ok(None) => None,
Err(_) => None, // No fallback - unified format is mandatory
}
}
/// Add assistant message with tool calls preserved
pub fn add_assistant_message_with_tool_calls(
chat_session: &mut ChatSession,
content: &str,
exchange: &ProviderExchange,
response_id: Option<String>,
) -> Result<()> {
// Extract the original tool_calls from the exchange response based on provider
let original_tool_calls = Self::extract_original_tool_calls(exchange);
// Create the assistant message directly with tool_calls preserved from the exchange
let assistant_message = crate::session::Message {
role: "assistant".to_string(),
content: content.to_string(),
timestamp: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
cached: false,
tool_call_id: None,
name: None,
tool_calls: original_tool_calls,
images: None,
videos: None,
thinking: None,
id: response_id,
};
// Add the assistant message to the session
chat_session.session.messages.push(assistant_message);
// Update last response
chat_session.last_response = content.to_string();
// CRITICAL FIX: Track API call and tokens (same logic as add_assistant_message)
// This was missing, causing api_calls=0 in compression analysis
if let Some(usage) = &exchange.usage {
// Track API time if available
if let Some(api_time_ms) = usage.request_time_ms {
chat_session.session.info.total_api_time_ms += api_time_ms;
}
// CACHE-AWARE COMPRESSION: Track API calls for amortized cost analysis
chat_session.session.info.total_api_calls += 1;
// Update token counts using cache manager with octolib data directly
// Update token counts using cache manager with octolib data directly
let cache_manager = crate::session::cache::CacheManager::new();
cache_manager.update_token_tracking(
&mut chat_session.session,
usage.input_tokens, // Non-cached input tokens from API
usage.output_tokens,
usage.cache_read_tokens,
usage.cache_write_tokens,
usage.reasoning_tokens,
);
// Track cost if available
if let Some(cost) = usage.cost {
chat_session.session.info.total_cost += cost;
}
}
// Save to session file to persist the message with id field
if let Some(session_file) = &chat_session.session.session_file {
let message_json =
serde_json::to_string(&chat_session.session.messages.last().unwrap())?;
crate::session::append_to_session_file(session_file, &message_json)?;
}
Ok(())
}
/// Log assistant response and exchange data
pub fn log_response_data(
session_name: &str,
content: &str,
exchange: &ProviderExchange,
) -> Result<()> {
let _ = crate::session::logger::log_assistant_response(session_name, content);
let _ = crate::session::logger::log_raw_exchange(exchange);
Ok(())
}
}