1use std::collections::{HashMap, HashSet};
37use std::fs;
38use std::path::{Path, PathBuf};
39
40use chrono::{DateTime, SecondsFormat, Utc};
41use serde::{Deserialize, Serialize};
42use serde_json::{Map, Value, json};
43use uuid::Uuid;
44
45use crate::common::{Block, ImageSource, Message, Meta, Role, StopReason, Tool, ToolOutput};
46use crate::error::Result;
47use crate::harness::jsonl;
48use crate::transcript::{Codec, Common, Discovered, Harness, Saved, Store, TextCodec, Transcript};
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub struct Grok;
53
54impl Harness for Grok {
55 const NAME: &'static str = "grok";
56 type Body = GrokSession;
57}
58
59#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
68pub struct GrokSession {
69 #[serde(default, skip_serializing_if = "Vec::is_empty")]
70 pub chat_history: Vec<ChatRecord>,
71 #[serde(default, skip_serializing_if = "Vec::is_empty")]
72 pub updates: Vec<Value>,
73 #[serde(default, skip_serializing_if = "Vec::is_empty")]
74 pub events: Vec<Value>,
75 #[serde(default, skip_serializing_if = "Vec::is_empty")]
76 pub rewind_points: Vec<Value>,
77 #[serde(default, skip_serializing_if = "Option::is_none")]
78 pub summary: Option<Value>,
79 #[serde(default, skip_serializing_if = "Option::is_none")]
80 pub prompt_context: Option<Value>,
81 #[serde(default, skip_serializing_if = "Option::is_none")]
82 pub resources_state: Option<Value>,
83 #[serde(default, skip_serializing_if = "Option::is_none")]
84 pub signals: Option<Value>,
85 #[serde(default, skip_serializing_if = "Option::is_none")]
86 pub system_prompt: Option<String>,
87}
88
89#[derive(Debug, Clone, PartialEq)]
92pub enum ChatRecord {
93 System(SystemLine),
94 User(UserLine),
95 Assistant(AssistantLine),
96 Reasoning(ReasoningLine),
97 ToolResult(ToolResultLine),
98 Other(Value),
99}
100
101#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
103pub struct SystemLine {
104 #[serde(default, skip_serializing_if = "Value::is_null")]
105 pub content: Value,
106 #[serde(flatten)]
107 pub extra: Map<String, Value>,
108}
109
110#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
113pub struct UserLine {
114 #[serde(default, skip_serializing_if = "Value::is_null")]
115 pub content: Value,
116 #[serde(default, skip_serializing_if = "Option::is_none")]
117 pub prior_turn_interrupt: Option<Value>,
118 #[serde(flatten)]
119 pub extra: Map<String, Value>,
120}
121
122#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
125pub struct AssistantLine {
126 #[serde(default, skip_serializing_if = "Value::is_null")]
127 pub content: Value,
128 #[serde(default, skip_serializing_if = "Option::is_none")]
129 pub tool_calls: Option<Value>,
130 #[serde(default, skip_serializing_if = "Option::is_none")]
131 pub model_id: Option<String>,
132 #[serde(default, skip_serializing_if = "Option::is_none")]
133 pub model_fingerprint: Option<String>,
134 #[serde(flatten)]
135 pub extra: Map<String, Value>,
136}
137
138#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
140pub struct ReasoningLine {
141 #[serde(default, skip_serializing_if = "Option::is_none")]
142 pub id: Option<String>,
143 #[serde(default, skip_serializing_if = "Value::is_null")]
144 pub summary: Value,
145 #[serde(default, skip_serializing_if = "Option::is_none")]
146 pub encrypted_content: Option<String>,
147 #[serde(default, skip_serializing_if = "Option::is_none")]
148 pub status: Option<String>,
149 #[serde(flatten)]
150 pub extra: Map<String, Value>,
151}
152
153#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
155pub struct ToolResultLine {
156 pub tool_call_id: String,
157 #[serde(default, skip_serializing_if = "Value::is_null")]
158 pub content: Value,
159 #[serde(flatten)]
160 pub extra: Map<String, Value>,
161}
162
163impl From<Value> for ChatRecord {
164 fn from(v: Value) -> Self {
165 fn typed<T: for<'de> Deserialize<'de>>(
170 v: &Value,
171 extra: impl Fn(&mut T) -> &mut Map<String, Value>,
172 f: impl Fn(T) -> ChatRecord,
173 ) -> Option<ChatRecord> {
174 T::deserialize(v).ok().map(|mut line| {
175 extra(&mut line).remove("type");
176 f(line)
177 })
178 }
179 let record = match v.get("type").and_then(Value::as_str) {
180 Some("system") => typed(&v, |l: &mut SystemLine| &mut l.extra, ChatRecord::System),
181 Some("user") => typed(&v, |l: &mut UserLine| &mut l.extra, ChatRecord::User),
182 Some("assistant") => typed(
183 &v,
184 |l: &mut AssistantLine| &mut l.extra,
185 ChatRecord::Assistant,
186 ),
187 Some("reasoning") => typed(
188 &v,
189 |l: &mut ReasoningLine| &mut l.extra,
190 ChatRecord::Reasoning,
191 ),
192 Some("tool_result") => typed(
193 &v,
194 |l: &mut ToolResultLine| &mut l.extra,
195 ChatRecord::ToolResult,
196 ),
197 _ => None,
199 };
200 record.unwrap_or(ChatRecord::Other(v))
201 }
202}
203
204impl From<ChatRecord> for Value {
205 fn from(r: ChatRecord) -> Self {
206 fn tagged(line: impl Serialize, ty: &str) -> Value {
207 let mut v = serde_json::to_value(line).unwrap_or(Value::Null);
208 if let Value::Object(obj) = &mut v {
209 obj.insert("type".into(), Value::String(ty.into()));
210 }
211 v
212 }
213 match r {
214 ChatRecord::System(l) => tagged(l, "system"),
215 ChatRecord::User(l) => tagged(l, "user"),
216 ChatRecord::Assistant(l) => tagged(l, "assistant"),
217 ChatRecord::Reasoning(l) => tagged(l, "reasoning"),
218 ChatRecord::ToolResult(l) => tagged(l, "tool_result"),
219 ChatRecord::Other(v) => v,
220 }
221 }
222}
223
224impl Serialize for ChatRecord {
225 fn serialize<S: serde::Serializer>(&self, s: S) -> std::result::Result<S::Ok, S::Error> {
226 Value::from(self.clone()).serialize(s)
227 }
228}
229
230impl<'de> Deserialize<'de> for ChatRecord {
231 fn deserialize<D: serde::Deserializer<'de>>(d: D) -> std::result::Result<Self, D::Error> {
232 Ok(ChatRecord::from(Value::deserialize(d)?))
233 }
234}
235
236impl Codec for Grok {
239 fn to_common(transcript: &Transcript<Self>) -> Result<Transcript<Common>> {
240 Ok(Transcript::new(
241 transcript.meta.clone(),
242 body_to_messages(&transcript.body, transcript.meta.timestamp),
243 ))
244 }
245
246 fn from_common(transcript: &Transcript<Common>) -> Result<Transcript<Self>> {
247 Ok(Transcript::new(
248 transcript.meta.clone(),
249 body_from_messages(&transcript.meta, &transcript.body),
250 ))
251 }
252}
253
254impl TextCodec for Grok {
255 fn from_text(text: &str) -> Result<Transcript<Self>> {
256 let body: GrokSession = serde_json::from_str(text)?;
257 let meta = meta_from_body(&body);
258 Ok(Transcript::new(meta, body))
259 }
260
261 fn to_text(transcript: &Transcript<Self>) -> Result<String> {
262 Ok(serde_json::to_string_pretty(&transcript.body)?)
263 }
264}
265
266#[derive(Default)]
271struct DisplayIndex {
272 prompts: Vec<PromptChunks>,
273 thought_ts: Vec<DateTime<Utc>>,
274 agent_ts: Vec<DateTime<Utc>>,
275 call_ts: HashMap<String, DateTime<Utc>>,
276 result_ts: HashMap<String, DateTime<Utc>>,
277 failed: HashSet<String>,
278 stop_reasons: Vec<StopReason>,
279}
280
281#[derive(Default)]
284struct PromptChunks {
285 ts: Option<DateTime<Utc>>,
286 images: Vec<ImageSource>,
287}
288
289fn index_updates(updates: &[Value]) -> DisplayIndex {
290 let mut idx = DisplayIndex::default();
291 let dated = updates
293 .iter()
294 .filter_map(|line| line.pointer("/params/update").zip(update_ts(line)));
295 for (update, ts) in dated {
296 let call_id = || {
297 update
298 .get("toolCallId")
299 .and_then(Value::as_str)
300 .map(String::from)
301 };
302 match update.get("sessionUpdate").and_then(Value::as_str) {
303 Some("user_message_chunk") => {
304 let content = update.get("content").unwrap_or(&Value::Null);
305 let is_image = content.get("type").and_then(Value::as_str) == Some("image");
306 let slot = update
309 .pointer("/_meta/promptIndex")
310 .and_then(Value::as_u64)
311 .and_then(|i| usize::try_from(i).ok())
312 .unwrap_or_else(|| {
313 let started = idx.prompts.len();
314 if is_image {
315 started.saturating_sub(1)
316 } else {
317 started
318 }
319 });
320 if idx.prompts.len() <= slot {
321 idx.prompts.resize_with(slot + 1, PromptChunks::default);
322 }
323 let prompt = &mut idx.prompts[slot];
324 prompt.ts.get_or_insert(ts);
325 if is_image && let Some(source) = image_from_chunk(content) {
326 prompt.images.push(source);
327 }
328 }
329 Some("agent_thought_chunk") => idx.thought_ts.push(ts),
330 Some("agent_message_chunk") => idx.agent_ts.push(ts),
331 Some("tool_call") => {
332 if let Some(id) = call_id() {
333 idx.call_ts.entry(id).or_insert(ts);
334 }
335 }
336 Some("tool_call_update") => {
337 let status = update.get("status").and_then(Value::as_str);
338 if let Some(id) = call_id() {
339 match status {
340 Some("completed") => {
341 idx.result_ts.insert(id, ts);
342 }
343 Some("failed") => {
344 idx.result_ts.insert(id.clone(), ts);
345 idx.failed.insert(id);
346 }
347 _ => {}
350 }
351 }
352 }
353 Some("turn_completed") => {
354 let reason = update
355 .get("stop_reason")
356 .and_then(Value::as_str)
357 .unwrap_or("end_turn");
358 idx.stop_reasons.push(parse_stop_reason(reason));
359 }
360 _ => {}
363 }
364 }
365 idx
366}
367
368fn update_ts(line: &Value) -> Option<DateTime<Utc>> {
372 line.pointer("/params/_meta/agentTimestampMs")
373 .and_then(Value::as_i64)
374 .and_then(DateTime::from_timestamp_millis)
375 .or_else(|| {
376 line.get("timestamp")
377 .and_then(Value::as_i64)
378 .and_then(|s| DateTime::from_timestamp(s, 0))
379 })
380}
381
382fn body_to_messages(body: &GrokSession, fallback_ts: DateTime<Utc>) -> Vec<Message> {
383 let idx = index_updates(&body.updates);
384 let mut prompts = idx.prompts.iter();
385 let mut thought_ts = idx.thought_ts.iter();
386 let mut agent_ts = idx.agent_ts.iter();
387
388 let mut messages: Vec<Message> = Vec::new();
389 let mut turn_last_assistant: Vec<Option<usize>> = Vec::new();
393 let mut interrupted_turns: HashSet<usize> = HashSet::new();
394 let mut current_last_assistant: Option<usize> = None;
395 let mut in_turn = false;
396
397 for record in &body.chat_history {
398 match record {
399 ChatRecord::User(line) if is_user_info(&line.content) => {}
401 ChatRecord::User(line) => {
402 let mut content = parse_user_blocks(&line.content);
403 if !content.is_empty() {
405 if in_turn {
406 turn_last_assistant.push(current_last_assistant.take());
407 }
408 if line.prior_turn_interrupt.is_some() && !turn_last_assistant.is_empty() {
409 interrupted_turns.insert(turn_last_assistant.len() - 1);
410 }
411 in_turn = true;
412 let prompt = prompts.next();
414 let images = prompt.map(|p| p.images.as_slice()).unwrap_or_default();
415 content.extend(images.iter().cloned().map(|source| Block::Image { source }));
416 let ts = prompt.and_then(|p| p.ts).unwrap_or(fallback_ts);
417 messages.push(plain_message(Role::User, content, ts));
418 }
419 }
420 ChatRecord::Assistant(line) => {
421 let (content, first_call_id) = assistant_blocks(line);
422 if !content.is_empty() {
424 let timestamp =
427 if line.content.as_str().is_some_and(|t| !t.trim().is_empty()) {
428 agent_ts.next().copied()
429 } else {
430 first_call_id.and_then(|id| idx.call_ts.get(&id).copied())
431 }
432 .unwrap_or(fallback_ts);
433 messages.push(Message {
434 role: Role::Assistant,
435 content,
436 timestamp,
437 model: line.model_id.clone(),
438 stop_reason: None,
439 usage: None,
440 });
441 current_last_assistant = Some(messages.len() - 1);
442 }
443 }
444 ChatRecord::Reasoning(line) => {
445 if let Some(text) = reasoning_summary_text(&line.summary) {
447 messages.push(plain_message(
448 Role::Assistant,
449 vec![Block::Thinking {
450 text,
451 signature: None,
452 encrypted: line.encrypted_content.clone(),
453 }],
454 thought_ts.next().copied().unwrap_or(fallback_ts),
455 ));
456 }
457 }
458 ChatRecord::ToolResult(line) => {
459 let content = match &line.content {
460 Value::String(s) => ToolOutput::Text(s.clone()),
461 Value::Null => ToolOutput::Text(String::new()),
462 other => ToolOutput::Json(other.clone()),
463 };
464 messages.push(plain_message(
465 Role::User,
466 vec![Block::ToolResult {
467 tool_use_id: line.tool_call_id.clone(),
468 content,
469 is_error: idx.failed.contains(&line.tool_call_id),
470 }],
471 idx.result_ts
472 .get(&line.tool_call_id)
473 .copied()
474 .unwrap_or(fallback_ts),
475 ));
476 }
477 ChatRecord::System(_) | ChatRecord::Other(_) => {}
479 }
480 }
481 if in_turn {
482 turn_last_assistant.push(current_last_assistant.take());
483 }
484 backfill_stop_reasons(
485 &mut messages,
486 &turn_last_assistant,
487 &interrupted_turns,
488 &idx,
489 );
490 messages
491}
492
493fn backfill_stop_reasons(
497 messages: &mut [Message],
498 turn_last_assistant: &[Option<usize>],
499 interrupted_turns: &HashSet<usize>,
500 idx: &DisplayIndex,
501) {
502 let stamped = turn_last_assistant
504 .iter()
505 .enumerate()
506 .filter_map(|(turn, last)| last.map(|msg_idx| (turn, msg_idx)));
507 for (turn, msg_idx) in stamped {
508 if let Some(reason) = idx.stop_reasons.get(turn) {
509 messages[msg_idx].stop_reason = Some(reason.clone());
510 } else if interrupted_turns.contains(&turn) {
511 messages[msg_idx].stop_reason = Some(StopReason::Aborted);
512 }
513 }
514}
515
516fn plain_message(role: Role, content: Vec<Block>, timestamp: DateTime<Utc>) -> Message {
517 Message {
518 role,
519 content,
520 timestamp,
521 model: None,
522 stop_reason: None,
523 usage: None,
524 }
525}
526
527fn assistant_blocks(line: &AssistantLine) -> (Vec<Block>, Option<String>) {
530 let mut content = Vec::new();
531 if let Some(text) = line.content.as_str().filter(|t| !t.trim().is_empty()) {
532 content.push(Block::Text {
533 text: text.to_string(),
534 });
535 }
536 let mut first_call_id = None;
537 let calls = line
539 .tool_calls
540 .as_ref()
541 .and_then(Value::as_array)
542 .into_iter()
543 .flatten()
544 .filter_map(|call| call.get("id").and_then(Value::as_str).map(|id| (id, call)));
545 for (id, call) in calls {
546 first_call_id.get_or_insert_with(|| id.to_string());
547 let name = call.get("name").and_then(Value::as_str).unwrap_or("tool");
548 let input = parse_arguments(call.get("arguments"));
549 content.push(Block::ToolUse {
550 id: id.to_string(),
551 tool: normalize_tool(name, input),
552 });
553 }
554 (content, first_call_id)
555}
556
557fn is_user_info(content: &Value) -> bool {
558 user_texts(content).any(|t| t.contains("<user_info>"))
559}
560
561fn user_texts(content: &Value) -> impl Iterator<Item = &str> {
562 let out: Vec<&str> = match content {
563 Value::String(s) => vec![s.as_str()],
564 Value::Array(arr) => arr
565 .iter()
566 .filter_map(|block| block.get("text").and_then(Value::as_str))
567 .collect(),
568 Value::Null | Value::Bool(_) | Value::Number(_) | Value::Object(_) => Vec::new(),
570 };
571 out.into_iter()
572}
573
574fn parse_user_blocks(content: &Value) -> Vec<Block> {
575 match content {
576 Value::String(s) => text_block(s).into_iter().collect(),
577 Value::Array(arr) => arr
578 .iter()
579 .filter_map(|v| match v.get("type").and_then(Value::as_str) {
580 Some("text") => text_block(v.get("text").and_then(Value::as_str).unwrap_or("")),
581 _ => None,
584 })
585 .collect(),
586 Value::Null | Value::Bool(_) | Value::Number(_) | Value::Object(_) => Vec::new(),
588 }
589}
590
591fn text_block(raw: &str) -> Option<Block> {
592 let text = strip_user_query(raw);
593 (!text.trim().is_empty()).then_some(Block::Text { text })
594}
595
596fn image_from_chunk(content: &Value) -> Option<ImageSource> {
599 Some(ImageSource {
600 source_type: "base64".to_string(),
601 media_type: content.get("mimeType")?.as_str()?.to_string(),
602 data: content.get("data")?.as_str()?.to_string(),
603 })
604}
605
606fn strip_user_query(text: &str) -> String {
607 let trimmed = text.trim();
608 if let Some(inner) = trimmed
609 .strip_prefix("<user_query>")
610 .and_then(|s| s.strip_suffix("</user_query>"))
611 {
612 inner.trim().to_string()
613 } else {
614 text.to_string()
615 }
616}
617
618fn wrap_user_query(text: &str) -> String {
619 if text.trim_start().starts_with("<user_query>") {
620 text.to_string()
621 } else {
622 format!("<user_query>\n{text}\n</user_query>")
623 }
624}
625
626fn reasoning_summary_text(summary: &Value) -> Option<String> {
627 let parts: Vec<&str> = summary
628 .as_array()?
629 .iter()
630 .filter_map(|e| {
631 (e.get("type").and_then(Value::as_str) == Some("summary_text"))
632 .then(|| e.get("text").and_then(Value::as_str))
633 .flatten()
634 })
635 .collect();
636 (!parts.is_empty()).then(|| parts.join("\n\n"))
637}
638
639fn parse_arguments(arguments: Option<&Value>) -> Value {
642 match arguments {
643 Some(Value::String(raw)) => {
644 serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.clone()))
645 }
646 Some(other) => other.clone(),
647 None => Value::Object(Map::new()),
648 }
649}
650
651fn parse_stop_reason(s: &str) -> StopReason {
652 match s {
653 "end_turn" => StopReason::EndTurn,
654 "tool_use" => StopReason::ToolUse,
655 "max_tokens" => StopReason::MaxTokens,
656 "stop_sequence" => StopReason::StopSequence,
657 "cancelled" => StopReason::Aborted,
658 "error" => StopReason::Error,
659 other => StopReason::Other(other.to_string()),
660 }
661}
662
663fn stop_reason_str(r: &StopReason) -> String {
664 match r {
665 StopReason::EndTurn => "end_turn".into(),
666 StopReason::ToolUse => "tool_use".into(),
667 StopReason::MaxTokens => "max_tokens".into(),
668 StopReason::StopSequence => "stop_sequence".into(),
669 StopReason::Aborted => "cancelled".into(),
670 StopReason::Error => "error".into(),
671 StopReason::Other(s) => s.clone(),
672 }
673}
674
675fn normalize_tool(name: &str, input: Value) -> Tool {
682 let canonical = match name {
683 "Shell" => "Bash",
684 "StrReplace" => "Edit",
685 other => other,
686 };
687 Tool::from_canonical(canonical, normalize_args(canonical, input))
688}
689
690fn normalize_args(tool: &str, args: Value) -> Value {
691 match args {
692 Value::Object(obj) => Value::Object(
693 obj.into_iter()
694 .map(|(k, v)| match (tool, k.as_str()) {
695 ("Read" | "Write" | "Edit" | "MultiEdit", "path") => ("file_path".into(), v),
696 ("Write", "contents") => ("content".into(), v),
697 ("Bash", "block_until_ms") => match integral(&v) {
698 Some(ms) => ("timeout_ms".into(), Value::from(ms)),
699 None => (k, v),
700 },
701 ("Glob", "glob_pattern") => ("pattern".into(), v),
702 ("Glob", "target_directory") => ("path".into(), v),
703 _ => (k, v),
704 })
705 .collect(),
706 ),
707 other => other,
709 }
710}
711
712fn denormalize_tool(tool: &Tool) -> (String, Value) {
714 let (name, input) = tool.to_canonical();
715 let grok_name = match name.as_str() {
716 "Bash" => "Shell",
717 "Edit" => "StrReplace",
718 other => other,
719 }
720 .to_string();
721 (grok_name, denormalize_args(&name, input))
722}
723
724fn denormalize_args(tool: &str, input: Value) -> Value {
725 match input {
726 Value::Object(obj) => Value::Object(
727 obj.into_iter()
728 .map(|(k, v)| {
729 let key = match (tool, k.as_str()) {
730 ("Read" | "Write" | "Edit" | "MultiEdit", "file_path") => "path".into(),
731 ("Write", "content") => "contents".into(),
732 ("Bash", "timeout_ms") => "block_until_ms".into(),
733 ("Glob", "pattern") => "glob_pattern".into(),
734 ("Glob", "path") => "target_directory".into(),
735 _ => k,
736 };
737 (key, v)
738 })
739 .collect(),
740 ),
741 other => other,
743 }
744}
745
746#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] fn integral(v: &Value) -> Option<u64> {
750 v.as_u64().or_else(|| {
751 v.as_f64()
752 .filter(|f| f.fract() == 0.0 && (0.0..=9_007_199_254_740_992.0).contains(f))
753 .map(|f| f as u64)
754 })
755}
756
757fn body_from_messages(meta: &Meta, messages: &[Message]) -> GrokSession {
760 let session_id = if meta.id.is_empty() {
761 Uuid::new_v4().to_string()
762 } else {
763 meta.id.clone()
764 };
765
766 let mut chat: Vec<ChatRecord> = Vec::new();
767 let mut updates = UpdateLog::new(&session_id, meta.model.as_deref());
768 let mut turn_open = false;
771 let mut last_stop: Option<StopReason> = None;
772
773 for (i, msg) in messages.iter().enumerate() {
774 match msg.role {
775 Role::User => {
776 let mut prompt_texts: Vec<String> = Vec::new();
777 let mut prompt_images: Vec<&ImageSource> = Vec::new();
778 for block in &msg.content {
779 match block {
780 Block::Text { text } => prompt_texts.push(text.clone()),
781 Block::Image { source } => prompt_images.push(source),
785 Block::ToolResult {
786 tool_use_id,
787 content,
788 is_error,
789 } => {
790 chat.push(ChatRecord::ToolResult(ToolResultLine {
791 tool_call_id: tool_use_id.clone(),
792 content: Value::String(tool_output_string(content)),
795 extra: Map::new(),
796 }));
797 updates.tool_result(msg.timestamp, tool_use_id, content, *is_error);
798 }
799 Block::Thinking { .. } | Block::ToolUse { .. } => {}
801 }
802 }
803 if !prompt_texts.is_empty() || !prompt_images.is_empty() {
804 if turn_open {
805 updates.turn_completed(msg.timestamp, last_stop.as_ref());
806 }
807 let interrupted = turn_open && matches!(last_stop, Some(StopReason::Aborted));
808 turn_open = true;
809 last_stop = None;
810 if prompt_texts.is_empty() {
812 prompt_texts.push("[image]".to_string());
813 }
814 let text = prompt_texts.join("\n\n");
815 chat.push(ChatRecord::User(UserLine {
816 content: json!([{"type": "text", "text": wrap_user_query(&text)}]),
817 prior_turn_interrupt: interrupted
818 .then(|| Value::String("mid_turn_abort".into())),
819 extra: Map::new(),
820 }));
821 updates.user_message(msg.timestamp, &text, &prompt_images);
822 }
823 }
824 Role::Assistant => {
825 if push_assistant(&mut chat, &mut updates, &session_id, i, msg) {
826 last_stop.clone_from(&msg.stop_reason);
827 }
828 }
829 }
830 }
831 if turn_open {
832 let final_ts = messages.last().map_or(meta.timestamp, |m| m.timestamp);
833 updates.turn_completed(final_ts, last_stop.as_ref());
834 }
835
836 let updates = updates.finish();
837 let summary = summary_value(meta, &session_id, chat.len(), updates.len());
838 GrokSession {
839 chat_history: chat,
840 updates,
841 events: Vec::new(),
842 rewind_points: Vec::new(),
843 summary: Some(summary),
844 prompt_context: None,
845 resources_state: None,
846 signals: None,
847 system_prompt: None,
848 }
849}
850
851fn push_assistant(
856 chat: &mut Vec<ChatRecord>,
857 updates: &mut UpdateLog,
858 session_id: &str,
859 msg_index: usize,
860 msg: &Message,
861) -> bool {
862 let mut texts: Vec<String> = Vec::new();
863 let mut tool_calls: Vec<Value> = Vec::new();
864 let mut pending_calls: Vec<(String, String, Value)> = Vec::new();
865 for (j, block) in msg.content.iter().enumerate() {
866 match block {
867 Block::Thinking {
868 text, encrypted, ..
869 } => {
870 chat.push(ChatRecord::Reasoning(ReasoningLine {
871 id: Some(format!("rs_{}", grok_uuid(session_id, msg_index, j))),
872 summary: json!([{ "type": "summary_text", "text": text }]),
873 encrypted_content: encrypted.clone(),
874 status: Some("completed".into()),
875 extra: Map::new(),
876 }));
877 updates.thought(msg.timestamp, text);
878 }
879 Block::Text { text } => texts.push(text.clone()),
880 Block::ToolUse { id, tool } => {
881 let (name, input) = denormalize_tool(tool);
882 let arguments = match &input {
883 Value::String(raw) => raw.clone(),
884 other => other.to_string(),
885 };
886 tool_calls.push(json!({
887 "id": id,
888 "name": name,
889 "arguments": arguments,
890 }));
891 pending_calls.push((id.clone(), name, input));
892 }
893 Block::Image { .. } | Block::ToolResult { .. } => {}
895 }
896 }
897 if texts.is_empty() && tool_calls.is_empty() {
898 false
900 } else {
901 let text = texts.join("\n\n");
902 chat.push(ChatRecord::Assistant(AssistantLine {
903 content: Value::String(text.clone()),
904 tool_calls: if tool_calls.is_empty() {
905 None
906 } else {
907 Some(Value::Array(tool_calls))
908 },
909 model_id: msg.model.clone(),
910 model_fingerprint: None,
911 extra: Map::new(),
912 }));
913 if !text.is_empty() {
914 updates.agent_message(msg.timestamp, &text);
915 }
916 for (id, name, input) in pending_calls {
917 updates.tool_call(msg.timestamp, &id, &name, &input);
918 }
919 true
920 }
921}
922
923struct UpdateLog {
925 session_id: String,
926 model: Option<String>,
927 lines: Vec<Value>,
928 prompt_index: u64,
929}
930
931impl UpdateLog {
932 fn new(session_id: &str, model: Option<&str>) -> Self {
933 Self {
934 session_id: session_id.to_string(),
935 model: model.map(String::from),
936 lines: Vec::new(),
937 prompt_index: 0,
938 }
939 }
940
941 fn push(&mut self, ts: DateTime<Utc>, method: &str, update: &Value) {
942 let seq = self.lines.len();
943 self.lines.push(json!({
944 "timestamp": ts.timestamp(),
945 "method": method,
946 "params": {
947 "sessionId": self.session_id,
948 "update": update,
949 "_meta": {
950 "eventId": format!("{}-{seq}", self.session_id),
951 "agentTimestampMs": ts.timestamp_millis(),
952 },
953 },
954 }));
955 }
956
957 fn user_message(&mut self, ts: DateTime<Utc>, text: &str, images: &[&ImageSource]) {
958 let meta = |log: &Self| {
959 let mut meta = Map::new();
960 if let Some(model) = &log.model {
961 meta.insert("modelId".into(), Value::String(model.clone()));
962 }
963 meta.insert("promptIndex".into(), Value::from(log.prompt_index));
964 meta
965 };
966 self.push(
967 ts,
968 "session/update",
969 &json!({
970 "sessionUpdate": "user_message_chunk",
971 "content": { "type": "text", "text": text },
972 "_meta": meta(self),
973 }),
974 );
975 for image in images {
978 self.push(
979 ts,
980 "session/update",
981 &json!({
982 "sessionUpdate": "user_message_chunk",
983 "content": {
984 "type": "image",
985 "data": image.data,
986 "mimeType": image.media_type,
987 },
988 "_meta": meta(self),
989 }),
990 );
991 }
992 }
993
994 fn agent_message(&mut self, ts: DateTime<Utc>, text: &str) {
995 self.push(
996 ts,
997 "session/update",
998 &json!({
999 "sessionUpdate": "agent_message_chunk",
1000 "content": { "type": "text", "text": text },
1001 }),
1002 );
1003 }
1004
1005 fn thought(&mut self, ts: DateTime<Utc>, text: &str) {
1006 self.push(
1007 ts,
1008 "session/update",
1009 &json!({
1010 "sessionUpdate": "agent_thought_chunk",
1011 "content": { "type": "text", "text": text },
1012 }),
1013 );
1014 }
1015
1016 fn tool_call(&mut self, ts: DateTime<Utc>, id: &str, name: &str, input: &Value) {
1017 self.push(
1018 ts,
1019 "session/update",
1020 &json!({
1021 "sessionUpdate": "tool_call",
1022 "toolCallId": id,
1023 "title": tool_title(name, input),
1024 "kind": tool_kind(name),
1025 "rawInput": input,
1026 }),
1027 );
1028 }
1029
1030 fn tool_result(&mut self, ts: DateTime<Utc>, id: &str, content: &ToolOutput, is_error: bool) {
1031 let text = tool_output_string(content);
1032 self.push(
1033 ts,
1034 "session/update",
1035 &json!({
1036 "sessionUpdate": "tool_call_update",
1037 "toolCallId": id,
1038 "status": if is_error { "failed" } else { "completed" },
1039 "content": [{ "type": "content", "content": { "type": "text", "text": text } }],
1040 }),
1041 );
1042 }
1043
1044 fn turn_completed(&mut self, ts: DateTime<Utc>, stop: Option<&StopReason>) {
1045 let prompt_id = grok_prompt_uuid(&self.session_id, self.prompt_index);
1046 let reason = stop.map_or_else(|| "end_turn".to_string(), stop_reason_str);
1047 self.push(
1048 ts,
1049 "_x.ai/session/update",
1050 &json!({
1051 "sessionUpdate": "turn_completed",
1052 "prompt_id": prompt_id,
1053 "stop_reason": reason,
1054 }),
1055 );
1056 self.prompt_index += 1;
1057 }
1058
1059 fn finish(self) -> Vec<Value> {
1060 self.lines
1061 }
1062}
1063
1064fn tool_output_string(content: &ToolOutput) -> String {
1067 match content {
1068 ToolOutput::Text(s) => s.clone(),
1069 ToolOutput::Json(v) => {
1070 let block_texts: Option<Vec<&str>> = v.as_array().and_then(|arr| {
1071 arr.iter()
1072 .map(|b| {
1073 (b.get("type").and_then(Value::as_str) == Some("text"))
1074 .then(|| b.get("text").and_then(Value::as_str))
1075 .flatten()
1076 })
1077 .collect()
1078 });
1079 match block_texts {
1080 Some(texts) if !texts.is_empty() => texts.join("\n\n"),
1081 Some(_) | None => v.to_string(),
1083 }
1084 }
1085 }
1086}
1087
1088fn tool_kind(name: &str) -> &'static str {
1090 match name {
1091 "Read" => "read",
1092 "Write" | "StrReplace" | "MultiEdit" => "edit",
1093 "Shell" => "execute",
1094 "Glob" | "Grep" => "search",
1095 _ => "other",
1096 }
1097}
1098
1099fn tool_title(name: &str, input: &Value) -> String {
1100 let get = |key: &str| input.get(key).and_then(Value::as_str);
1101 match name {
1102 "Read" => get("path").map(|p| format!("Read `{p}`")),
1103 "Write" | "StrReplace" => get("path").map(|p| format!("Edit `{p}`")),
1104 "Shell" => get("command").map(|c| format!("Execute `{c}`")),
1105 "Glob" => get("glob_pattern").map(|p| format!("Glob `{p}`")),
1106 "Grep" => get("pattern").map(String::from),
1107 _ => None,
1108 }
1109 .unwrap_or_else(|| name.to_string())
1110}
1111
1112fn summary_value(meta: &Meta, session_id: &str, num_chat: usize, num_updates: usize) -> Value {
1113 let ts = meta.timestamp.to_rfc3339_opts(SecondsFormat::Micros, true);
1114 let title = meta.title.clone().unwrap_or_default();
1115 let mut summary = json!({
1116 "info": {
1117 "id": session_id,
1118 "cwd": meta.cwd.clone().unwrap_or_default(),
1119 },
1120 "session_summary": title,
1121 "generated_title": title,
1122 "created_at": ts,
1123 "updated_at": ts,
1124 "num_messages": num_updates,
1125 "num_chat_messages": num_chat,
1126 "current_model_id": meta.model.clone().unwrap_or_default(),
1127 "chat_format_version": 1,
1128 });
1129 if let Some(branch) = meta.git_branch.as_deref()
1130 && let Value::Object(obj) = &mut summary
1131 {
1132 obj.insert("head_branch".into(), Value::String(branch.into()));
1133 }
1134 summary
1135}
1136
1137const NS: Uuid = Uuid::from_bytes([
1139 0x6b, 0x2e, 0x41, 0x7d, 0x35, 0x0a, 0x4f, 0x91, 0x8c, 0x27, 0xd4, 0x5b, 0x9e, 0x63, 0x18, 0x2f,
1140]);
1141
1142fn grok_uuid(session_id: &str, i: usize, j: usize) -> String {
1143 Uuid::new_v5(&NS, format!("{session_id}:{i}:{j}").as_bytes()).to_string()
1144}
1145
1146fn grok_prompt_uuid(session_id: &str, turn: u64) -> String {
1147 Uuid::new_v5(&NS, format!("{session_id}:prompt:{turn}").as_bytes()).to_string()
1148}
1149
1150fn meta_from_body(body: &GrokSession) -> Meta {
1153 let summary = body.summary.as_ref();
1154 let get = |key: &str| {
1155 summary
1156 .and_then(|s| s.get(key))
1157 .and_then(Value::as_str)
1158 .filter(|v| !v.is_empty())
1159 .map(String::from)
1160 };
1161
1162 let id = summary
1163 .and_then(|s| s.pointer("/info/id"))
1164 .and_then(Value::as_str)
1165 .unwrap_or_default()
1166 .to_string();
1167 let cwd = summary
1168 .and_then(|s| s.pointer("/info/cwd"))
1169 .and_then(Value::as_str)
1170 .filter(|v| !v.is_empty())
1171 .map(String::from)
1172 .or_else(|| workspace_from_chat(&body.chat_history));
1173 let timestamp = get("created_at")
1174 .and_then(|s| s.parse::<DateTime<Utc>>().ok())
1175 .or_else(|| body.updates.iter().find_map(update_ts))
1176 .unwrap_or_else(Utc::now);
1177 let model = get("current_model_id").or_else(|| {
1178 body.chat_history.iter().find_map(|r| match r {
1179 ChatRecord::Assistant(line) => line.model_id.clone(),
1180 ChatRecord::System(_)
1182 | ChatRecord::User(_)
1183 | ChatRecord::Reasoning(_)
1184 | ChatRecord::ToolResult(_)
1185 | ChatRecord::Other(_) => None,
1186 })
1187 });
1188
1189 Meta {
1190 id,
1191 timestamp,
1192 cwd,
1193 git_branch: get("head_branch"),
1194 title: get("generated_title").or_else(|| get("session_summary")),
1195 cli_version: None,
1196 model,
1197 }
1198}
1199
1200fn workspace_from_chat(records: &[ChatRecord]) -> Option<String> {
1203 records.iter().find_map(|record| match record {
1204 ChatRecord::User(line) => user_texts(&line.content).find_map(|text| {
1205 text.lines()
1206 .find_map(|l| l.strip_prefix("Workspace Path: ").map(String::from))
1207 }),
1208 ChatRecord::System(_)
1210 | ChatRecord::Assistant(_)
1211 | ChatRecord::Reasoning(_)
1212 | ChatRecord::ToolResult(_)
1213 | ChatRecord::Other(_) => None,
1214 })
1215}
1216
1217#[derive(Debug, Clone)]
1222pub struct GrokStore {
1223 pub sessions_dir: PathBuf,
1224}
1225
1226impl GrokStore {
1227 pub fn new(sessions_dir: impl Into<PathBuf>) -> Self {
1228 Self {
1229 sessions_dir: sessions_dir.into(),
1230 }
1231 }
1232
1233 #[must_use]
1236 pub fn default_root() -> Option<Self> {
1237 std::env::var_os("GROK_HOME")
1238 .filter(|v| !v.is_empty())
1239 .map(|home| Self::new(PathBuf::from(home).join("sessions")))
1240 .or_else(|| super::home_dir().map(|h| Self::new(h.join(".grok").join("sessions"))))
1241 }
1242}
1243
1244impl GrokStore {
1245 fn discover_meta(&self, dir: &Path) -> Option<Meta> {
1251 let summary: Option<Value> = fs::read_to_string(dir.join("summary.json"))
1252 .ok()
1253 .and_then(|text| serde_json::from_str(&text).ok());
1254 if summary.as_ref().is_some_and(summary_answers_meta) {
1255 let mut meta = meta_from_body(&GrokSession {
1256 chat_history: Vec::new(),
1257 updates: Vec::new(),
1258 events: Vec::new(),
1259 rewind_points: Vec::new(),
1260 summary,
1261 prompt_context: None,
1262 resources_state: None,
1263 signals: None,
1264 system_prompt: None,
1265 });
1266 if meta.id.is_empty() {
1267 meta.id = jsonl::file_id(dir);
1268 }
1269 Some(meta)
1270 } else {
1271 self.load(&dir.to_path_buf()).ok().map(|t| t.meta)
1272 }
1273 }
1274}
1275
1276fn summary_answers_meta(summary: &Value) -> bool {
1282 summary
1283 .get("created_at")
1284 .and_then(Value::as_str)
1285 .is_some_and(|s| s.parse::<DateTime<Utc>>().is_ok())
1286 && summary
1287 .pointer("/info/cwd")
1288 .and_then(Value::as_str)
1289 .is_some_and(|cwd| !cwd.is_empty())
1290 && summary
1291 .get("current_model_id")
1292 .and_then(Value::as_str)
1293 .is_some_and(|model| !model.is_empty())
1294}
1295
1296impl Store for GrokStore {
1297 type H = Grok;
1298 type Ref = PathBuf;
1299
1300 fn discover(&self) -> Result<Vec<Discovered<PathBuf>>> {
1301 match fs::read_dir(&self.sessions_dir) {
1302 Err(_) => Ok(Vec::new()),
1304 Ok(projects) => Ok(projects
1305 .flatten()
1306 .map(|project| project.path())
1307 .filter(|project_dir| project_dir.is_dir())
1309 .filter_map(|project_dir| fs::read_dir(project_dir).ok())
1311 .flat_map(Iterator::flatten)
1312 .map(|session| session.path())
1313 .filter(|dir| {
1316 dir.join("updates.jsonl").is_file() || dir.join("chat_history.jsonl").is_file()
1317 })
1318 .filter_map(|dir| {
1319 let meta = self.discover_meta(&dir)?;
1320 Some(Discovered {
1321 meta,
1322 reference: dir,
1323 })
1324 })
1325 .collect()),
1326 }
1327 }
1328
1329 fn load(&self, reference: &PathBuf) -> Result<Transcript<Grok>> {
1330 if reference.is_dir() {
1331 let text = |name: &str| fs::read_to_string(reference.join(name)).ok();
1332 let lines = |name: &str| -> Vec<Value> {
1333 text(name).map(|t| jsonl::parse(&t)).unwrap_or_default()
1334 };
1335 let value = |name: &str| -> Option<Value> {
1336 text(name).and_then(|t| serde_json::from_str(&t).ok())
1337 };
1338
1339 let body = GrokSession {
1340 chat_history: text("chat_history.jsonl")
1341 .map(|t| jsonl::parse(&t))
1342 .unwrap_or_default(),
1343 updates: lines("updates.jsonl"),
1344 events: lines("events.jsonl"),
1345 rewind_points: lines("rewind_points.jsonl"),
1346 summary: value("summary.json"),
1347 prompt_context: value("prompt_context.json"),
1348 resources_state: value("resources_state.json"),
1349 signals: value("signals.json"),
1350 system_prompt: text("system_prompt.txt"),
1351 };
1352 let mut meta = meta_from_body(&body);
1353 if meta.id.is_empty() {
1354 meta.id = jsonl::file_id(reference);
1355 }
1356 Ok(Transcript::new(meta, body))
1357 } else {
1358 Err(std::io::Error::new(
1359 std::io::ErrorKind::NotFound,
1360 format!("no session directory at {}", reference.display()),
1361 )
1362 .into())
1363 }
1364 }
1365
1366 fn save(&self, transcript: &Transcript<Grok>) -> Result<Saved<PathBuf>> {
1367 let meta = &transcript.meta;
1368 let id = if meta.id.is_empty() {
1369 Uuid::new_v4().to_string()
1370 } else {
1371 meta.id.clone()
1372 };
1373 let cwd = meta.cwd.clone().unwrap_or_default();
1374 let dir = self.sessions_dir.join(encode_cwd(&cwd)).join(&id);
1375 fs::create_dir_all(&dir)?;
1376
1377 let body = &transcript.body;
1378 fs::write(
1379 dir.join("chat_history.jsonl"),
1380 jsonl::render(&body.chat_history)?,
1381 )?;
1382 fs::write(dir.join("updates.jsonl"), jsonl::render(&body.updates)?)?;
1383 if !body.events.is_empty() {
1384 fs::write(dir.join("events.jsonl"), jsonl::render(&body.events)?)?;
1385 }
1386 if !body.rewind_points.is_empty() {
1387 fs::write(
1388 dir.join("rewind_points.jsonl"),
1389 jsonl::render(&body.rewind_points)?,
1390 )?;
1391 }
1392 if let Some(summary) = &body.summary {
1393 fs::write(
1394 dir.join("summary.json"),
1395 serde_json::to_string_pretty(summary)?,
1396 )?;
1397 }
1398 if let Some(v) = &body.prompt_context {
1399 fs::write(dir.join("prompt_context.json"), serde_json::to_string(v)?)?;
1400 }
1401 if let Some(v) = &body.resources_state {
1402 fs::write(dir.join("resources_state.json"), serde_json::to_string(v)?)?;
1403 }
1404 if let Some(v) = &body.signals {
1405 fs::write(dir.join("signals.json"), serde_json::to_string(v)?)?;
1406 }
1407 if let Some(prompt) = &body.system_prompt {
1408 fs::write(dir.join("system_prompt.txt"), prompt)?;
1409 }
1410
1411 Ok(Saved { id, reference: dir })
1412 }
1413
1414 fn delete(&self, reference: &PathBuf) -> Result<()> {
1416 Ok(fs::remove_dir_all(reference)?)
1417 }
1418
1419 fn fingerprints(&self, refs: &[PathBuf]) -> Result<HashMap<String, String>> {
1420 let mut out = HashMap::with_capacity(refs.len());
1421 for dir in refs {
1422 let file = ["updates.jsonl", "chat_history.jsonl"]
1423 .iter()
1424 .map(|n| dir.join(n))
1425 .find(|p| p.is_file());
1426 let fingerprint = file.map(|p| file_fingerprint(&p)).unwrap_or_default();
1427 out.insert(dir.to_string_lossy().into_owned(), fingerprint);
1428 }
1429 Ok(out)
1430 }
1431}
1432
1433fn encode_cwd(cwd: &str) -> String {
1436 let mut out = String::with_capacity(cwd.len() * 3);
1437 for byte in cwd.bytes() {
1438 match byte {
1439 b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' => {
1440 out.push(char::from(byte));
1441 }
1442 _ => {
1443 const HEX: &[u8; 16] = b"0123456789ABCDEF";
1444 out.push('%');
1445 out.push(char::from(HEX[usize::from(byte >> 4)]));
1446 out.push(char::from(HEX[usize::from(byte & 0x0F)]));
1447 }
1448 }
1449 }
1450 out
1451}
1452
1453fn file_fingerprint(path: &Path) -> String {
1454 fs::metadata(path)
1455 .ok()
1456 .and_then(|m| {
1457 let len = m.len();
1458 m.modified()
1459 .ok()
1460 .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
1461 .map(|d| format!("{}:{len}", d.as_nanos()))
1462 })
1463 .unwrap_or_default()
1464}