packages/kap-server/src/protocol/ws-control.ts

import { z } from 'zod';

import { isoDateTimeSchema } from '@moonshot-ai/agent-core-v2/_base/utils/isoDateTime';
import { transcriptGradeSpecSchema, transcriptSeqSchema } from '@moonshot-ai/transcript';

import { eventSchema } from './events-zod';

export const WS_PROTOCOL_VERSION = 2;

export const sessionCursorSchema = z.object({
  seq: z.number().int().nonnegative(),
  epoch: z.string().min(1).optional(),
});

export type SessionCursor = z.infer<typeof sessionCursorSchema>;

export const cursorsBySessionSchema = z.record(z.string(), sessionCursorSchema);

export type CursorsBySession = z.infer<typeof cursorsBySessionSchema>;

export const wsEventEnvelopeSchema = <T extends z.ZodTypeAny>(payload: T) =>
  z.object({
    type: z.string(),
    seq: z.number().int().nonnegative(),
    epoch: z.string().optional(),
    volatile: z.boolean().optional(),
    offset: z.number().int().nonnegative().optional(),
    session_id: z.string().optional(),
    timestamp: isoDateTimeSchema,
    payload,
  });

export const wsControlEnvelopeSchema = <T extends z.ZodTypeAny>(payload: T) =>
  z.object({
    type: z.string(),
    id: z.string().optional(),
    payload,
  });

export const wsAckEnvelopeSchema = <T extends z.ZodTypeAny>(payload: T) =>
  z.object({
    type: z.literal('ack'),
    id: z.string(),
    code: z.number().int(),
    msg: z.string(),
    payload,
  });

export const serverHelloPayloadSchema = z.object({
  ws_connection_id: z.string(),
  protocol_version: z.number().int().positive(),
  heartbeat_ms: z.number().int().positive().optional(),
  max_event_buffer_size: z.number().int().positive(),
  capabilities: z.object({
    event_batching: z.boolean(),
    compression: z.boolean(),
  }),
});

export const serverHelloMessageSchema = z.object({
  type: z.literal('server_hello'),
  timestamp: isoDateTimeSchema,
  payload: serverHelloPayloadSchema,
});

export type ServerHelloMessage = z.infer<typeof serverHelloMessageSchema>;

export const agentFilterSchema = z.record(z.string(), z.array(z.string()).min(1));

export type AgentFilter = z.infer<typeof agentFilterSchema>;

export const clientHelloPayloadSchema = z.object({
  client_id: z.string(),
  subscriptions: z.array(z.string()).optional(),
  cursors: cursorsBySessionSchema.optional(),
  agent_filter: agentFilterSchema.optional(),
});

export const clientHelloMessageSchema = z.object({
  type: z.literal('client_hello'),
  id: z.string(),
  payload: clientHelloPayloadSchema,
});

export type ClientHelloMessage = z.infer<typeof clientHelloMessageSchema>;

export const clientHelloAckPayloadSchema = z.object({
  accepted_subscriptions: z.array(z.string()),
  resync_required: z.array(z.string()),
  cursors: cursorsBySessionSchema.optional(),
});

export const helloAckPayloadSchema = clientHelloAckPayloadSchema;

export const clientHelloAckMessageSchema = wsAckEnvelopeSchema(clientHelloAckPayloadSchema);

export const subscribePayloadSchema = z.object({
  session_ids: z.array(z.string()),
  cursors: cursorsBySessionSchema.optional(),
  agent_filter: agentFilterSchema.optional(),
});

export const subscribeMessageSchema = z.object({
  type: z.literal('subscribe'),
  id: z.string(),
  payload: subscribePayloadSchema,
});

export type SubscribeMessage = z.infer<typeof subscribeMessageSchema>;

export const subscribeV2PayloadSchema = z.object({
  session_id: z.string().min(1),
  transcript: transcriptGradeSpecSchema,
  transcript_since: z.record(z.string(), transcriptSeqSchema).optional(),
});

export const subscribeV2MessageSchema = z.object({
  type: z.literal('subscribe_v2'),
  id: z.string(),
  payload: subscribeV2PayloadSchema,
});

export type SubscribeV2Message = z.infer<typeof subscribeV2MessageSchema>;

export const unsubscribeV2PayloadSchema = z.object({
  session_id: z.string().min(1),
  agent_ids: z.array(z.string().min(1)).min(1).optional(),
});

export const unsubscribeV2MessageSchema = z.object({
  type: z.literal('unsubscribe_v2'),
  id: z.string(),
  payload: unsubscribeV2PayloadSchema,
});

export type UnsubscribeV2Message = z.infer<typeof unsubscribeV2MessageSchema>;

export const subscribeAckPayloadSchema = z.object({
  accepted: z.array(z.string()),
  not_found: z.array(z.string()),
  resync_required: z.array(z.string()),
  cursors: cursorsBySessionSchema.optional(),
});

export const subscribeAckMessageSchema = wsAckEnvelopeSchema(subscribeAckPayloadSchema);

export const subscribeV2AckMessageSchema = wsAckEnvelopeSchema(subscribeAckPayloadSchema);

export const unsubscribeV2AckMessageSchema = wsAckEnvelopeSchema(subscribeAckPayloadSchema);

export const unsubscribePayloadSchema = z.object({
  session_ids: z.array(z.string()),
});

export const unsubscribeMessageSchema = z.object({
  type: z.literal('unsubscribe'),
  id: z.string(),
  payload: unsubscribePayloadSchema,
});

export type UnsubscribeMessage = z.infer<typeof unsubscribeMessageSchema>;

export const unsubscribeAckPayloadSchema = subscribeAckPayloadSchema;

export const unsubscribeAckMessageSchema = wsAckEnvelopeSchema(unsubscribeAckPayloadSchema);

export const abortPayloadSchema = z.object({
  session_id: z.string(),
  prompt_id: z.string(),
});

export const abortMessageSchema = z.object({
  type: z.literal('abort'),
  id: z.string(),
  payload: abortPayloadSchema,
});

export type AbortMessage = z.infer<typeof abortMessageSchema>;

export const abortAckPayloadSchema = z.object({
  aborted: z.boolean().optional(),
  at_seq: z.number().int().nonnegative().optional(),
});

export const abortAckMessageSchema = wsAckEnvelopeSchema(abortAckPayloadSchema);

export const terminalAttachPayloadSchema = z.object({
  session_id: z.string().min(1),
  terminal_id: z.string().min(1),
  since_seq: z.number().int().nonnegative().optional(),
});

export const terminalAttachMessageSchema = z.object({
  type: z.literal('terminal_attach'),
  id: z.string(),
  payload: terminalAttachPayloadSchema,
});

export type TerminalAttachMessage = z.infer<typeof terminalAttachMessageSchema>;

export const terminalAttachAckPayloadSchema = z.object({
  attached: z.literal(true),
  replayed: z.number().int().nonnegative(),
});

export const terminalAttachAckMessageSchema = wsAckEnvelopeSchema(
  terminalAttachAckPayloadSchema,
);

export const terminalDetachPayloadSchema = z.object({
  session_id: z.string().min(1),
  terminal_id: z.string().min(1),
});

export const terminalDetachMessageSchema = z.object({
  type: z.literal('terminal_detach'),
  id: z.string(),
  payload: terminalDetachPayloadSchema,
});

export type TerminalDetachMessage = z.infer<typeof terminalDetachMessageSchema>;

export const terminalDetachAckPayloadSchema = z.object({
  detached: z.literal(true),
});

export const terminalDetachAckMessageSchema = wsAckEnvelopeSchema(
  terminalDetachAckPayloadSchema,
);

export const terminalInputPayloadSchema = z.object({
  session_id: z.string().min(1),
  terminal_id: z.string().min(1),
  data: z.string(),
});

export const terminalInputMessageSchema = z.object({
  type: z.literal('terminal_input'),
  id: z.string(),
  payload: terminalInputPayloadSchema,
});

export type TerminalInputMessage = z.infer<typeof terminalInputMessageSchema>;

export const terminalInputAckPayloadSchema = z.object({
  accepted: z.literal(true),
});

export const terminalInputAckMessageSchema = wsAckEnvelopeSchema(
  terminalInputAckPayloadSchema,
);

export const terminalResizePayloadSchema = z.object({
  session_id: z.string().min(1),
  terminal_id: z.string().min(1),
  cols: z.number().int().positive(),
  rows: z.number().int().positive(),
});

export const terminalResizeMessageSchema = z.object({
  type: z.literal('terminal_resize'),
  id: z.string(),
  payload: terminalResizePayloadSchema,
});

export type TerminalResizeMessage = z.infer<typeof terminalResizeMessageSchema>;

export const terminalResizeAckPayloadSchema = z.object({
  resized: z.literal(true),
});

export const terminalResizeAckMessageSchema = wsAckEnvelopeSchema(
  terminalResizeAckPayloadSchema,
);

export const terminalClosePayloadSchema = z.object({
  session_id: z.string().min(1),
  terminal_id: z.string().min(1),
});

export const terminalCloseMessageSchema = z.object({
  type: z.literal('terminal_close'),
  id: z.string(),
  payload: terminalClosePayloadSchema,
});

export type TerminalCloseMessage = z.infer<typeof terminalCloseMessageSchema>;

export const terminalCloseAckPayloadSchema = z.object({
  closed: z.literal(true),
});

export const terminalCloseAckMessageSchema = wsAckEnvelopeSchema(
  terminalCloseAckPayloadSchema,
);

export const pingPayloadSchema = z.object({
  nonce: z.string(),
});

export const pingMessageSchema = z.object({
  type: z.literal('ping'),
  timestamp: isoDateTimeSchema,
  payload: pingPayloadSchema,
});

export type PingMessage = z.infer<typeof pingMessageSchema>;

export const pongPayloadSchema = z.object({
  nonce: z.string(),
});

export const pongMessageSchema = z.object({
  type: z.literal('pong'),
  payload: pongPayloadSchema,
});

export type PongMessage = z.infer<typeof pongMessageSchema>;

export const resyncRequiredPayloadSchema = z.object({
  session_id: z.string(),
  reason: z.enum(['buffer_overflow', 'session_recreated', 'epoch_changed']),
  current_seq: z.number().int().nonnegative(),
  epoch: z.string().min(1).optional(),
});

export const resyncRequiredMessageSchema = z.object({
  type: z.literal('resync_required'),
  timestamp: isoDateTimeSchema,
  payload: resyncRequiredPayloadSchema,
});

export type ResyncRequiredMessage = z.infer<typeof resyncRequiredMessageSchema>;

export const wsErrorPayloadSchema = z.object({
  code: z.number().int(),
  msg: z.string(),
  fatal: z.boolean(),
  request_id: z.string().optional(),
  details: z.unknown().optional(),
});

export const wsErrorMessageSchema = z.object({
  type: z.literal('error'),
  timestamp: isoDateTimeSchema,
  payload: wsErrorPayloadSchema,
});

export type WsErrorMessage = z.infer<typeof wsErrorMessageSchema>;

export const sessionEventMessageSchema = wsEventEnvelopeSchema(eventSchema);

export const terminalOutputPayloadSchema = z.object({
  data: z.string(),
});

export const terminalOutputMessageSchema = z.object({
  type: z.literal('terminal_output'),
  seq: z.number().int().positive(),
  session_id: z.string().min(1),
  terminal_id: z.string().min(1),
  timestamp: isoDateTimeSchema,
  payload: terminalOutputPayloadSchema,
});

export type TerminalOutputMessage = z.infer<typeof terminalOutputMessageSchema>;

export const terminalExitPayloadSchema = z.object({
  exit_code: z.number().int().nullable().optional(),
});

export const terminalExitMessageSchema = z.object({
  type: z.literal('terminal_exit'),
  session_id: z.string().min(1),
  terminal_id: z.string().min(1),
  timestamp: isoDateTimeSchema,
  payload: terminalExitPayloadSchema,
});

export type TerminalExitMessage = z.infer<typeof terminalExitMessageSchema>;

export const clientControlMessageSchema = z.discriminatedUnion('type', [
  clientHelloMessageSchema,
  subscribeMessageSchema,
  subscribeV2MessageSchema,
  unsubscribeMessageSchema,
  unsubscribeV2MessageSchema,
  abortMessageSchema,
  terminalAttachMessageSchema,
  terminalDetachMessageSchema,
  terminalInputMessageSchema,
  terminalResizeMessageSchema,
  terminalCloseMessageSchema,
  pongMessageSchema,
]);

export type ClientControlMessage = z.infer<typeof clientControlMessageSchema>;

export const serverSystemMessageSchema = z.discriminatedUnion('type', [
  serverHelloMessageSchema,
  pingMessageSchema,
  resyncRequiredMessageSchema,
  wsErrorMessageSchema,
]);

export type ServerSystemMessage = z.infer<typeof serverSystemMessageSchema>;

export type WsOperationDirection = 'client_to_server' | 'server_to_client';

export type WsOperationKind = 'control' | 'system' | 'event';

export interface WsOperationDefinition {
  readonly type: string;
  readonly direction: WsOperationDirection;
  readonly kind: WsOperationKind;
  readonly messageSchema: z.ZodTypeAny;
  readonly ackSchema?: z.ZodTypeAny;
  readonly description: string;
}

export const clientControlOperations = [
  {
    type: 'client_hello',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: clientHelloMessageSchema,
    ackSchema: clientHelloAckMessageSchema,
    description: 'Start a client session and optionally subscribe to existing daemon sessions.',
  },
  {
    type: 'subscribe',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: subscribeMessageSchema,
    ackSchema: subscribeAckMessageSchema,
    description: 'Subscribe the connection to one or more session event streams.',
  },
  {
    type: 'subscribe_v2',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: subscribeV2MessageSchema,
    ackSchema: subscribeV2AckMessageSchema,
    description:
      "Attach or update this connection's per-agent transcript grade stream for one session.",
  },
  {
    type: 'unsubscribe_v2',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: unsubscribeV2MessageSchema,
    ackSchema: unsubscribeV2AckMessageSchema,
    description:
      "Detach this connection's transcript grade stream for one session, optionally per agent.",
  },
  {
    type: 'unsubscribe',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: unsubscribeMessageSchema,
    ackSchema: unsubscribeAckMessageSchema,
    description: 'Remove one or more session event stream subscriptions.',
  },
  {
    type: 'abort',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: abortMessageSchema,
    ackSchema: abortAckMessageSchema,
    description: 'Abort a running prompt in a session.',
  },
  {
    type: 'terminal_attach',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: terminalAttachMessageSchema,
    ackSchema: terminalAttachAckMessageSchema,
    description: 'Attach this connection to a terminal stream.',
  },
  {
    type: 'terminal_detach',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: terminalDetachMessageSchema,
    ackSchema: terminalDetachAckMessageSchema,
    description: 'Detach this connection from a terminal stream.',
  },
  {
    type: 'terminal_input',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: terminalInputMessageSchema,
    ackSchema: terminalInputAckMessageSchema,
    description: 'Write raw input bytes to a terminal.',
  },
  {
    type: 'terminal_resize',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: terminalResizeMessageSchema,
    ackSchema: terminalResizeAckMessageSchema,
    description: 'Resize a terminal.',
  },
  {
    type: 'terminal_close',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: terminalCloseMessageSchema,
    ackSchema: terminalCloseAckMessageSchema,
    description: 'Close a terminal.',
  },
  {
    type: 'pong',
    direction: 'client_to_server',
    kind: 'control',
    messageSchema: pongMessageSchema,
    description: 'Reply to a server ping with the same nonce.',
  },
] as const satisfies readonly WsOperationDefinition[];

export const serverSystemOperations = [
  {
    type: 'server_hello',
    direction: 'server_to_client',
    kind: 'system',
    messageSchema: serverHelloMessageSchema,
    description: 'Initial server greeting sent immediately after the socket opens.',
  },
  {
    type: 'ping',
    direction: 'server_to_client',
    kind: 'system',
    messageSchema: pingMessageSchema,
    description: 'Heartbeat ping sent by the server; clients must answer with pong.',
  },
  {
    type: 'resync_required',
    direction: 'server_to_client',
    kind: 'system',
    messageSchema: resyncRequiredMessageSchema,
    description: 'Signals that a client must rebuild local session state from REST history.',
  },
  {
    type: 'error',
    direction: 'server_to_client',
    kind: 'system',
    messageSchema: wsErrorMessageSchema,
    description: 'Server-side WebSocket protocol or runtime error.',
  },
] as const satisfies readonly WsOperationDefinition[];

export const sessionEventOperation = {
  type: 'session_event',
  direction: 'server_to_client',
  kind: 'event',
  messageSchema: sessionEventMessageSchema,
  description: 'Session-scoped agent event envelope; frame type is the payload event type.',
} as const satisfies WsOperationDefinition;

export const wsOperations = [
  ...clientControlOperations,
  ...serverSystemOperations,
  sessionEventOperation,
] as const satisfies readonly WsOperationDefinition[];

export function getClientControlOperation(
  type: string,
): (typeof clientControlOperations)[number] | undefined {
  return clientControlOperations.find((operation) => operation.type === type);
}