5315e88494
- Enrich Session with Workspace, Label, WaitingSince fields - Protocol types: Request, Response, SessionInfo, SwitchArgs, LabelArgs - SessionRegistry with WaitingSince transition tracking - LabelStore with JSON file persistence
142 lines
3.1 KiB
Go
142 lines
3.1 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"os"
|
|
"path/filepath"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// TrackedSession holds a session's info alongside state-transition tracking.
|
|
type TrackedSession struct {
|
|
Info SessionInfo
|
|
PrevState string
|
|
}
|
|
|
|
// SessionRegistry maintains an in-memory index of active sessions.
|
|
// It tracks WaitingSince transitions: when a session moves to "Needs Input",
|
|
// the timestamp is recorded; when it leaves that state, it is cleared.
|
|
type SessionRegistry struct {
|
|
mu sync.RWMutex
|
|
sessions map[string]*TrackedSession
|
|
}
|
|
|
|
func NewRegistry() *SessionRegistry {
|
|
return &SessionRegistry{
|
|
sessions: make(map[string]*TrackedSession),
|
|
}
|
|
}
|
|
|
|
// Update adds or refreshes a session in the registry.
|
|
// WaitingSince is set when transitioning to "Needs Input" and cleared otherwise.
|
|
func (r *SessionRegistry) Update(info SessionInfo) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
existing, ok := r.sessions[info.SessionID]
|
|
if !ok {
|
|
existing = &TrackedSession{}
|
|
r.sessions[info.SessionID] = existing
|
|
}
|
|
|
|
isWaiting := info.State == "Needs Input"
|
|
wasWaiting := existing.PrevState == "Needs Input"
|
|
|
|
if isWaiting && !wasWaiting {
|
|
now := time.Now()
|
|
info.WaitingSince = &now
|
|
} else if isWaiting && wasWaiting {
|
|
// Keep existing timestamp
|
|
info.WaitingSince = existing.Info.WaitingSince
|
|
}
|
|
// If not waiting, WaitingSince stays nil (default)
|
|
|
|
existing.Info = info
|
|
existing.PrevState = info.State
|
|
}
|
|
|
|
// List returns a snapshot of all tracked sessions.
|
|
func (r *SessionRegistry) List() []SessionInfo {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
|
|
result := make([]SessionInfo, 0, len(r.sessions))
|
|
for _, ts := range r.sessions {
|
|
result = append(result, ts.Info)
|
|
}
|
|
return result
|
|
}
|
|
|
|
// RemoveStale removes sessions whose SessionID is not in activeIDs.
|
|
func (r *SessionRegistry) RemoveStale(activeIDs map[string]bool) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
for id := range r.sessions {
|
|
if !activeIDs[id] {
|
|
delete(r.sessions, id)
|
|
}
|
|
}
|
|
}
|
|
|
|
// LabelStore persists user-assigned labels in a JSON file.
|
|
type LabelStore struct {
|
|
mu sync.RWMutex
|
|
labels map[string]string
|
|
path string
|
|
}
|
|
|
|
// NewLabelStore loads labels from path. Returns an empty store if the file does not exist.
|
|
func NewLabelStore(path string) (*LabelStore, error) {
|
|
ls := &LabelStore{
|
|
labels: make(map[string]string),
|
|
path: path,
|
|
}
|
|
|
|
data, err := os.ReadFile(path)
|
|
if err != nil {
|
|
if os.IsNotExist(err) {
|
|
return ls, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
|
|
if err := json.Unmarshal(data, &ls.labels); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return ls, nil
|
|
}
|
|
|
|
// Set assigns a label to a session and persists to disk.
|
|
func (ls *LabelStore) Set(sessionID, label string) error {
|
|
ls.mu.Lock()
|
|
defer ls.mu.Unlock()
|
|
|
|
ls.labels[sessionID] = label
|
|
return ls.save()
|
|
}
|
|
|
|
// Get returns the label for a session, or empty string if not found.
|
|
func (ls *LabelStore) Get(sessionID string) string {
|
|
ls.mu.RLock()
|
|
defer ls.mu.RUnlock()
|
|
|
|
return ls.labels[sessionID]
|
|
}
|
|
|
|
func (ls *LabelStore) save() error {
|
|
data, err := json.MarshalIndent(ls.labels, "", " ")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
dir := filepath.Dir(ls.path)
|
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
|
return err
|
|
}
|
|
|
|
return os.WriteFile(ls.path, data, 0o644)
|
|
}
|