package agent
import (
"context"
"fmt"
"strings"
"time"
"github.com/charmbracelet/bubbles/textarea"
"github.com/charmbracelet/bubbles/viewport"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/glamour"
"github.com/charmbracelet/lipgloss"
)
// Message types for Bubbletea updates
type statusUpdateMsg struct {
message string
}
type agentResponseMsg struct {
content string
}
type agentThinkingMsg struct {
thinking bool
}
type agentToolCallMsg struct {
toolName string
toolArgs string
}
type thinkingTickMsg struct{}
var (
headerStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color("240")).
Padding(0, 1)
inputPromptStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color("240"))
inputCursorStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color("255"))
inputBorderStyle = lipgloss.NewStyle().
Border(lipgloss.NormalBorder()).
BorderForeground(lipgloss.Color("240")).
Padding(0, 1)
thinkingStyle = lipgloss.NewStyle().
Foreground(lipgloss.Color("33")).
Bold(true)
)
// AgentUI represents the main TUI model
type agentUI struct {
viewport viewport.Model
textInput textarea.Model
width int
height int
statusMessage string
isThinking bool
messages []uiMessage
ready bool
agent Agent
session Session
config AgentUIConfig
thinkingFrame int
inputHistory []string
historyIndex int
currentInput string
}
// Message represents a chat message
type uiMessage struct {
Role string // "user", "agent", "system"
Content string
Timestamp time.Time
}
// AgentUIConfig defines the configuration for the UI
type AgentUIConfig struct {
Width int
Height int
InitialSystemMessage string
TextInputPlaceholder string
TitleText string
MaxHistory int
// Only for informational purposes.
ModelName string
ModelVendor string
ModelFast bool
}
// DefaultAgentUIConfig returns the opinionated default configuration for the UI
func DefaultAgentUIConfig() AgentUIConfig {
return AgentUIConfig{
Width: 80,
Height: 20,
MaxHistory: 50,
InitialSystemMessage: "Security Agent initialized",
TextInputPlaceholder: "Ask me anything...",
TitleText: "Security Agent",
}
}
// NewAgentUI creates a new agent UI instance
func NewAgentUI(agent Agent, session Session, config AgentUIConfig) *agentUI {
vp := viewport.New(config.Width, config.Height)
ta := textarea.New()
ta.Placeholder = ""
ta.Focus()
ta.SetHeight(1)
ta.SetWidth(80)
ta.CharLimit = 1000
ta.ShowLineNumbers = false
ui := &agentUI{
viewport: vp,
textInput: ta,
statusMessage: "",
messages: []uiMessage{},
agent: agent,
session: session,
config: config,
thinkingFrame: 0,
inputHistory: []string{},
historyIndex: -1,
currentInput: "",
}
ui.addSystemMessage(config.InitialSystemMessage)
return ui
}
// Init implements the tea.Model interface
func (m *agentUI) Init() tea.Cmd {
return tea.Batch(
textarea.Blink,
m.tickThinking(),
)
}
// Update implements the tea.Model interface
func (m *agentUI) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
var cmds []tea.Cmd
var cmd tea.Cmd
switch msg := msg.(type) {
case tea.KeyMsg:
switch msg.Type {
case tea.KeyCtrlC, tea.KeyEsc:
return m, tea.Quit
case tea.KeyEnter:
if m.textInput.Focused() && !m.isThinking {
// Handle user input only if agent is not in thinking mode
userInput := strings.TrimSpace(m.textInput.Value())
if userInput != "" {
// Add to history and reset navigation
m.addToHistory(userInput)
// Add the input to the message list and reset user input field
m.addUserMessage(userInput)
m.resetInputField()
// Check if it's a slash command
if strings.HasPrefix(userInput, "/") {
// Handle slash command
cmd := m.handleSlashCommand(userInput)
if cmd != nil {
cmds = append(cmds, cmd)
}
} else {
// Execute agent query
cmds = append(cmds,
m.setThinking(true),
m.executeAgentQuery(userInput),
)
}
}
}
case tea.KeyTab:
// Switch focus between input and viewport, but not while agent is thinking
if !m.isThinking {
if m.textInput.Focused() {
m.textInput.Blur()
} else {
m.textInput.Focus()
cmds = append(cmds, textarea.Blink)
}
}
case tea.KeyUp, tea.KeyDown:
if m.textInput.Focused() && !m.isThinking {
// Navigate input history when text input is focused
var direction int
if msg.Type == tea.KeyUp {
direction = 1 // Go back in history
} else {
direction = -1 // Go forward in history
}
historyEntry := m.navigateHistory(direction)
m.textInput.SetValue(historyEntry)
m.textInput.CursorEnd()
} else if !m.textInput.Focused() {
// Allow scrolling in viewport when not focused on text input
m.viewport, cmd = m.viewport.Update(msg)
cmds = append(cmds, cmd)
}
case tea.KeyPgUp, tea.KeyPgDown:
// Allow scrolling in viewport when not focused on text input
if !m.textInput.Focused() {
m.viewport, cmd = m.viewport.Update(msg)
cmds = append(cmds, cmd)
}
case tea.KeyHome:
if !m.textInput.Focused() {
m.viewport.GotoTop()
}
case tea.KeyEnd:
if !m.textInput.Focused() {
m.viewport.GotoBottom()
}
}
case tea.WindowSizeMsg:
// Handle window resize
m.width = msg.Width
m.height = msg.Height
// Calculate dimensions for minimal UI
headerHeight := 2 // Header + blank line
inputHeight := 2 // Input area + status
spacing := 1 // Bottom spacing
// Calculate viewport dimensions to maximize output area
viewportHeight := m.height - headerHeight - inputHeight - spacing
// Ensure minimum height
if viewportHeight < 10 {
viewportHeight = 10
}
// Full width utilization
viewportWidth := m.width
// Ensure minimum width
if viewportWidth < 50 {
viewportWidth = 50
}
m.viewport.Width = viewportWidth
m.viewport.Height = viewportHeight
m.textInput.SetWidth(m.width - 3)
// Update content when dimensions change
m.viewport.SetContent(m.renderMessages())
if !m.ready {
m.ready = true
}
case statusUpdateMsg:
m.statusMessage = msg.message
case agentThinkingMsg:
m.isThinking = msg.thinking
// When agent starts thinking, blur the input
if m.isThinking {
m.resetInputField()
m.textInput.Blur()
m.thinkingFrame = 0
cmds = append(cmds, m.tickThinking())
} else {
// Re-focus input when thinking stops
m.textInput.Focus()
cmds = append(cmds, textarea.Blink)
}
case agentResponseMsg:
m.addAgentMessage(msg.content)
cmds = append(cmds, m.setThinking(false))
case agentToolCallMsg:
m.addToolCallMessage(fmt.Sprintf("🔧 %s", msg.toolName), msg.toolArgs)
case thinkingTickMsg:
if m.isThinking {
m.thinkingFrame = (m.thinkingFrame + 1) % 4
cmds = append(cmds, m.tickThinking())
}
}
// Update child components
m.viewport, cmd = m.viewport.Update(msg)
cmds = append(cmds, cmd)
// Only update text input if not thinking
if !m.isThinking {
m.textInput, cmd = m.textInput.Update(msg)
cmds = append(cmds, cmd)
}
return m, tea.Batch(cmds...)
}
// View implements the tea.Model interface
func (m *agentUI) View() string {
if !m.ready {
return "Loading..."
}
if m.width == 0 || m.height == 0 {
return "Initializing..."
}
modelAbility := "fast"
if !m.config.ModelFast {
modelAbility = "slow"
}
modelStatusLine := fmt.Sprintf("%s/%s (%s)", m.config.ModelVendor, m.config.ModelName, modelAbility)
header := headerStyle.Render(fmt.Sprintf("%s %s", m.config.TitleText, modelStatusLine))
content := m.viewport.View()
var thinkingIndicator string
if m.isThinking {
thinkingFrames := []string{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"}
spinner := thinkingFrames[m.thinkingFrame%len(thinkingFrames)]
thinkingIndicator = thinkingStyle.Render(fmt.Sprintf("%s thinking...", spinner))
}
var inputArea string
userInput := m.textInput.Value()
cursor := ""
if m.textInput.Focused() && !m.isThinking {
cursor = inputCursorStyle.Render("▊")
}
inputContent := fmt.Sprintf("%s%s%s", inputPromptStyle.Render("> "), userInput, cursor)
inputArea = inputBorderStyle.Width(m.width - 2).Render(inputContent)
statusLine := inputPromptStyle.Render(fmt.Sprintf("** %s | ctrl+c to exit", modelStatusLine))
var components []string
components = append(components, header, "", content, "")
if thinkingIndicator != "" {
components = append(components, thinkingIndicator)
}
components = append(components, inputArea, statusLine)
return lipgloss.JoinVertical(lipgloss.Left, components...)
}
func (m *agentUI) resetInputField() {
m.textInput.Reset()
m.textInput.SetValue("")
m.textInput.CursorStart()
}
func (m *agentUI) addUserMessage(content string) {
m.messages = append(m.messages, uiMessage{
Role: "user",
Content: content,
Timestamp: time.Now(),
})
m.viewport.SetContent(m.renderMessages())
m.viewport.GotoBottom()
}
func (m *agentUI) addAgentMessage(content string) {
m.messages = append(m.messages, uiMessage{
Role: "agent",
Content: content,
Timestamp: time.Now(),
})
m.viewport.SetContent(m.renderMessages())
m.viewport.GotoBottom()
}
func (m *agentUI) addSystemMessage(content string) {
m.messages = append(m.messages, uiMessage{
Role: "system",
Content: content,
Timestamp: time.Now(),
})
m.viewport.SetContent(m.renderMessages())
m.viewport.GotoBottom()
}
func (m *agentUI) addToolCallMessage(toolName string, toolArgs string) {
content := fmt.Sprintf(" %s", toolName)
if toolArgs != "" && toolArgs != "{}" {
content += fmt.Sprintf("\n └─ %s", toolArgs)
}
m.messages = append(m.messages, uiMessage{
Role: "tool",
Content: content,
Timestamp: time.Now(),
})
m.viewport.SetContent(m.renderMessages())
m.viewport.GotoBottom()
}
// renderMessages formats all messages for display
func (m *agentUI) renderMessages() string {
var rendered []string
rendered = append(rendered, "", "")
contentWidth := m.viewport.Width - 2 // Account for internal padding
if contentWidth < 40 {
contentWidth = 40
}
r, err := glamour.NewTermRenderer(
glamour.WithStandardStyle("notty"),
glamour.WithWordWrap(contentWidth),
)
if err != nil {
r = nil
}
for _, msg := range m.messages {
timestamp := msg.Timestamp.Format("15:04:05")
switch msg.Role {
case "user":
userHeaderStyle := lipgloss.NewStyle().
Foreground(lipgloss.Color("86")).
Bold(true).
Border(lipgloss.NormalBorder(), false, false, false, true).
BorderForeground(lipgloss.Color("86")).
Padding(0, 1)
userContentStyle := lipgloss.NewStyle().
Foreground(lipgloss.Color("255")).
Padding(0, 2)
rendered = append(rendered,
userHeaderStyle.Render(fmt.Sprintf("[%s] → You:", timestamp)),
userContentStyle.Render(msg.Content),
"",
)
case "agent":
var content string
if r != nil {
renderedMarkdown, err := r.Render(msg.Content)
if err == nil {
content = strings.TrimSpace(renderedMarkdown)
} else {
content = msg.Content // Fallback to plain text
}
} else {
content = msg.Content
}
agentHeaderStyle := lipgloss.NewStyle().
Foreground(lipgloss.Color("39")).
Bold(true).
Border(lipgloss.NormalBorder(), false, false, false, true).
BorderForeground(lipgloss.Color("39")).
Padding(0, 1)
agentContentStyle := lipgloss.NewStyle().
Foreground(lipgloss.Color("255")).
Padding(0, 2)
rendered = append(rendered,
agentHeaderStyle.Render(fmt.Sprintf("[%s] ← Agent:", timestamp)),
agentContentStyle.Render(content),
"",
)
case "system":
systemStyle := lipgloss.NewStyle().
Foreground(lipgloss.Color("241")).
Italic(true).
Border(lipgloss.NormalBorder(), false, false, false, true).
BorderForeground(lipgloss.Color("241")).
Padding(0, 1)
rendered = append(rendered,
systemStyle.Render(fmt.Sprintf("[%s] ℹ %s", timestamp, msg.Content)),
"",
)
case "tool":
toolStyle := lipgloss.NewStyle().
Foreground(lipgloss.Color("245")).
Italic(true).
Faint(true).
Border(lipgloss.NormalBorder(), false, false, false, true).
BorderForeground(lipgloss.Color("245")).
Padding(0, 1)
rendered = append(rendered,
toolStyle.Render(fmt.Sprintf("[%s] %s", timestamp, msg.Content)),
"",
)
}
}
rendered = append(rendered, "", "")
return strings.Join(rendered, "\n")
}
func (m *agentUI) updateStatus(message string) tea.Cmd {
return func() tea.Msg {
return statusUpdateMsg{message: message}
}
}
func (m *agentUI) setThinking(thinking bool) tea.Cmd {
return func() tea.Msg {
return agentThinkingMsg{thinking: thinking}
}
}
func (m *agentUI) executeAgentQuery(userInput string) tea.Cmd {
return func() tea.Msg {
ctx := context.Background()
input := Input{
Query: userInput,
}
toolCallHook := func(_ context.Context, _ Session, _ Input, toolName string, toolArgs string) error {
m.Update(agentToolCallMsg{toolName: toolName, toolArgs: toolArgs})
return nil
}
output, err := m.agent.Execute(ctx, m.session, input, WithToolCallHook(toolCallHook))
if err != nil {
return agentResponseMsg{
content: fmt.Sprintf("❌ **Error**\n\nSorry, I encountered an error while processing your query:\n\n%s", err.Error()),
}
}
return agentResponseMsg{content: output.Answer}
}
}
// StartUI starts the TUI application with the default configuration
func StartUI(agent Agent, session Session) error {
config := DefaultAgentUIConfig()
config.InitialSystemMessage = ""
return StartUIWithConfig(agent, session, config)
}
// StartUIWithConfig starts the TUI application with the provided configuration
func StartUIWithConfig(agent Agent, session Session, config AgentUIConfig) error {
ui := NewAgentUI(agent, session, config)
p := tea.NewProgram(
ui,
tea.WithAltScreen(),
tea.WithMouseCellMotion(),
)
_, err := p.Run()
return err
}
func (m *agentUI) tickThinking() tea.Cmd {
return tea.Tick(150*time.Millisecond, func(time.Time) tea.Msg {
return thinkingTickMsg{}
})
}
// handleSlashCommand processes commands that start with '/'
func (m *agentUI) handleSlashCommand(command string) tea.Cmd {
switch command {
case "/exit":
m.addSystemMessage("Goodbye! Exiting gracefully...")
return tea.Quit
default:
m.addSystemMessage(fmt.Sprintf("Unknown command: %s", command))
return nil
}
}
// addToHistory adds input to history buffer with a maximum of 50 entries
func (m *agentUI) addToHistory(input string) {
// Don't add empty strings or duplicates of the last entry
if input == "" || (len(m.inputHistory) > 0 && m.inputHistory[len(m.inputHistory)-1] == input) {
return
}
m.inputHistory = append(m.inputHistory, input)
// Keep only the last maxHistory entries
if len(m.inputHistory) > m.config.MaxHistory {
m.inputHistory = m.inputHistory[len(m.inputHistory)-m.config.MaxHistory:]
}
// Reset history navigation
m.historyIndex = -1
m.currentInput = ""
}
// navigateHistory moves through input history and returns the selected entry
func (m *agentUI) navigateHistory(direction int) string {
if len(m.inputHistory) == 0 {
return ""
}
// Save current input when starting navigation
if m.historyIndex == -1 {
m.currentInput = m.textInput.Value()
}
// Calculate new index
newIndex := m.historyIndex + direction
// Handle boundaries
if newIndex < -1 {
newIndex = -1
} else if newIndex >= len(m.inputHistory) {
newIndex = len(m.inputHistory) - 1
}
m.historyIndex = newIndex
// Return the appropriate entry
if m.historyIndex == -1 {
return m.currentInput
}
return m.inputHistory[len(m.inputHistory)-1-m.historyIndex]
}
