Add agent color styling and session persistence

Author: chriswritescode-dev

frontend/src/components/agent/AgentQuickSelect.tsx

@@ -9,6 +9,7 @@ import {
   DropdownMenuTrigger,
 } from '@/components/ui/dropdown-menu'
 import { useAgents } from '@/hooks/useOpenCode'
+import { getAgentStyleVars } from '@/lib/agent-colors'
 
 interface AgentQuickSelectProps {
   opcodeUrl: string | null | undefined
@@ -19,33 +20,33 @@ interface AgentQuickSelectProps {
   disabled?: boolean
 }
 
-const getAgentStyles = (agent: string) => {
-  const lowerAgent = agent.toLowerCase()
-  if (lowerAgent === 'plan') {
-    return {
-      color: 'text-yellow-600 dark:text-yellow-500',
-      bg: 'bg-yellow-500/20 border-yellow-400 hover:bg-yellow-500/30 hover:border-yellow-300',
-      shadow: 'shadow-yellow-500/20 hover:shadow-yellow-500/30',
-    }
-  }
-  if (lowerAgent === 'build') {
-    return {
-      color: 'text-green-600 dark:text-green-500',
-      bg: 'bg-green-500/20 border-green-400 hover:bg-green-500/30 hover:border-green-300',
-      shadow: 'shadow-green-500/20 hover:shadow-green-500/30',
-    }
-  }
-  return {
-    color: 'text-blue-600 dark:text-blue-500',
-    bg: 'bg-blue-500/20 border-blue-400 hover:bg-blue-500/30 hover:border-blue-300',
-    shadow: 'shadow-blue-500/20 hover:shadow-blue-500/30',
-  }
+interface AgentInfo {
+  name: string
+  color?: string
+  description?: string
+  mode?: string
+  hidden?: boolean
+}
+
+const findAgentColor = (agents: AgentInfo[], agentName: string): string | undefined => {
+  return agents.find(a => a.name.toLowerCase() === agentName.toLowerCase())?.color
 }
 
-const bashStyles = {
-  color: 'text-purple-700 dark:text-purple-300',
-  bg: 'bg-purple-500/20 border-purple-400',
-  shadow: 'shadow-purple-500/20 hover:shadow-purple-500/30',
+const bashStyleVars: Record<string, string> = {
+  '--agent-color-light': '#a753ae',
+  '--agent-color-dark': '#edb2f1',
+  '--agent-bg-light': 'rgba(167, 83, 174, 0.2)',
+  '--agent-bg-dark': 'rgba(237, 178, 241, 0.2)',
+  '--agent-bg-hover-light': 'rgba(167, 83, 174, 0.3)',
+  '--agent-bg-hover-dark': 'rgba(237, 178, 241, 0.3)',
+  '--agent-border-light': 'rgba(167, 83, 174, 0.6)',
+  '--agent-border-dark': 'rgba(237, 178, 241, 0.6)',
+  '--agent-border-hover-light': 'rgba(167, 83, 174, 0.5)',
+  '--agent-border-hover-dark': 'rgba(237, 178, 241, 0.5)',
+  '--agent-shadow-light': 'rgba(167, 83, 174, 0.2)',
+  '--agent-shadow-dark': 'rgba(237, 178, 241, 0.2)',
+  '--agent-shadow-hover-light': 'rgba(167, 83, 174, 0.3)',
+  '--agent-shadow-hover-dark': 'rgba(237, 178, 241, 0.3)',
 }
 
 export function AgentQuickSelect({
@@ -66,60 +67,52 @@ export function AgentQuickSelect({
     )
   }, [agents])
 
-  const handleToggle = () => {
-    if (primaryAgents.length === 0) return
-    
-    const currentIndex = primaryAgents.findIndex(
-      (a) => a.name.toLowerCase() === currentAgent.toLowerCase()
-    )
-    const nextIndex = (currentIndex + 1) % primaryAgents.length
-    onAgentChange(primaryAgents[nextIndex].name)
-  }
-
   const handleSelect = (agentName: string) => {
     onAgentChange(agentName)
   }
 
-  const styles = isBashMode ? bashStyles : getAgentStyles(currentAgent)
+  const styleVars = isBashMode 
+    ? bashStyleVars 
+    : getAgentStyleVars(currentAgent, findAgentColor(agents, currentAgent))
   const displayName = isBashMode ? 'Bash' : capitalize(currentAgent)
 
   const buttonContent = (
     <button
       data-toggle-mode
-      onClick={primaryAgents.length <= 2 ? handleToggle : undefined}
       disabled={disabled}
-      className={`px-2 md:px-3.5 py-1 h-[36px] rounded-lg text-sm font-medium border min-w-[48px] max-w-[80px] md:max-w-[100px] flex items-center justify-center transition-all duration-200 active:scale-95 hover:scale-105 shadow-md ${styles.bg} ${styles.color} ${styles.shadow}`}
+      style={styleVars as React.CSSProperties}
+      className="px-2 md:px-3.5 py-1 h-[36px] rounded-lg text-sm font-medium border min-w-[56px] max-w-[80px] md:max-w-[100px] flex-shrink-0 flex items-center justify-center transition-all duration-200 active:scale-95 hover:scale-105 shadow-md text-[var(--agent-color-light)] dark:text-[var(--agent-color-dark)] bg-[var(--agent-bg-light)] dark:bg-[var(--agent-bg-dark)] border-[var(--agent-border-light)] dark:border-[var(--agent-border-dark)] hover:bg-[var(--agent-bg-hover-light)] dark:hover:bg-[var(--agent-bg-hover-dark)] hover:border-[var(--agent-border-hover-light)] dark:hover:border-[var(--agent-border-hover-dark)] shadow-[var(--agent-shadow-light)] dark:shadow-[var(--agent-shadow-dark)] hover:shadow-[var(--agent-shadow-hover-light)] dark:hover:shadow-[var(--agent-shadow-hover-dark)]"
     >
       <span className="truncate">{displayName}</span>
     </button>
   )
 
-  if (primaryAgents.length <= 2) {
-    return buttonContent
-  }
-
   return (
     <DropdownMenu>
       <DropdownMenuTrigger asChild disabled={disabled}>
         {buttonContent}
       </DropdownMenuTrigger>
-      <DropdownMenuContent align="start" className="w-48">
+      <DropdownMenuContent align="start" className="w-64">
         {primaryAgents.map((agent) => {
-          const agentStyles = getAgentStyles(agent.name)
+          const apiColor = agent.color
+          const itemStyleVars = getAgentStyleVars(agent.name, apiColor)
           const isSelected = agent.name.toLowerCase() === currentAgent.toLowerCase()
 
           return (
             <DropdownMenuItem
               key={agent.name}
               onClick={() => handleSelect(agent.name)}
-              className="flex items-center justify-between"
+              className="group flex items-center justify-between"
             >
-              <div className="flex flex-col">
-                <span className={`font-medium ${agentStyles.color}`}>
+              <div className="flex flex-col min-w-0">
+                <span 
+                  className="font-medium"
+                  style={{ color: itemStyleVars['--agent-color-light'] }}
+                >
                   {capitalize(agent.name)}
                 </span>
                 {agent.description && (
-                  <span className="text-xs text-muted-foreground truncate max-w-[160px]">
+                  <span className="text-xs text-muted-foreground line-clamp-2 group-hover:line-clamp-none">
                     {agent.description}
                   </span>
                 )}

frontend/src/components/message/PromptInput.tsx

@@ -9,6 +9,7 @@ import { useSessionAgent } from '@/hooks/useSessionAgent'
 import { useSTT } from '@/hooks/useSTT'
 
 import { useUserBash } from '@/stores/userBashStore'
+import { useSessionAgentStore } from '@/stores/sessionAgentStore'
 import { useMobile } from '@/hooks/useMobile'
 
 import { usePermissions } from '@/contexts/EventContext'
@@ -193,13 +194,15 @@ export const PromptInput = memo(forwardRef<PromptInputHandle, PromptInputProps>(
 
     if (hasActiveStream) {
       const parts = parsePromptToParts(prompt, attachedFiles, imageAttachments)
+      const agentUsed = selectedAgent || currentMode
       sendPrompt.mutate({
         sessionID,
         parts,
         model: currentModel,
-        agent: selectedAgent || currentMode,
+        agent: agentUsed,
         variant: currentVariant
       })
+      setStoredAgent(sessionID, agentUsed)
       setPrompt('')
       setAttachedFiles(new Map())
       revokeBlobUrls(imageAttachments)
@@ -217,6 +220,7 @@ export const PromptInput = memo(forwardRef<PromptInputHandle, PromptInputProps>(
         command,
         agent: currentMode
       })
+      setStoredAgent(sessionID, currentMode)
       setPrompt('')
       setIsBashMode(false)
       clearSTT()
@@ -239,15 +243,17 @@ export const PromptInput = memo(forwardRef<PromptInputHandle, PromptInputProps>(
     }
 
     const parts = parsePromptToParts(prompt, attachedFiles, imageAttachments)
+    const agentUsed = selectedAgent || currentMode
 
     sendPrompt.mutate({
       sessionID,
       parts,
       model: currentModel,
-      agent: selectedAgent || currentMode,
+      agent: agentUsed,
       variant: currentVariant
     })
 
+    setStoredAgent(sessionID, agentUsed)
     setPrompt('')
     setAttachedFiles(new Map())
     revokeBlobUrls(imageAttachments)
@@ -356,6 +362,7 @@ export const PromptInput = memo(forwardRef<PromptInputHandle, PromptInputProps>(
 
   const handleAgentChange = (agent: string) => {
     setLocalMode(agent)
+    setStoredAgent(sessionID, agent)
   }
 
   const handleVoiceToggle = async () => {
@@ -718,6 +725,7 @@ if (isIOS && isSecureContext && navigator.clipboard && navigator.clipboard.read)
 
   const sessionAgent = useSessionAgent(opcodeUrl, sessionID, directory)
   const currentMode = localMode ?? sessionAgent.agent
+  const setStoredAgent = useSessionAgentStore((s) => s.setAgent)
 
 const { model, modelString } = useModelSelection(opcodeUrl, directory)
   const currentModel = modelString || ''
@@ -806,7 +814,7 @@ return (
       )}
 
       <div className="flex gap-1.5 md:gap-2 items-center justify-between">
-        <div className="flex gap-1.5 md:gap-2 items-center">
+        <div className="flex gap-1.5 md:gap-2 items-center min-w-0">
           <AgentQuickSelect
             opcodeUrl={opcodeUrl}
             directory={directory}

frontend/src/hooks/useSessionAgent.ts

@@ -1,10 +1,11 @@
-import { useMemo, useRef } from 'react'
+import { useMemo, useRef, useEffect } from 'react'
 import { useMessages } from './useOpenCode'
+import { useSessionAgentStore } from '@/stores/sessionAgentStore'
 import type { components } from '@/api/opencode-types'
 
 type UserMessage = components['schemas']['UserMessage']
 
-const DEFAULT_AGENT = 'plan'
+const DEFAULT_AGENT = 'build'
 
 interface SessionAgentResult {
   agent: string
@@ -17,56 +18,102 @@ export function useSessionAgent(
   sessionID: string | undefined,
   directory?: string
 ) {
-  const { data: messages } = useMessages(opcodeUrl, sessionID, directory)
+  const { data: messages, isLoading: messagesLoading } = useMessages(opcodeUrl, sessionID, directory)
+  const storedAgent = useSessionAgentStore((s) => s.agents[sessionID ?? ''] ?? null)
+  const setAgent = useSessionAgentStore((s) => s.setAgent)
   const prevRef = useRef<SessionAgentResult>({ agent: DEFAULT_AGENT, model: undefined, variant: undefined })
 
-  return useMemo(() => {
-    let agent = DEFAULT_AGENT
-    let model: { providerID: string; modelID: string } | undefined
-    let variant: string | undefined
-
-    if (messages && messages.length > 0) {
-      for (let i = messages.length - 1; i >= 0; i--) {
-        const msgWithParts = messages[i]
-        if (msgWithParts.info.role === 'user') {
-          const userInfo = msgWithParts.info as UserMessage
-          agent = userInfo.agent || DEFAULT_AGENT
-          model = userInfo.model
-          variant = userInfo.variant
-          break
+  const result = useMemo(() => {
+    if (storedAgent) {
+      let model: { providerID: string; modelID: string } | undefined
+      let variant: string | undefined
+
+      if (messages && messages.length > 0) {
+        for (let i = messages.length - 1; i >= 0; i--) {
+          const msgWithParts = messages[i]
+          if (msgWithParts.info.role === 'user') {
+            const userInfo = msgWithParts.info as UserMessage
+            model = userInfo.model
+            variant = userInfo.variant
+            break
+          }
+        }
+      }
+
+      const prev = prevRef.current
+      if (
+        prev.agent === storedAgent &&
+        prev.variant === variant &&
+        prev.model?.providerID === model?.providerID &&
+        prev.model?.modelID === model?.modelID
+      ) {
+        return prev
+      }
+
+      const next: SessionAgentResult = { agent: storedAgent, model, variant }
+      prevRef.current = next
+      return next
+    }
+
+    if (messagesLoading) {
+      return { agent: DEFAULT_AGENT, model: undefined, variant: undefined }
+    }
+
+    if (!messages || messages.length === 0) {
+      return { agent: DEFAULT_AGENT, model: undefined, variant: undefined }
+    }
+
+    for (let i = messages.length - 1; i >= 0; i--) {
+      const msgWithParts = messages[i]
+      if (msgWithParts.info.role === 'user') {
+        const userInfo = msgWithParts.info as UserMessage
+        if (userInfo.agent) {
+          const prev = prevRef.current
+          if (
+            prev.agent === userInfo.agent &&
+            prev.variant === userInfo.variant &&
+            prev.model?.providerID === userInfo.model?.providerID &&
+            prev.model?.modelID === userInfo.model?.modelID
+          ) {
+            return prev
+          }
+
+          const next: SessionAgentResult = {
+            agent: userInfo.agent,
+            model: userInfo.model,
+            variant: userInfo.variant,
+          }
+          prevRef.current = next
+          return next
         }
       }
     }
 
-    const prev = prevRef.current
-    if (
-      prev.agent === agent &&
-      prev.variant === variant &&
-      prev.model?.providerID === model?.providerID &&
-      prev.model?.modelID === model?.modelID
-    ) {
-      return prev
+    return { agent: DEFAULT_AGENT, model: undefined, variant: undefined }
+  }, [messages, messagesLoading, storedAgent])
+
+  useEffect(() => {
+    if (result.agent && sessionID) {
+      setAgent(sessionID, result.agent)
     }
+  }, [result.agent, sessionID, setAgent])
 
-    const next: SessionAgentResult = { agent, model, variant }
-    prevRef.current = next
-    return next
-  }, [messages])
+  return result
 }
 
 export function getSessionAgentFromMessages(
   messages: Array<{ role: string; agent?: string }> | undefined
-): string {
+): string | undefined {
   if (!messages || messages.length === 0) {
-    return DEFAULT_AGENT
+    return undefined
   }
 
   for (let i = messages.length - 1; i >= 0; i--) {
     const msg = messages[i]
-    if (msg.role === 'user' && 'agent' in msg) {
-      return (msg.agent as string) || DEFAULT_AGENT
+    if (msg.role === 'user' && 'agent' in msg && msg.agent) {
+      return msg.agent
     }
   }
 
-  return DEFAULT_AGENT
+  return undefined
 }