Skip to content

Repository navigation

Commit fe75641

Browse files
tanmayvdanitanmayvdani
committed
fix(gemini): order SQLite metadata queries, synchronize thoughts total, and expand docstrings
1 parent 04bc2a4 commit fe75641

3 files changed

Lines changed: 139 additions & 85 deletions

File tree

‎rust/adapters/gemini/src/loader.rs‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use super::{
77
paths::discover_log_files,
88
};
99

10+
/// Loads Gemini CLI and Antigravity entries with progress tracking.
1011
pub fn load_entries(shared: &SharedArgs, pricing: &PricingMap) -> Result<Vec<LoadedEntry>> {
1112
crate::progress::track_usage_load(
1213
crate::progress::UsageLoadAgent("Gemini CLI"),
@@ -15,6 +16,7 @@ pub fn load_entries(shared: &SharedArgs, pricing: &PricingMap) -> Result<Vec<Loa
1516
)
1617
}
1718

19+
/// Internal loader function that discovers and parses Gemini and Antigravity log files in parallel.
1820
fn load_entries_inner(shared: &SharedArgs, pricing: &PricingMap) -> Result<Vec<LoadedEntry>> {
1921
let tz = parse_tz(shared.timezone.as_deref());
2022
let files = discover_log_files()?;
@@ -118,12 +120,6 @@ mod tests {
118120
blob.push(field1_payload.len() as u8);
119121
blob.extend_from_slice(&field1_payload);
120122

121-
let mut statement = connection
122-
.prepare("INSERT INTO gen_metadata (idx, data) VALUES (1, ?)")
123-
.unwrap();
124-
statement.bind((1, blob.as_slice())).unwrap();
125-
statement.next().unwrap();
126-
127123
// Row 2: Continuation row with tokens only (no model field)
128124
let cont_tokens = vec![
129125
0x10, 0x90, 0x03, // 2: varint 400
@@ -143,12 +139,19 @@ mod tests {
143139
cont_blob.push(cont_field1.len() as u8);
144140
cont_blob.extend_from_slice(&cont_field1);
145141

142+
// Insert Row 2 first and Row 1 second to verify ORDER BY idx ASC
146143
let mut stmt2 = connection
147144
.prepare("INSERT INTO gen_metadata (idx, data) VALUES (2, ?)")
148145
.unwrap();
149146
stmt2.bind((1, cont_blob.as_slice())).unwrap();
150147
stmt2.next().unwrap();
151148

149+
let mut statement = connection
150+
.prepare("INSERT INTO gen_metadata (idx, data) VALUES (1, ?)")
151+
.unwrap();
152+
statement.bind((1, blob.as_slice())).unwrap();
153+
statement.next().unwrap();
154+
152155
let _env_guard = super::super::GeminiDataDirEnvGuard::set(fixture.root());
153156
let shared = SharedArgs {
154157
timezone: Some("UTC".to_string()),
@@ -157,7 +160,7 @@ mod tests {
157160
let entries = load_entries(&shared, &PricingMap::load_embedded()).unwrap();
158161

159162
assert_eq!(entries.len(), 2);
160-
// Row 1: Normalized from "Gemini 2.5 Flash (High)"
163+
// Row 1 (idx=1): Normalized from "Gemini 2.5 Flash (High)"
161164
assert_eq!(entries[0].session_id.as_ref(), "session-db");
162165
assert_eq!(entries[0].model.as_deref(), Some("gemini-2.5-flash"));
163166
assert_eq!(entries[0].data.message.usage.input_tokens, 1000);
@@ -166,7 +169,7 @@ mod tests {
166169
assert_eq!(entries[0].extra_total_tokens, 50);
167170
assert!(entries[0].cost > 0.0);
168171

169-
// Row 2: Inherited from active session model
172+
// Row 2 (idx=2): Inherited from active session model
170173
assert_eq!(entries[1].model.as_deref(), Some("gemini-2.5-flash"));
171174
assert_eq!(entries[1].data.message.usage.input_tokens, 400);
172175
assert_eq!(entries[1].data.message.usage.output_tokens, 100);

‎rust/adapters/gemini/src/parser.rs‎

Lines changed: 94 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -218,14 +218,21 @@ pub(super) fn parse_jsonl_file(path: &Path) -> Result<Vec<GeminiUsageEvent>> {
218218
Ok(events)
219219
}
220220

221+
/// Parses an Antigravity SQLite conversation database file.
222+
///
223+
/// Reads generation metadata records from the `gen_metadata` table ordered by `idx ASC`
224+
/// to guarantee deterministic chronological sequence and ensure model names propagate
225+
/// accurately to continuation rows.
221226
pub(super) fn parse_sqlite_file(path: &Path) -> Result<Vec<GeminiUsageEvent>> {
222227
let fallback_timestamp = file_modified_timestamp(path);
223228
let Ok(connection) =
224229
sqlite::Connection::open_with_flags(path, sqlite::OpenFlags::new().with_read_only())
225230
else {
226231
return Ok(Vec::new());
227232
};
228-
let Ok(mut statement) = connection.prepare("SELECT idx, data FROM gen_metadata") else {
233+
let Ok(mut statement) =
234+
connection.prepare("SELECT idx, data FROM gen_metadata ORDER BY idx ASC")
235+
else {
229236
return Ok(Vec::new());
230237
};
231238

@@ -237,83 +244,91 @@ pub(super) fn parse_sqlite_file(path: &Path) -> Result<Vec<GeminiUsageEvent>> {
237244

238245
let mut current_model = None::<String>;
239246
let mut events = Vec::new();
240-
while let Ok(sqlite::State::Row) = statement.next() {
241-
let idx: i64 = statement.read(0).unwrap_or_default();
242-
let blob: Vec<u8> = statement.read(1).unwrap_or_default();
243-
if blob.is_empty() {
244-
continue;
245-
}
246-
let parsed = parse_antigravity_protobuf(&blob);
247-
let row_model = parsed
248-
.strings
249-
.get("1.19")
250-
.or_else(|| parsed.strings.get("1.21"))
251-
.or_else(|| parsed.strings.get("1.3"))
252-
.map(|s| normalize_antigravity_model(s));
253-
254-
if let Some(ref m) = row_model {
255-
current_model = Some(m.clone());
256-
}
247+
loop {
248+
match statement.next() {
249+
Ok(sqlite::State::Row) => {
250+
let idx: i64 = statement.read(0).unwrap_or_default();
251+
let blob: Vec<u8> = statement.read(1).unwrap_or_default();
252+
if blob.is_empty() {
253+
continue;
254+
}
255+
let parsed = parse_antigravity_protobuf(&blob);
256+
let row_model = parsed
257+
.strings
258+
.get("1.19")
259+
.or_else(|| parsed.strings.get("1.21"))
260+
.or_else(|| parsed.strings.get("1.3"))
261+
.map(|s| normalize_antigravity_model(s));
262+
263+
if let Some(ref m) = row_model {
264+
current_model = Some(m.clone());
265+
}
257266

258-
let model = row_model
259-
.or_else(|| current_model.clone())
260-
.unwrap_or_else(|| "gemini-internal-model".to_string());
261-
let input_tokens = parsed.numbers.get("1.4.2").copied().unwrap_or(0);
262-
let output_tokens = parsed.numbers.get("1.4.3").copied().unwrap_or(0);
263-
let cache_read_tokens = parsed.numbers.get("1.4.5").copied().unwrap_or(0);
264-
let reasoning_tokens = parsed.numbers.get("1.4.9").copied().unwrap_or(0);
265-
let visible_tokens = parsed.numbers.get("1.4.10").copied().unwrap_or(0);
266-
let timestamp_seconds = parsed.numbers.get("1.9.4.1").copied();
267-
let timestamp_nanos = parsed
268-
.numbers
269-
.get("1.9.4.2")
270-
.copied()
271-
.unwrap_or(0)
272-
.min(999_999_999);
273-
let timestamp = timestamp_seconds
274-
.and_then(|value| i64::try_from(value).ok())
275-
.and_then(|secs| {
276-
if secs > 0 {
277-
let ms = secs
278-
.saturating_mul(1000)
279-
.saturating_add((timestamp_nanos / 1_000_000) as i64);
280-
Some(TimestampMs::from_millis(ms))
281-
} else {
282-
None
267+
let model = row_model
268+
.or_else(|| current_model.clone())
269+
.unwrap_or_else(|| "gemini-internal-model".to_string());
270+
let input_tokens = parsed.numbers.get("1.4.2").copied().unwrap_or(0);
271+
let output_tokens = parsed.numbers.get("1.4.3").copied().unwrap_or(0);
272+
let cache_read_tokens = parsed.numbers.get("1.4.5").copied().unwrap_or(0);
273+
let reasoning_tokens = parsed.numbers.get("1.4.9").copied().unwrap_or(0);
274+
let visible_tokens = parsed.numbers.get("1.4.10").copied().unwrap_or(0);
275+
let timestamp_seconds = parsed.numbers.get("1.9.4.1").copied();
276+
let timestamp_nanos = parsed
277+
.numbers
278+
.get("1.9.4.2")
279+
.copied()
280+
.unwrap_or(0)
281+
.min(999_999_999);
282+
let timestamp = timestamp_seconds
283+
.and_then(|value| i64::try_from(value).ok())
284+
.and_then(|secs| {
285+
if secs > 0 {
286+
let ms = secs
287+
.saturating_mul(1000)
288+
.saturating_add((timestamp_nanos / 1_000_000) as i64);
289+
Some(TimestampMs::from_millis(ms))
290+
} else {
291+
None
292+
}
293+
})
294+
.unwrap_or(fallback_timestamp);
295+
if input_tokens == 0
296+
&& output_tokens == 0
297+
&& cache_read_tokens == 0
298+
&& reasoning_tokens == 0
299+
&& visible_tokens == 0
300+
{
301+
continue;
283302
}
284-
})
285-
.unwrap_or(fallback_timestamp);
286-
if input_tokens == 0
287-
&& output_tokens == 0
288-
&& cache_read_tokens == 0
289-
&& reasoning_tokens == 0
290-
&& visible_tokens == 0
291-
{
292-
continue;
293-
}
294-
let event = build_event(
295-
Some(&model),
296-
&session_id,
297-
timestamp,
298-
GeminiTokens {
299-
input: input_tokens,
300-
output: output_tokens,
301-
cached: cache_read_tokens,
302-
thoughts: reasoning_tokens.max(visible_tokens.saturating_sub(output_tokens)),
303-
tool: 0,
304-
total: Some(input_tokens + output_tokens + cache_read_tokens + reasoning_tokens),
305-
},
306-
normalize_session_input,
307-
Some(format!("antigravity-gen-{idx}")),
308-
);
309-
if let Some(event) = event {
310-
events.push(event);
303+
let thoughts = reasoning_tokens.max(visible_tokens.saturating_sub(output_tokens));
304+
let event = build_event(
305+
Some(&model),
306+
&session_id,
307+
timestamp,
308+
GeminiTokens {
309+
input: input_tokens,
310+
output: output_tokens,
311+
cached: cache_read_tokens,
312+
thoughts,
313+
tool: 0,
314+
total: Some(input_tokens + output_tokens + cache_read_tokens + thoughts),
315+
},
316+
normalize_session_input,
317+
Some(format!("antigravity-gen-{idx}")),
318+
);
319+
if let Some(event) = event {
320+
events.push(event);
321+
}
322+
}
323+
Ok(sqlite::State::Done) => break,
324+
Err(_) => break,
311325
}
312326
}
313327
Ok(events)
314328
}
315329

316-
fn normalize_antigravity_model(raw: &str) -> String {
330+
/// Normalizes Antigravity model identifiers and display strings into standard model IDs.
331+
pub(super) fn normalize_antigravity_model(raw: &str) -> String {
317332
let trimmed = raw.trim();
318333
if trimmed.is_empty() {
319334
return "gemini-internal-model".to_string();
@@ -327,11 +342,14 @@ fn normalize_antigravity_model(raw: &str) -> String {
327342
};
328343

329344
match base {
345+
"gemini 3.7 flash" | "gemini 3.7 flash thinking" => "gemini-3.7-flash".to_string(),
346+
"gemini 3.7 pro" | "gemini 3.7 pro thinking" => "gemini-3.7-pro".to_string(),
330347
"gemini 3.6 flash" | "gemini 3 flash" => "gemini-3.6-flash".to_string(),
331348
"gemini 3.6 pro" | "gemini 3 pro" => "gemini-3.6-pro".to_string(),
332349
"gemini 2.5 flash" => "gemini-2.5-flash".to_string(),
333350
"gemini 2.5 pro" => "gemini-2.5-pro".to_string(),
334351
"gemini 2.0 flash" | "gemini 2 flash" => "gemini-2.0-flash".to_string(),
352+
"gemini 2.0 pro" => "gemini-2.0-pro".to_string(),
335353
"gemini 1.5 flash" => "gemini-1.5-flash".to_string(),
336354
"gemini 1.5 pro" => "gemini-1.5-pro".to_string(),
337355
"claude 3.7 sonnet" | "claude 3.7 sonnet thinking" => "claude-3-7-sonnet".to_string(),
@@ -352,20 +370,24 @@ fn normalize_antigravity_model(raw: &str) -> String {
352370
}
353371
}
354372

373+
/// Maximum recursion depth allowed during nested protobuf message parsing.
355374
const MAX_PROTOBUF_DEPTH: usize = 16;
356375

376+
/// Intermediate storage for decoded protobuf field numbers and values.
357377
#[derive(Default)]
358378
struct AntigravityProtobuf {
359379
numbers: HashMap<String, u64>,
360380
strings: HashMap<String, String>,
361381
}
362382

383+
/// Decodes raw protobuf payload into numeric and string field maps.
363384
fn parse_antigravity_protobuf(blob: &[u8]) -> AntigravityProtobuf {
364385
let mut out = AntigravityProtobuf::default();
365386
parse_antigravity_protobuf_fields(blob, &mut out, "", 0);
366387
out
367388
}
368389

390+
/// Recursively parses protobuf wire-format fields up to `MAX_PROTOBUF_DEPTH`.
369391
fn parse_antigravity_protobuf_fields(
370392
blob: &[u8],
371393
out: &mut AntigravityProtobuf,
@@ -427,6 +449,7 @@ fn parse_antigravity_protobuf_fields(
427449
}
428450
}
429451

452+
/// Reads a single varint from `blob` starting at `offset`, checking for 10th-byte overflow.
430453
fn read_varint(blob: &[u8], mut offset: usize) -> Option<(u64, usize)> {
431454
let mut value = 0u64;
432455
let mut shift = 0u32;

‎rust/adapters/gemini/src/paths.rs‎

Lines changed: 34 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@ use std::{collections::HashSet, env, path::PathBuf};
33
use crate::{Result, collect_files_with_extension};
44

55
pub(super) const GEMINI_DATA_DIR_ENV: &str = "GEMINI_DATA_DIR";
6+
pub(super) const ANTIGRAVITY_DATA_DIR_ENV: &str = "ANTIGRAVITY_DATA_DIR";
67

8+
/// Returns all discovery candidate directories for Gemini and Antigravity logs.
79
fn paths() -> Result<Vec<PathBuf>> {
810
let mut paths = Vec::new();
911
let mut seen = HashSet::new();
@@ -21,22 +23,48 @@ fn paths() -> Result<Vec<PathBuf>> {
2123
return Ok(paths);
2224
}
2325

26+
if let Ok(env_paths) = env::var(ANTIGRAVITY_DATA_DIR_ENV) {
27+
for raw in env_paths
28+
.split(',')
29+
.map(str::trim)
30+
.filter(|path| !path.is_empty())
31+
{
32+
let path = PathBuf::from(raw);
33+
if path.is_dir() && seen.insert(path.clone()) {
34+
paths.push(path);
35+
}
36+
}
37+
}
38+
2439
if let Some(home) = crate::home::home_dir() {
2540
let path = home.join(".gemini").join("tmp");
2641
if path.is_dir() && seen.insert(path.clone()) {
2742
paths.push(path);
2843
}
29-
let antigravity_dir = home
30-
.join(".gemini")
31-
.join("antigravity")
32-
.join("conversations");
33-
if antigravity_dir.is_dir() && seen.insert(antigravity_dir.clone()) {
34-
paths.push(antigravity_dir);
44+
let antigravity_dirs = [
45+
home.join(".gemini")
46+
.join("antigravity")
47+
.join("conversations"),
48+
home.join(".gemini")
49+
.join("antigravity-cli")
50+
.join("conversations"),
51+
home.join(".gemini")
52+
.join("antigravity-ide")
53+
.join("conversations"),
54+
home.join(".gemini")
55+
.join("antigravity-backup")
56+
.join("conversations"),
57+
];
58+
for dir in antigravity_dirs {
59+
if dir.is_dir() && seen.insert(dir.clone()) {
60+
paths.push(dir);
61+
}
3562
}
3663
}
3764
Ok(paths)
3865
}
3966

67+
/// Discovers all `.json`, `.jsonl`, and SQLite `.db` log files across known directories.
4068
pub(super) fn discover_log_files() -> Result<Vec<PathBuf>> {
4169
let mut files = Vec::new();
4270
for path in paths()? {

0 commit comments

Comments
 (0)