next-tjwater-drainage-frontend/features/agent/api/client.ts

export type AgentRunStatus = "running" | "completed" | "error" | "aborted";

export type AgentPermissionReply = "once" | "always" | "reject";

export type AgentModelOption = {
  id: string;
  label: string;
  description: string;
  icon: "bolt" | "sparkle";
};

export type AgentModelsResponse = {
  defaultModel: string;
  models: AgentModelOption[];
};

export type AgentChatSession = {
  id?: string;
  session_id: string;
  title?: string;
  status?: string;
  run_status?: AgentRunStatus;
  is_streaming?: boolean;
  messages?: unknown[];
  created_at?: number | string;
  updated_at?: number | string;
  is_title_manually_edited?: boolean;
  parent_session_id?: string;
};

export type AgentChatSessionSummary = {
  id: string;
  title: string;
  createdAt: number;
  updatedAt: number;
  status?: string;
  parentSessionId?: string;
  isStreaming?: boolean;
  runStatus?: AgentRunStatus;
};

export type AgentLoadedChatSession = AgentChatSessionSummary & {
  sessionId: string;
  isTitleManuallyEdited: boolean;
  messages: unknown[];
};

export type AgentSessionStreamDataEventType =
  | "token"
  | "progress"
  | "todo_update"
  | "permission_request"
  | "permission_response"
  | "question_request"
  | "question_response"
  | "tool_call"
  | "ui_envelope"
  | "session_title"
  | "done"
  | "error"
  | "auth_required";

export type AgentSessionStreamEvent =
  | {
      type: "state";
      sessionId: string;
      messages: unknown[];
      isStreaming: boolean;
      runStatus?: AgentRunStatus;
    }
  | {
      type: AgentSessionStreamDataEventType;
      sessionId?: string;
      data: Record<string, unknown>;
    };

export type AgentApiClient = {
  createSession: () => Promise<AgentChatSession>;
  listSessions: () => Promise<AgentChatSessionSummary[]>;
  loadSession: (sessionId: string) => Promise<AgentLoadedChatSession | null>;
  streamSession: (
    sessionId: string,
    options: {
      signal?: AbortSignal;
      onEvent: (event: AgentSessionStreamEvent) => void;
    }
  ) => Promise<void>;
  updateSessionTitle: (
    sessionId: string,
    title: string,
    options?: { isTitleManuallyEdited?: boolean }
  ) => Promise<void>;
  deleteSession: (sessionId: string) => Promise<void>;
  getModels: () => Promise<AgentModelsResponse>;
  getUiRegistry: () => Promise<unknown>;
  resolveRenderRef: (renderRef: string, sessionId: string) => Promise<unknown>;
  replyPermission: (
    requestId: string,
    options: { sessionId: string; reply: AgentPermissionReply; message?: string }
  ) => Promise<void>;
  replyQuestion: (
    requestId: string,
    options: { sessionId: string; answers: string[][] }
  ) => Promise<void>;
  rejectQuestion: (requestId: string, options: { sessionId: string }) => Promise<void>;
  abort: (sessionId: string) => Promise<void>;
};

const AGENT_API_BASE_URLS = process.env.NEXT_PUBLIC_AGENT_API_BASE_URL
  ? [process.env.NEXT_PUBLIC_AGENT_API_BASE_URL.replace(/\/$/, "")]
  : ["", "http://127.0.0.1:8787"];

const CHAT_PATH = "/api/v1/agent/chat";

export function createAgentApiClient(baseUrls: string | string[] = AGENT_API_BASE_URLS): AgentApiClient {
  const candidates = (Array.isArray(baseUrls) ? baseUrls : [baseUrls]).map((item) => item.replace(/\/$/, ""));
  let activeBaseUrl = candidates[0] ?? "";

  return {
    async createSession() {
      return requestJsonWithFallback<AgentChatSession>(candidates, activeBaseUrl, (nextBaseUrl) => {
        activeBaseUrl = nextBaseUrl;
      }, "/session", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({})
      });
    },

    async listSessions() {
      const payload = await requestJsonWithFallback<{ sessions?: unknown[] }>(
        candidates,
        activeBaseUrl,
        (nextBaseUrl) => {
          activeBaseUrl = nextBaseUrl;
        },
        "/sessions"
      );
      return (payload.sessions ?? []).map(toSessionSummary).filter(isPresent).sort(compareSessionSummaries);
    },

    async loadSession(sessionId) {
      const response = await fetchWithFallback(candidates, activeBaseUrl, (nextBaseUrl) => {
        activeBaseUrl = nextBaseUrl;
      }, `/session/${encodeURIComponent(sessionId)}`);
      const text = await response.text();
      const data = text ? JSON.parse(text) : null;

      if (response.status === 404) {
        return null;
      }
      if (!response.ok) {
        throw new Error(getResponseErrorMessage(data, response.status));
      }

      return toLoadedSession(data);
    },

    async streamSession(sessionId, options) {
      const response = await fetchWithFallback(
        candidates,
        activeBaseUrl,
        (nextBaseUrl) => {
          activeBaseUrl = nextBaseUrl;
        },
        `/session/${encodeURIComponent(sessionId)}/stream`,
        { signal: options.signal }
      );

      if (!response.ok) {
        const text = await response.text();
        const data = text ? JSON.parse(text) : null;
        throw new Error(getResponseErrorMessage(data, response.status));
      }

      await readSessionSseStream(response, options.onEvent);
    },

    async updateSessionTitle(sessionId, title, options) {
      const normalizedTitle = title.trim();
      if (!normalizedTitle) {
        return;
      }

      await requestJsonWithFallback<unknown>(
        candidates,
        activeBaseUrl,
        (nextBaseUrl) => {
          activeBaseUrl = nextBaseUrl;
        },
        `/session/${encodeURIComponent(sessionId)}/title`,
        {
          method: "PATCH",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({
            title: normalizedTitle,
            is_title_manually_edited: options?.isTitleManuallyEdited
          })
        }
      );
    },

    async deleteSession(sessionId) {
      await requestJsonWithFallback<unknown>(
        candidates,
        activeBaseUrl,
        (nextBaseUrl) => {
          activeBaseUrl = nextBaseUrl;
        },
        `/session/${encodeURIComponent(sessionId)}`,
        { method: "DELETE" }
      );
    },

    async getModels() {
      const payload = await requestJsonWithFallback<unknown>(candidates, activeBaseUrl, (nextBaseUrl) => {
        activeBaseUrl = nextBaseUrl;
      }, "/models");
      return toModelsResponse(payload);
    },

    async getUiRegistry() {
      return requestJsonWithFallback<unknown>(candidates, activeBaseUrl, (nextBaseUrl) => {
        activeBaseUrl = nextBaseUrl;
      }, "/ui-registry");
    },

    async resolveRenderRef(renderRef, sessionId) {
      const params = new URLSearchParams({ session_id: sessionId });
      return requestJsonWithFallback<unknown>(
        candidates,
        activeBaseUrl,
        (nextBaseUrl) => {
          activeBaseUrl = nextBaseUrl;
        },
        `/render-ref/${encodeURIComponent(renderRef)}?${params.toString()}`
      );
    },

    async replyPermission(requestId, options) {
      await requestJsonWithFallback<unknown>(candidates, activeBaseUrl, (nextBaseUrl) => {
        activeBaseUrl = nextBaseUrl;
      }, `/permission/${encodeURIComponent(requestId)}/reply`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          session_id: options.sessionId,
          reply: options.reply,
          message: options.message
        })
      });
    },

    async replyQuestion(requestId, options) {
      await requestJsonWithFallback<unknown>(candidates, activeBaseUrl, (nextBaseUrl) => {
        activeBaseUrl = nextBaseUrl;
      }, `/question/${encodeURIComponent(requestId)}/reply`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          session_id: options.sessionId,
          answers: options.answers
        })
      });
    },

    async rejectQuestion(requestId, options) {
      await requestJsonWithFallback<unknown>(candidates, activeBaseUrl, (nextBaseUrl) => {
        activeBaseUrl = nextBaseUrl;
      }, `/question/${encodeURIComponent(requestId)}/reject`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          session_id: options.sessionId
        })
      });
    },

    async abort(sessionId) {
      await requestJsonWithFallback<unknown>(candidates, activeBaseUrl, (nextBaseUrl) => {
        activeBaseUrl = nextBaseUrl;
      }, "/abort", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ session_id: sessionId })
      });
    }
  };
}

async function requestJsonWithFallback<T>(
  baseUrls: string[],
  activeBaseUrl: string,
  setActiveBaseUrl: (baseUrl: string) => void,
  path: string,
  init?: RequestInit
) {
  const response = await fetchWithFallback(baseUrls, activeBaseUrl, setActiveBaseUrl, path, init);
  const text = await response.text();
  const data = text ? JSON.parse(text) : null;

  if (!response.ok) {
    throw new Error(getResponseErrorMessage(data, response.status));
  }

  return data as T;
}

async function fetchWithFallback(
  baseUrls: string[],
  activeBaseUrl: string,
  setActiveBaseUrl: (baseUrl: string) => void,
  path: string,
  init?: RequestInit
) {
  const orderedBaseUrls = [activeBaseUrl, ...baseUrls.filter((item) => item !== activeBaseUrl)];
  let lastError: unknown;
  let lastResponse: Response | null = null;

  for (const baseUrl of orderedBaseUrls) {
    try {
      const response = await fetch(`${baseUrl}${CHAT_PATH}${path}`, init);
      if (response.ok) {
        setActiveBaseUrl(baseUrl);
        return response;
      }

      lastResponse = response;
      if (shouldFallbackOnHttpStatus(response.status) && baseUrl !== orderedBaseUrls[orderedBaseUrls.length - 1]) {
        await response.body?.cancel();
        continue;
      }
      return response;
    } catch (error) {
      lastError = error;
    }
  }

  if (lastResponse) {
    return lastResponse;
  }

  throw lastError instanceof Error ? lastError : new Error("Agent API unavailable");
}

function shouldFallbackOnHttpStatus(status: number) {
  return status === 404 || status === 405 || status === 502 || status === 503 || status === 504;
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null;
}

function isPresent<T>(value: T | null | undefined): value is T {
  return value !== null && value !== undefined;
}

function toModelsResponse(value: unknown): AgentModelsResponse {
  if (!isRecord(value)) {
    return { defaultModel: "", models: [] };
  }

  const models = Array.isArray(value.models) ? value.models.map(toModelOption).filter(isPresent) : [];
  const defaultModel = typeof value.default_model === "string" ? value.default_model : models[0]?.id ?? "";
  return {
    defaultModel,
    models
  };
}

function toModelOption(value: unknown): AgentModelOption | null {
  if (!isRecord(value) || typeof value.id !== "string") {
    return null;
  }

  const label = typeof value.label === "string" && value.label.trim() ? value.label : value.id;
  const description =
    typeof value.description === "string" && value.description.trim() ? value.description : label;
  const icon = value.icon === "bolt" || value.icon === "sparkle" ? value.icon : "sparkle";

  return {
    id: value.id,
    label,
    description,
    icon
  };
}

function getResponseErrorMessage(data: unknown, status: number) {
  return isRecord(data) && typeof data.message === "string"
    ? data.message
    : `Agent API failed with HTTP ${status}`;
}

function toMillis(value: unknown) {
  if (typeof value === "number" && Number.isFinite(value)) {
    return value;
  }
  if (typeof value === "string" && value.trim()) {
    const parsed = new Date(value).getTime();
    return Number.isFinite(parsed) ? parsed : Date.now();
  }
  return Date.now();
}

function normalizeTitle(value: unknown) {
  return typeof value === "string" && value.trim() ? value.trim() : "新对话";
}

function readRunStatus(value: unknown): AgentRunStatus | undefined {
  return value === "running" || value === "completed" || value === "error" || value === "aborted"
    ? value
    : undefined;
}

function toSessionSummary(value: unknown): AgentChatSessionSummary | null {
  if (!isRecord(value)) {
    return null;
  }

  const id = typeof value.id === "string" ? value.id : typeof value.session_id === "string" ? value.session_id : "";
  if (!id) {
    return null;
  }

  return {
    id,
    title: normalizeTitle(value.title),
    createdAt: toMillis(value.created_at),
    updatedAt: toMillis(value.updated_at),
    status: typeof value.status === "string" ? value.status : undefined,
    parentSessionId: typeof value.parent_session_id === "string" ? value.parent_session_id : undefined,
    isStreaming: typeof value.is_streaming === "boolean" ? value.is_streaming : undefined,
    runStatus: readRunStatus(value.run_status)
  };
}

function toLoadedSession(value: unknown): AgentLoadedChatSession {
  const summary = toSessionSummary(value);
  if (!summary || !isRecord(value)) {
    throw new Error("Invalid agent session payload");
  }

  return {
    ...summary,
    sessionId: typeof value.session_id === "string" ? value.session_id : summary.id,
    isTitleManuallyEdited:
      typeof value.is_title_manually_edited === "boolean" ? value.is_title_manually_edited : false,
    messages: Array.isArray(value.messages) ? value.messages : []
  };
}

async function readSessionSseStream(
  response: Response,
  onEvent: (event: AgentSessionStreamEvent) => void
) {
  if (!response.body) {
    return;
  }

  const reader = response.body.getReader();
  const decoder = new TextDecoder();
  let buffer = "";

  while (true) {
    const { done, value } = await reader.read();
    buffer += decoder.decode(value, { stream: !done });

    let delimiterIndex = buffer.indexOf("\n\n");
    while (delimiterIndex !== -1) {
      const rawEvent = buffer.slice(0, delimiterIndex);
      buffer = buffer.slice(delimiterIndex + 2);
      emitSessionSseEvent(rawEvent, onEvent);
      delimiterIndex = buffer.indexOf("\n\n");
    }

    if (done) {
      break;
    }
  }

  if (buffer.trim()) {
    emitSessionSseEvent(buffer, onEvent);
  }
}

function emitSessionSseEvent(
  rawEvent: string,
  onEvent: (event: AgentSessionStreamEvent) => void
) {
  const lines = rawEvent.split(/\r?\n/);
  const eventType = lines
    .find((line) => line.startsWith("event:"))
    ?.slice("event:".length)
    .trim();
  const dataText = lines
    .filter((line) => line.startsWith("data:"))
    .map((line) => line.slice("data:".length).trimStart())
    .join("\n");

  if (!eventType || !dataText) {
    return;
  }

  const data = JSON.parse(dataText) as unknown;
  if (!isRecord(data)) {
    return;
  }

  if (eventType === "state") {
    const sessionId = typeof data.session_id === "string" ? data.session_id : "";
    onEvent({
      type: "state",
      sessionId,
      messages: Array.isArray(data.messages) ? data.messages : [],
      isStreaming: data.is_streaming === true,
      runStatus: readRunStatus(data.run_status)
    });
    return;
  }

  if (!isSessionStreamDataEventType(eventType)) {
    return;
  }

  onEvent({
    type: eventType,
    sessionId: typeof data.session_id === "string" ? data.session_id : undefined,
    data
  });
}

function isSessionStreamDataEventType(value: string): value is AgentSessionStreamDataEventType {
  return (
    value === "token" ||
    value === "progress" ||
    value === "todo_update" ||
    value === "permission_request" ||
    value === "permission_response" ||
    value === "question_request" ||
    value === "question_response" ||
    value === "tool_call" ||
    value === "ui_envelope" ||
    value === "session_title" ||
    value === "done" ||
    value === "error" ||
    value === "auth_required"
  );
}

function compareSessionSummaries(left: AgentChatSessionSummary, right: AgentChatSessionSummary) {
  const createdAtDiff = right.createdAt - left.createdAt;
  if (createdAtDiff !== 0) {
    return createdAtDiff;
  }

  const updatedAtDiff = right.updatedAt - left.updatedAt;
  if (updatedAtDiff !== 0) {
    return updatedAtDiff;
  }

  return right.id.localeCompare(left.id);
}