Skip to content
Source
Sign inSign up
Open History to see the last commit

AgentSessionJson.ts

192 lines6.3 KB

1import * as Schema from "effect/Schema";
2import * as SchemaAST from "effect/SchemaAST";
3import { none, type Many } from "stream-chain/defs.js";
4import { Assembler } from "stream-json/core/assembler.js";
5import * as StreamJson from "stream-json/core/parser.js";
6import type { ParserOptions, Token } from "stream-json/core/parser.js";
7​
8type JsonPath = ReadonlyArray<string | number | null>;
9​
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 if (node.indexSignatures.length > 0) return true;
17 return node.propertySignatures.some(
18 (property) =>
19 String(property.name) === path[index] && includes(property.type, path, index + 1),
20 );
21 case "Arrays": {
22 const key = path[index];
23 if (typeof key !== "number") return true;
24 const element = node.elements[key];
25 if (element) return includes(element, path, index + 1);
26 return node.rest.length === 0 || node.rest.some((item) => includes(item, path, index + 1));
27 }
28 case "Union":
29 return node.types.some((type) => includes(type, path, index));
30 case "Suspend":
31 return includes(node.thunk(), path, index);
32 case "Unknown":
33 case "Any":
34 case "ObjectKeyword":
35 case "Declaration":
36 return true;
37 default:
38 return false;
39 }
40 };
41 return (path: JsonPath) => includes(ast, path, 0);
43​
44export class TranscriptJsonLimitError extends Schema.TaggedError<TranscriptJsonLimitError>()(
45 "TranscriptJsonLimitError",
46 { message: Schema.String },
47) {}
48​
49export function createTranscriptJsonReader(
50 reserve: (bytes: number) => void,
51 selectPath: (path: JsonPath) => boolean,
52 options?: { readonly maxDepth?: number },
53) {
54 // 3.6.0 omits the synchronous tokenizer from types; parser() adds an unnecessary async generator.
55 const { jsonParser } = StreamJson as typeof StreamJson & {
56 jsonParser: (
57 options: ParserOptions,
58 ) => (input: string | typeof none) => Many<Token> | typeof none;
59 };
60 const tokenize = jsonParser({ packValues: false });
61 const assembler = new Assembler();
62 let key: string | null = null;
63 let value = "";
64 let depth = 0;
65 let complete = false;
66 let malformed = false;
67​
68 const assemble = (token: Token) => {
69 reserve(
70 64 + ("value" in token && typeof token.value === "string" ? token.value.length * 2 : 0),
71 );
72 switch (token.name) {
73 case "startString":
74 case "startNumber":
75 value = "";
76 break;
77 case "stringChunk":
78 case "numberChunk":
79 value += token.value;
80 break;
81 case "endString":
82 assembler.consume({ name: "stringValue", value });
83 value = "";
84 break;
85 case "endNumber":
86 assembler.consume({ name: "numberValue", value });
87 value = "";
88 break;
89 default:
90 assembler.consume(token);
91 }
92 };
93 // Preserve actual keys; adjacent duplicates share a path, but JSON.parse keeps the last value.
94 const stack: Array<{ path: JsonPath; key: string | number | null; selected: boolean }> = [];
95 let selectedValue = false;
96 const startValue = () => {
97 const parent = stack.at(-1);
98 const path = parent?.selected ? [...parent.path, parent.key] : [];
99 const selected = (parent?.selected ?? true) && selectPath(path);
100 if (selected && typeof parent?.key === "string") {
101 assemble({ name: "keyValue", value: parent.key });
102 }
103 return { path, selected };
104 };
105 const endValue = () => {
106 const parent = stack.at(-1);
107 if (parent && typeof parent.key === "number") parent.key++;
108 };
109 const selectToken = (token: Token | typeof none) => {
110 if (token === none) return;
111 switch (token.name) {
112 case "keyValue": {
113 const parent = stack.at(-1);
114 if (parent) parent.key = token.value;
115 return;
116 }
117 case "startObject":
118 case "startArray": {
119 const frame = startValue();
120 stack.push({ ...frame, key: token.name === "startArray" ? 0 : null });
121 if (frame.selected) assemble(token);
122 return;
123 }
124 case "endObject":
125 case "endArray":
126 if (stack.pop()?.selected) assemble(token);
127 endValue();
128 return;
129 case "startString":
130 case "startNumber":
131 selectedValue = startValue().selected;
132 if (selectedValue) assemble(token);
133 return;
134 case "endString":
135 case "endNumber":
136 if (selectedValue) assemble(token);
137 endValue();
138 return;
139 case "nullValue":
140 case "trueValue":
141 case "falseValue":
142 if (startValue().selected) assemble(token);
143 endValue();
144 return;
145 default:
146 if (selectedValue) assemble(token);
147 }
148 };
149 const consume = (input: string | typeof none) => {
150 if (malformed) return;
151 try {
152 const tokens = tokenize(input);
153 if (tokens === none) return;
154 for (const token of tokens.values) {
155 if (token.name === "startObject" || token.name === "startArray") {
156 if (++depth > (options?.maxDepth ?? 128))
157 throw new TranscriptJsonLimitError({
158 message: "Transcript JSON nesting exceeds the depth limit",
159 });
160 } else if (token.name === "endObject" || token.name === "endArray") {
161 if (--depth === 0) complete = true;
162 }
163 if (token.name === "startKey") {
164 key = "";
165 } else if (token.name === "stringChunk" && key !== null) {
166 reserve(token.value.length * 2);
167 key += token.value;
168 } else if (token.name === "endKey") {
169 selectToken({ name: "keyValue", value: key ?? "" });
170 key = null;
171 } else {
172 selectToken(token);
173 }
174 }
175 } catch (cause) {
176 if (cause instanceof Error && cause.message.startsWith("Parser ")) {
177 malformed = true;
178 } else {
179 throw cause;
180 }
181 }
182 };
183 return {
184 write: (chunk: string) => consume(chunk),
185 finish: (): unknown => {
186 consume(none);
187 if (malformed || !complete) return undefined;
188 selectToken(none);
189 return assembler.done ? assembler.current : undefined;
190 },
191 };