572 lines · 18.6 KB · TypeScript
RawDownload
import * as Cause from "effect/Cause";import * as Effect from "effect/Effect";import * as Deferred from "effect/Deferred";import * as Queue from "effect/Queue";import * as Ref from "effect/Ref";import * as Schema from "effect/Schema";import * as Scope from "effect/Scope";import * as Stream from "effect/Stream";import * as Stdio from "effect/Stdio";import * as RpcClient from "effect/unstable/rpc/RpcClient";import * as RpcClientError from "effect/unstable/rpc/RpcClientError";import * as RpcMessage from "effect/unstable/rpc/RpcMessage";import * as RpcSerialization from "effect/unstable/rpc/RpcSerialization";import * as RpcServer from "effect/unstable/rpc/RpcServer"; import * as AcpSchema from "./_generated/schema.gen.ts";import { CLIENT_METHODS } from "./_generated/meta.gen.ts";import * as AcpError from "./errors.ts";const isAcpError = Schema.is(AcpError.AcpError); export interface AcpProtocolLogEvent {  readonly direction: "incoming" | "outgoing";  readonly stage: "raw" | "decoded" | "decode_failed";  readonly payload: unknown;} export type AcpIncomingNotification =  | {      readonly _tag: "SessionUpdate";      readonly method: typeof CLIENT_METHODS.session_update;      readonly params: AcpSchema.SessionNotification;    }  | {      readonly _tag: "ElicitationComplete";      readonly method: typeof CLIENT_METHODS.session_elicitation_complete;      readonly params: AcpSchema.ElicitationCompleteNotification;    }  | {      readonly _tag: "ExtNotification";      readonly method: string;      readonly params: unknown;    }; export interface AcpPatchedProtocolOptions {  readonly stdio: Stdio.Stdio;  readonly terminationError?: Effect.Effect<AcpError.AcpError>;  readonly serverRequestMethods: ReadonlySet<string>;  readonly logIncoming?: boolean;  readonly logOutgoing?: boolean;  readonly logger?: (event: AcpProtocolLogEvent) => Effect.Effect<void, never>;  readonly onNotification?: (    notification: AcpIncomingNotification,  ) => Effect.Effect<void, AcpError.AcpError, never>;  readonly onExtRequest?: (    method: string,    params: unknown,  ) => Effect.Effect<unknown, AcpError.AcpError, never>;  readonly onTermination?: (error: AcpError.AcpError) => Effect.Effect<void, never, never>;} export interface AcpPatchedProtocol {  readonly clientProtocol: RpcClient.Protocol["Service"];  readonly serverProtocol: RpcServer.Protocol["Service"];  readonly incoming: Stream.Stream<AcpIncomingNotification>;  readonly request: (method: string, payload: unknown) => Effect.Effect<unknown, AcpError.AcpError>;  readonly notify: (method: string, payload: unknown) => Effect.Effect<void, AcpError.AcpError>;} interface AcpPendingRequest {  readonly deferred: Deferred.Deferred<unknown, AcpError.AcpError>;  readonly method: string;} const decodeSessionUpdate = Schema.decodeUnknownEffect(AcpSchema.SessionNotification);const decodeElicitationComplete = Schema.decodeUnknownEffect(  AcpSchema.ElicitationCompleteNotification,);const parserFactory = RpcSerialization.ndJsonRpc(); export const makeAcpPatchedProtocol = Effect.fn("makeAcpPatchedProtocol")(function* (  options: AcpPatchedProtocolOptions,): Effect.fn.Return<AcpPatchedProtocol, never, Scope.Scope> {  const parser = parserFactory.makeUnsafe();  const serverQueue = yield* Queue.unbounded<RpcMessage.FromClientEncoded>();  const clientQueue = yield* Queue.unbounded<RpcMessage.FromServerEncoded>();  const notificationQueue = yield* Queue.unbounded<AcpIncomingNotification>();  const disconnects = yield* Queue.unbounded<number>();  const outgoing = yield* Queue.unbounded<string | Uint8Array, Cause.Done<void>>();  const nextRequestId = yield* Ref.make(1n);  const terminationHandled = yield* Ref.make(false);  const extPending = yield* Ref.make(new Map<string, AcpPendingRequest>());   const logProtocol = (event: AcpProtocolLogEvent) => {    if (event.direction === "incoming" && !options.logIncoming) {      return Effect.void;    }    if (event.direction === "outgoing" && !options.logOutgoing) {      return Effect.void;    }    return (      options.logger?.(event) ??      Effect.logDebug("ACP protocol event").pipe(Effect.annotateLogs({ event }))    );  };   const offerOutgoing = Effect.fn("offerOutgoing")(function* (    message: RpcMessage.FromClientEncoded | RpcMessage.FromServerEncoded,  ) {    yield* logProtocol({      direction: "outgoing",      stage: "decoded",      payload: message,    });     const method = message._tag === "Request" ? message.tag : undefined;    const encodedRequestId =      message._tag === "Request"        ? message.id        : "requestId" in message          ? message.requestId          : undefined;    const requestId = encodedRequestId === "" ? undefined : encodedRequestId;    const encoded = yield* Effect.try({      try: () => parser.encode(message),      catch: (cause) => AcpError.AcpProtocolParseError.fromEncodingError(method, requestId, cause),    });     if (encoded) {      yield* logProtocol({        direction: "outgoing",        stage: "raw",        payload: typeof encoded === "string" ? encoded : new TextDecoder().decode(encoded),      });       yield* Queue.offer(outgoing, encoded).pipe(Effect.asVoid);    }  });   const resolveExtPending = (    requestId: string,    onFound: (pendingRequest: AcpPendingRequest) => Effect.Effect<void>,  ) =>    Ref.modify(extPending, (pending) => {      const pendingRequest = pending.get(requestId);      if (!pendingRequest) {        return [Effect.void, pending] as const;      }      const next = new Map(pending);      next.delete(requestId);      return [onFound(pendingRequest), next] as const;    }).pipe(Effect.flatten);   const removeExtPending = (requestId: string) =>    Ref.update(extPending, (pending) => {      if (!pending.has(requestId)) {        return pending;      }      const next = new Map(pending);      next.delete(requestId);      return next;    });   const completeExtPendingFailure = (requestId: string, error: AcpError.AcpError) =>    resolveExtPending(requestId, ({ deferred }) => Deferred.fail(deferred, error));   const completeExtPendingSuccess = (requestId: string, value: unknown) =>    resolveExtPending(requestId, ({ deferred }) => Deferred.succeed(deferred, value));   const failAllExtPending = (error: AcpError.AcpError) =>    Ref.getAndSet(extPending, new Map()).pipe(      Effect.flatMap((pending) =>        Effect.forEach([...pending.values()], ({ deferred }) => Deferred.fail(deferred, error), {          discard: true,        }),      ),    );   const dispatchNotification = (notification: AcpIncomingNotification) =>    Queue.offer(notificationQueue, notification).pipe(      Effect.andThen(        options.onNotification          ? options.onNotification(notification).pipe(Effect.catch(() => Effect.void))          : Effect.void,      ),      Effect.asVoid,    );   const emitClientProtocolError = (error: AcpError.AcpError) =>    Queue.offer(clientQueue, {      _tag: "ClientProtocolError",      error: new RpcClientError.RpcClientError({        reason: new RpcClientError.RpcClientDefect({          message: "ACP protocol terminated.",          cause: error,        }),      }),    }).pipe(Effect.asVoid);   const handleTermination = (classify: () => Effect.Effect<AcpError.AcpError | undefined>) =>    Ref.modify(terminationHandled, (handled) => {      if (handled) {        return [Effect.void, true] as const;      }      return [        Effect.gen(function* () {          yield* Queue.offer(disconnects, 0);          const error = yield* classify();          if (!error) {            return;          }          yield* failAllExtPending(error);          yield* emitClientProtocolError(error);          if (options.onTermination) {            yield* options.onTermination(error);          }        }),        true,      ] as const;    }).pipe(Effect.flatten);   const respondWithSuccess = (requestId: string, value: unknown) =>    offerOutgoing({      _tag: "Exit",      requestId,      exit: {        _tag: "Success",        value,      },    });   const respondWithError = (requestId: string, error: AcpError.AcpRequestError) =>    offerOutgoing({      _tag: "Exit",      requestId,      exit: {        _tag: "Failure",        cause: [          {            _tag: "Fail",            error: error.toProtocolError(),          },        ],      },    });   const handleExtRequest = (message: RpcMessage.RequestEncoded) => {    if (!options.onExtRequest) {      return respondWithError(message.id, AcpError.AcpRequestError.methodNotFound(message.tag));    }    return options.onExtRequest(message.tag, message.payload).pipe(      Effect.matchEffect({        onFailure: (error) =>          respondWithError(            message.id,            AcpError.AcpRequestError.fromExtensionHandlerError(error, message.tag),          ),        onSuccess: (value) => respondWithSuccess(message.id, value),      }),    );  };   const handleRequestEncoded = (message: RpcMessage.RequestEncoded) => {    if (message.id === "") {      if (message.tag === CLIENT_METHODS.session_update) {        return decodeSessionUpdate(message.payload).pipe(          Effect.map(            (params) =>              ({                _tag: "SessionUpdate",                method: CLIENT_METHODS.session_update,                params,              }) satisfies AcpIncomingNotification,          ),          Effect.mapError((cause) =>            AcpError.AcpProtocolParseError.fromSchemaError(              "decode-notification-payload",              CLIENT_METHODS.session_update,              cause,            ),          ),          Effect.flatMap(dispatchNotification),        );      }      if (message.tag === CLIENT_METHODS.session_elicitation_complete) {        return decodeElicitationComplete(message.payload).pipe(          Effect.map(            (params) =>              ({                _tag: "ElicitationComplete",                method: CLIENT_METHODS.session_elicitation_complete,                params,              }) satisfies AcpIncomingNotification,          ),          Effect.mapError((cause) =>            AcpError.AcpProtocolParseError.fromSchemaError(              "decode-notification-payload",              CLIENT_METHODS.session_elicitation_complete,              cause,            ),          ),          Effect.flatMap(dispatchNotification),        );      }      return dispatchNotification({        _tag: "ExtNotification",        method: message.tag,        params: message.payload,      });    }     if (!options.serverRequestMethods.has(message.tag)) {      return handleExtRequest(message).pipe(        Effect.catchTags({          AcpProtocolParseError: (error) =>            Effect.logWarning(error).pipe(              Effect.annotateLogs({                method: message.tag,                requestId: message.id,                operation: error.operation,              }),              Effect.andThen(                respondWithError(                  message.id,                  AcpError.AcpRequestError.fromExtensionResponseEncodingError(                    message.tag,                    message.id,                    error,                  ),                ),              ),            ),        }),        Effect.asVoid,      );    }     return Queue.offer(serverQueue, message).pipe(Effect.asVoid);  };   const handleExitEncoded = (message: RpcMessage.ResponseExitEncoded) =>    Ref.get(extPending).pipe(      Effect.flatMap((pending) => {        const pendingRequest = pending.get(message.requestId);        if (!pendingRequest) {          return Queue.offer(clientQueue, message).pipe(Effect.asVoid);        }        if (message.exit._tag === "Success") {          return completeExtPendingSuccess(message.requestId, message.exit.value);        }        const failure = message.exit.cause.find((entry) => entry._tag === "Fail");        if (failure && isProtocolError(failure.error)) {          return completeExtPendingFailure(            message.requestId,            AcpError.AcpRequestError.fromProtocolError(failure.error, {              method: pendingRequest.method,              requestId: message.requestId,              cause: message.exit.cause,            }),          );        }        return completeExtPendingFailure(          message.requestId,          AcpError.AcpRequestError.fromExtensionResponseFailure(            pendingRequest.method,            message.requestId,            message.exit.cause,          ),        );      }),    );   const routeDecodedMessage = (    message: RpcMessage.FromClientEncoded | RpcMessage.FromServerEncoded,  ): Effect.Effect<void, AcpError.AcpError> => {    switch (message._tag) {      case "Request":        return handleRequestEncoded(message);      case "Exit":        return handleExitEncoded(message);      case "Chunk":        return Ref.get(extPending).pipe(          Effect.flatMap((pending) => {            const pendingRequest = pending.get(message.requestId);            return pendingRequest              ? completeExtPendingFailure(                  message.requestId,                  AcpError.AcpRequestError.unsupportedStreamingResponse(                    pendingRequest.method,                    message.requestId,                  ),                )              : Queue.offer(clientQueue, message).pipe(Effect.asVoid);          }),        );      case "Defect":      case "ClientProtocolError":      case "Pong":        return Queue.offer(clientQueue, message).pipe(Effect.asVoid);      case "Ack":      case "Interrupt":      case "Ping":      case "Eof":        return Queue.offer(serverQueue, message).pipe(Effect.asVoid);    }  };   yield* options.stdio.stdin.pipe(    Stream.runForEach((data) =>      logProtocol({        direction: "incoming",        stage: "raw",        payload: typeof data === "string" ? data : new TextDecoder().decode(data),      }).pipe(        Effect.flatMap(() =>          Effect.try({            try: () =>              parser.decode(data) as ReadonlyArray<                RpcMessage.FromClientEncoded | RpcMessage.FromServerEncoded              >,            catch: (cause) =>              new AcpError.AcpProtocolParseError({                operation: "decode-wire-message",                cause,              }),          }),        ),        Effect.tap((messages) =>          logProtocol({            direction: "incoming",            stage: "decoded",            payload: messages,          }),        ),        Effect.tapErrorTag("AcpProtocolParseError", (error) =>          logProtocol({            direction: "incoming",            stage: "decode_failed",            payload: {              operation: error.operation,              ...(error.method === undefined ? {} : { method: error.method }),              ...(error.requestId === undefined ? {} : { requestId: error.requestId }),              ...(error.issueCount === undefined ? {} : { issueCount: error.issueCount }),              ...(error.issueKinds === undefined ? {} : { issueKinds: error.issueKinds }),              ...(error.maximumPathDepth === undefined                ? {}                : { maximumPathDepth: error.maximumPathDepth }),            },          }),        ),        Effect.flatMap((messages) =>          Effect.forEach(messages, routeDecodedMessage, {            discard: true,          }),        ),      ),    ),    Effect.matchEffect({      onFailure: (error) => {        const normalized: AcpError.AcpError = isAcpError(error)          ? error          : new AcpError.AcpTransportError({              operation: "read-input-stream",              cause: error,            });        return handleTermination(() => Effect.succeed(normalized));      },      onSuccess: () =>        handleTermination(          () =>            options.terminationError ?? Effect.succeed(new AcpError.AcpInputStreamEndedError({})),        ),    }),    Effect.forkScoped,  );   yield* Stream.fromQueue(outgoing).pipe(Stream.run(options.stdio.stdout()), Effect.forkScoped);   const clientProtocol = RpcClient.Protocol.of({    run: (_clientId, f) =>      Stream.fromQueue(clientQueue).pipe(        Stream.runForEach((message) => f(message)),        Effect.forever,      ),    send: (_clientId, request) =>      offerOutgoing(request).pipe(        Effect.mapError(          (error) =>            new RpcClientError.RpcClientError({              reason: new RpcClientError.RpcClientDefect({                message: "Failed to send ACP protocol message.",                cause: error,              }),            }),        ),      ),    supportsAck: true,    supportsTransferables: false,  });   const serverProtocol = RpcServer.Protocol.of({

Lines 1–500 of 572.