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