@bytespell/shella 0.2.4 → 0.2.6

package/bundled-plugins/agent/src/store.ts

@@ -0,0 +1,379 @@
+import { create } from 'zustand'
+import type {
+  ContentBlock,
+  AgentMessage,
+  ActiveToolCall,
+  SessionState,
+  SessionStats,
+  AvailableModel,
+  ConnectionStatus,
+  ToolResultMessage,
+  Session,
+} from './types'
+
+interface AgentStore {
+  // Connection state
+  status: ConnectionStatus
+  error: string | null
+  currentCwd: string | null
+
+  // Session state
+  sessionState: SessionState | null
+  sessionStats: SessionStats | null
+  availableModels: AvailableModel[]
+  sessions: Session[]
+
+  // Messages
+  messages: AgentMessage[]
+  isStreaming: boolean
+  streamingMessageIndex: number | null
+  activeToolCalls: Map<string, ActiveToolCall>
+
+  // Actions
+  connect: () => void
+  disconnect: () => void
+  send: (data: unknown) => void
+  sendPrompt: (message: string) => void
+  abort: () => void
+  changeCwd: (path: string) => void
+  setError: (error: string | null) => void
+
+  // Commands
+  getState: () => void
+  getMessages: () => void
+  getAvailableModels: () => void
+  getSessionStats: () => void
+  getSessions: () => void
+  setModel: (provider: string, modelId: string) => void
+  setThinkingLevel: (level: string) => void
+  setAutoCompaction: (enabled: boolean) => void
+  compact: () => void
+  newSession: () => void
+  switchSession: (sessionId: string) => void
+}
+
+// WebSocket singleton - lives outside React
+let ws: WebSocket | null = null
+let reconnectTimeout: number | null = null
+let reconnectAttempts = 0
+
+export const useAgentStore = create<AgentStore>((set, get) => ({
+  // Initial state
+  status: 'connecting',
+  error: null,
+  currentCwd: null,
+  sessionState: null,
+  sessionStats: null,
+  availableModels: [],
+  sessions: [],
+  messages: [],
+  isStreaming: false,
+  streamingMessageIndex: null,
+  activeToolCalls: new Map(),
+
+  // Connect to WebSocket
+  connect: () => {
+    // Don't reconnect if already open or connecting
+    if (ws && (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING)) return
+
+    if (ws) {
+      ws.onclose = null // Prevent reconnect on intentional close
+      ws.close()
+    }
+
+    set({ status: 'connecting' })
+    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
+    ws = new WebSocket(`${protocol}//${window.location.host}/ws`)
+
+    ws.onopen = () => {
+      set({ status: 'connected', error: null })
+      reconnectAttempts = 0
+      ws?.send(JSON.stringify({ type: 'start' }))
+    }
+
+    ws.onclose = () => {
+      set({ status: 'disconnected' })
+      ws = null
+      const delay = Math.min(1000 * Math.pow(2, reconnectAttempts), 10000)
+      reconnectAttempts++
+      reconnectTimeout = window.setTimeout(() => {
+        get().connect()
+      }, delay)
+    }
+
+    ws.onerror = () => {
+      set({ status: 'error', error: 'WebSocket connection failed' })
+    }
+
+    ws.onmessage = (event) => {
+      try {
+        const data = JSON.parse(event.data)
+        handleMessage(data, set, get)
+      } catch (err) {
+        console.error('Parse error:', err)
+      }
+    }
+  },
+
+  disconnect: () => {
+    if (reconnectTimeout) {
+      clearTimeout(reconnectTimeout)
+      reconnectTimeout = null
+    }
+    if (ws) {
+      ws.onclose = null
+      ws.close()
+      ws = null
+    }
+    set({ status: 'disconnected' })
+  },
+
+  send: (data) => {
+    ws?.send(JSON.stringify(data))
+  },
+
+  sendPrompt: (message) => {
+    const state = get()
+    if (!message.trim() || !ws || state.sessionState?.isStreaming) return
+
+    // Optimistically add user message
+    set({
+      messages: [...state.messages, { role: 'user', content: [{ type: 'text', text: message }] }]
+    })
+    ws.send(JSON.stringify({ type: 'prompt', message }))
+  },
+
+  abort: () => {
+    ws?.send(JSON.stringify({ type: 'abort' }))
+  },
+
+  changeCwd: (path) => {
+    ws?.send(JSON.stringify({ type: 'change_cwd', path }))
+  },
+
+  setError: (error) => set({ error }),
+
+  // Commands
+  getState: () => {
+    ws?.send(JSON.stringify({ type: 'command', command: { type: 'get_state' } }))
+  },
+
+  getMessages: () => {
+    ws?.send(JSON.stringify({ type: 'command', command: { type: 'get_messages' } }))
+  },
+
+  getAvailableModels: () => {
+    ws?.send(JSON.stringify({ type: 'command', command: { type: 'get_available_models' } }))
+  },
+
+  getSessionStats: () => {
+    ws?.send(JSON.stringify({ type: 'command', command: { type: 'get_session_stats' } }))
+  },
+
+  getSessions: () => {
+    ws?.send(JSON.stringify({ type: 'command', command: { type: 'get_sessions' } }))
+  },
+
+  setModel: (provider, modelId) => {
+    ws?.send(JSON.stringify({ type: 'command', command: { type: 'set_model', provider, modelId } }))
+  },
+
+  setThinkingLevel: (level) => {
+    ws?.send(JSON.stringify({ type: 'command', command: { type: 'set_thinking_level', level } }))
+  },
+
+  setAutoCompaction: (enabled) => {
+    ws?.send(JSON.stringify({ type: 'command', command: { type: 'set_auto_compaction', enabled } }))
+  },
+
+  compact: () => {
+    ws?.send(JSON.stringify({ type: 'command', command: { type: 'compact' } }))
+  },
+
+  newSession: () => {
+    ws?.send(JSON.stringify({ type: 'command', command: { type: 'new_session' } }))
+  },
+
+  switchSession: (sessionPath) => {
+    ws?.send(JSON.stringify({ type: 'command', command: { type: 'switch_session', sessionPath } }))
+  },
+}))
+
+// Handle incoming WebSocket messages
+function handleMessage(
+  data: { type: string; event?: Record<string, unknown>; message?: string; cwd?: string },
+  set: (partial: Partial<AgentStore>) => void,
+  get: () => AgentStore
+) {
+  if (data.type === 'ready') {
+    set({ status: 'connected' })
+    if (data.cwd) set({ currentCwd: data.cwd })
+  } else if (data.type === 'cwd_changed') {
+    if (data.cwd) set({ currentCwd: data.cwd })
+    set({ messages: [], sessionStats: null, activeToolCalls: new Map() })
+  } else if (data.type === 'error') {
+    set({ error: data.message || 'Unknown error' })
+  } else if (data.type === 'rpc_event' && data.event) {
+    handleRpcEvent(data.event, set, get)
+  }
+}
+
+// Handle RPC events from pi-agent
+function handleRpcEvent(
+  event: Record<string, unknown>,
+  set: (partial: Partial<AgentStore>) => void,
+  get: () => AgentStore
+) {
+  const eventType = event.type as string
+
+  // Response events
+  if (eventType === 'response') {
+    const command = event.command as string
+    if (command === 'get_state' && event.success) {
+      set({ sessionState: event.data as SessionState })
+    } else if (command === 'get_messages' && event.success) {
+      const data = event.data as { messages: AgentMessage[] }
+      console.log('[store] get_messages response:', data.messages?.length, 'messages', data)
+      set({ messages: data.messages || [] })
+    } else if (command === 'get_available_models' && event.success) {
+      const data = event.data as { models: AvailableModel[] }
+      set({ availableModels: data.models || [] })
+    } else if (command === 'get_session_stats' && event.success) {
+      set({ sessionStats: event.data as SessionStats })
+    } else if (command === 'set_model' && event.success) {
+      get().getState()
+    } else if (command === 'set_thinking_level' && event.success) {
+      get().getState()
+    } else if (command === 'set_auto_compaction' && event.success) {
+      get().getState()
+    } else if (command === 'compact' && event.success) {
+      get().getState()
+      get().getSessionStats()
+    } else if (command === 'new_session' && event.success) {
+      set({ messages: [], sessionStats: null })
+      get().getState()
+      get().getSessions() // Refresh sessions list
+    } else if (command === 'get_sessions' && event.success) {
+      const data = event.data as { sessions: Session[] }
+      set({ sessions: data.sessions || [] })
+    } else if (command === 'switch_session' && event.success) {
+      // Session switched - clear current state and fetch new messages
+      set({ messages: [], sessionStats: null, activeToolCalls: new Map() })
+      get().getState()
+      get().getMessages()
+      get().getSessionStats()
+      get().getSessions() // Refresh to update current session indicator
+    }
+  }
+
+  // Agent lifecycle events
+  else if (eventType === 'agent_start') {
+    set({ isStreaming: true, activeToolCalls: new Map() })
+  }
+  else if (eventType === 'agent_end') {
+    set({ isStreaming: false, activeToolCalls: new Map(), streamingMessageIndex: null })
+    get().getState()
+    get().getSessionStats()
+  }
+
+  // Message events
+  else if (eventType === 'message_start') {
+    const msg = event.message as AgentMessage
+    console.log('[store] message_start:', msg?.role, msg)
+    if (msg) {
+      // Skip user messages - we add them optimistically in sendPrompt
+      if (msg.role === 'user') {
+        return
+      }
+
+      const messages = get().messages
+      const newMessages = [...messages, msg]
+
+      const updates: Partial<AgentStore> = { messages: newMessages }
+
+      if (msg.role === 'assistant') {
+        updates.streamingMessageIndex = newMessages.length - 1
+      }
+
+      if (msg.role === 'toolResult') {
+        const toolCallId = (msg as ToolResultMessage).toolCallId
+        const activeToolCalls = new Map(get().activeToolCalls)
+        activeToolCalls.delete(toolCallId)
+        updates.activeToolCalls = activeToolCalls
+      }
+
+      set(updates)
+    }
+  }
+  else if (eventType === 'message_update') {
+    const msg = event.message as AgentMessage
+    if (msg) {
+      const messages = get().messages
+      if (messages.length === 0) {
+        set({ messages: [msg] })
+      } else {
+        const updated = [...messages]
+        updated[updated.length - 1] = msg
+        set({ messages: updated })
+      }
+    }
+  }
+  else if (eventType === 'message_end') {
+    set({ streamingMessageIndex: null })
+  }
+
+  // Tool execution events
+  else if (eventType === 'tool_execution_start') {
+    const toolCallId = event.toolCallId as string
+    const activeToolCalls = new Map(get().activeToolCalls)
+    activeToolCalls.set(toolCallId, {
+      toolCallId,
+      toolName: event.toolName as string,
+      args: (event.args as Record<string, unknown>) || {},
+      status: 'running',
+      partialOutput: null,
+    })
+    set({ activeToolCalls })
+  }
+  else if (eventType === 'tool_execution_update') {
+    const toolCallId = event.toolCallId as string
+    const partialResult = event.partialResult as { content?: ContentBlock[] } | undefined
+    const activeToolCalls = new Map(get().activeToolCalls)
+    const existing = activeToolCalls.get(toolCallId)
+    if (existing) {
+      const text = partialResult?.content?.find((c): c is { type: 'text'; text: string } => c.type === 'text')?.text
+      activeToolCalls.set(toolCallId, { ...existing, partialOutput: text || null })
+      set({ activeToolCalls })
+    }
+  }
+  else if (eventType === 'tool_execution_end') {
+    const toolCallId = event.toolCallId as string
+    const result = event.result as { content?: ContentBlock[] } | undefined
+    const activeToolCalls = new Map(get().activeToolCalls)
+    const existing = activeToolCalls.get(toolCallId)
+    if (existing) {
+      const text = result?.content?.find((c): c is { type: 'text'; text: string } => c.type === 'text')?.text
+      activeToolCalls.set(toolCallId, {
+        ...existing,
+        status: event.isError ? 'error' : 'done',
+        partialOutput: text || null,
+      })
+      set({ activeToolCalls })
+    }
+  }
+}
+
+// Reconnect on visibility change
+if (typeof document !== 'undefined') {
+  document.addEventListener('visibilitychange', () => {
+    const { status, connect } = useAgentStore.getState()
+    if (document.visibilityState === 'visible' && (status === 'disconnected' || status === 'error')) {
+      if (reconnectTimeout) {
+        clearTimeout(reconnectTimeout)
+        reconnectTimeout = null
+      }
+      reconnectAttempts = 0
+      connect()
+    }
+  })
+}

package/bundled-plugins/agent/src/types.ts

@@ -0,0 +1,98 @@
+// =============================================================================
+// Types - Following AGENT_SPEC.md
+// =============================================================================
+
+export type ContentBlock =
+  | { type: 'text'; text: string }
+  | { type: 'thinking'; thinking: string }
+  | { type: 'toolCall'; id: string; name: string; arguments: Record<string, unknown> }
+  | { type: 'image'; data: string; mimeType: string }
+
+export interface UserMessage {
+  role: 'user'
+  content: ContentBlock[]
+  timestamp?: number
+}
+
+export interface AssistantMessage {
+  role: 'assistant'
+  content: ContentBlock[]
+  model?: string
+  stopReason?: string
+  timestamp?: number
+}
+
+export interface ToolResultMessage {
+  role: 'toolResult'
+  toolCallId: string
+  toolName: string
+  content: ContentBlock[]
+  isError?: boolean
+  timestamp?: number
+  details?: ToolResultDetails
+}
+
+// Tool-specific details
+export type ToolResultDetails = EditToolDetails | BashToolDetails | unknown
+
+export interface EditToolDetails {
+  diff: string
+  firstChangedLine?: number
+}
+
+export interface BashToolDetails {
+  exitCode?: number
+  truncated?: boolean
+}
+
+export type AgentMessage = UserMessage | AssistantMessage | ToolResultMessage
+
+export interface ActiveToolCall {
+  toolCallId: string
+  toolName: string
+  args: Record<string, unknown>
+  status: 'running' | 'done' | 'error'
+  partialOutput: string | null
+}
+
+export interface SessionState {
+  model?: { provider: string; id: string; name?: string; reasoning?: boolean }
+  thinkingLevel?: string
+  isStreaming: boolean
+  isCompacting?: boolean
+  autoCompactionEnabled?: boolean
+}
+
+export interface SessionStats {
+  tokens?: { input: number; output: number; cacheRead?: number; cacheWrite?: number; total: number }
+  cost?: number
+}
+
+export interface AvailableModel {
+  provider: string
+  id: string
+  name?: string
+  reasoning?: boolean
+}
+
+export type ConnectionStatus = 'connecting' | 'connected' | 'disconnected' | 'error'
+
+export interface Session {
+  id: string
+  name?: string | null
+  file?: string
+  messageCount?: number
+  created?: number | null
+  lastModified?: number
+  isCurrent?: boolean
+}
+
+// =============================================================================
+// Render Item Types - For splitting assistant messages
+// =============================================================================
+
+export type RenderItem =
+  | { type: 'user-message'; content: ContentBlock[] }
+  | { type: 'text-chunk'; blocks: ContentBlock[]; isStreaming: boolean }
+  | { type: 'tool-call'; id: string; name: string; args: Record<string, unknown> }
+  | { type: 'tool-result'; toolCallId: string; toolName: string; command?: string; output: string; isError: boolean; diff?: string; filePath?: string }