Shade/packages/shade-files/src/client/client.ts

import type { ShadeBridge } from '../integration/shade-bridge.js';
import {
  KIND_CUSTOM_V1,
  KIND_DELETE_V1,
  KIND_GET_THUMBNAIL_V1,
  KIND_LIST_V1,
  KIND_MKDIR_V1,
  KIND_MOVE_V1,
  KIND_READ_V1,
  KIND_STAT_V1,
  KIND_WRITE_V1,
  MUTATION_OPS,
  type StandardOp,
} from '../protocol/kinds.js';
import type { CustomOpsMap } from '../server/custom-ops.js';
import { generateIdempotencyKey, generateRequestId } from '../protocol/correlate.js';
import {
  base64ToBytes,
  bytesToBase64,
  canonicalRpcBytes,
  hashArgs,
} from '../protocol/canonical.js';
import { ShadeFileRpcChannel } from '../rpc/channel.js';
import { PendingRpcRegistry, type RegisterOptions } from '../rpc/pending.js';
import type { RpcRequest } from '../schemas/envelope.js';
import {
  CustomArgsSchema,
  CustomResultSchema,
  DeleteArgsSchema,
  DeleteResultSchema,
  GetThumbnailArgsSchema,
  GetThumbnailResultSchema,
  ListArgsSchema,
  ListResultSchema,
  MkdirArgsSchema,
  MkdirResultSchema,
  MoveArgsSchema,
  MoveResultSchema,
  ReadArgsSchema,
  ReadResultSchema,
  StatArgsSchema,
  StatResultSchema,
  WriteArgsSchema,
  WriteResultSchema,
  type DeleteArgs,
  type DeleteResult,
  type GetThumbnailArgs,
  type ListArgs,
  type ListResult,
  type MkdirArgs,
  type MkdirResult,
  type MoveArgs,
  type MoveResult,
  type ReadArgs,
  type StatResult,
  type ThumbnailSize,
  type WriteResult,
} from '../schemas/ops.js';
import { ConflictError, InternalFileError } from '../schemas/errors.js';
import { decideInline, INLINE_THRESHOLD, type WriteSource } from './inline-threshold.js';
import type { ClientStreamsBridge } from './streams-bridge.js';

export interface BaseOpts {
  signal?: AbortSignal;
  /** Auto-generated for mutations if not provided. */
  idempotencyKey?: string;
  /** Per-call timeout. Default 30_000 ms. */
  timeoutMs?: number;
}

// ─── read/write public types ─────────────────────────────────

export interface ReadInlineOutput {
  kind: 'inline';
  bytes: Uint8Array;
  size: number;
  sha256: string;
  contentType?: string;
}
export interface ReadStreamsOutput {
  kind: 'streams';
  stream: ReadableStream<Uint8Array>;
  size: number;
  sha256: string;
  contentType?: string;
  /** Resolves once the entire transfer has been received and verified. */
  done(): Promise<void>;
}
export type ReadOutput = ReadInlineOutput | ReadStreamsOutput;

export interface WriteOpts extends BaseOpts {
  contentType?: string;
  overwrite?: boolean;
  /** Force inline even if size > 256 KiB. Throws if input is too big. */
  forceInline?: boolean;
}

export interface ReadOpts extends BaseOpts {
  range?: { start: number; end: number };
  preferInline?: boolean;
}

export interface ThumbnailResult {
  bytes: Uint8Array;
  format: 'png' | 'webp' | 'jpeg';
  width: number;
  height: number;
  sha256: string;
}

// ─── FileClient interface ────────────────────────────────────

/**
 * Untyped fallback for `FileClient.custom()` — used when the consumer
 * hasn't extended `CustomOpsMap` for a given op name.
 */
type CustomOpArgs<K> = K extends keyof CustomOpsMap
  ? CustomOpsMap[K] extends { args: infer A }
    ? A
    : unknown
  : unknown;
type CustomOpResponse<K> = K extends keyof CustomOpsMap
  ? CustomOpsMap[K] extends { response: infer R }
    ? R
    : unknown
  : unknown;

export interface FileClient {
  list(path: string, opts?: BaseOpts & Partial<Omit<ListArgs, 'path'>>): Promise<ListResult>;
  stat(path: string, opts?: BaseOpts): Promise<StatResult>;
  mkdir(path: string, opts?: BaseOpts & Partial<Omit<MkdirArgs, 'path'>>): Promise<MkdirResult>;
  delete(path: string, opts?: BaseOpts & Partial<Omit<DeleteArgs, 'path'>>): Promise<DeleteResult>;
  move(src: string, dst: string, opts?: BaseOpts & Partial<Omit<MoveArgs, 'src' | 'dst'>>): Promise<MoveResult>;
  read(path: string, opts?: ReadOpts): Promise<ReadOutput>;
  write(path: string, input: WriteSource, opts?: WriteOpts): Promise<WriteResult>;
  getThumbnail(
    path: string,
    size: ThumbnailSize,
    opts?: BaseOpts & { format?: 'png' | 'webp' | 'jpeg' },
  ): Promise<ThumbnailResult>;
  /**
   * Invoke a custom op registered on the server. Args/response types are
   * pulled from `CustomOpsMap` via TypeScript declaration merging — see
   * `server/custom-ops.ts` for the registration pattern.
   */
  custom<K extends keyof CustomOpsMap | string>(
    name: K & string,
    args: CustomOpArgs<K>,
    opts?: BaseOpts,
  ): Promise<CustomOpResponse<K>>;
  close(): void;
}

export interface CreateFileClientOptions {
  /** Default per-call timeout. Default 30_000. */
  defaultTimeoutMs?: number;
  /** Hard deadline for incoming-read awaits. Default 60_000. */
  ioTimeoutMs?: number;
  /**
   * Required for read/write `streams` ops. Coordinates inbound/outbound
   * `@shade/transfer` transfers via `userMetadata.shadeFiles*Id` keys.
   */
  streamsBridge?: ClientStreamsBridge;
  /**
   * Optional: sign the canonical bytes of every outgoing RPC envelope.
   * Pluggable so apps can plug their own signing-key store (e.g.,
   * Ed25519-as-a-service, browser SubtleCrypto). When omitted, ships
   * `'unsigned'` — the server's `verifySender` should also be unset, or
   * be configured to accept the placeholder.
   */
  signRequest?: (canonicalBytes: Uint8Array) => Promise<string> | string;
}

/**
 * Client-side proxy for `@shade/files` ops. Each method ships an
 * `RpcRequest` over `Shade.send`/`Shade.receive` and awaits the matching
 * response (or error/timeout) from `PendingRpcRegistry`.
 *
 * Mutations auto-generate an idempotency key per logical call (not per
 * attempt) so transparent retries under the SDK don't produce duplicates.
 *
 * Read/write content I/O over 256 KiB requires a `streamsBridge` to be
 * passed via options — it coordinates the inbound/outbound `@shade/transfer`
 * transfers that carry the actual bytes.
 */
export function createFileClient(
  shade: ShadeBridge,
  channel: ShadeFileRpcChannel,
  pending: PendingRpcRegistry,
  peerAddress: string,
  options: CreateFileClientOptions = {},
): FileClient {
  const defaultTimeout = options.defaultTimeoutMs ?? 30_000;
  const ioTimeoutMs = options.ioTimeoutMs ?? 60_000;
  const streamsBridge = options.streamsBridge;
  const signRequest = options.signRequest;
  const senderAddress = shade.myAddress;

  async function request<TResult>(
    kind: string,
    op: StandardOp | 'custom',
    args: unknown,
    opts: BaseOpts | undefined,
  ): Promise<TResult> {
    const requestId = generateRequestId();
    const isMutation = MUTATION_OPS.has(op);
    const idempotencyKey =
      opts?.idempotencyKey ?? (isMutation ? generateIdempotencyKey() : undefined);
    const signedAt = Date.now();
    let sig = 'unsigned';
    if (signRequest !== undefined) {
      // Server reconstructs canonical bytes using `address = from`, which
      // is OUR own address as authenticated by the ratchet. So we sign
      // over the same identifier here.
      const canonical = canonicalRpcBytes({
        address: senderAddress,
        signedAt,
        kind,
        id: requestId,
        argsHash: hashArgs(args),
      });
      sig = await signRequest(canonical);
    }
    const env: RpcRequest = {
      kind,
      id: requestId,
      args,
      ...(idempotencyKey !== undefined ? { idempotencyKey } : {}),
      sig,
      signedAt,
    };

    const registerOpts: RegisterOptions = {
      timeoutMs: opts?.timeoutMs ?? defaultTimeout,
      onCancel: (reason) => {
        // Fire-and-forget cancel envelope so server can release resources.
        void channel
          .send(peerAddress, {
            kind: 'shade.fs.cancel/v1',
            id: requestId,
            reason,
          })
          .catch(() => {
            /* swallow — cancellation is best-effort */
          });
      },
    };
    if (opts?.signal !== undefined) registerOpts.signal = opts.signal;

    const pendingPromise = pending.register<TResult>(requestId, registerOpts);
    try {
      await channel.send(peerAddress, env);
    } catch (err) {
      // If the send itself fails, the pending entry will never resolve;
      // reject it directly.
      pending.rejectAll(err);
      throw err;
    }
    return pendingPromise;
  }

  return {
    async list(path, opts) {
      const args: ListArgs = ListArgsSchema.parse({
        path,
        ...(opts?.cursor !== undefined ? { cursor: opts.cursor } : {}),
        ...(opts?.pageSize !== undefined ? { pageSize: opts.pageSize } : {}),
        ...(opts?.filter !== undefined ? { filter: opts.filter } : {}),
      });
      const raw = await request<unknown>(KIND_LIST_V1, 'list', args, opts);
      return ListResultSchema.parse(raw);
    },

    async stat(path, opts) {
      const args = StatArgsSchema.parse({ path });
      const raw = await request<unknown>(KIND_STAT_V1, 'stat', args, opts);
      return StatResultSchema.parse(raw);
    },

    async mkdir(path, opts) {
      const args = MkdirArgsSchema.parse({
        path,
        ...(opts?.recursive !== undefined ? { recursive: opts.recursive } : {}),
      });
      const raw = await request<unknown>(KIND_MKDIR_V1, 'mkdir', args, opts);
      return MkdirResultSchema.parse(raw);
    },

    async delete(path, opts) {
      const args = DeleteArgsSchema.parse({
        path,
        ...(opts?.recursive !== undefined ? { recursive: opts.recursive } : {}),
      });
      const raw = await request<unknown>(KIND_DELETE_V1, 'delete', args, opts);
      return DeleteResultSchema.parse(raw);
    },

    async move(src, dst, opts) {
      const args = MoveArgsSchema.parse({
        src,
        dst,
        ...(opts?.overwrite !== undefined ? { overwrite: opts.overwrite } : {}),
      });
      const raw = await request<unknown>(KIND_MOVE_V1, 'move', args, opts);
      return MoveResultSchema.parse(raw);
    },

    async read(path, opts) {
      const args: ReadArgs = ReadArgsSchema.parse({
        path,
        ...(opts?.range !== undefined ? { range: opts.range } : {}),
        ...(opts?.preferInline !== undefined ? { preferInline: opts.preferInline } : {}),
      });
      const raw = await request<unknown>(KIND_READ_V1, 'read', args, opts);
      const wire = ReadResultSchema.parse(raw);
      if (wire.kind === 'inline') {
        const bytes = base64ToBytes(wire.bytesB64);
        const out: ReadInlineOutput = {
          kind: 'inline',
          bytes,
          size: wire.size,
          sha256: wire.sha256,
          ...(wire.contentType !== undefined ? { contentType: wire.contentType } : {}),
        };
        return out;
      }
      // streams — wait for the matching incoming transfer via the bridge.
      if (streamsBridge === undefined) {
        throw new InternalFileError(
          'streams-bridge not configured: cannot consume streamed read',
        );
      }
      const bridgeSignal = opts?.signal ?? new AbortController().signal;
      const parked = await streamsBridge.awaitRead(wire.streamId, {
        expectedFrom: peerAddress,
        signal: bridgeSignal,
        timeoutMs: ioTimeoutMs,
      });
      const out: ReadStreamsOutput = {
        kind: 'streams',
        stream: parked.readable,
        size: wire.size,
        sha256: wire.sha256,
        ...(wire.contentType !== undefined ? { contentType: wire.contentType } : {}),
        done: async () => {
          await parked.done;
        },
      };
      return out;
    },

    async write(path, input, opts) {
      const decision = await decideInline(input);
      const overwrite = opts?.overwrite ?? false;
      const contentType = opts?.contentType ?? decision.contentType;

      if (decision.kind === 'inline' || opts?.forceInline === true) {
        // Inline path — base64 in the RPC envelope.
        const bytes =
          decision.kind === 'inline'
            ? decision.bytes
            : await drainToUint8Array(decision.stream, decision.size ?? Number.POSITIVE_INFINITY);
        if (bytes.byteLength > INLINE_THRESHOLD && opts?.forceInline !== true) {
          throw new ConflictError(
            `inline write exceeds ${INLINE_THRESHOLD}-byte threshold (got ${bytes.byteLength}); pass forceInline=true to override`,
          );
        }
        const args = WriteArgsSchema.parse({
          kind: 'inline',
          path,
          bytesB64: bytesToBase64(bytes),
          ...(contentType !== undefined ? { contentType } : {}),
          overwrite,
        });
        const raw = await request<unknown>(KIND_WRITE_V1, 'write', args, opts);
        return WriteResultSchema.parse(raw);
      }

      // Streams path — kick the upload, then ship the RPC.
      if (streamsBridge === undefined) {
        throw new InternalFileError(
          'streams-bridge not configured: cannot ship streamed write',
        );
      }
      const size = decision.size;
      if (size === undefined) {
        throw new ConflictError(
          'streams write requires a known plaintext size; pass `{ stream, size }` instead of a bare ReadableStream',
        );
      }
      const { writeId, handle } = await streamsBridge.initiateWrite({
        peer: peerAddress,
        stream: decision.stream,
        size,
        ...(contentType !== undefined ? { contentType } : {}),
        name: path,
        ...(opts?.signal !== undefined ? { signal: opts.signal } : {}),
      });
      const args = WriteArgsSchema.parse({
        kind: 'streams',
        path,
        size,
        ...(contentType !== undefined ? { contentType } : {}),
        overwrite,
        writeId,
      });
      try {
        const [raw] = await Promise.all([
          request<unknown>(KIND_WRITE_V1, 'write', args, opts),
          handle.done(),
        ]);
        return WriteResultSchema.parse(raw);
      } catch (err) {
        // Best-effort cancel of the transfer on RPC failure.
        await handle.abort('rpc-failed').catch(() => undefined);
        throw err;
      }
    },

    async getThumbnail(path, size, opts) {
      const args: GetThumbnailArgs = GetThumbnailArgsSchema.parse({
        path,
        size,
        ...(opts?.format !== undefined ? { format: opts.format } : {}),
      });
      const raw = await request<unknown>(KIND_GET_THUMBNAIL_V1, 'getThumbnail', args, opts);
      const wire = GetThumbnailResultSchema.parse(raw);
      return {
        bytes: base64ToBytes(wire.bytesB64),
        format: wire.format,
        width: wire.width,
        height: wire.height,
        sha256: wire.sha256,
      };
    },

    async custom(name, args, opts) {
      const wireArgs = CustomArgsSchema.parse({ name, payload: args });
      // `custom` is a mutation in the rate-limit sense; auto-key for retries.
      const raw = await request<unknown>(KIND_CUSTOM_V1, 'custom' as StandardOp, wireArgs, opts);
      const wire = CustomResultSchema.parse(raw);
      // The result schema is `{ result: unknown }` — the inner `result` is
      // already validated against the consumer's response schema on the
      // server side, so we trust it here.
      return wire.result as never;
    },

    close() {
      pending.rejectAll(new Error('FileClient closed'));
    },
  };
}

/** Drain a stream into a single buffer; used for the inline-write fallback. */
async function drainToUint8Array(
  stream: ReadableStream<Uint8Array>,
  cap: number,
): Promise<Uint8Array> {
  const reader = stream.getReader();
  const chunks: Uint8Array[] = [];
  let total = 0;
  try {
    while (true) {
      const { value, done } = await reader.read();
      if (done) break;
      if (value === undefined) continue;
      chunks.push(value);
      total += value.byteLength;
      if (total > cap) {
        throw new Error(`stream produced more than declared size cap (${cap})`);
      }
    }
  } finally {
    reader.releaseLock();
  }
  const out = new Uint8Array(total);
  let offset = 0;
  for (const c of chunks) {
    out.set(c, offset);
    offset += c.byteLength;
  }
  return out;
}