Atlas/web/src/composables/useInvoke.ts

import { onMounted, onUnmounted, ref } from 'vue'
import api from '@/api/client'

export interface TaskItem {
  id: number
  status: string
  progress: number
  result_json: string
  error: string
}

const terminalStatuses = new Set(['succeeded', 'failed', 'canceled'])

export function useInvoke(serviceType: string) {
  const activeTasks = ref<TaskItem[]>([])
  const streams = new Map<number, EventSource>()

  function upsertTask(task: TaskItem) {
    const idx = activeTasks.value.findIndex((t) => t.id === task.id)
    if (idx >= 0) activeTasks.value[idx] = task
    else activeTasks.value.push(task)
  }

  function closeStream(taskId: number) {
    const es = streams.get(taskId)
    if (es) {
      es.close()
      streams.delete(taskId)
    }
  }

  function watchTask(taskId: number) {
    if (streams.has(taskId)) return
    const es = new EventSource(`/api/v1/user/tasks/${taskId}/events`, { withCredentials: true })
    streams.set(taskId, es)

    es.addEventListener('task', (ev) => {
      const task = JSON.parse(ev.data) as TaskItem
      upsertTask(task)
      if (terminalStatuses.has(task.status)) {
        closeStream(taskId)
      }
    })
    es.addEventListener('done', () => closeStream(taskId))
    es.onerror = () => closeStream(taskId)
  }

  async function loadActiveTasks() {
    const res = await api.get<{ items: TaskItem[] }>('/user/tasks/active', {
      params: { service_type: serviceType },
    })
    activeTasks.value = res.data.items
    for (const t of res.data.items) {
      watchTask(t.id)
    }
  }

  async function generate(params: Record<string, unknown>, mode: 'sync' | 'async', model: string) {
    const res = await api.post('/user/generate', {
      service_type: serviceType,
      params,
      mode,
      model,
    })
    if (mode === 'async' && res.data.task_id) {
      watchTask(res.data.task_id)
    }
    return res.data
  }

  async function cancelTask(taskId: number) {
    await api.post(`/user/tasks/${taskId}/cancel`)
  }

  onMounted(loadActiveTasks)
  onUnmounted(() => {
    streams.forEach((es) => es.close())
    streams.clear()
  })

  return { activeTasks, generate, loadActiveTasks, cancelTask, watchTask }
}