1import * as Cause from "effect/Cause";
2import * as Deferred from "effect/Deferred";
3import * as Effect from "effect/Effect";
4import * as Exit from "effect/Exit";
5import * as Option from "effect/Option";
6import * as Queue from "effect/Queue";
7import * as Ref from "effect/Ref";
8import * as Scope from "effect/Scope";
9import * as Schema from "effect/Schema";
10import * as Stdio from "effect/Stdio";
11import * as Stream from "effect/Stream";
12
13import * as CodexError from "./errors.ts";
14import { JsonRpcId, JsonRpcResponseEnvelope } from "./_internal/shared.ts";
15const isJsonRpcId = Schema.is(JsonRpcId);
16const isJsonRpcResponseEnvelope = Schema.is(JsonRpcResponseEnvelope);
17const isCodexAppServerError = Schema.is(CodexError.CodexAppServerError);
18const MAX_BUFFERED_RAW_MESSAGES = 32;
19
20export interface CodexAppServerProtocolLogEvent {
21 readonly direction: "incoming" | "outgoing";
22 readonly stage: "raw" | "decoded" | "decode_failed";
23 readonly payload: unknown;
24}
25
26export interface CodexAppServerIncomingNotification {
27 readonly method: string;
28 readonly params?: unknown;
29}
30
31export interface CodexAppServerIncomingRequest {
32 readonly id: string | number;
33 readonly method: string;
34 readonly params?: unknown;
35}
36
37export interface CodexAppServerPatchedProtocolOptions {
38 readonly stdio: Stdio.Stdio;
39 readonly terminationError?: Effect.Effect<CodexError.CodexAppServerError>;
40 readonly logIncoming?: boolean;
41 readonly logOutgoing?: boolean;
42 readonly logger?: (event: CodexAppServerProtocolLogEvent) => Effect.Effect<void, never>;
43 readonly onNotification?: (
44 notification: CodexAppServerIncomingNotification,
45 ) => Effect.Effect<void, never>;
46 readonly onRequest?: (
47 request: CodexAppServerIncomingRequest,
48 ) => Effect.Effect<unknown, CodexError.CodexAppServerError>;
49 readonly onTermination?: (error: CodexError.CodexAppServerError) => Effect.Effect<void, never>;
50}
51
52export interface CodexAppServerPatchedProtocol {
53 readonly incomingNotifications: Stream.Stream<CodexAppServerIncomingNotification>;
54 readonly incomingRequests: Stream.Stream<CodexAppServerIncomingRequest>;
55 readonly request: (
56 method: string,
57 payload?: unknown,
58 ) => Effect.Effect<unknown, CodexError.CodexAppServerError>;
59 readonly notify: (
60 method: string,
61 payload?: unknown,
62 ) => Effect.Effect<void, CodexError.CodexAppServerError>;
63 readonly respond: (
64 requestId: string | number,
65 result: unknown,
66 ) => Effect.Effect<void, CodexError.CodexAppServerError>;
67 readonly respondError: (
68 requestId: string | number,
69 error: CodexError.CodexAppServerRequestError,
70 ) => Effect.Effect<void, CodexError.CodexAppServerError>;
71}
72
73interface CodexAppServerPendingRequest {
74 readonly deferred: Deferred.Deferred<unknown, CodexError.CodexAppServerError>;
75 readonly method: string;
76}
77
78function isObject(value: unknown): value is Record<string, unknown> {
79 return typeof value === "object" && value !== null;
80}
81
82function isIncomingRequest(value: unknown): value is CodexAppServerIncomingRequest {
83 if (!isObject(value) || typeof value.method !== "string") {
84 return false;
85 }
86 return isJsonRpcId(value.id);
87}
88
89function isIncomingNotification(value: unknown): value is CodexAppServerIncomingNotification {
90 return isObject(value) && typeof value.method === "string" && !("id" in value);
91}
92
93function isIncomingResponse(value: unknown): value is typeof JsonRpcResponseEnvelope.Type {
94 return isJsonRpcResponseEnvelope(value);
95}
96
97const encodeJsonString = Schema.encodeUnknownEffect(Schema.fromJsonString(Schema.Unknown));
98const decodeJsonString = Schema.decodeUnknownEffect(Schema.fromJsonString(Schema.Unknown));
99
100const encodeWireMessage = (
101 message: Record<string, unknown>,
102): Effect.Effect<string, CodexError.CodexAppServerProtocolParseError> =>
103 encodeJsonString(message).pipe(
104 Effect.map((encoded) => `${encoded}\n`),
105 Effect.mapError((cause) => {
106 const method = typeof message.method === "string" ? message.method : undefined;
107 const requestId =
108 typeof message.id === "string" || typeof message.id === "number"
109 ? String(message.id)
110 : undefined;
111 return CodexError.CodexAppServerProtocolParseError.fromSchemaError(
112 "encode-wire-message",
113 cause,
114 {
115 ...(method === undefined ? {} : { method }),
116 ...(requestId === undefined ? {} : { requestId }),
117 },
118 );
119 }),
120 );
121
122const decodeWireMessage = (
123 line: string,
124): Effect.Effect<unknown, CodexError.CodexAppServerProtocolParseError> =>
125 decodeJsonString(line).pipe(
126 Effect.mapError((cause) =>
127 CodexError.CodexAppServerProtocolParseError.fromSchemaError("decode-wire-message", cause),
128 ),
129 );
130
131const normalizeIncomingError = (
132 error: unknown,
133 operation: CodexError.CodexAppServerTransportOperation,
134): CodexError.CodexAppServerError =>
135 isCodexAppServerError(error)
136 ? error
137 : new CodexError.CodexAppServerTransportError({
138 operation,
139 cause: error,
140 });
141
142const toProtocolMessage = (
143 requestId: string | number,
144 fields: {
145 readonly result?: unknown;
146 readonly error?: CodexError.CodexAppServerProtocolErrorShape;
147 },
148): { readonly [key: string]: unknown } => ({
149 id: requestId,
150 ...(fields.result !== undefined ? { result: fields.result } : {}),
151 ...(fields.error !== undefined ? { error: fields.error } : {}),
152});
153
154export const makeCodexAppServerPatchedProtocol = Effect.fn("makeCodexAppServerPatchedProtocol")(
155 function* (
156 options: CodexAppServerPatchedProtocolOptions,
157 ): Effect.fn.Return<CodexAppServerPatchedProtocol, never, Scope.Scope> {
158 const protocolScope = yield* Scope.Scope;
159 const requestHandlerScope = yield* Scope.fork(protocolScope, "parallel");
160 const outgoing = yield* Queue.unbounded<string, Cause.Done<void>>();
161 const incomingNotifications =
162 yield* Queue.sliding<CodexAppServerIncomingNotification>(MAX_BUFFERED_RAW_MESSAGES);
163 const incomingRequests =
164 yield* Queue.sliding<CodexAppServerIncomingRequest>(MAX_BUFFERED_RAW_MESSAGES);
165 const pending = yield* Ref.make(new Map<string, CodexAppServerPendingRequest>());
166 const nextRequestId = yield* Ref.make(1);
167 const remainder: Array<string> = [];
168 const terminationHandled = yield* Ref.make(false);
169 const terminationFailure = yield* Ref.make(Option.none<CodexError.CodexAppServerError>());
170 const terminationSignal = yield* Deferred.make<void>();
171 const activeRequestHandlers = yield* Ref.make(0);
172
173 const logProtocol = (event: CodexAppServerProtocolLogEvent) => {
174 if (event.direction === "incoming" && !options.logIncoming) {
175 return Effect.void;
176 }
177 if (event.direction === "outgoing" && !options.logOutgoing) {
178 return Effect.void;
179 }
180 return (
181 options.logger?.(event) ??
182 Effect.logDebug("Codex App Server protocol event").pipe(Effect.annotateLogs({ event }))
183 );
184 };
185
186 const failAllPending = (error: CodexError.CodexAppServerError) =>
187 Ref.get(pending).pipe(
188 Effect.flatMap((current) =>
189 Effect.forEach([...current.values()], ({ deferred }) => Deferred.fail(deferred, error), {
190 discard: true,
191 }),
192 ),
193 Effect.andThen(Ref.set(pending, new Map())),
194 );
195
196 const handleTermination = (classify: () => Effect.Effect<CodexError.CodexAppServerError>) =>
197 Ref.modify(terminationHandled, (handled) => {
198 if (handled) {
199 return [Effect.void, true] as const;
200 }
201 return [
202 Effect.gen(function* () {
203 const error = yield* classify();
204 yield* Ref.set(terminationFailure, Option.some(error));
205 yield* failAllPending(error);
206 yield* Queue.end(outgoing);
207 yield* Deferred.succeed(terminationSignal, undefined);
208 yield* Scope.close(requestHandlerScope, Exit.void).pipe(
209 Effect.forkIn(protocolScope, { startImmediately: true }),
210 Effect.asVoid,
211 );
212 if (options.onTermination) {
213 yield* options.onTermination(error);
214 }
215 }),
216 true,
217 ] as const;
218 }).pipe(Effect.flatten);
219
220 const offerOutgoing = (message: Record<string, unknown>) =>
221 Effect.gen(function* () {
222 const failure = yield* Ref.get(terminationFailure);
223 if (Option.isSome(failure)) return yield* failure.value;
224
225 yield* logProtocol({
226 direction: "outgoing",
227 stage: "decoded",
228 payload: message,
229 });
230 const encoded = yield* encodeWireMessage(message);
231 yield* logProtocol({
232 direction: "outgoing",
233 stage: "raw",
234 payload: encoded,
235 });
236 const accepted = yield* Queue.offer(outgoing, encoded);
237 if (!accepted) {
238 const closed = yield* Ref.get(terminationFailure);
239 return yield* Option.getOrElse(
240 closed,
241 () => new CodexError.CodexAppServerInputStreamEndedError({}),
242 );
243 }
244 });
245
246 const removePending = (requestId: string) =>
247 Ref.update(pending, (current) => {
248 if (!current.has(requestId)) {
249 return current;
250 }
251 const next = new Map(current);
252 next.delete(requestId);
253 return next;
254 });
255
256 const resolvePending = (
257 requestId: string,
258 handler: (pendingRequest: CodexAppServerPendingRequest) => Effect.Effect<void>,
259 ) =>
260 Ref.modify(pending, (current) => {
261 const pendingRequest = current.get(requestId);
262 if (!pendingRequest) {
263 return [Effect.void, current] as const;
264 }
265 const next = new Map(current);
266 next.delete(requestId);
267 return [handler(pendingRequest), next] as const;
268 }).pipe(Effect.flatten);
269
270 const respond = (requestId: string | number, result: unknown) =>
271 offerOutgoing(toProtocolMessage(requestId, { result }));
272
273 const respondError = (
274 requestId: string | number,
275 error: CodexError.CodexAppServerRequestError,
276 ) => offerOutgoing(toProtocolMessage(requestId, { error: error.toProtocolError() }));
277
278 const handleResponse = (response: typeof JsonRpcResponseEnvelope.Type) => {
279 const requestId = String(response.id);
280 const protocolError = response.error;
281 if (protocolError !== undefined) {
282 return resolvePending(requestId, ({ deferred, method }) =>
283 Deferred.fail(
284 deferred,
285 CodexError.CodexAppServerRequestError.fromProtocolError(
286 protocolError,
287 method,
288 requestId,
289 ),
290 ),
291 );
292 }
293 return resolvePending(requestId, ({ deferred }) =>
294 Deferred.succeed(deferred, response.result),
295 );
296 };
297
298 const handleRequest = (request: CodexAppServerIncomingRequest) =>
299 Queue.offer(incomingRequests, request).pipe(
300 Effect.flatMap(() => {
301 const handler = options.onRequest;
302 if (!handler) return Effect.void;
303
304 return Ref.modify(activeRequestHandlers, (count) =>
305 count >= MAX_BUFFERED_RAW_MESSAGES ? [false, count] : [true, count + 1],
306 ).pipe(
307 Effect.flatMap((accepted) => {
308 if (!accepted) {
309 return respondError(
310 request.id,
311 CodexError.CodexAppServerRequestError.overloaded(
312 "Too many Codex requests are already active.",
313 ),
314 );
315 }
316
317 return handler(request).pipe(
318 Effect.matchEffect({
319 onFailure: (error) =>
320 respondError(
321 request.id,
322 CodexError.CodexAppServerRequestError.fromAppServerError(
323 error,
324 request.method,
325 ),
326 ),
327 onSuccess: (result) => respond(request.id, result),
328 }),
329 Effect.ensuring(
330 Ref.update(activeRequestHandlers, (count) => Math.max(0, count - 1)),
331 ),
332 Effect.catch((error) =>
333 handleTermination(() => Effect.succeed(error)).pipe(
334 Effect.forkIn(protocolScope),
335 Effect.asVoid,
336 ),
337 ),
338 Effect.forkIn(requestHandlerScope, { startImmediately: true }),
339 Effect.asVoid,
340 );
341 }),
342 );
343 }),
344 Effect.asVoid,
345 );
346
347 const handleNotification = (notification: CodexAppServerIncomingNotification) =>
348 Queue.offer(incomingNotifications, notification).pipe(
349 Effect.andThen(options.onNotification ? options.onNotification(notification) : Effect.void),
350 Effect.asVoid,
351 );
352
353 const routeMessage = Effect.fnUntraced(function* (message: unknown) {
354 if (Option.isSome(yield* Ref.get(terminationFailure))) return;
355 if (isIncomingRequest(message)) return yield* handleRequest(message);
356 if (isIncomingNotification(message)) return yield* handleNotification(message);
357 if (isIncomingResponse(message)) return yield* handleResponse(message);
358 return yield* CodexError.CodexAppServerProtocolParseError.fromUnroutableMessage(message);
359 });
360
361 const handleLine = (line: string): Effect.Effect<void, CodexError.CodexAppServerError> => {
362 if (line.trim().length === 0) {
363 return Effect.void;
364 }
365 return logProtocol({
366 direction: "incoming",
367 stage: "raw",
368 payload: line,
369 }).pipe(
370 Effect.flatMap(() => decodeWireMessage(line)),
371 Effect.tap((decoded) =>
372 logProtocol({
373 direction: "incoming",
374 stage: "decoded",
375 payload: decoded,
376 }),
377 ),
378 Effect.tapErrorTag("CodexAppServerProtocolParseError", (error) =>
379 logProtocol({
380 direction: "incoming",
381 stage: "decode_failed",
382 payload: {
383 operation: error.operation,
384 ...(error.method === undefined ? {} : { method: error.method }),
385 ...(error.requestId === undefined ? {} : { requestId: error.requestId }),
386 ...(error.issueCount === undefined ? {} : { issueCount: error.issueCount }),
387 ...(error.issueKinds === undefined ? {} : { issueKinds: error.issueKinds }),
388 ...(error.maximumPathDepth === undefined
389 ? {}
390 : { maximumPathDepth: error.maximumPathDepth }),
391 },
392 }),
393 ),
394 Effect.flatMap(routeMessage),
395 );
396 };
397
398 yield* options.stdio.stdin.pipe(
399 Stream.interruptWhen(Deferred.await(terminationSignal)),
400 Stream.decodeText(),
401 Stream.runForEach((chunk) =>
402 Effect.sync(() => {
403 const lines: Array<string> = [];
404 let start = 0;
405 for (
406 let newline = chunk.indexOf("\n");
407 newline !== -1;
408 newline = chunk.indexOf("\n", start)
409 ) {
410 remainder.push(chunk.slice(start, newline));
411 lines.push(remainder.join("").replace(/\r$/, ""));
412 remainder.length = 0;
413 start = newline + 1;
414 }
415 // Keep unfinished lines in fragments so each chunk is scanned only once.
416 if (start < chunk.length) {
417 remainder.push(chunk.slice(start));
418 }
419 return lines;
420 }).pipe(Effect.flatMap((lines) => Effect.forEach(lines, handleLine, { discard: true }))),
421 ),
422 Effect.matchEffect({
423 onFailure: (error) =>
424 handleTermination(() =>
425 Effect.succeed(normalizeIncomingError(error, "read-input-stream")),
426 ),
427 onSuccess: () =>
428 Effect.sync(() => {
429 const line = remainder.join("");
430 remainder.length = 0;
431 return line;
432 }).pipe(
433 Effect.flatMap(handleLine),
434 Effect.matchEffect({
435 onFailure: (error) => handleTermination(() => Effect.succeed(error)),
436 onSuccess: () =>
437 handleTermination(
438 () =>
439 options.terminationError ??
440 Effect.succeed(new CodexError.CodexAppServerInputStreamEndedError({})),
441 ),
442 }),
443 ),
444 }),
445 Effect.forkScoped,
446 );
447
448 yield* Stream.fromQueue(outgoing).pipe(Stream.run(options.stdio.stdout()), Effect.forkScoped);
449
450 const request = (method: string, payload?: unknown) =>
451 Effect.gen(function* () {
452 const requestId = yield* Ref.modify(
453 nextRequestId,
454 (current) => [current, current + 1] as const,
455 );
456 const deferred = yield* Deferred.make<unknown, CodexError.CodexAppServerError>();
457 yield* Ref.update(pending, (current) =>
458 new Map(current).set(String(requestId), { deferred, method }),
459 );
460 yield* offerOutgoing({
461 id: requestId,
462 method,
463 ...(payload !== undefined ? { params: payload } : {}),
464 }).pipe(Effect.tapError(() => removePending(String(requestId))));
465 return yield* Deferred.await(deferred).pipe(
466 Effect.onInterrupt(() => removePending(String(requestId))),
467 );
468 });
469
470 const notify = (method: string, payload?: unknown) =>
471 offerOutgoing({
472 method,
473 ...(payload !== undefined ? { params: payload } : {}),
474 });
475
476 return {
477 incomingNotifications: Stream.fromQueue(incomingNotifications),
478 incomingRequests: Stream.fromQueue(incomingRequests),
479 request,
480 notify,
481 respond,
482 respondError,
483 } satisfies CodexAppServerPatchedProtocol;
484 },
485);