1import * as SchemaAST from "effect/SchemaAST";
2import { none, type Many } from "stream-chain/defs.js";
3import { Assembler } from "stream-json/core/assembler.js";
4import * as StreamJson from "stream-json/core/parser.js";
5import type { ParserOptions, Token } from "stream-json/core/parser.js";
6
7type JsonPath = ReadonlyArray<string | number | null>;
8
9/** Select schema fields before assembling their values, without a second field list. */
10export function createTranscriptJsonSelector(schema: { readonly ast: SchemaAST.AST }) {
11 const ast = SchemaAST.toEncoded(schema.ast);
12 const includes = (node: SchemaAST.AST, path: JsonPath, index: number): boolean => {
13 if (index === path.length) return true;
14 switch (node._tag) {
15 case "Objects":
16 // Records have dynamic keys. Keep their values for the decoder to validate.
17 if (node.indexSignatures.length > 0) return true;
18 return node.propertySignatures.some(
19 (property) =>
20 String(property.name) === path[index] && includes(property.type, path, index + 1),
21 );
22 case "Arrays": {
23 const key = path[index];
24 if (typeof key !== "number") return true;
25 const element = node.elements[key];
26 if (element) return includes(element, path, index + 1);
27 return node.rest.length === 0 || node.rest.some((item) => includes(item, path, index + 1));
28 }
29 case "Union":
30 return node.types.some((type) => includes(type, path, index));
31 case "Suspend":
32 return includes(node.thunk(), path, index);
33 case "Unknown":
34 case "Any":
35 case "ObjectKeyword":
36 case "Declaration":
37 // Unstructured/custom schemas must reach the decoder intact. The
38 // shared budget still bounds their allocations.
39 return true;
40 default:
41 return false;
42 }
43 };
44 return (path: JsonPath) => includes(ast, path, 0);
45}
46
47export class TranscriptJsonLimitError extends Error {}
48
49/**
50 * Project a single JSONL record without materializing unselected string values.
51 * The caller supplies a shared allocation budget for the entire transcript.
52 * Budget exhaustion rejects the transcript, never a message within it.
53 */
54export function createTranscriptJsonReader(
55 reserve: (bytes: number) => void,
56 selectPath: (path: JsonPath) => boolean,
57 options?: { readonly maxDepth?: number },
58) {
59 // The synchronous tokenizer is exported at runtime in 3.6.0, but omitted
60 // from its bundled types. Unlike parser(), it does not wrap tokens in an
61 // async generator; the file reader already supplies backpressure and UTF-8.
62 const { jsonParser } = StreamJson as typeof StreamJson & {
63 jsonParser: (
64 options: ParserOptions,
65 ) => (input: string | typeof none) => Many<Token> | typeof none;
66 };
67 const tokenize = jsonParser({ packValues: false });
68 const assembler = new Assembler();
69 let key: string | null = null;
70 let value = "";
71 let depth = 0;
72 let complete = false;
73 let malformed = false;
74
75 const assemble = (token: Token) => {
76 reserve(
77 64 + ("value" in token && typeof token.value === "string" ? token.value.length * 2 : 0),
78 );
79 switch (token.name) {
80 case "startString":
81 case "startNumber":
82 value = "";
83 break;
84 case "stringChunk":
85 case "numberChunk":
86 value += token.value;
87 break;
88 case "endString":
89 assembler.consume({ name: "stringValue", value });
90 value = "";
91 break;
92 case "endNumber":
93 assembler.consume({ name: "numberValue", value });
94 value = "";
95 break;
96 default:
97 assembler.consume(token);
98 }
99 };
100 // Forward actual selected keys instead of reconstructing them from path
101 // changes: adjacent duplicate keys have the same path but JSON.parse keeps
102 // the last value. Reconstructing paths can silently retain the first value.
103 const stack: Array<{ path: JsonPath; key: string | number | null; selected: boolean }> = [];
104 let selectedValue = false;
105 const startValue = () => {
106 const parent = stack.at(-1);
107 const path = parent?.selected ? [...parent.path, parent.key] : [];
108 const selected = (parent?.selected ?? true) && selectPath(path);
109 if (selected && typeof parent?.key === "string") {
110 assemble({ name: "keyValue", value: parent.key });
111 }
112 return { path, selected };
113 };
114 const endValue = () => {
115 const parent = stack.at(-1);
116 if (parent && typeof parent.key === "number") parent.key++;
117 };
118 const selectToken = (token: Token | typeof none) => {
119 if (token === none) return;
120 switch (token.name) {
121 case "keyValue": {
122 const parent = stack.at(-1);
123 if (parent) parent.key = token.value;
124 return;
125 }
126 case "startObject":
127 case "startArray": {
128 const frame = startValue();
129 stack.push({ ...frame, key: token.name === "startArray" ? 0 : null });
130 if (frame.selected) assemble(token);
131 return;
132 }
133 case "endObject":
134 case "endArray":
135 if (stack.pop()?.selected) assemble(token);
136 endValue();
137 return;
138 case "startString":
139 case "startNumber":
140 selectedValue = startValue().selected;
141 if (selectedValue) assemble(token);
142 return;
143 case "endString":
144 case "endNumber":
145 if (selectedValue) assemble(token);
146 endValue();
147 return;
148 case "nullValue":
149 case "trueValue":
150 case "falseValue":
151 if (startValue().selected) assemble(token);
152 endValue();
153 return;
154 default:
155 if (selectedValue) assemble(token);
156 }
157 };
158 const consume = (input: string | typeof none) => {
159 if (malformed) return;
160 try {
161 const tokens = tokenize(input);
162 if (tokens === none) return;
163 for (const token of tokens.values) {
164 if (token.name === "startObject" || token.name === "startArray") {
165 if (++depth > (options?.maxDepth ?? 128))
166 throw new TranscriptJsonLimitError("Transcript JSON nesting exceeds the depth limit");
167 } else if (token.name === "endObject" || token.name === "endArray") {
168 if (--depth === 0) complete = true;
169 }
170 // Charge keys before assembling them, including unknown names. Reject
171 // the transcript on exhaustion instead of silently shortening a key.
172 if (token.name === "startKey") {
173 key = "";
174 } else if (token.name === "stringChunk" && key !== null) {
175 reserve(token.value.length * 2);
176 key += token.value;
177 } else if (token.name === "endKey") {
178 selectToken({ name: "keyValue", value: key ?? "" });
179 key = null;
180 } else {
181 selectToken(token);
182 }
183 }
184 } catch (cause) {
185 if (cause instanceof Error && cause.message.startsWith("Parser ")) {
186 malformed = true;
187 } else {
188 throw cause;
189 }
190 }
191 };
192 return {
193 write: (chunk: string) => consume(chunk),
194 finish: (): unknown => {
195 consume(none);
196 if (malformed || !complete) return undefined;
197 selectToken(none);
198 return assembler.done ? assembler.current : undefined;
199 },
200 };
201}