863 lines · 29.8 KB · TypeScript
RawDownload
import * as Cause from "effect/Cause";import * as Effect from "effect/Effect";import * as Deferred from "effect/Deferred";import * as Exit from "effect/Exit";import * as Queue from "effect/Queue";import * as Ref from "effect/Ref";import type * as PlatformError from "effect/PlatformError";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;    }; /** Standard I/O whose input can report provider-specific ACP failures. */export interface AcpStdio extends Omit<Stdio.Stdio, "stdin"> {  readonly stdin: Stream.Stream<Uint8Array, PlatformError.PlatformError | AcpError.AcpError>;} export interface AcpPatchedProtocolOptions {  readonly stdio: AcpStdio;  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 onIncomingRequest?: (    requestId: string,    method: string,    payload: unknown,  ) => Effect.Effect<void, never>;  readonly transformSessionUpdate?: (    notification: AcpSchema.SessionNotification,  ) => AcpSchema.SessionNotification;  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>;  readonly onOutgoingResponseFailure?: (    requestId: string,    error: AcpError.AcpError,  ) => Effect.Effect<void, never>;  readonly onOutgoingResponse?: (requestId: string) => Effect.Effect<void, never>;  readonly testHooks?: {    readonly onOutgoingWriteAdmitted?: () => Effect.Effect<void>;    /** Invoked from the stdout writer exit path (after intentional/unexpected handling). */    readonly onOutgoingWriterExit?: (input: {      readonly intentionalEnd: boolean;    }) => Effect.Effect<void>;  };} 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;} interface AcpOutgoingWrite {  readonly acknowledgement?: Deferred.Deferred<void, AcpError.AcpError>;  readonly payload: string | Uint8Array;} interface AcpOutgoingWriterState {  readonly outstandingAcknowledgements: ReadonlySet<Deferred.Deferred<void, AcpError.AcpError>>;  readonly intentionalEnd: boolean;  readonly terminalError?: AcpError.AcpError;} const decodeSessionUpdate = Schema.decodeUnknownEffect(AcpSchema.SessionNotification);const decodeElicitationComplete = Schema.decodeUnknownEffect(  AcpSchema.ElicitationCompleteNotification,);const parserFactory = RpcSerialization.ndJsonRpc();const MAX_BUFFERED_RAW_NOTIFICATIONS = 32;// Outbound JSON-RPC notification: no `id`, so peers never treat it as a request.const encodeJsonRpcNotification = Schema.encodeUnknownExit(  Schema.fromJsonString(    Schema.Struct({      jsonrpc: Schema.Literal("2.0"),      method: Schema.String,      params: Schema.Unknown,    }),  ),); const isEffectRpcRequestId = (requestId: AcpError.AcpRequestId): boolean =>  typeof requestId === "number" && Number.isSafeInteger(requestId); 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.sliding<AcpIncomingNotification>(    MAX_BUFFERED_RAW_NOTIFICATIONS,  );  const disconnects = yield* Queue.unbounded<number>();  const outgoing = yield* Queue.unbounded<AcpOutgoingWrite, AcpError.AcpError | Cause.Done<void>>();  const outgoingWriterState = yield* Ref.make<AcpOutgoingWriterState>({    intentionalEnd: false,    outstandingAcknowledgements: new Set(),  });  const nextRequestId = yield* Ref.make(1);  const terminationHandled = yield* Ref.make(false);  const terminationFailure = yield* Deferred.make<never, AcpError.AcpError>();  const extPending = yield* Ref.make(new Map<string, AcpPendingRequest>());   const ensureActive = Ref.get(terminationHandled).pipe(    Effect.flatMap((terminated) => (terminated ? Deferred.await(terminationFailure) : Effect.void)),  );   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,  ) {    // RpcClient emits `@effect/rpc/Interrupt` when a pending request's fiber is interrupted.    // ACP has no such method; agents log it as an error and cannot act on it, so drop it.    if (message._tag === "Interrupt") {      return;    }    yield* ensureActive;    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* ensureActive;      const acknowledgement =        message._tag === "Exit" ? yield* Deferred.make<void, AcpError.AcpError>() : undefined;      const admissionError = yield* Ref.modify(outgoingWriterState, (state) => {        if (state.terminalError !== undefined) {          return [state.terminalError, state] as const;        }        if (acknowledgement === undefined) {          return [undefined, state] as const;        }        return [          undefined,          {            ...state,            outstandingAcknowledgements: new Set(state.outstandingAcknowledgements).add(              acknowledgement,            ),          },        ] as const;      });      if (admissionError !== undefined) {        return yield* admissionError;      }      if (options.testHooks?.onOutgoingWriteAdmitted !== undefined) {        yield* options.testHooks.onOutgoingWriteAdmitted();      }      yield* Queue.offer(outgoing, {        payload: encoded,        ...(acknowledgement === undefined ? {} : { acknowledgement }),      }).pipe(        Effect.mapError(          (cause) =>            new AcpError.AcpTransportError({              detail: "Failed to queue an outgoing ACP message",              cause,            }),        ),        Effect.filterOrFail(          (accepted) => accepted,          () =>            new AcpError.AcpTransportError({              detail: "The ACP output queue was closed before accepting a message",              cause: "Output queue closed",            }),        ),        Effect.tapError((error) =>          acknowledgement === undefined            ? Effect.void            : Ref.update(outgoingWriterState, (state) => {                const updated = new Set(state.outstandingAcknowledgements);                updated.delete(acknowledgement);                return { ...state, outstandingAcknowledgements: updated };              }).pipe(Effect.andThen(Deferred.fail(acknowledgement, error)), Effect.asVoid),        ),      );      if (acknowledgement !== undefined) {        yield* Deferred.await(acknowledgement);      }      if (message._tag === "Exit" && options.onOutgoingResponse !== undefined) {        yield* options.onOutgoingResponse(String(message.requestId));      }    }  });   const resolveExtPending = (    requestId: AcpError.AcpRequestId,    onFound: (pendingRequest: AcpPendingRequest) => Effect.Effect<void>,  ) =>    Ref.modify(extPending, (pending) => {      const pendingKey = String(requestId);      const pendingRequest = pending.get(pendingKey);      if (!pendingRequest) {        return [Effect.void, pending] as const;      }      const next = new Map(pending);      next.delete(pendingKey);      return [onFound(pendingRequest), next] as const;    }).pipe(Effect.flatten);   const removeExtPending = (requestId: AcpError.AcpRequestId) =>    Ref.update(extPending, (pending) => {      const pendingKey = String(requestId);      if (!pending.has(pendingKey)) {        return pending;      }      const next = new Map(pending);      next.delete(pendingKey);      return next;    });   const completeExtPendingFailure = (requestId: AcpError.AcpRequestId, error: AcpError.AcpError) =>    resolveExtPending(requestId, ({ deferred }) => Deferred.fail(deferred, error));   const completeExtPendingSuccess = (requestId: AcpError.AcpRequestId, 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>) =>    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();          yield* Deferred.fail(terminationFailure, error);          yield* failAllExtPending(error);          yield* emitClientProtocolError(error);          if (options.onTermination) {            yield* options.onTermination(error);          }        }),        true,      ] as const;    }).pipe(Effect.flatten);   const respondWithSuccess = (requestId: AcpError.AcpRequestId, value: unknown) =>    offerOutgoing({      _tag: "Exit",      requestId,      exit: {        _tag: "Success",        value,      },    });   const respondWithError = (requestId: AcpError.AcpRequestId, 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: options.transformSessionUpdate?.(params) ?? 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,      });    }     const observeIncoming =      options.onIncomingRequest?.(String(message.id), message.tag, message.payload) ?? Effect.void;     if (!options.serverRequestMethods.has(message.tag)) {      return observeIncoming.pipe(Effect.andThen(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 observeIncoming.pipe(Effect.andThen(Queue.offer(serverQueue, message)), Effect.asVoid);  };   const forwardToRpcClient = (    message: RpcMessage.ResponseChunkEncoded | RpcMessage.ResponseExitEncoded,  ) =>    isEffectRpcRequestId(message.requestId)      ? Queue.offer(clientQueue, message).pipe(Effect.asVoid)      : Effect.void;   const handleExitEncoded = (message: RpcMessage.ResponseExitEncoded) =>    Ref.get(extPending).pipe(      Effect.flatMap((pending) => {        const pendingRequest = pending.get(String(message.requestId));        if (!pendingRequest) {          return forwardToRpcClient(message);        }        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(

Lines 1–500 of 863.