1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182

package main

import (
	"sync"
	"time"

	"github.com/google/uuid"
)

// TokenUsageCounter tracks usage of a specific token within a session.
// This allows enforcing limits like "max_full_reads: 5".
type TokenUsageCounter struct {
	MessagesRead int
	SearchesDone int
}

// SessionInfo holds information about an active session
type SessionInfo struct {
	CreatedAt  time.Time
	ExpiresAt  time.Time
	ClientInfo string // Optional: IP address, user-agent, etc.
	TokenUsage map[string]*TokenUsageCounter // key: token signature
}

// SessionManager manages active sessions locally (no network calls needed)
type SessionManager struct {
	sessions map[string]SessionInfo
	mu       sync.RWMutex

	// Optional: cleanup ticker
	cleanupTicker *time.Ticker
	stopCleanup   chan bool
}

// NewSessionManager creates a new session manager
func NewSessionManager() *SessionManager {
	sm := &SessionManager{
		sessions:    make(map[string]SessionInfo),
		stopCleanup: make(chan bool),
	}

	// Start background cleanup goroutine
	sm.cleanupTicker = time.NewTicker(5 * time.Minute)
	go sm.cleanupLoop()

	return sm
}

// CreateSession creates a new session with the given TTL
// Returns the session UUID and expiration time
func (sm *SessionManager) CreateSession(ttl time.Duration) (string, time.Time) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	sessionID := uuid.New().String()
	expiresAt := time.Now().Add(ttl)

	sm.sessions[sessionID] = SessionInfo{
		CreatedAt: time.Now(),
		ExpiresAt: expiresAt,
	}

	return sessionID, expiresAt
}

// IsActive checks if a session is currently active (exists and not expired)
func (sm *SessionManager) IsActive(sessionID string) bool {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	info, exists := sm.sessions[sessionID]
	if !exists {
		return false
	}

	// Check if expired
	return time.Now().Before(info.ExpiresAt)
}

// GetExpiry returns the expiration time for a session
func (sm *SessionManager) GetExpiry(sessionID string) (time.Time, bool) {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	info, exists := sm.sessions[sessionID]
	if !exists {
		return time.Time{}, false
	}

	return info.ExpiresAt, true
}

// EndSession removes a session, immediately invalidating all its tokens
func (sm *SessionManager) EndSession(sessionID string) bool {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	_, existed := sm.sessions[sessionID]
	delete(sm.sessions, sessionID)
	return existed
}

// cleanupLoop runs in the background, removing expired sessions
func (sm *SessionManager) cleanupLoop() {
	for {
		select {
		case <-sm.cleanupTicker.C:
			sm.cleanupExpired()
		case <-sm.stopCleanup:
			return
		}
	}
}

// cleanupExpired removes all expired sessions
func (sm *SessionManager) cleanupExpired() {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	now := time.Now()
	for id, info := range sm.sessions {
		if now.After(info.ExpiresAt) {
			delete(sm.sessions, id)
		}
	}
}

// IncrementMessageReads increments the message read counter for a specific token
// within a session. Returns the new count after incrementing.
// This is used to enforce max_full_reads limits on tokens.
func (sm *SessionManager) IncrementMessageReads(sessionID, tokenSig string) int {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	session, exists := sm.sessions[sessionID]
	if !exists {
		return 0
	}

	// Initialize TokenUsage map if needed
	if session.TokenUsage == nil {
		session.TokenUsage = make(map[string]*TokenUsageCounter)
	}

	// Initialize counter for this token if needed
	if session.TokenUsage[tokenSig] == nil {
		session.TokenUsage[tokenSig] = &TokenUsageCounter{}
	}

	// Increment and save back to map
	session.TokenUsage[tokenSig].MessagesRead++
	sm.sessions[sessionID] = session

	return session.TokenUsage[tokenSig].MessagesRead
}

// GetMessageReads returns the current message read count for a token in a session.
// Returns 0 if the session or token doesn't exist.
func (sm *SessionManager) GetMessageReads(sessionID, tokenSig string) int {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	session, exists := sm.sessions[sessionID]
	if !exists || session.TokenUsage == nil {
		return 0
	}

	counter := session.TokenUsage[tokenSig]
	if counter == nil {
		return 0
	}

	return counter.MessagesRead
}

// Stop stops the cleanup goroutine
func (sm *SessionManager) Stop() {
	if sm.cleanupTicker != nil {
		sm.cleanupTicker.Stop()
	}
	close(sm.stopCleanup)
}