client.ts

import {
  EnvironmentAuthorizationError,
  clientRpcRequiredScopes,
  authScopeRequiredResponse,
  type EnvironmentId,
  type ClientGuardedRpcTag,
  ORCHESTRATION_V2_WS_METHODS,
  WS_METHODS,
} from "@t3tools/contracts";
import * as Cause from "effect/Cause";
import * as Context from "effect/Context";
import * as Duration from "effect/Duration";
import * as Effect from "effect/Effect";
import * as Option from "effect/Option";
import * as Schema from "effect/Schema";
import * as Stream from "effect/Stream";
import * as SubscriptionRef from "effect/SubscriptionRef";
import { RpcClientError } from "effect/rpc";

import * as EnvironmentSupervisor from "../connection/supervisor.ts";
import type { WsRpcProtocolClient } from "../rpc/protocol.ts";
import type { RpcSession } from "../rpc/session.ts";

export class EnvironmentRpcUnavailableError extends Schema.TaggedError<EnvironmentRpcUnavailableError>()(
  "EnvironmentRpcUnavailableError",
  {
    environmentId: Schema.String,
    message: Schema.String,
  },
) {}

export interface EnvironmentRpcRequestObservation {
  readonly environmentId: string;
  readonly environmentLabel: string;
  readonly method: string;
}

export class EnvironmentRpcRequestObserver extends Context.Reference<{
  readonly observe: (
    request: EnvironmentRpcRequestObservation,
  ) => Effect.Effect<Effect.Effect<void>>;
}>("@t3tools/client-runtime/rpc/EnvironmentRpcRequestObserver", {
  defaultValue: () => ({
    observe: () => Effect.succeed(Effect.void),
  }),
}) {}

export type EnvironmentRpcTag = keyof WsRpcProtocolClient & string;
type RpcMethod<TTag extends EnvironmentRpcTag> = WsRpcProtocolClient[TTag];

export type EnvironmentSubscriptionRpcTag =
  | typeof WS_METHODS.providerAuthSubscribe
  | typeof WS_METHODS.providerInstallSubscribe
  | typeof ORCHESTRATION_V2_WS_METHODS.subscribeShell
  | typeof ORCHESTRATION_V2_WS_METHODS.subscribeThread
  | typeof ORCHESTRATION_V2_WS_METHODS.subscribeCommandOutput
  | typeof ORCHESTRATION_V2_WS_METHODS.subscribePromptSuggestion
  | typeof ORCHESTRATION_V2_WS_METHODS.searchThreadStream
  | typeof WS_METHODS.subscribeAuthAccess
  | typeof WS_METHODS.subscribeServerConfig
  | typeof WS_METHODS.subscribeServerLifecycle
  | typeof WS_METHODS.scheduledTasksSubscribe
  | typeof WS_METHODS.serverGetStorageCleanupReport
  | typeof WS_METHODS.serverGetDatabaseStorage
  | typeof WS_METHODS.subscribeTerminalMetadata
  | typeof WS_METHODS.subscribePreviewEvents
  | typeof WS_METHODS.subscribeDiscoveredLocalServers
  | typeof WS_METHODS.subscribeDeviceState
  | typeof WS_METHODS.subscribeResourceTelemetry
  | typeof WS_METHODS.pullRequestsSubscribeRefreshes
  | typeof WS_METHODS.pullRequestsSubscribeLinkedThreadRefreshes
  | typeof WS_METHODS.subscribeVcsStatus
  | typeof WS_METHODS.subscribeVcsBranch
  | typeof WS_METHODS.subscribeWorktreeSetup
  | typeof WS_METHODS.subscribeProjectClones
  | typeof WS_METHODS.terminalAttach
  | typeof WS_METHODS.terminalObserve
  // Fork: kitshn client state.
  | typeof WS_METHODS.machineMapSubscribe
  // Fork: preview local server groups (forkPreview.ts).
  | typeof WS_METHODS.forkPreviewSubscribeLocalListeners
  // Fork: desktop "Install when idle" waits for terminal thread programs (forkTerminalThreads.ts).
  | typeof WS_METHODS.forkUpdatesTerminalHold;

export type EnvironmentStreamCommandRpcTag =
  | typeof WS_METHODS.chatGptHandoffSubscribe
  | typeof WS_METHODS.cloudInstallRelayClient
  | typeof WS_METHODS.serverUpdateServerWithProgress
  | typeof WS_METHODS.gitRunStackedAction;

export type EnvironmentStreamRpcTag =
  | EnvironmentSubscriptionRpcTag
  | EnvironmentStreamCommandRpcTag;

export type EnvironmentUnaryRpcTag = Exclude<EnvironmentRpcTag, EnvironmentStreamRpcTag>;

export interface EnvironmentRpcSubscriptionObservation {
  readonly environmentId: string;
  readonly method: EnvironmentSubscriptionRpcTag;
  readonly input: unknown;
}

export class EnvironmentRpcSubscriptionObserver extends Context.Reference<{
  readonly observe: (
    subscription: EnvironmentRpcSubscriptionObservation,
  ) => Effect.Effect<Effect.Effect<void>>;
}>("@t3tools/client-runtime/rpc/EnvironmentRpcSubscriptionObserver", {
  defaultValue: () => ({
    observe: () => Effect.succeed(Effect.void),
  }),
}) {}

export const isRpcClientError = Schema.is(RpcClientError.RpcClientError);
const isEnvironmentAuthorizationError = Schema.is(EnvironmentAuthorizationError);

/** A subscription never retries an authorization failure on its own. */
export const isAuthorizationFailure = (cause: Cause.Cause<unknown>) =>
  cause.reasons.some(
    (reason) => reason._tag === "Fail" && isEnvironmentAuthorizationError(reason.error),
  );

const MAX_EXPECTED_FAILURE_RETRY_DELAY_MS = 30_000;
const STREAM_FAILURE_RETRY_DELAY = Duration.seconds(1);

export const subscriptionRetryDelay = (base: Duration.Input, retries: number) =>
  Duration.millis(
    Math.min(Duration.toMillis(base) * 2 ** retries, MAX_EXPECTED_FAILURE_RETRY_DELAY_MS),
  );

export type EnvironmentRpcInput<TTag extends EnvironmentRpcTag> = Parameters<RpcMethod<TTag>>[0];

export type EnvironmentRpcSuccess<TTag extends EnvironmentUnaryRpcTag> =
  RpcMethod<TTag> extends (input: any, options?: any) => Effect.Effect<infer A, any, any>
    ? A
    : never;

export type EnvironmentRpcFailure<TTag extends EnvironmentUnaryRpcTag> =
  RpcMethod<TTag> extends (input: any, options?: any) => Effect.Effect<any, infer E, any>
    ? E
    : never;

export type EnvironmentRpcStreamValue<TTag extends EnvironmentStreamRpcTag> =
  RpcMethod<TTag> extends (input: any, options?: any) => Stream.Stream<infer A, any, any>
    ? A
    : never;

export type EnvironmentRpcStreamFailure<TTag extends EnvironmentStreamRpcTag> =
  RpcMethod<TTag> extends (input: any, options?: any) => Stream.Stream<any, infer E, any>
    ? E
    : never;

/**
 * Fork: how long a request waits for a session that is opening. After the app returns, the phase
 * can still say "connected" while the old socket is gone and a new one opens. A request in that
 * moment, such as a home screen shortcut, waits for it instead of failing with "not connected".
 */
const OPENING_SESSION_WAIT = Duration.seconds(10);

export const currentSession = Effect.fn("EnvironmentRpc.currentSession")(function* () {
  const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
  const live = yield* SubscriptionRef.get(supervisor.session);
  if (Option.isSome(live)) return live.value;
  // The wait ends with the first session, or when the phase turns offline, available (not
  // wanted) or blocked: then no session is about to open, so the request fails as it did before.
  // A backoff keeps the wait: after the app returns, the old socket fails once and the next
  // attempt opens a session within a second.
  const opened = yield* Stream.merge(
    SubscriptionRef.changes(supervisor.session).pipe(Stream.filter(Option.isSome)),
    SubscriptionRef.changes(supervisor.state).pipe(
      Stream.filter(
        ({ phase }) => phase === "offline" || phase === "available" || phase === "blocked",
      ),
      Stream.as(Option.none<RpcSession>()),
    ),
  ).pipe(
    Stream.runHead,
    Effect.timeoutOption(OPENING_SESSION_WAIT),
    Effect.map((result) => Option.flatten(Option.flatten(result))),
  );
  if (Option.isSome(opened)) return opened.value;
  return yield* new EnvironmentRpcUnavailableError({
    environmentId: supervisor.target.environmentId,
    message: `${supervisor.target.label} is not connected.`,
  });
});

export const getInitialServerConfig = Effect.fn("EnvironmentRpc.getInitialServerConfig")(
  function* () {
    const session = yield* currentSession();
    return yield* session.initialConfig;
  },
);

/** Installed by the shared command boundary. A missing boundary denies protected writes. */
export class RpcPermissionGuard extends Context.Reference<{
  readonly authorize: (
    environmentId: EnvironmentId,
    method: string,
    input: unknown,
  ) => Effect.Effect<void, EnvironmentAuthorizationError>;
}>("@t3tools/client-runtime/rpc/RpcPermissionGuard", {
  defaultValue: () => ({
    authorize: (_environmentId, method, input) => {
      const scope = clientRpcRequiredScopes(method, input)[0];
      return scope === undefined
        ? Effect.void
        : Effect.fail(
            new EnvironmentAuthorizationError({
              ...authScopeRequiredResponse(scope),
              message: `This connection requires ${scope}.`,
            }),
          );
    },
  }),
}) {}

const authorizeRequest = Effect.fn("EnvironmentRpc.authorize")(function* (
  method: string,
  input: unknown,
) {
  if (clientRpcRequiredScopes(method, input).length === 0) return;
  const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
  const guard = yield* RpcPermissionGuard;
  yield* guard.authorize(supervisor.target.environmentId, method, input);
});

export const requestInSession = Effect.fn("EnvironmentRpc.request")(function* <
  TTag extends EnvironmentUnaryRpcTag,
>(session: RpcSession, tag: TTag, input: EnvironmentRpcInput<TTag>) {
  const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
  yield* Effect.annotateCurrentSpan({
    "environment.id": supervisor.target.environmentId,
    "rpc.method": tag,
  });
  yield* authorizeRequest(tag, input);
  const observer = yield* EnvironmentRpcRequestObserver;
  const method = session.client[tag] as (
    input: EnvironmentRpcInput<TTag>,
  ) => Effect.Effect<EnvironmentRpcSuccess<TTag>, EnvironmentRpcFailure<TTag>>;
  const completeObservation = yield* observer.observe({
    environmentId: supervisor.target.environmentId,
    environmentLabel: supervisor.target.label,
    method: tag,
  });
  return yield* method(input).pipe(Effect.ensuring(completeObservation));
});

// The session-bound request owns the existing EnvironmentRpc.request span.
export const requestGuarded = Effect.fnUntraced(function* <TTag extends EnvironmentUnaryRpcTag>(
  tag: TTag,
  input: EnvironmentRpcInput<TTag>,
) {
  return yield* requestInSession(yield* currentSession(), tag, input);
});

export function runStreamGuarded<TTag extends EnvironmentStreamCommandRpcTag>(
  tag: TTag,
  input: EnvironmentRpcInput<TTag>,
): Stream.Stream<
  EnvironmentRpcStreamValue<TTag>,
  | EnvironmentRpcStreamFailure<TTag>
  | EnvironmentRpcUnavailableError
  | EnvironmentAuthorizationError,
  EnvironmentSupervisor.EnvironmentSupervisor
> {
  return Stream.unwrap(
    currentSession().pipe(Effect.map((session) => runStreamInSession(session, tag, input))),
  );
}

export function runStreamInSession<TTag extends EnvironmentStreamCommandRpcTag>(
  session: RpcSession,
  tag: TTag,
  input: EnvironmentRpcInput<TTag>,
): Stream.Stream<
  EnvironmentRpcStreamValue<TTag>,
  EnvironmentRpcStreamFailure<TTag> | EnvironmentAuthorizationError,
  EnvironmentSupervisor.EnvironmentSupervisor
> {
  const method = session.client[tag] as (
    input: EnvironmentRpcInput<TTag>,
  ) => Stream.Stream<EnvironmentRpcStreamValue<TTag>, EnvironmentRpcStreamFailure<TTag>>;
  return Stream.unwrap(authorizeRequest(tag, input).pipe(Effect.map(() => method(input)))).pipe(
    Stream.withSpan("EnvironmentRpc.runStream", {
      attributes: { "rpc.method": tag },
    }),
  );
}

interface SubscriptionOptions<TTag extends EnvironmentSubscriptionRpcTag> {
  /** Reports defects; the subscription then waits for the next session or resubscribe. */
  readonly onDefect?: (
    cause: Cause.Cause<EnvironmentRpcStreamFailure<TTag>>,
  ) => Effect.Effect<void, never, never>;
  readonly onExpectedFailure?: (
    cause: Cause.Cause<EnvironmentRpcStreamFailure<TTag>>,
  ) => Effect.Effect<void, never, never>;
  readonly retryExpectedFailureAfter?: Duration.Input;
  readonly resubscribe?: Stream.Stream<unknown, never, never>;
}

function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
  tag: TTag,
  makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
  mapStream: (
    session: RpcSession,
    stream: Stream.Stream<EnvironmentRpcStreamValue<TTag>, EnvironmentRpcStreamFailure<TTag>>,
  ) => Stream.Stream<A, EnvironmentRpcStreamFailure<TTag>>,
  options?: SubscriptionOptions<TTag>,
): Stream.Stream<
  A,
  EnvironmentRpcStreamFailure<TTag>,
  EnvironmentSupervisor.EnvironmentSupervisor
> {
  return Stream.unwrap(
    Effect.gen(function* () {
      const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
      const observer = yield* EnvironmentRpcSubscriptionObserver;
      const sessionChanges = SubscriptionRef.changes(supervisor.session);
      const sessions =
        options?.resubscribe === undefined
          ? sessionChanges
          : Stream.merge(
              sessionChanges,
              options.resubscribe.pipe(
                Stream.mapEffect(() => SubscriptionRef.get(supervisor.session)),
              ),
            );
      return sessions.pipe(
        Stream.switchMap(
          Option.match({
            onNone: () => Stream.empty,
            onSome: (session) => {
              const method = (
                tag === WS_METHODS.subscribeServerConfig
                  ? session.subscribeServerConfig
                  : session.client[tag]
              ) as (
                input: EnvironmentRpcInput<TTag>,
              ) => Stream.Stream<
                EnvironmentRpcStreamValue<TTag>,
                EnvironmentRpcStreamFailure<TTag>
              >;
              let expectedFailureRetries = 0;
              // Not reset by a first value: a stream can emit before every chunk failure.
              let streamFailureRetries = 0;
              const subscribeToSession = (): Stream.Stream<A, EnvironmentRpcStreamFailure<TTag>> =>
                Stream.suspend(() =>
                  Stream.unwrap(
                    Effect.gen(function* () {
                      const input = yield* makeInput(session);
                      const completeObservation = yield* observer.observe({
                        environmentId: supervisor.target.environmentId,
                        method: tag,
                        input,
                      });
                      const stream = mapStream(
                        session,
                        Stream.suspend(() => method(input)),
                      ).pipe(
                        Stream.onFirst(() =>
                          Effect.sync(() => {
                            expectedFailureRetries = 0;
                          }),
                        ),
                      );
                      return stream.pipe(Stream.ensuring(completeObservation));
                    }),
                  ).pipe(
                    Stream.tapCause((cause) =>
                      options?.onDefect !== undefined &&
                      cause.reasons.some(
                        (reason) =>
                          reason._tag === "Die" ||
                          (reason._tag === "Fail" &&
                            isRpcClientError(reason.error) &&
                            reason.error.reason._tag === "RpcClientDefect"),
                      )
                        ? options.onDefect(cause)
                        : Effect.void,
                    ),
                    Stream.catchCause((cause) => {
                      const retryAfter = (base: Duration.Input, retries: number) =>
                        Stream.fromEffect(Effect.sleep(subscriptionRetryDelay(base, retries))).pipe(
                          Stream.drain,
                          Stream.concat(subscribeToSession()),
                        );
                      const hasOnlyExpectedFailures =
                        cause.reasons.length > 0 &&
                        cause.reasons.every((reason) => reason._tag === "Fail");
                      const isTransportFailure =
                        hasOnlyExpectedFailures &&
                        cause.reasons.every(
                          (reason) => reason._tag === "Fail" && isRpcClientError(reason.error),
                        );
                      // The RPC client dies with a SchemaError when one stream chunk fails to
                      // decode, ending only that request while the socket stays open. A caller
                      // with `onDefect` owns its defect recovery.
                      const isChunkDecodeDefect =
                        options?.onDefect === undefined &&
                        cause.reasons.length > 0 &&
                        cause.reasons.every(
                          (reason) => reason._tag === "Die" && Schema.isSchemaError(reason.defect),
                        );
                      // Both end one stream while the session can stay healthy, so nothing else
                      // would replace it. A dead session is replaced by the supervisor, which
                      // interrupts this retry through switchMap. A finite search must settle
                      // rather than restart, so it keeps waiting for the next session.
                      const resubscribesOnSameSession =
                        tag !== ORCHESTRATION_V2_WS_METHODS.searchThreadStream;
                      if (isTransportFailure && !resubscribesOnSameSession) {
                        return Stream.fromEffect(
                          Effect.logWarning(
                            "Durable RPC subscription lost its transport; waiting for the next session.",
                            {
                              cause: Cause.pretty(cause),
                              method: tag,
                              environmentId: supervisor.target.environmentId,
                            },
                          ),
                        ).pipe(Stream.drain);
                      }
                      if (
                        resubscribesOnSameSession &&
                        (isTransportFailure || isChunkDecodeDefect)
                      ) {
                        return Stream.fromEffect(
                          Effect.logWarning(
                            "Durable RPC subscription stream ended; resubscribing.",
                            {
                              cause: Cause.pretty(cause),
                              method: tag,
                              environmentId: supervisor.target.environmentId,
                            },
                          ),
                        ).pipe(
                          Stream.drain,
                          Stream.concat(
                            retryAfter(
                              options?.retryExpectedFailureAfter ?? STREAM_FAILURE_RETRY_DELAY,
                              streamFailureRetries++,
                            ),
                          ),
                        );
                      }
                      if (hasOnlyExpectedFailures && options?.onExpectedFailure !== undefined) {
                        const handled = Stream.fromEffect(options.onExpectedFailure(cause)).pipe(
                          Stream.drain,
                        );
                        if (
                          options.retryExpectedFailureAfter === undefined ||
                          isAuthorizationFailure(cause)
                        ) {
                          return handled;
                        }
                        return handled.pipe(
                          Stream.concat(
                            retryAfter(options.retryExpectedFailureAfter, expectedFailureRetries++),
                          ),
                        );
                      }
                      // A reported defect parks until the next session or resubscribe. Failing would
                      // end the whole subscription, so a mounted screen never syncs again.
                      if (options?.onDefect !== undefined && Cause.hasDies(cause)) {
                        return Stream.empty;
                      }
                      return Stream.failCause(cause);
                    }),
                  ),
                );
              return subscribeToSession();
            },
          }),
        ),
      );
    }),
  ).pipe(
    Stream.withSpan("EnvironmentRpc.subscribe", {
      attributes: { "rpc.method": tag },
    }),
  );
}

export function subscribeDynamic<TTag extends EnvironmentSubscriptionRpcTag>(
  tag: TTag,
  makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
  options?: SubscriptionOptions<TTag>,
): Stream.Stream<
  EnvironmentRpcStreamValue<TTag>,
  EnvironmentRpcStreamFailure<TTag>,
  EnvironmentSupervisor.EnvironmentSupervisor
> {
  return subscribeDynamicMapped(tag, makeInput, (_session, stream) => stream, options);
}

/** Tags each value before `switchMap` can buffer it across a session change. */
export function subscribeDynamicWithSession<TTag extends EnvironmentSubscriptionRpcTag>(
  tag: TTag,
  makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
  options?: SubscriptionOptions<TTag>,
): Stream.Stream<
  readonly [session: RpcSession, value: EnvironmentRpcStreamValue<TTag>],
  EnvironmentRpcStreamFailure<TTag>,
  EnvironmentSupervisor.EnvironmentSupervisor
> {
  return subscribeDynamicMapped(
    tag,
    makeInput,
    (session, stream) => stream.pipe(Stream.map((value) => [session, value] as const)),
    options,
  );
}

export function subscribe<TTag extends EnvironmentSubscriptionRpcTag>(
  tag: TTag,
  input: EnvironmentRpcInput<TTag>,
  options?: SubscriptionOptions<TTag>,
): Stream.Stream<
  EnvironmentRpcStreamValue<TTag>,
  EnvironmentRpcStreamFailure<TTag>,
  EnvironmentSupervisor.EnvironmentSupervisor
> {
  return subscribeDynamic(tag, () => Effect.succeed(input), options);
}

export const subscribeVcsBranch = (
  input: EnvironmentRpcInput<typeof WS_METHODS.subscribeVcsBranch>,
) =>
  subscribeDynamicMapped(
    WS_METHODS.subscribeVcsBranch,
    () => Effect.succeed(input),
    (session, stream) =>
      Stream.unwrap(
        session.initialConfig.pipe(
          Effect.mapError(
            (error) =>
              new RpcClientError.RpcClientError({
                reason: new RpcClientError.RpcClientDefect({
                  message: error.message,
                  cause: error,
                }),
              }),
          ),
          Effect.map((config) =>
            config.environment.capabilities.vcsBranchStatus === true
              ? stream
              : session.client[WS_METHODS.subscribeVcsStatus]({
                  ...input,
                  includeRemote: false,
                }).pipe(
                  Stream.filter((event) => event._tag !== "remoteUpdated"),
                  Stream.map(({ local }) => ({ refName: local.refName, isRepo: local.isRepo })),
                ),
          ),
        ),
      ),
  );

export const refreshVcsBranch = Effect.fn("EnvironmentRpc.refreshVcsBranch")(function* (
  input: EnvironmentRpcInput<typeof WS_METHODS.vcsRefreshBranch>,
) {
  const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
  const session = yield* currentSession();
  const config = yield* session.initialConfig.pipe(
    Effect.mapError(
      (error) =>
        new EnvironmentRpcUnavailableError({
          environmentId: supervisor.target.environmentId,
          message: error.message,
        }),
    ),
  );
  if (config.environment.capabilities.vcsBranchStatus === true) {
    return yield* requestInSession(session, WS_METHODS.vcsRefreshBranch, input);
  }
  // Older hosts retain their existing status stream without adding per-row wake scans.
});

export const request = <TTag extends Exclude<EnvironmentUnaryRpcTag, ClientGuardedRpcTag>>(
  tag: TTag,
  input: EnvironmentRpcInput<TTag>,
) => requestGuarded(tag, input);

export const runStream = <
  TTag extends Exclude<EnvironmentStreamCommandRpcTag, ClientGuardedRpcTag>,
>(
  tag: TTag,
  input: EnvironmentRpcInput<TTag>,
) => runStreamGuarded(tag, input);