server/runtime/session.ts

/**
 * Session serialization module — matches original rust/crates/runtime/src/session.rs exactly.
 *
 * Provides:
 * - MessageRole enum
 * - ContentBlock discriminated union (Text, ToolUse, ToolResult)
 * - ConversationMessage with role, blocks, usage
 * - Session with version + messages
 * - to_json / from_json serialization with VALIDATED deserialization (matches Rust Result<>)
 * - user_text(), assistant(), assistant_with_usage(), tool_result() constructors
 */

import type { TokenUsage } from "./usage";

// ─── SessionError (matches original SessionError enum) ──────────────────────
export class SessionError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "SessionError";
  }
}

// ─── MessageRole ────────────────────────────────────────────────────────────
export type MessageRole = "system" | "user" | "assistant" | "tool";

const VALID_ROLES: Set<string> = new Set(["system", "user", "assistant", "tool"]);

// ─── ContentBlock ───────────────────────────────────────────────────────────
export type ContentBlock =
  | { type: "text"; text: string }
  | { type: "tool_use"; id: string; name: string; input: string }
  | { type: "tool_result"; tool_use_id: string; tool_name: string; output: string; is_error: boolean };

// ─── ConversationMessage ────────────────────────────────────────────────────
export interface ConversationMessage {
  role: MessageRole;
  blocks: ContentBlock[];
  usage?: TokenUsage | null;
}

// ─── Session ────────────────────────────────────────────────────────────────
export interface Session {
  version: number;
  messages: ConversationMessage[];
}

// ─── Constructors (match original impl ConversationMessage) ─────────────────
export function userText(text: string): ConversationMessage {
  return {
    role: "user",
    blocks: [{ type: "text", text }],
    usage: null,
  };
}

export function assistant(blocks: ContentBlock[]): ConversationMessage {
  return {
    role: "assistant",
    blocks,
    usage: null,
  };
}

export function assistantWithUsage(
  blocks: ContentBlock[],
  usage: TokenUsage | null
): ConversationMessage {
  return {
    role: "assistant",
    blocks,
    usage,
  };
}

export function toolResult(
  toolUseId: string,
  toolName: string,
  output: string,
  isError: boolean
): ConversationMessage {
  return {
    role: "tool",
    blocks: [
      {
        type: "tool_result",
        tool_use_id: toolUseId,
        tool_name: toolName,
        output,
        is_error: isError,
      },
    ],
    usage: null,
  };
}

export function newSession(): Session {
  return { version: 1, messages: [] };
}

// ─── Validation helpers (match original required_string / required_u32) ─────

function requiredString(obj: Record<string, unknown>, key: string): string {
  const value = obj[key];
  if (typeof value !== "string") {
    throw new SessionError(`missing or invalid string field "${key}"`);
  }
  return value;
}

function requiredBoolean(obj: Record<string, unknown>, key: string): boolean {
  const value = obj[key];
  if (typeof value !== "boolean") {
    throw new SessionError(`missing or invalid boolean field "${key}"`);
  }
  return value;
}

function requiredNumber(obj: Record<string, unknown>, key: string): number {
  const value = obj[key];
  if (typeof value !== "number") {
    throw new SessionError(`missing or invalid number field "${key}"`);
  }
  return value;
}

function requiredArray(obj: Record<string, unknown>, key: string): unknown[] {
  const value = obj[key];
  if (!Array.isArray(value)) {
    throw new SessionError(`missing or invalid array field "${key}"`);
  }
  return value;
}

function requiredObject(value: unknown, context: string): Record<string, unknown> {
  if (value === null || typeof value !== "object" || Array.isArray(value)) {
    throw new SessionError(`expected object for ${context}`);
  }
  return value as Record<string, unknown>;
}

// ─── Serialization (matches original to_json / from_json) ───────────────────

function contentBlockToJson(block: ContentBlock): Record<string, unknown> {
  switch (block.type) {
    case "text":
      return { type: "text", text: block.text };
    case "tool_use":
      return { type: "tool_use", id: block.id, name: block.name, input: block.input };
    case "tool_result":
      return {
        type: "tool_result",
        tool_use_id: block.tool_use_id,
        tool_name: block.tool_name,
        output: block.output,
        is_error: block.is_error,
      };
  }
}

/**
 * Validated deserialization — matches original ContentBlock::from_json() with Result<>.
 * Throws SessionError on invalid data instead of silently passing through.
 */
function contentBlockFromJson(raw: unknown): ContentBlock {
  const obj = requiredObject(raw, "content block");
  const type = requiredString(obj, "type");

  switch (type) {
    case "text":
      return {
        type: "text",
        text: requiredString(obj, "text"),
      };
    case "tool_use":
      return {
        type: "tool_use",
        id: requiredString(obj, "id"),
        name: requiredString(obj, "name"),
        input: requiredString(obj, "input"),
      };
    case "tool_result":
      return {
        type: "tool_result",
        tool_use_id: requiredString(obj, "tool_use_id"),
        tool_name: requiredString(obj, "tool_name"),
        output: requiredString(obj, "output"),
        is_error: requiredBoolean(obj, "is_error"),
      };
    default:
      throw new SessionError(`unknown content block type: "${type}"`);
  }
}

/**
 * Validated usage deserialization — matches original usage_from_json().
 */
function usageFromJson(raw: unknown): TokenUsage | null {
  if (raw === null || raw === undefined) return null;
  const obj = requiredObject(raw, "usage");
  return {
    inputTokens: typeof obj.inputTokens === "number" ? obj.inputTokens : (typeof obj.input_tokens === "number" ? obj.input_tokens : 0),
    outputTokens: typeof obj.outputTokens === "number" ? obj.outputTokens : (typeof obj.output_tokens === "number" ? obj.output_tokens : 0),
    cacheReadInputTokens: typeof obj.cacheReadInputTokens === "number" ? obj.cacheReadInputTokens : (typeof obj.cache_read_input_tokens === "number" ? obj.cache_read_input_tokens : 0),
    cacheCreationInputTokens: typeof obj.cacheCreationInputTokens === "number" ? obj.cacheCreationInputTokens : (typeof obj.cache_creation_input_tokens === "number" ? obj.cache_creation_input_tokens : 0),
  } as TokenUsage;
}

function conversationMessageToJson(msg: ConversationMessage): Record<string, unknown> {
  const result: Record<string, unknown> = {
    role: msg.role,
    blocks: msg.blocks.map(contentBlockToJson),
  };
  if (msg.usage) {
    result.usage = msg.usage;
  }
  return result;
}

/**
 * Validated deserialization — matches original ConversationMessage::from_json() with Result<>.
 */
function conversationMessageFromJson(raw: unknown): ConversationMessage {
  const obj = requiredObject(raw, "conversation message");
  const role = requiredString(obj, "role");
  if (!VALID_ROLES.has(role)) {
    throw new SessionError(`invalid message role: "${role}"`);
  }
  const blocksRaw = requiredArray(obj, "blocks");
  const blocks = blocksRaw.map(contentBlockFromJson);
  const usage = usageFromJson(obj.usage);
  return {
    role: role as MessageRole,
    blocks,
    usage,
  };
}

export function sessionToJson(session: Session): Record<string, unknown> {
  return {
    version: session.version,
    messages: session.messages.map(conversationMessageToJson),
  };
}

/**
 * Validated deserialization — matches original Session::from_json() with Result<>.
 * Returns structured SessionError instead of generic throws.
 */
export function sessionFromJson(raw: unknown): Session {
  const obj = requiredObject(raw, "session");
  const version = requiredNumber(obj, "version");
  const messagesRaw = requiredArray(obj, "messages");
  const messages = messagesRaw.map(conversationMessageFromJson);
  return { version, messages };
}

/**
 * Matches original Session::save_to_path() — writes session JSON to file.
 */
export function saveSessionToPath(session: Session, path: string): void {
  const fs = require("fs");
  fs.writeFileSync(path, JSON.stringify(sessionToJson(session), null, 2));
}

/**
 * Matches original Session::load_from_path() — reads session JSON from file.
 * Throws SessionError on invalid data.
 */
export function loadSessionFromPath(path: string): Session {
  const fs = require("fs");
  let contents: string;
  try {
    contents = fs.readFileSync(path, "utf-8");
  } catch (err: any) {
    throw new SessionError(`failed to read session file: ${err.message}`);
  }
  let parsed: unknown;
  try {
    parsed = JSON.parse(contents);
  } catch (err: any) {
    throw new SessionError(`invalid JSON in session file: ${err.message}`);
  }
  return sessionFromJson(parsed);
}