claude-web-chat/sessions.go

package main

import (
	"fmt"
	"log"
	"sync"
	"time"
)

const (
	maxIdleTime     = 30 * time.Minute
	maxBufferSize   = 500
	cleanupInterval = 5 * time.Minute
)

// ChatMessage represents a single message in the chat history.
type ChatMessage struct {
	Role      string    // "user", "assistant", "system", "tool"
	Content   string    // HTML fragment
	Timestamp time.Time
}

// Subscriber receives broadcast messages via a channel.
type Subscriber struct {
	Ch chan string
	ID string
}

// ChatSession holds the state for a single chat with a project.
type ChatSession struct {
	ID         string
	ProjectDir string
	Process    *CLIProcess
	Buffer     []ChatMessage
	LastActive time.Time
	Streaming  bool // true while Claude is generating

	subscribers map[string]*Subscriber
	mu          sync.Mutex
}

// AddMessage appends a message to the buffer and broadcasts to subscribers.
func (cs *ChatSession) AddMessage(msg ChatMessage) {
	cs.mu.Lock()
	defer cs.mu.Unlock()

	cs.Buffer = append(cs.Buffer, msg)
	if len(cs.Buffer) > maxBufferSize {
		cs.Buffer = cs.Buffer[len(cs.Buffer)-maxBufferSize:]
	}
	cs.LastActive = time.Now()

	// Broadcast to all subscribers
	for id, sub := range cs.subscribers {
		select {
		case sub.Ch <- msg.Content:
		default:
			log.Printf("Subscriber %s buffer full, skipping", id)
		}
	}
}

// GetBuffer returns a copy of the message buffer.
func (cs *ChatSession) GetBuffer() []ChatMessage {
	cs.mu.Lock()
	defer cs.mu.Unlock()

	buf := make([]ChatMessage, len(cs.Buffer))
	copy(buf, cs.Buffer)
	return buf
}

// Subscribe adds a new subscriber and returns it.
func (cs *ChatSession) Subscribe(id string) *Subscriber {
	cs.mu.Lock()
	defer cs.mu.Unlock()

	sub := &Subscriber{
		Ch: make(chan string, 100),
		ID: id,
	}
	if cs.subscribers == nil {
		cs.subscribers = make(map[string]*Subscriber)
	}
	cs.subscribers[id] = sub
	return sub
}

// Unsubscribe removes a subscriber.
func (cs *ChatSession) Unsubscribe(id string) {
	cs.mu.Lock()
	defer cs.mu.Unlock()

	if sub, ok := cs.subscribers[id]; ok {
		close(sub.Ch)
		delete(cs.subscribers, id)
	}
}

// SubscriberCount returns the number of active subscribers.
func (cs *ChatSession) SubscriberCount() int {
	cs.mu.Lock()
	defer cs.mu.Unlock()
	return len(cs.subscribers)
}

// ChatSessionManager manages all active chat sessions.
type ChatSessionManager struct {
	sessions map[string]*ChatSession
	mu       sync.RWMutex
	stopCh   chan struct{}
}

func NewChatSessionManager() *ChatSessionManager {
	csm := &ChatSessionManager{
		sessions: make(map[string]*ChatSession),
		stopCh:   make(chan struct{}),
	}
	go csm.cleanup()
	return csm
}

// GetOrCreate returns an existing session or creates a new one.
func (csm *ChatSessionManager) GetOrCreate(sessionID, projectDir string) (*ChatSession, bool, error) {
	csm.mu.Lock()
	defer csm.mu.Unlock()

	if sess, ok := csm.sessions[sessionID]; ok {
		sess.mu.Lock()
		sess.LastActive = time.Now()
		sess.mu.Unlock()
		return sess, false, nil
	}

	// Spawn new CLI process
	proc, err := SpawnCLI(projectDir)
	if err != nil {
		return nil, false, fmt.Errorf("spawn CLI: %w", err)
	}

	sess := &ChatSession{
		ID:          sessionID,
		ProjectDir:  projectDir,
		Process:     proc,
		Buffer:      make([]ChatMessage, 0),
		LastActive:  time.Now(),
		subscribers: make(map[string]*Subscriber),
	}

	csm.sessions[sessionID] = sess
	return sess, true, nil
}

// Get returns an existing session or nil.
func (csm *ChatSessionManager) Get(sessionID string) *ChatSession {
	csm.mu.RLock()
	defer csm.mu.RUnlock()
	return csm.sessions[sessionID]
}

// Remove terminates and removes a session.
func (csm *ChatSessionManager) Remove(sessionID string) {
	csm.mu.Lock()
	sess, ok := csm.sessions[sessionID]
	if ok {
		delete(csm.sessions, sessionID)
	}
	csm.mu.Unlock()

	if ok && sess.Process != nil {
		sess.Process.Close()
	}
}

// Stop shuts down the manager and all sessions.
func (csm *ChatSessionManager) Stop() {
	close(csm.stopCh)

	csm.mu.Lock()
	for id, sess := range csm.sessions {
		if sess.Process != nil {
			sess.Process.Close()
		}
		delete(csm.sessions, id)
	}
	csm.mu.Unlock()
}

// Count returns the number of active sessions.
func (csm *ChatSessionManager) Count() int {
	csm.mu.RLock()
	defer csm.mu.RUnlock()
	return len(csm.sessions)
}

func (csm *ChatSessionManager) cleanup() {
	ticker := time.NewTicker(cleanupInterval)
	defer ticker.Stop()

	for {
		select {
		case <-csm.stopCh:
			return
		case <-ticker.C:
			csm.cleanupIdle()
		}
	}
}

func (csm *ChatSessionManager) cleanupIdle() {
	csm.mu.Lock()
	var toRemove []string
	for id, sess := range csm.sessions {
		sess.mu.Lock()
		idle := time.Since(sess.LastActive) > maxIdleTime
		sess.mu.Unlock()
		if idle {
			toRemove = append(toRemove, id)
		}
	}
	csm.mu.Unlock()

	for _, id := range toRemove {
		log.Printf("Cleaning up idle session: %s", id)
		csm.Remove(id)
	}
}