Dynavera/site/src/stores/agentStore.ts

import { defineStore } from 'pinia'
import { ref } from 'vue'
import type {
    AgentEvent,
    AgentExecutionStatus,
    AgentSocketEventPayload,
    AgentStartPayload,
} from '../types/agent'
import { BACKOFF_BASE_MS, BACKOFF_MAX_MS, BACKOFF_MAX_ATTEMPTS } from './agentBackoff'

export const useAgentStore = defineStore('agent', () => {
    const isConnected = ref(false)
    const executionStatus = ref<AgentExecutionStatus>('idle')
    const eventLog = ref<AgentEvent[]>([])
    const lastExecutionId = ref<string | null>(null)
    const socket = ref<WebSocket | null>(null)
    const streamBuffer = ref('')

    let currentUrl = ''
    let reconnectAttempts = 0
    let reconnectTimer: ReturnType<typeof setTimeout> | null = null
    let intentionalClose = false

    const pushEvent = (evt: AgentSocketEventPayload) => {
        eventLog.value.unshift({
            type: evt.type,
            message: evt.message,
            content: evt.content,
            timestamp: evt.timestamp ? new Date(evt.timestamp) : new Date(),
        })
    }

    const clearReconnectTimer = () => {
        if (reconnectTimer !== null) {
            clearTimeout(reconnectTimer)
            reconnectTimer = null
        }
    }

    const scheduleReconnect = () => {
        if (reconnectAttempts >= BACKOFF_MAX_ATTEMPTS) {
            pushEvent({ type: 'error', message: 'Connection lost. Please refresh the page.' })
            executionStatus.value = 'idle'
            return
        }
        const delay = Math.min(BACKOFF_BASE_MS * 2 ** reconnectAttempts, BACKOFF_MAX_MS)
        reconnectAttempts++
        pushEvent({ type: 'status', message: `Reconnecting in ${Math.round(delay / 1000)}s (attempt ${reconnectAttempts}/${BACKOFF_MAX_ATTEMPTS})...` })
        reconnectTimer = setTimeout(() => {
            if (!intentionalClose) openSocket(currentUrl)
        }, delay)
    }

    const openSocket = (url: string) => {
        socket.value = new WebSocket(url)

        socket.value.onopen = () => {
            reconnectAttempts = 0
            isConnected.value = true
            pushEvent({ type: 'status', message: 'Connected to Orchestrator' })
        }

        socket.value.onmessage = (event) => {
            try {
                const payload = JSON.parse(event.data) as AgentSocketEventPayload
                const type = payload.type

                if (payload.execution_id) {
                    lastExecutionId.value = String(payload.execution_id)
                }

                if (type === 'stream_chunk') {
                    executionStatus.value = 'running'
                    streamBuffer.value += payload.message || ''
                } else if (type === 'status' || type === 'thought' || type === 'tool_start') {
                    executionStatus.value = 'running'
                    pushEvent({
                        type,
                        message: payload.message || payload.thought,
                        content: payload.content,
                    })
                } else if (
                    type === 'tool_call' ||
                    type === 'tool_result' ||
                    type === 'tool_complete'
                ) {
                    pushEvent({
                        type,
                        message: payload.message,
                        content: payload.content || payload,
                    })
                } else if (type === 'completed') {
                    executionStatus.value = 'completed'
                    streamBuffer.value = ''
                    pushEvent({
                        type: 'completed',
                        message: 'Generation loop finished successfully',
                        content: payload.content,
                        timestamp: payload.timestamp,
                    })
                } else if (type === 'error') {
                    executionStatus.value = 'failed'
                    pushEvent({ type: 'error', message: payload.message })
                }
            } catch (e) {
                console.error('Store message error', e)
            }
        }

        socket.value.onclose = (event) => {
            isConnected.value = false
            // code 1000 = clean close (server finished normally); don't reconnect
            if (!intentionalClose && event.code !== 1000) {
                scheduleReconnect()
            } else {
                executionStatus.value = 'idle'
            }
        }
    }

    const connect = (id: string) => {
        intentionalClose = false
        clearReconnectTimer()
        reconnectAttempts = 0
        streamBuffer.value = ''

        if (socket.value) {
            socket.value.close()
            socket.value = null
        }

        const wsProtocol = window.location.protocol === 'https:' ? 'wss' : 'ws'
        currentUrl = `${wsProtocol}://${window.location.host}/ws/onboarding/chat/${id}/`
        openSocket(currentUrl)
    }

    const disconnect = () => {
        intentionalClose = true
        clearReconnectTimer()
        if (socket.value) {
            socket.value.close()
            socket.value = null
        }
        isConnected.value = false
        executionStatus.value = 'idle'
    }

    const startAgent = (data: AgentStartPayload) => {
        if (!socket.value || socket.value.readyState !== WebSocket.OPEN) return

        executionStatus.value = 'running'
        socket.value.send(
            JSON.stringify({
                action: 'message',
                query: data.query,
                max_tokens: data.max_tokens,
            }),
        )
    }

    const stopAgent = (executionId?: string) => {
        if (!socket.value || socket.value.readyState !== WebSocket.OPEN) return

        socket.value.send(
            JSON.stringify({
                action: 'stop_agent',
                execution_id: executionId ?? lastExecutionId.value,
            }),
        )
    }

    const clearLog = () => {
        eventLog.value = []
    }

    return {
        isConnected,
        executionStatus,
        eventLog,
        socket,
        streamBuffer,
        connect,
        disconnect,
        startAgent,
        stopAgent,
        clearLog,
        lastExecutionId,
    }
})