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::{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
312 .pointer("/_meta/promptIndex")
313 .and_then(Value::as_u64)
314 .and_then(|i| usize::try_from(i).ok())
315 .filter(|&i| i <= updates.len())
316 .unwrap_or_else(|| {
317 let started = idx.prompts.len();
318 if is_image {
319 started.saturating_sub(1)
320 } else {
321 started
322 }
323 });
324 if idx.prompts.len() <= slot {
325 idx.prompts.resize_with(slot + 1, PromptChunks::default);
326 }
327 let prompt = &mut idx.prompts[slot];
328 prompt.ts.get_or_insert(ts);
329 if is_image && let Some(source) = image_from_chunk(content) {
330 prompt.images.push(source);
331 }
332 }
333 Some("agent_thought_chunk") => idx.thought_ts.push(ts),
334 Some("agent_message_chunk") => idx.agent_ts.push(ts),
335 Some("tool_call") => {
336 if let Some(id) = call_id() {
337 idx.call_ts.entry(id).or_insert(ts);
338 }
339 }
340 Some("tool_call_update") => {
341 let status = update.get("status").and_then(Value::as_str);
342 if let Some(id) = call_id() {
343 match status {
344 Some("completed") => {
345 idx.result_ts.insert(id, ts);
346 }
347 Some("failed") => {
348 idx.result_ts.insert(id.clone(), ts);
349 idx.failed.insert(id);
350 }
351 _ => {}
354 }
355 }
356 }
357 Some("turn_completed") => {
358 let reason = update
359 .get("stop_reason")
360 .and_then(Value::as_str)
361 .unwrap_or("end_turn");
362 idx.stop_reasons.push(parse_stop_reason(reason));
363 }
364 _ => {}
367 }
368 }
369 idx
370}
371
372fn update_ts(line: &Value) -> Option<DateTime<Utc>> {
376 line.pointer("/params/_meta/agentTimestampMs")
377 .and_then(Value::as_i64)
378 .and_then(DateTime::from_timestamp_millis)
379 .or_else(|| {
380 line.get("timestamp")
381 .and_then(Value::as_i64)
382 .and_then(|s| DateTime::from_timestamp(s, 0))
383 })
384}
385
386fn body_to_messages(body: &GrokSession, fallback_ts: DateTime<Utc>) -> Vec<Message> {
387 let idx = index_updates(&body.updates);
388 let mut prompts = idx.prompts.iter();
389 let mut thought_ts = idx.thought_ts.iter();
390 let mut agent_ts = idx.agent_ts.iter();
391
392 let mut messages: Vec<Message> = Vec::new();
393 let mut turn_last_assistant: Vec<Option<usize>> = Vec::new();
397 let mut interrupted_turns: HashSet<usize> = HashSet::new();
398 let mut current_last_assistant: Option<usize> = None;
399 let mut in_turn = false;
400
401 for record in &body.chat_history {
402 match record {
403 ChatRecord::User(line) if is_user_info(&line.content) => {}
405 ChatRecord::User(line) => {
406 let mut content = parse_user_blocks(&line.content);
407 if !content.is_empty() {
409 if in_turn {
410 turn_last_assistant.push(current_last_assistant.take());
411 }
412 if line.prior_turn_interrupt.is_some() && !turn_last_assistant.is_empty() {
413 interrupted_turns.insert(turn_last_assistant.len() - 1);
414 }
415 in_turn = true;
416 let prompt = prompts.next();
418 let images = prompt.map(|p| p.images.as_slice()).unwrap_or_default();
419 content.extend(images.iter().cloned().map(|source| Block::Image { source }));
420 let ts = prompt.and_then(|p| p.ts).unwrap_or(fallback_ts);
421 messages.push(plain_message(Role::User, content, ts));
422 }
423 }
424 ChatRecord::Assistant(line) => {
425 let (content, first_call_id) = assistant_blocks(line);
426 if !content.is_empty() {
428 let timestamp =
431 if line.content.as_str().is_some_and(|t| !t.trim().is_empty()) {
432 agent_ts.next().copied()
433 } else {
434 first_call_id.and_then(|id| idx.call_ts.get(&id).copied())
435 }
436 .unwrap_or(fallback_ts);
437 messages.push(Message {
438 role: Role::Assistant,
439 content,
440 timestamp,
441 model: line.model_id.clone(),
442 stop_reason: None,
443 usage: None,
444 });
445 current_last_assistant = Some(messages.len() - 1);
446 }
447 }
448 ChatRecord::Reasoning(line) => {
449 if let Some(text) = reasoning_summary_text(&line.summary) {
451 messages.push(plain_message(
452 Role::Assistant,
453 vec![Block::Thinking {
454 text,
455 signature: None,
456 encrypted: line.encrypted_content.clone(),
457 }],
458 thought_ts.next().copied().unwrap_or(fallback_ts),
459 ));
460 }
461 }
462 ChatRecord::ToolResult(line) => {
463 let content = match &line.content {
464 Value::String(s) => ToolOutput::Text(s.clone()),
465 Value::Null => ToolOutput::Text(String::new()),
466 other => ToolOutput::Json(other.clone()),
467 };
468 messages.push(plain_message(
469 Role::User,
470 vec![Block::ToolResult {
471 tool_use_id: line.tool_call_id.clone(),
472 content,
473 is_error: idx.failed.contains(&line.tool_call_id),
474 }],
475 idx.result_ts
476 .get(&line.tool_call_id)
477 .copied()
478 .unwrap_or(fallback_ts),
479 ));
480 }
481 ChatRecord::System(_) | ChatRecord::Other(_) => {}
483 }
484 }
485 if in_turn {
486 turn_last_assistant.push(current_last_assistant.take());
487 }
488 backfill_stop_reasons(
489 &mut messages,
490 &turn_last_assistant,
491 &interrupted_turns,
492 &idx,
493 );
494 messages
495}
496
497fn backfill_stop_reasons(
501 messages: &mut [Message],
502 turn_last_assistant: &[Option<usize>],
503 interrupted_turns: &HashSet<usize>,
504 idx: &DisplayIndex,
505) {
506 let stamped = turn_last_assistant
508 .iter()
509 .enumerate()
510 .filter_map(|(turn, last)| last.map(|msg_idx| (turn, msg_idx)));
511 for (turn, msg_idx) in stamped {
512 if let Some(reason) = idx.stop_reasons.get(turn) {
513 messages[msg_idx].stop_reason = Some(reason.clone());
514 } else if interrupted_turns.contains(&turn) {
515 messages[msg_idx].stop_reason = Some(StopReason::Aborted);
516 }
517 }
518}
519
520fn plain_message(role: Role, content: Vec<Block>, timestamp: DateTime<Utc>) -> Message {
521 Message {
522 role,
523 content,
524 timestamp,
525 model: None,
526 stop_reason: None,
527 usage: None,
528 }
529}
530
531fn assistant_blocks(line: &AssistantLine) -> (Vec<Block>, Option<String>) {
534 let mut content = Vec::new();
535 if let Some(text) = line.content.as_str().filter(|t| !t.trim().is_empty()) {
536 content.push(Block::Text {
537 text: text.to_string(),
538 });
539 }
540 let mut first_call_id = None;
541 let calls = line
543 .tool_calls
544 .as_ref()
545 .and_then(Value::as_array)
546 .into_iter()
547 .flatten()
548 .filter_map(|call| call.get("id").and_then(Value::as_str).map(|id| (id, call)));
549 for (id, call) in calls {
550 first_call_id.get_or_insert_with(|| id.to_string());
551 let name = call.get("name").and_then(Value::as_str).unwrap_or("tool");
552 let input = parse_arguments(call.get("arguments"));
553 content.push(Block::ToolUse {
554 id: id.to_string(),
555 tool: normalize_tool(name, input),
556 });
557 }
558 (content, first_call_id)
559}
560
561fn is_user_info(content: &Value) -> bool {
562 user_texts(content).any(|t| t.contains("<user_info>"))
563}
564
565fn user_texts(content: &Value) -> impl Iterator<Item = &str> {
566 let out: Vec<&str> = match content {
567 Value::String(s) => vec![s.as_str()],
568 Value::Array(arr) => arr
569 .iter()
570 .filter_map(|block| block.get("text").and_then(Value::as_str))
571 .collect(),
572 Value::Null | Value::Bool(_) | Value::Number(_) | Value::Object(_) => Vec::new(),
574 };
575 out.into_iter()
576}
577
578fn parse_user_blocks(content: &Value) -> Vec<Block> {
579 match content {
580 Value::String(s) => text_block(s).into_iter().collect(),
581 Value::Array(arr) => arr
582 .iter()
583 .filter_map(|v| match v.get("type").and_then(Value::as_str) {
584 Some("text") => text_block(v.get("text").and_then(Value::as_str).unwrap_or("")),
585 _ => None,
588 })
589 .collect(),
590 Value::Null | Value::Bool(_) | Value::Number(_) | Value::Object(_) => Vec::new(),
592 }
593}
594
595fn text_block(raw: &str) -> Option<Block> {
596 let text = strip_user_query(raw);
597 (!text.trim().is_empty()).then_some(Block::Text { text })
598}
599
600fn image_from_chunk(content: &Value) -> Option<ImageSource> {
603 Some(ImageSource {
604 source_type: "base64".to_string(),
605 media_type: content.get("mimeType")?.as_str()?.to_string(),
606 data: content.get("data")?.as_str()?.to_string(),
607 })
608}
609
610fn strip_user_query(text: &str) -> String {
611 let trimmed = text.trim();
612 if let Some(inner) = trimmed
613 .strip_prefix("<user_query>")
614 .and_then(|s| s.strip_suffix("</user_query>"))
615 {
616 inner.trim().to_string()
617 } else {
618 text.to_string()
619 }
620}
621
622fn wrap_user_query(text: &str) -> String {
623 if text.trim_start().starts_with("<user_query>") {
624 text.to_string()
625 } else {
626 format!("<user_query>\n{text}\n</user_query>")
627 }
628}
629
630fn reasoning_summary_text(summary: &Value) -> Option<String> {
631 let parts: Vec<&str> = summary
632 .as_array()?
633 .iter()
634 .filter_map(|e| {
635 (e.get("type").and_then(Value::as_str) == Some("summary_text"))
636 .then(|| e.get("text").and_then(Value::as_str))
637 .flatten()
638 })
639 .collect();
640 (!parts.is_empty()).then(|| parts.join("\n\n"))
641}
642
643fn parse_arguments(arguments: Option<&Value>) -> Value {
646 match arguments {
647 Some(Value::String(raw)) => {
648 serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.clone()))
649 }
650 Some(other) => other.clone(),
651 None => Value::Object(Map::new()),
652 }
653}
654
655fn parse_stop_reason(s: &str) -> StopReason {
656 match s {
657 "end_turn" => StopReason::EndTurn,
658 "tool_use" => StopReason::ToolUse,
659 "max_tokens" => StopReason::MaxTokens,
660 "stop_sequence" => StopReason::StopSequence,
661 "cancelled" => StopReason::Aborted,
662 "error" => StopReason::Error,
663 other => StopReason::Other(other.to_string()),
664 }
665}
666
667fn stop_reason_str(r: &StopReason) -> String {
668 match r {
669 StopReason::EndTurn => "end_turn".into(),
670 StopReason::ToolUse => "tool_use".into(),
671 StopReason::MaxTokens => "max_tokens".into(),
672 StopReason::StopSequence => "stop_sequence".into(),
673 StopReason::Aborted => "cancelled".into(),
674 StopReason::Error => "error".into(),
675 StopReason::Other(s) => s.clone(),
676 }
677}
678
679fn normalize_tool(name: &str, input: Value) -> Tool {
686 let canonical = match name {
687 "Shell" => "Bash",
688 "StrReplace" => "Edit",
689 other => other,
690 };
691 Tool::from_canonical(canonical, normalize_args(canonical, input))
692}
693
694fn normalize_args(tool: &str, args: Value) -> Value {
695 match args {
696 Value::Object(obj) => Value::Object(
697 obj.into_iter()
698 .map(|(k, v)| match (tool, k.as_str()) {
699 ("Read" | "Write" | "Edit" | "MultiEdit", "path") => ("file_path".into(), v),
700 ("Write", "contents") => ("content".into(), v),
701 ("Bash", "block_until_ms") => match integral(&v) {
702 Some(ms) => ("timeout_ms".into(), Value::from(ms)),
703 None => (k, v),
704 },
705 ("Glob", "glob_pattern") => ("pattern".into(), v),
706 ("Glob", "target_directory") => ("path".into(), v),
707 _ => (k, v),
708 })
709 .collect(),
710 ),
711 other => other,
713 }
714}
715
716fn denormalize_tool(tool: &Tool) -> (String, Value) {
718 let (name, input) = tool.to_canonical();
719 let grok_name = match name.as_str() {
720 "Bash" => "Shell",
721 "Edit" => "StrReplace",
722 other => other,
723 }
724 .to_string();
725 (grok_name, denormalize_args(&name, input))
726}
727
728fn denormalize_args(tool: &str, input: Value) -> Value {
729 match input {
730 Value::Object(obj) => Value::Object(
731 obj.into_iter()
732 .map(|(k, v)| {
733 let key = match (tool, k.as_str()) {
734 ("Read" | "Write" | "Edit" | "MultiEdit", "file_path") => "path".into(),
735 ("Write", "content") => "contents".into(),
736 ("Bash", "timeout_ms") => "block_until_ms".into(),
737 ("Glob", "pattern") => "glob_pattern".into(),
738 ("Glob", "path") => "target_directory".into(),
739 _ => k,
740 };
741 (key, v)
742 })
743 .collect(),
744 ),
745 other => other,
747 }
748}
749
750#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] fn integral(v: &Value) -> Option<u64> {
754 v.as_u64().or_else(|| {
755 v.as_f64()
756 .filter(|f| f.fract() == 0.0 && (0.0..=9_007_199_254_740_992.0).contains(f))
757 .map(|f| f as u64)
758 })
759}
760
761fn body_from_messages(meta: &Meta, messages: &[Message]) -> GrokSession {
764 let session_id = if meta.id.is_empty() {
765 Uuid::new_v4().to_string()
766 } else {
767 meta.id.clone()
768 };
769
770 let mut chat: Vec<ChatRecord> = Vec::new();
771 let mut updates = UpdateLog::new(&session_id, meta.model.as_deref());
772 let mut turn_open = false;
775 let mut last_stop: Option<StopReason> = None;
776
777 for (i, msg) in messages.iter().enumerate() {
778 match msg.role {
779 Role::User => {
780 let mut prompt_texts: Vec<String> = Vec::new();
781 let mut prompt_images: Vec<&ImageSource> = Vec::new();
782 for block in &msg.content {
783 match block {
784 Block::Text { text } => prompt_texts.push(text.clone()),
785 Block::Image { source } => prompt_images.push(source),
789 Block::Artifact { artifact } => {
790 prompt_texts.push(artifact.display_text());
791 }
792 Block::ToolResult {
793 tool_use_id,
794 content,
795 is_error,
796 } => {
797 chat.push(ChatRecord::ToolResult(ToolResultLine {
798 tool_call_id: tool_use_id.clone(),
799 content: Value::String(tool_output_string(content)),
802 extra: Map::new(),
803 }));
804 updates.tool_result(msg.timestamp, tool_use_id, content, *is_error);
805 }
806 Block::Thinking { .. } | Block::ToolUse { .. } => {}
808 }
809 }
810 if !prompt_texts.is_empty() || !prompt_images.is_empty() {
811 if turn_open {
812 updates.turn_completed(msg.timestamp, last_stop.as_ref());
813 }
814 let interrupted = turn_open && matches!(last_stop, Some(StopReason::Aborted));
815 turn_open = true;
816 last_stop = None;
817 if prompt_texts.is_empty() {
819 prompt_texts.push("[image]".to_string());
820 }
821 let text = prompt_texts.join("\n\n");
822 chat.push(ChatRecord::User(UserLine {
823 content: json!([{"type": "text", "text": wrap_user_query(&text)}]),
824 prior_turn_interrupt: interrupted
825 .then(|| Value::String("mid_turn_abort".into())),
826 extra: Map::new(),
827 }));
828 updates.user_message(msg.timestamp, &text, &prompt_images);
829 }
830 }
831 Role::Assistant => {
832 if push_assistant(&mut chat, &mut updates, &session_id, i, msg) {
833 last_stop.clone_from(&msg.stop_reason);
834 }
835 }
836 }
837 }
838 if turn_open {
839 let final_ts = messages.last().map_or(meta.timestamp, |m| m.timestamp);
840 updates.turn_completed(final_ts, last_stop.as_ref());
841 }
842
843 let updates = updates.finish();
844 let summary = summary_value(meta, &session_id, chat.len(), updates.len());
845 GrokSession {
846 chat_history: chat,
847 updates,
848 events: Vec::new(),
849 rewind_points: Vec::new(),
850 summary: Some(summary),
851 prompt_context: None,
852 resources_state: None,
853 signals: None,
854 system_prompt: None,
855 }
856}
857
858fn push_assistant(
863 chat: &mut Vec<ChatRecord>,
864 updates: &mut UpdateLog,
865 session_id: &str,
866 msg_index: usize,
867 msg: &Message,
868) -> bool {
869 let mut texts: Vec<String> = Vec::new();
870 let mut tool_calls: Vec<Value> = Vec::new();
871 let mut pending_calls: Vec<(String, String, Value)> = Vec::new();
872 for (j, block) in msg.content.iter().enumerate() {
873 match block {
874 Block::Thinking {
875 text, encrypted, ..
876 } => {
877 chat.push(ChatRecord::Reasoning(ReasoningLine {
878 id: Some(format!("rs_{}", grok_uuid(session_id, msg_index, j))),
879 summary: json!([{ "type": "summary_text", "text": text }]),
880 encrypted_content: encrypted.clone(),
881 status: Some("completed".into()),
882 extra: Map::new(),
883 }));
884 updates.thought(msg.timestamp, text);
885 }
886 Block::Text { text } => texts.push(text.clone()),
887 Block::Artifact { artifact } => texts.push(artifact.display_text()),
888 Block::ToolUse { id, tool } => {
889 let (name, input) = denormalize_tool(tool);
890 let arguments = match &input {
891 Value::String(raw) => raw.clone(),
892 other => other.to_string(),
893 };
894 tool_calls.push(json!({
895 "id": id,
896 "name": name,
897 "arguments": arguments,
898 }));
899 pending_calls.push((id.clone(), name, input));
900 }
901 Block::Image { .. } | Block::ToolResult { .. } => {}
903 }
904 }
905 if texts.is_empty() && tool_calls.is_empty() {
906 false
908 } else {
909 let text = texts.join("\n\n");
910 chat.push(ChatRecord::Assistant(AssistantLine {
911 content: Value::String(text.clone()),
912 tool_calls: if tool_calls.is_empty() {
913 None
914 } else {
915 Some(Value::Array(tool_calls))
916 },
917 model_id: msg.model.clone(),
918 model_fingerprint: None,
919 extra: Map::new(),
920 }));
921 if !text.is_empty() {
922 updates.agent_message(msg.timestamp, &text);
923 }
924 for (id, name, input) in pending_calls {
925 updates.tool_call(msg.timestamp, &id, &name, &input);
926 }
927 true
928 }
929}
930
931struct UpdateLog {
933 session_id: String,
934 model: Option<String>,
935 lines: Vec<Value>,
936 prompt_index: u64,
937}
938
939impl UpdateLog {
940 fn new(session_id: &str, model: Option<&str>) -> Self {
941 Self {
942 session_id: session_id.to_string(),
943 model: model.map(String::from),
944 lines: Vec::new(),
945 prompt_index: 0,
946 }
947 }
948
949 fn push(&mut self, ts: DateTime<Utc>, method: &str, update: &Value) {
950 let seq = self.lines.len();
951 self.lines.push(json!({
952 "timestamp": ts.timestamp(),
953 "method": method,
954 "params": {
955 "sessionId": self.session_id,
956 "update": update,
957 "_meta": {
958 "eventId": format!("{}-{seq}", self.session_id),
959 "agentTimestampMs": ts.timestamp_millis(),
960 },
961 },
962 }));
963 }
964
965 fn user_message(&mut self, ts: DateTime<Utc>, text: &str, images: &[&ImageSource]) {
966 let meta = |log: &Self| {
967 let mut meta = Map::new();
968 if let Some(model) = &log.model {
969 meta.insert("modelId".into(), Value::String(model.clone()));
970 }
971 meta.insert("promptIndex".into(), Value::from(log.prompt_index));
972 meta
973 };
974 self.push(
975 ts,
976 "session/update",
977 &json!({
978 "sessionUpdate": "user_message_chunk",
979 "content": { "type": "text", "text": text },
980 "_meta": meta(self),
981 }),
982 );
983 for image in images {
986 self.push(
987 ts,
988 "session/update",
989 &json!({
990 "sessionUpdate": "user_message_chunk",
991 "content": {
992 "type": "image",
993 "data": image.data,
994 "mimeType": image.media_type,
995 },
996 "_meta": meta(self),
997 }),
998 );
999 }
1000 }
1001
1002 fn agent_message(&mut self, ts: DateTime<Utc>, text: &str) {
1003 self.push(
1004 ts,
1005 "session/update",
1006 &json!({
1007 "sessionUpdate": "agent_message_chunk",
1008 "content": { "type": "text", "text": text },
1009 }),
1010 );
1011 }
1012
1013 fn thought(&mut self, ts: DateTime<Utc>, text: &str) {
1014 self.push(
1015 ts,
1016 "session/update",
1017 &json!({
1018 "sessionUpdate": "agent_thought_chunk",
1019 "content": { "type": "text", "text": text },
1020 }),
1021 );
1022 }
1023
1024 fn tool_call(&mut self, ts: DateTime<Utc>, id: &str, name: &str, input: &Value) {
1025 self.push(
1026 ts,
1027 "session/update",
1028 &json!({
1029 "sessionUpdate": "tool_call",
1030 "toolCallId": id,
1031 "title": tool_title(name, input),
1032 "kind": tool_kind(name),
1033 "rawInput": input,
1034 }),
1035 );
1036 }
1037
1038 fn tool_result(&mut self, ts: DateTime<Utc>, id: &str, content: &ToolOutput, is_error: bool) {
1039 let text = tool_output_string(content);
1040 self.push(
1041 ts,
1042 "session/update",
1043 &json!({
1044 "sessionUpdate": "tool_call_update",
1045 "toolCallId": id,
1046 "status": if is_error { "failed" } else { "completed" },
1047 "content": [{ "type": "content", "content": { "type": "text", "text": text } }],
1048 }),
1049 );
1050 }
1051
1052 fn turn_completed(&mut self, ts: DateTime<Utc>, stop: Option<&StopReason>) {
1053 let prompt_id = grok_prompt_uuid(&self.session_id, self.prompt_index);
1054 let reason = stop.map_or_else(|| "end_turn".to_string(), stop_reason_str);
1055 self.push(
1056 ts,
1057 "_x.ai/session/update",
1058 &json!({
1059 "sessionUpdate": "turn_completed",
1060 "prompt_id": prompt_id,
1061 "stop_reason": reason,
1062 }),
1063 );
1064 self.prompt_index += 1;
1065 }
1066
1067 fn finish(self) -> Vec<Value> {
1068 self.lines
1069 }
1070}
1071
1072fn tool_output_string(content: &ToolOutput) -> String {
1075 match content {
1076 ToolOutput::Text(s) => s.clone(),
1077 ToolOutput::Json(v) => {
1078 let block_texts: Option<Vec<&str>> = v.as_array().and_then(|arr| {
1079 arr.iter()
1080 .map(|b| {
1081 (b.get("type").and_then(Value::as_str) == Some("text"))
1082 .then(|| b.get("text").and_then(Value::as_str))
1083 .flatten()
1084 })
1085 .collect()
1086 });
1087 match block_texts {
1088 Some(texts) if !texts.is_empty() => texts.join("\n\n"),
1089 Some(_) | None => v.to_string(),
1091 }
1092 }
1093 }
1094}
1095
1096fn tool_kind(name: &str) -> &'static str {
1098 match name {
1099 "Read" => "read",
1100 "Write" | "StrReplace" | "MultiEdit" => "edit",
1101 "Shell" => "execute",
1102 "Glob" | "Grep" => "search",
1103 _ => "other",
1104 }
1105}
1106
1107fn tool_title(name: &str, input: &Value) -> String {
1108 let get = |key: &str| input.get(key).and_then(Value::as_str);
1109 match name {
1110 "Read" => get("path").map(|p| format!("Read `{p}`")),
1111 "Write" | "StrReplace" => get("path").map(|p| format!("Edit `{p}`")),
1112 "Shell" => get("command").map(|c| format!("Execute `{c}`")),
1113 "Glob" => get("glob_pattern").map(|p| format!("Glob `{p}`")),
1114 "Grep" => get("pattern").map(String::from),
1115 _ => None,
1116 }
1117 .unwrap_or_else(|| name.to_string())
1118}
1119
1120fn summary_value(meta: &Meta, session_id: &str, num_chat: usize, num_updates: usize) -> Value {
1121 let ts = meta.timestamp.to_rfc3339_opts(SecondsFormat::Micros, true);
1122 let title = meta.title.clone().unwrap_or_default();
1123 let mut summary = json!({
1124 "info": {
1125 "id": session_id,
1126 "cwd": meta.cwd.clone().unwrap_or_default(),
1127 },
1128 "session_summary": title,
1129 "generated_title": title,
1130 "created_at": ts,
1131 "updated_at": ts,
1132 "num_messages": num_updates,
1133 "num_chat_messages": num_chat,
1134 "current_model_id": meta.model.clone().unwrap_or_default(),
1135 "chat_format_version": 1,
1136 });
1137 if let Some(branch) = meta.git_branch.as_deref()
1138 && let Value::Object(obj) = &mut summary
1139 {
1140 obj.insert("head_branch".into(), Value::String(branch.into()));
1141 }
1142 summary
1143}
1144
1145const NS: Uuid = Uuid::from_bytes([
1147 0x6b, 0x2e, 0x41, 0x7d, 0x35, 0x0a, 0x4f, 0x91, 0x8c, 0x27, 0xd4, 0x5b, 0x9e, 0x63, 0x18, 0x2f,
1148]);
1149
1150fn grok_uuid(session_id: &str, i: usize, j: usize) -> String {
1151 Uuid::new_v5(&NS, format!("{session_id}:{i}:{j}").as_bytes()).to_string()
1152}
1153
1154fn grok_prompt_uuid(session_id: &str, turn: u64) -> String {
1155 Uuid::new_v5(&NS, format!("{session_id}:prompt:{turn}").as_bytes()).to_string()
1156}
1157
1158fn meta_from_body(body: &GrokSession) -> Meta {
1161 let summary = body.summary.as_ref();
1162 let get = |key: &str| {
1163 summary
1164 .and_then(|s| s.get(key))
1165 .and_then(Value::as_str)
1166 .filter(|v| !v.is_empty())
1167 .map(String::from)
1168 };
1169
1170 let id = summary
1171 .and_then(|s| s.pointer("/info/id"))
1172 .and_then(Value::as_str)
1173 .unwrap_or_default()
1174 .to_string();
1175 let cwd = summary
1176 .and_then(|s| s.pointer("/info/cwd"))
1177 .and_then(Value::as_str)
1178 .filter(|v| !v.is_empty())
1179 .map(String::from)
1180 .or_else(|| workspace_from_chat(&body.chat_history));
1181 let timestamp = get("created_at")
1182 .and_then(|s| s.parse::<DateTime<Utc>>().ok())
1183 .or_else(|| body.updates.iter().find_map(update_ts))
1184 .unwrap_or_else(Utc::now);
1185 let model = get("current_model_id").or_else(|| {
1186 body.chat_history.iter().find_map(|r| match r {
1187 ChatRecord::Assistant(line) => line.model_id.clone(),
1188 ChatRecord::System(_)
1190 | ChatRecord::User(_)
1191 | ChatRecord::Reasoning(_)
1192 | ChatRecord::ToolResult(_)
1193 | ChatRecord::Other(_) => None,
1194 })
1195 });
1196
1197 Meta {
1198 id,
1199 timestamp,
1200 cwd,
1201 git_branch: get("head_branch"),
1202 title: get("generated_title").or_else(|| get("session_summary")),
1203 cli_version: None,
1204 model,
1205 }
1206}
1207
1208fn workspace_from_chat(records: &[ChatRecord]) -> Option<String> {
1211 records.iter().find_map(|record| match record {
1212 ChatRecord::User(line) => user_texts(&line.content).find_map(|text| {
1213 text.lines()
1214 .find_map(|l| l.strip_prefix("Workspace Path: ").map(String::from))
1215 }),
1216 ChatRecord::System(_)
1218 | ChatRecord::Assistant(_)
1219 | ChatRecord::Reasoning(_)
1220 | ChatRecord::ToolResult(_)
1221 | ChatRecord::Other(_) => None,
1222 })
1223}
1224
1225#[derive(Debug, Clone)]
1230pub struct GrokStore {
1231 pub sessions_dir: PathBuf,
1232}
1233
1234impl GrokStore {
1235 pub fn new(sessions_dir: impl Into<PathBuf>) -> Self {
1236 Self {
1237 sessions_dir: sessions_dir.into(),
1238 }
1239 }
1240
1241 #[must_use]
1244 pub fn default_root() -> Option<Self> {
1245 std::env::var_os("GROK_HOME")
1246 .filter(|v| !v.is_empty())
1247 .map(|home| Self::new(PathBuf::from(home).join("sessions")))
1248 .or_else(|| super::home_dir().map(|h| Self::new(h.join(".grok").join("sessions"))))
1249 }
1250}
1251
1252impl GrokStore {
1253 fn discover_meta(&self, dir: &Path) -> Option<Meta> {
1259 let summary: Option<Value> = fs::read_to_string(dir.join("summary.json"))
1260 .ok()
1261 .and_then(|text| serde_json::from_str(&text).ok());
1262 if summary.as_ref().is_some_and(summary_answers_meta) {
1263 let mut meta = meta_from_body(&GrokSession {
1264 chat_history: Vec::new(),
1265 updates: Vec::new(),
1266 events: Vec::new(),
1267 rewind_points: Vec::new(),
1268 summary,
1269 prompt_context: None,
1270 resources_state: None,
1271 signals: None,
1272 system_prompt: None,
1273 });
1274 if meta.id.is_empty() {
1275 meta.id = jsonl::file_id(dir);
1276 }
1277 Some(meta)
1278 } else {
1279 self.load(&dir.to_path_buf()).ok().map(|t| t.meta)
1280 }
1281 }
1282}
1283
1284fn summary_answers_meta(summary: &Value) -> bool {
1290 summary
1291 .get("created_at")
1292 .and_then(Value::as_str)
1293 .is_some_and(|s| s.parse::<DateTime<Utc>>().is_ok())
1294 && summary
1295 .pointer("/info/cwd")
1296 .and_then(Value::as_str)
1297 .is_some_and(|cwd| !cwd.is_empty())
1298 && summary
1299 .get("current_model_id")
1300 .and_then(Value::as_str)
1301 .is_some_and(|model| !model.is_empty())
1302}
1303
1304impl Store for GrokStore {
1305 type H = Grok;
1306 type Ref = PathBuf;
1307
1308 fn discover(&self) -> Result<Vec<Discovered<PathBuf>>> {
1309 match fs::read_dir(&self.sessions_dir) {
1310 Err(_) => Ok(Vec::new()),
1312 Ok(projects) => Ok(projects
1313 .flatten()
1314 .map(|project| project.path())
1315 .filter(|project_dir| project_dir.is_dir())
1317 .filter_map(|project_dir| fs::read_dir(project_dir).ok())
1319 .flat_map(Iterator::flatten)
1320 .map(|session| session.path())
1321 .filter(|dir| {
1324 dir.join("updates.jsonl").is_file() || dir.join("chat_history.jsonl").is_file()
1325 })
1326 .filter_map(|dir| {
1327 let meta = self.discover_meta(&dir)?;
1328 Some(Discovered {
1329 meta,
1330 reference: dir,
1331 })
1332 })
1333 .collect()),
1334 }
1335 }
1336
1337 fn load(&self, reference: &PathBuf) -> Result<Transcript<Grok>> {
1338 if reference.is_dir() {
1339 let text = |name: &str| fs::read_to_string(reference.join(name)).ok();
1340 let lines = |name: &str| -> Vec<Value> {
1341 text(name).map(|t| jsonl::parse(&t)).unwrap_or_default()
1342 };
1343 let value = |name: &str| -> Option<Value> {
1344 text(name).and_then(|t| serde_json::from_str(&t).ok())
1345 };
1346
1347 let body = GrokSession {
1348 chat_history: text("chat_history.jsonl")
1349 .map(|t| jsonl::parse(&t))
1350 .unwrap_or_default(),
1351 updates: lines("updates.jsonl"),
1352 events: lines("events.jsonl"),
1353 rewind_points: lines("rewind_points.jsonl"),
1354 summary: value("summary.json"),
1355 prompt_context: value("prompt_context.json"),
1356 resources_state: value("resources_state.json"),
1357 signals: value("signals.json"),
1358 system_prompt: text("system_prompt.txt"),
1359 };
1360 let mut meta = meta_from_body(&body);
1361 if meta.id.is_empty() {
1362 meta.id = jsonl::file_id(reference);
1363 }
1364 Ok(Transcript::new(meta, body))
1365 } else {
1366 Err(std::io::Error::new(
1367 std::io::ErrorKind::NotFound,
1368 format!("no session directory at {}", reference.display()),
1369 )
1370 .into())
1371 }
1372 }
1373
1374 fn save(&self, transcript: &Transcript<Grok>) -> Result<Saved<PathBuf>> {
1375 let meta = &transcript.meta;
1376 let id = if meta.id.is_empty() {
1377 Uuid::new_v4().to_string()
1378 } else {
1379 meta.id.clone()
1380 };
1381 super::checked_id_component(Grok::NAME, &id)?;
1382 let cwd = meta.cwd.clone().unwrap_or_default();
1383 let dir = self.sessions_dir.join(encode_cwd(&cwd)).join(&id);
1384 fs::create_dir_all(&dir)?;
1385
1386 let body = &transcript.body;
1387 fs::write(
1388 dir.join("chat_history.jsonl"),
1389 jsonl::render(&body.chat_history)?,
1390 )?;
1391 fs::write(dir.join("updates.jsonl"), jsonl::render(&body.updates)?)?;
1392 if !body.events.is_empty() {
1393 fs::write(dir.join("events.jsonl"), jsonl::render(&body.events)?)?;
1394 }
1395 if !body.rewind_points.is_empty() {
1396 fs::write(
1397 dir.join("rewind_points.jsonl"),
1398 jsonl::render(&body.rewind_points)?,
1399 )?;
1400 }
1401 if let Some(summary) = &body.summary {
1402 fs::write(
1403 dir.join("summary.json"),
1404 serde_json::to_string_pretty(summary)?,
1405 )?;
1406 }
1407 if let Some(v) = &body.prompt_context {
1408 fs::write(dir.join("prompt_context.json"), serde_json::to_string(v)?)?;
1409 }
1410 if let Some(v) = &body.resources_state {
1411 fs::write(dir.join("resources_state.json"), serde_json::to_string(v)?)?;
1412 }
1413 if let Some(v) = &body.signals {
1414 fs::write(dir.join("signals.json"), serde_json::to_string(v)?)?;
1415 }
1416 if let Some(prompt) = &body.system_prompt {
1417 fs::write(dir.join("system_prompt.txt"), prompt)?;
1418 }
1419
1420 Ok(Saved { id, reference: dir })
1421 }
1422
1423 fn delete(&self, reference: &PathBuf) -> Result<()> {
1429 if !(reference.join("updates.jsonl").is_file()
1430 || reference.join("chat_history.jsonl").is_file())
1431 {
1432 return Err(Error::Malformed {
1433 harness: Grok::NAME,
1434 detail: format!("not a grok session directory: {}", reference.display()),
1435 });
1436 }
1437 let canon = reference.canonicalize()?;
1438 let root = self.sessions_dir.canonicalize()?;
1439 let contained = canon
1440 .strip_prefix(&root)
1441 .is_ok_and(|rest| rest.components().count() == 2);
1442 if !contained {
1443 return Err(Error::Malformed {
1444 harness: Grok::NAME,
1445 detail: format!(
1446 "refusing to delete outside the sessions root: {}",
1447 reference.display()
1448 ),
1449 });
1450 }
1451 Ok(fs::remove_dir_all(canon)?)
1452 }
1453
1454 fn fingerprints(&self, refs: &[PathBuf]) -> Result<HashMap<String, String>> {
1455 let mut out = HashMap::with_capacity(refs.len());
1456 for dir in refs {
1457 let file = ["updates.jsonl", "chat_history.jsonl"]
1458 .iter()
1459 .map(|n| dir.join(n))
1460 .find(|p| p.is_file());
1461 let fingerprint = file.map(|p| file_fingerprint(&p)).unwrap_or_default();
1462 out.insert(dir.to_string_lossy().into_owned(), fingerprint);
1463 }
1464 Ok(out)
1465 }
1466}
1467
1468fn encode_cwd(cwd: &str) -> String {
1471 let mut out = String::with_capacity(cwd.len() * 3);
1472 for byte in cwd.bytes() {
1473 match byte {
1474 b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' => {
1475 out.push(char::from(byte));
1476 }
1477 _ => {
1478 const HEX: &[u8; 16] = b"0123456789ABCDEF";
1479 out.push('%');
1480 out.push(char::from(HEX[usize::from(byte >> 4)]));
1481 out.push(char::from(HEX[usize::from(byte & 0x0F)]));
1482 }
1483 }
1484 }
1485 out
1486}
1487
1488fn file_fingerprint(path: &Path) -> String {
1489 fs::metadata(path)
1490 .ok()
1491 .and_then(|m| {
1492 let len = m.len();
1493 m.modified()
1494 .ok()
1495 .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
1496 .map(|d| format!("{}:{len}", d.as_nanos()))
1497 })
1498 .unwrap_or_default()
1499}