Dynavera/site/src/stores/agentStore.ts

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

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 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 connect = (id: string) => {
        if (socket.value) {
            socket.value.close()
            socket.value = null
        }

        const wsProtocol = window.location.protocol === 'https:' ? 'wss' : 'ws'
        const wsUrl = `${wsProtocol}://${window.location.host}/ws/onboarding/${id}/`
        socket.value = new WebSocket(wsUrl)

        socket.value.onopen = () => {
            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 === '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'
                    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 = () => {
            isConnected.value = false
            executionStatus.value = 'idle'
        }
    }

    const disconnect = () => {
        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({
                query: data.query,
                role_uuid: data.role_uuid,
                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,
        connect,
        disconnect,
        startAgent,
        stopAgent,
        clearLog,
        lastExecutionId,
    }
})