1use std::collections::BTreeMap;
19
20use crate::drift::optionality;
21use crate::drift::text::{
22 find_top_level, match_delim, read_ident, screaming_snake_to_pascal, split_top_level,
23 strip_comments,
24};
25
26#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct Field {
29 pub name: String,
30 pub required: bool,
31}
32
33#[derive(Debug, Clone)]
35pub struct UpstreamEvent {
36 pub event_type: String,
37 pub schema: Option<String>,
39 pub fields: Vec<Field>,
42 pub unparsed: Option<String>,
44}
45
46#[derive(Debug, Clone)]
48pub struct Upstream {
49 pub event_types: Vec<String>,
51 pub base_fields: Vec<Field>,
53 pub events: Vec<UpstreamEvent>,
55 pub notes: Vec<String>,
59}
60
61pub fn extract(source: &str) -> Result<Upstream, String> {
63 let src = strip_comments(source);
64 let (enum_idents, event_types) = parse_event_type_enum(&src)?;
65 let decls = collect_const_decls(&src);
66
67 let mut parsed: BTreeMap<&str, ParsedSchema> = BTreeMap::new();
68 for name in decls.keys() {
69 let schema = resolve(name, &decls, &mut Vec::new());
70 parsed.insert(name, schema);
71 }
72
73 let discriminator = ["type".to_string()];
75 let mut notes = Vec::new();
76 let base_fields = parsed
77 .get("BaseEventSchema")
78 .map(|s| to_fields(&s.fields, &discriminator, &decls, &mut notes))
79 .unwrap_or_default();
80 let inherited: Vec<String> = discriminator
81 .iter()
82 .cloned()
83 .chain(base_fields.iter().map(|f| f.name.clone()))
84 .collect();
85
86 let mut by_event: BTreeMap<String, (&str, &ParsedSchema)> = BTreeMap::new();
90 for (name, schema) in &parsed {
91 let Some(literal) = schema.fields.iter().find(|(k, _)| k == "type") else {
92 continue;
93 };
94 let Some(member) = event_type_literal(&literal.1) else {
95 continue;
96 };
97 let value = enum_idents.get(&member).cloned().unwrap_or(member);
98 if !event_types.contains(&value) {
99 continue;
100 }
101 by_event.entry(value).or_insert((name, schema));
102 }
103
104 let events = event_types
105 .iter()
106 .map(|event_type| {
107 build_event(
108 event_type, &by_event, &parsed, &inherited, &decls, &mut notes,
109 )
110 })
111 .collect();
112
113 Ok(Upstream {
114 event_types,
115 base_fields,
116 events,
117 notes,
118 })
119}
120
121fn build_event(
122 event_type: &str,
123 by_event: &BTreeMap<String, (&str, &ParsedSchema)>,
124 parsed: &BTreeMap<&str, ParsedSchema>,
125 inherited: &[String],
126 decls: &BTreeMap<String, String>,
127 notes: &mut Vec<String>,
128) -> UpstreamEvent {
129 if let Some((name, schema)) = by_event.get(event_type) {
130 return UpstreamEvent {
131 event_type: event_type.to_string(),
132 schema: Some((*name).to_string()),
133 fields: to_fields(&schema.fields, inherited, decls, notes),
134 unparsed: schema.error.clone(),
135 };
136 }
137
138 let guess = format!("{}EventSchema", screaming_snake_to_pascal(event_type));
141 match parsed.get(guess.as_str()) {
142 Some(schema) => UpstreamEvent {
143 event_type: event_type.to_string(),
144 schema: Some(guess.clone()),
145 fields: to_fields(&schema.fields, inherited, decls, notes),
146 unparsed: Some(schema.error.clone().unwrap_or_else(|| {
147 format!("{guess} does not declare a recognisable `type` literal")
148 })),
149 },
150 None => UpstreamEvent {
151 event_type: event_type.to_string(),
152 schema: None,
153 fields: Vec::new(),
154 unparsed: Some("no Zod schema found for this event type".to_string()),
155 },
156 }
157}
158
159fn parse_event_type_enum(src: &str) -> Result<(BTreeMap<String, String>, Vec<String>), String> {
164 let at = src
165 .find("enum EventType")
166 .ok_or("`enum EventType` not found in the upstream source")?;
167 let open = src[at..]
168 .find('{')
169 .map(|i| at + i)
170 .ok_or("`enum EventType` has no body")?;
171 let close = match_delim(src, open).ok_or("`enum EventType` body is unbalanced")?;
172
173 let mut idents = BTreeMap::new();
174 let mut values = Vec::new();
175 for entry in split_top_level(&src[open + 1..close], ',') {
176 let Some((ident, literal)) = entry.split_once('=') else {
177 return Err(format!("unrecognised `EventType` member: `{entry}`"));
178 };
179 let ident = ident.trim();
180 let literal = literal.trim();
181 let value = literal
182 .strip_prefix('"')
183 .and_then(|v| v.strip_suffix('"'))
184 .ok_or_else(|| format!("`EventType.{ident}` is not a string literal: `{literal}`"))?;
185 idents.insert(ident.to_string(), value.to_string());
186 values.push(value.to_string());
187 }
188 if values.is_empty() {
189 return Err("`enum EventType` is empty".to_string());
190 }
191 Ok((idents, values))
192}
193
194fn collect_const_decls(src: &str) -> BTreeMap<String, String> {
196 let mut out = BTreeMap::new();
197 let mut at = 0usize;
198 while let Some(found) = src[at..].find("const ") {
199 let kw = at + found;
200 at = kw + "const ".len();
201 if src[..kw]
203 .chars()
204 .next_back()
205 .is_some_and(|c| c.is_alphanumeric() || c == '_' || c == '.')
206 {
207 continue;
208 }
209 let name = read_ident(src, at);
210 if name.is_empty() {
211 continue;
212 }
213 let after = at + name.len();
214 let Some(eq) = src[after..].find('=').map(|i| after + i) else {
215 continue;
216 };
217 if !src[after..eq].trim().is_empty() {
218 continue;
219 }
220 let expr_start = eq + 1;
221 let end = find_top_level(&src[expr_start..], ";")
222 .map(|i| expr_start + i)
223 .unwrap_or(src.len());
224 out.insert(name.to_string(), src[expr_start..end].trim().to_string());
225 at = end;
226 }
227 out
228}
229
230#[derive(Debug, Clone, Default)]
232struct ParsedSchema {
233 fields: Vec<(String, String)>,
234 error: Option<String>,
237}
238
239fn resolve(name: &str, decls: &BTreeMap<String, String>, stack: &mut Vec<String>) -> ParsedSchema {
242 if stack.iter().any(|n| n == name) {
243 return ParsedSchema {
244 fields: Vec::new(),
245 error: Some(format!("`{name}` is defined in terms of itself")),
246 };
247 }
248 let Some(expr) = decls.get(name) else {
249 return ParsedSchema {
250 fields: Vec::new(),
251 error: Some(format!("`{name}` is not declared in this file")),
252 };
253 };
254 stack.push(name.to_string());
255 let parsed = parse_expr(expr, decls, stack);
256 stack.pop();
257 parsed
258}
259
260fn parse_expr(
263 expr: &str,
264 decls: &BTreeMap<String, String>,
265 stack: &mut Vec<String>,
266) -> ParsedSchema {
267 let expr = expr.trim();
268 let (mut schema, mut rest, head_is_z) = match parse_primary(expr, decls, stack) {
269 Ok(parts) => parts,
270 Err(error) => {
271 return ParsedSchema {
272 fields: Vec::new(),
273 error: Some(error),
274 };
275 }
276 };
277 let mut first_call = true;
278
279 loop {
280 rest = rest.trim_start();
281 if rest.is_empty() {
282 return schema;
283 }
284 let Some(after_dot) = rest.strip_prefix('.') else {
285 schema.error = Some(format!("unexpected trailing expression `{}`", brief(rest)));
286 return schema;
287 };
288 let method = read_ident(after_dot, 0);
289 let Some(open) = after_dot[method.len()..]
290 .find('(')
291 .map(|i| method.len() + i)
292 else {
293 schema.error = Some(format!("`.{method}` is not a call"));
294 return schema;
295 };
296 let Some(close) = match_delim(after_dot, open) else {
297 schema.error = Some(format!("`.{method}(` is unbalanced"));
298 return schema;
299 };
300 let arg = after_dot[open + 1..close].trim();
301 rest = &after_dot[close + 1..];
302 let is_head_call = first_call && head_is_z;
303 first_call = false;
304
305 match method {
306 "object" if is_head_call => match parse_object_literal(arg) {
308 Ok(fields) => schema.fields = fields,
309 Err(error) => {
310 schema.error = Some(error);
311 return schema;
312 }
313 },
314 _ if is_head_call => {
315 schema.error = Some(format!("unsupported schema head `z.{method}()`"));
316 return schema;
317 }
318 "extend" => match parse_object_literal(arg) {
319 Ok(fields) => {
320 for (key, value) in fields {
321 match schema.fields.iter_mut().find(|(k, _)| *k == key) {
322 Some(slot) => slot.1 = value,
323 None => schema.fields.push((key, value)),
324 }
325 }
326 }
327 Err(error) => {
328 schema.error = Some(error);
329 return schema;
330 }
331 },
332 "omit" | "pick" => match parse_object_literal(arg) {
333 Ok(keys) => {
334 let keys: Vec<String> = keys.into_iter().map(|(k, _)| k).collect();
335 let keep = method == "pick";
336 schema.fields.retain(|(k, _)| keys.contains(k) == keep);
337 }
338 Err(error) => {
339 schema.error = Some(error);
340 return schema;
341 }
342 },
343 "passthrough" | "strict" | "strip" | "describe" | "catchall" | "readonly" => {}
345 _ => {
346 schema.error = Some(format!("unsupported schema modifier `.{method}()`"));
347 return schema;
348 }
349 }
350 }
351}
352
353fn parse_primary<'a>(
360 expr: &'a str,
361 decls: &BTreeMap<String, String>,
362 stack: &mut Vec<String>,
363) -> Result<(ParsedSchema, &'a str, bool), String> {
364 let ident = read_ident(expr, 0);
365 if ident.is_empty() {
366 return Err(format!("unsupported schema expression `{}`", brief(expr)));
367 }
368 if ident == "z" {
369 return Ok((ParsedSchema::default(), &expr[1..], true));
370 }
371 Ok((resolve(ident, decls, stack), &expr[ident.len()..], false))
372}
373
374fn parse_object_literal(arg: &str) -> Result<Vec<(String, String)>, String> {
376 let arg = arg.trim();
377 if !arg.starts_with('{') {
378 return Err(format!(
379 "expected an object literal, found `{}`",
380 brief(arg)
381 ));
382 }
383 let close = match_delim(arg, 0).ok_or_else(|| "object literal is unbalanced".to_string())?;
384 let mut out = Vec::new();
385 for entry in split_top_level(&arg[1..close], ',') {
386 if entry.starts_with("...") {
387 return Err(format!(
388 "object spread is not supported: `{}`",
389 brief(entry)
390 ));
391 }
392 let colon = find_top_level(entry, ":")
393 .ok_or_else(|| format!("field `{}` has no value", brief(entry)))?;
394 let key = entry[..colon].trim().trim_matches(['"', '\''].as_slice());
395 out.push((key.to_string(), entry[colon + 1..].trim().to_string()));
396 }
397 Ok(out)
398}
399
400fn event_type_literal(value: &str) -> Option<String> {
402 let inner = value.strip_prefix("z.literal(")?.trim_end();
403 let inner = inner.strip_suffix(')')?.trim();
404 if let Some(member) = inner.strip_prefix("EventType.") {
405 return Some(read_ident(member, 0).to_string());
406 }
407 inner
408 .strip_prefix('"')
409 .and_then(|v| v.strip_suffix('"'))
410 .map(str::to_string)
411}
412
413fn to_fields(
416 pairs: &[(String, String)],
417 inherited: &[String],
418 decls: &BTreeMap<String, String>,
419 notes: &mut Vec<String>,
420) -> Vec<Field> {
421 pairs
422 .iter()
423 .filter(|(name, _)| !inherited.iter().any(|i| i == name))
424 .map(|(name, value)| {
425 let (required, note) = field_optionality(value, decls);
426 if let Some(note) = note.filter(|n| !notes.contains(n)) {
427 notes.push(note);
428 }
429 Field {
430 name: name.clone(),
431 required,
432 }
433 })
434 .collect()
435}
436
437fn field_optionality(value: &str, decls: &BTreeMap<String, String>) -> (bool, Option<String>) {
450 let mut expr = value.trim().to_string();
451 let mut followed: Vec<String> = Vec::new();
452 loop {
453 if !is_required(&expr) {
454 return (false, None);
455 }
456 let Some(ident) = bare_ident(&expr) else {
457 return (true, None);
458 };
459 match decls.get(ident) {
460 Some(declared) if !followed.iter().any(|f| f == ident) => {
461 followed.push(ident.to_string());
462 expr = declared.clone();
463 }
464 _ => {
465 let optional = ident.starts_with("Optional");
466 return (
467 !optional,
468 Some(format!(
469 "`{ident}` is not declared in this file, so it was read as {} from \
470 its name",
471 optionality(!optional)
472 )),
473 );
474 }
475 }
476 }
477}
478
479fn bare_ident(value: &str) -> Option<&str> {
482 let value = value.trim();
483 let ident = read_ident(value, 0);
484 (!ident.is_empty() && ident.len() == value.len()).then_some(ident)
485}
486
487fn is_required(value: &str) -> bool {
492 find_top_level(value, ".optional()").is_none() && find_top_level(value, ".default(").is_none()
493}
494
495fn brief(s: &str) -> String {
497 let flat = s.split_whitespace().collect::<Vec<_>>().join(" ");
498 if flat.chars().count() > 60 {
499 let cut: String = flat.chars().take(57).collect();
500 format!("{cut}...")
501 } else {
502 flat
503 }
504}
505
506#[cfg(test)]
507mod tests {
508 use super::*;
509
510 const SAMPLE: &str = r#"
511import { z } from "zod";
512
513const RoleSchema = z.union([z.literal("user"), z.literal("assistant")]);
514
515export enum EventType {
516 TEXT_MESSAGE_START = "TEXT_MESSAGE_START",
517 /** @deprecated */
518 THINKING_END = "THINKING_END",
519 RAW = "RAW",
520 ORPHAN = "ORPHAN",
521}
522
523export const BaseEventSchema = z
524 .object({
525 type: z.nativeEnum(EventType),
526 timestamp: z.number().optional(),
527 rawEvent: z.any().optional(),
528 })
529 .passthrough();
530
531export const TextMessageStartEventSchema = BaseEventSchema.extend({
532 type: z.literal(EventType.TEXT_MESSAGE_START),
533 messageId: z.string(),
534 role: RoleSchema.default("assistant"),
535 name: z.string().optional(),
536});
537
538export const ThinkingEndEventSchema = TextMessageStartEventSchema.omit({
539 role: true,
540 name: true,
541}).extend({
542 type: z.literal(EventType.THINKING_END),
543});
544
545export const RawEventSchema = BaseEventSchema.extend({
546 type: z.literal(EventType.RAW),
547 event: z.any(),
548}).superRefine((v) => v);
549"#;
550
551 fn event<'a>(up: &'a Upstream, ty: &str) -> &'a UpstreamEvent {
552 up.events.iter().find(|e| e.event_type == ty).unwrap()
553 }
554
555 #[test]
556 fn reads_event_types_in_order() {
557 let up = extract(SAMPLE).unwrap();
558 assert_eq!(
559 up.event_types,
560 ["TEXT_MESSAGE_START", "THINKING_END", "RAW", "ORPHAN"]
561 );
562 }
563
564 #[test]
565 fn base_fields_are_inherited_not_repeated() {
566 let up = extract(SAMPLE).unwrap();
567 assert_eq!(
568 up.base_fields,
569 [
570 Field {
571 name: "timestamp".into(),
572 required: false
573 },
574 Field {
575 name: "rawEvent".into(),
576 required: false
577 },
578 ]
579 );
580 let start = event(&up, "TEXT_MESSAGE_START");
581 assert_eq!(
582 start.fields,
583 [
584 Field {
585 name: "messageId".into(),
586 required: true
587 },
588 Field {
589 name: "role".into(),
590 required: false
591 },
592 Field {
593 name: "name".into(),
594 required: false
595 },
596 ]
597 );
598 assert_eq!(start.schema.as_deref(), Some("TextMessageStartEventSchema"));
599 assert!(start.unparsed.is_none());
600 }
601
602 #[test]
603 fn follows_omit_and_extend_through_another_schema() {
604 let up = extract(SAMPLE).unwrap();
605 let thinking = event(&up, "THINKING_END");
606 assert_eq!(
607 thinking.fields,
608 [Field {
609 name: "messageId".into(),
610 required: true
611 }]
612 );
613 assert!(thinking.unparsed.is_none());
614 }
615
616 #[test]
617 fn unknown_modifier_is_a_warning_not_a_loss_of_the_event() {
618 let up = extract(SAMPLE).unwrap();
619 let raw = event(&up, "RAW");
620 assert_eq!(raw.schema.as_deref(), Some("RawEventSchema"));
621 assert!(raw.unparsed.as_deref().unwrap().contains("superRefine"));
622 }
623
624 #[test]
625 fn event_type_without_a_schema_is_a_warning() {
626 let up = extract(SAMPLE).unwrap();
627 let orphan = event(&up, "ORPHAN");
628 assert_eq!(orphan.schema, None);
629 assert!(orphan.unparsed.is_some());
630 }
631
632 #[test]
633 fn missing_enum_is_a_hard_error() {
634 assert!(extract("export const x = 1;").is_err());
635 }
636
637 const IMPORTED: &str = r#"
641import { z } from "zod";
642import { OptionalMetadataSchema } from "./metadata";
643import { StateSchema } from "./types";
644
645const MaybeName = z.string().optional();
646
647export enum EventType {
648 STATE_SNAPSHOT = "STATE_SNAPSHOT",
649}
650
651export const BaseEventSchema = z
652 .object({
653 type: z.nativeEnum(EventType),
654 metadata: OptionalMetadataSchema,
655 })
656 .passthrough();
657
658export const StateSnapshotEventSchema = BaseEventSchema.extend({
659 type: z.literal(EventType.STATE_SNAPSHOT),
660 snapshot: StateSchema,
661 name: MaybeName,
662});
663"#;
664
665 #[test]
666 fn an_imported_optional_schema_is_read_as_optional() {
667 let up = extract(IMPORTED).unwrap();
668 assert_eq!(
669 up.base_fields,
670 [Field {
671 name: "metadata".into(),
672 required: false
673 }]
674 );
675 assert!(
676 up.notes
677 .iter()
678 .any(|n| n.contains("OptionalMetadataSchema") && n.contains("read as optional")),
679 "{:?}",
680 up.notes
681 );
682 }
683
684 #[test]
685 fn an_imported_schema_without_the_prefix_stays_required() {
686 let up = extract(IMPORTED).unwrap();
687 let snapshot = event(&up, "STATE_SNAPSHOT");
688 assert_eq!(
689 snapshot.fields[0],
690 Field {
691 name: "snapshot".into(),
692 required: true
693 }
694 );
695 assert!(
696 up.notes
697 .iter()
698 .any(|n| n.contains("StateSchema") && n.contains("read as required")),
699 "{:?}",
700 up.notes
701 );
702 }
703
704 #[test]
707 fn a_locally_declared_schema_is_followed_not_guessed() {
708 let up = extract(IMPORTED).unwrap();
709 let snapshot = event(&up, "STATE_SNAPSHOT");
710 assert_eq!(
711 snapshot.fields[1],
712 Field {
713 name: "name".into(),
714 required: false
715 }
716 );
717 assert!(!up.notes.iter().any(|n| n.contains("MaybeName")));
718 }
719
720 #[test]
721 fn a_note_is_made_once_however_many_fields_share_the_schema() {
722 let up = extract(IMPORTED).unwrap();
723 assert_eq!(up.notes.len(), 2, "{:?}", up.notes);
724 }
725}