CodeNomad/src/stores/session-status.ts

import type { Session, SessionStatus } from "../types/session"
import type { Message, MessageInfo } from "../types/message"
import { sessions } from "./sessions"
import { isSessionCompactionActive } from "./session-compaction"

function getSession(instanceId: string, sessionId: string): Session | null {
  const instanceSessions = sessions().get(instanceId)
  return instanceSessions?.get(sessionId) ?? null
}

function isSessionCompacting(session: Session): boolean {
  const time = (session.time as (Session["time"] & { compacting?: number }) | undefined)
  const compactingFlag = time?.compacting
  if (typeof compactingFlag === "number") {
    return compactingFlag > 0
  }
  return Boolean(compactingFlag)
}

function getMessageTimestamp(session: Session, message?: Message): number {
  if (!message) return Number.NEGATIVE_INFINITY
  if (typeof message.timestamp === "number" && Number.isFinite(message.timestamp)) {
    return message.timestamp
  }
  const info = session.messagesInfo.get(message.id)
  return info?.time?.created ?? Number.NEGATIVE_INFINITY
}

function getLastMessage(session: Session): Message | undefined {
  let latest: Message | undefined
  let latestTimestamp = Number.NEGATIVE_INFINITY
  for (const message of session.messages) {
    if (!message) continue
    const timestamp = getMessageTimestamp(session, message)
    if (timestamp >= latestTimestamp) {
      latest = message
      latestTimestamp = timestamp
    }
  }
  return latest
}

function getLastMessageInfo(session: Session, role?: MessageInfo["role"]): MessageInfo | undefined {
  if (session.messagesInfo.size === 0) {
    return undefined
  }
  let latest: MessageInfo | undefined
  let latestTimestamp = Number.NEGATIVE_INFINITY
  for (const info of session.messagesInfo.values()) {
    if (!info) continue
    if (role && info.role !== role) continue
    const timestamp = info.time?.created ?? 0
    if (timestamp >= latestTimestamp) {
      latest = info
      latestTimestamp = timestamp
    }
  }
  return latest
}

function getInfoCreatedTimestamp(info?: MessageInfo): number {
  if (!info) {
    return Number.NEGATIVE_INFINITY
  }
  const created = info.time?.created
  if (typeof created === "number" && Number.isFinite(created)) {
    return created
  }
  return Number.NEGATIVE_INFINITY
}

function getAssistantCompletionTimestamp(info?: MessageInfo): number {
  if (!info) {
    return Number.NEGATIVE_INFINITY
  }
  const completed = (info.time as { completed?: number } | undefined)?.completed
  if (typeof completed === "number" && Number.isFinite(completed)) {
    return completed
  }
  return Number.NEGATIVE_INFINITY
}

function isAssistantInfoPending(info?: MessageInfo): boolean {
  if (!info) {
    return false
  }
  const completed = (info.time as { completed?: number } | undefined)?.completed
  if (completed === undefined || completed === null) {
    return true
  }
  const created = getInfoCreatedTimestamp(info)
  return completed < created
}

function isAssistantStillGenerating(message: Message, info?: MessageInfo): boolean {
  if (message.type !== "assistant") {
    return false
  }

  if (message.status === "error") {
    return false
  }

  if (message.status === "streaming" || message.status === "sending") {
    return true
  }

  const completedAt = (info?.time as { completed?: number } | undefined)?.completed
  if (completedAt !== undefined && completedAt !== null) {
    return false
  }

  return !(message.status === "complete" || message.status === "sent")
}

export function getSessionStatus(instanceId: string, sessionId: string): SessionStatus {
  const session = getSession(instanceId, sessionId)
  if (!session) {
    return "idle"
  }

  if (isSessionCompactionActive(instanceId, sessionId) || isSessionCompacting(session)) {
    return "compacting"
  }

  const latestUserInfo = getLastMessageInfo(session, "user")
  const latestAssistantInfo = getLastMessageInfo(session, "assistant")
  const lastMessage = getLastMessage(session)
  if (!lastMessage) {
    const latestInfo = getLastMessageInfo(session)
    if (!latestInfo) {
      return "idle"
    }
    if (latestInfo.role === "user") {
      return "working"
    }
    const infoCompleted = latestInfo.time?.completed
    return infoCompleted ? "idle" : "working"
  }

  if (lastMessage.type === "user") {
    return "working"
  }

  const infoForMessage = session.messagesInfo.get(lastMessage.id) ?? latestAssistantInfo
  if (isAssistantStillGenerating(lastMessage, infoForMessage)) {
    return "working"
  }

  if (isAssistantInfoPending(latestAssistantInfo)) {
    return "working"
  }

  const userTimestamp = getInfoCreatedTimestamp(latestUserInfo)
  const assistantCompletedAt = getAssistantCompletionTimestamp(latestAssistantInfo)
  if (userTimestamp > assistantCompletedAt) {
    return "working"
  }

  return "idle"
}

export function isSessionBusy(instanceId: string, sessionId: string): boolean {
  const status = getSessionStatus(instanceId, sessionId)
  return status === "working" || status === "compacting"
}