Copilot commented on code in PR #1540: URL: https://github.com/apache/dubbo-admin/pull/1540#discussion_r3886732363
########## ai/store/store.go: ########## @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package store + +import ( + "context" + "time" + + "github.com/firebase/genkit/go/ai" +) + +// SessionStore persists Session metadata. +type SessionStore interface { + Create(ctx context.Context, session *Session) error + Get(ctx context.Context, sessionID string) (*Session, error) + List(ctx context.Context) ([]*Session, error) + Touch(ctx context.Context, sessionID string, updatedAt time.Time) error + Delete(ctx context.Context, sessionID string) error + DeleteExpired(ctx context.Context, now time.Time) (int, error) +} + +// MessageStore persists conversation Turns and Genkit messages. +type MessageStore interface { + AddHistory(ctx context.Context, sessionID string, messages ...*ai.Message) error + IsEmpty(ctx context.Context, sessionID string) (bool, error) + WindowMemory(ctx context.Context, sessionID string) ([]*ai.Message, error) + AllMemory(ctx context.Context, sessionID string) ([]*ai.Message, error) + NextTurn(ctx context.Context, sessionID string) error Review Comment: The contract addresses the active turn only by `sessionID`, so two concurrent requests (including requests handled by different instances) append to the same turn. The first request to call `NextTurn` finalizes their combined history, and the other then receives `ErrNoActiveTurn`; prompts and responses can therefore be mixed across users' requests. Introduce an interaction/turn identifier (or a session lease covering the whole interaction) and require writes/finalization to target that identifier. ########## ai/component/agent/react/react.go: ########## @@ -101,36 +117,34 @@ func (ra *ReActAgent) Interact(input *schema.UserInput, sessionID string) *agent chans.Send(schema.StreamFinal(final)) chans.Close() - history.NextTurn(sessionID) }() return chans } // newInteraction records the user input into history and returns a session-scoped // context plus a fresh state. -func (ra *ReActAgent) newInteraction(input *schema.UserInput, sessionID string) (context.Context, *state, *memory.HistoryMemory, error) { - history, err := memory.GetHistoryMemory(ra.memoryCtx, memory.ChatHistoryKey) - if err != nil { - return nil, nil, nil, fmt.Errorf("failed to get history from context: %w", err) +func (ra *ReActAgent) newInteraction(parent context.Context, input *schema.UserInput, sessionID string) (context.Context, *state, error) { + if ra.messageStore == nil { + return nil, nil, fmt.Errorf("message store is not configured") + } + if input == nil { + return nil, nil, fmt.Errorf("user input is nil") } // Record the user's message as plain text. The session id travels via - // context (memory.SessionIDKey), so there is no need to wrap the input in a + // context, so there is no need to wrap the input in a // JSON envelope the model would otherwise have to read through. - history.AddHistory(sessionID, ai.NewUserMessage(ai.NewTextPart(input.Content))) + if err := ra.messageStore.AddHistory(parent, sessionID, ai.NewUserMessage(ai.NewTextPart(input.Content))); err != nil { + return nil, nil, fmt.Errorf("failed to record user message: %w", err) + } - ctx := context.WithValue(ra.memoryCtx, memory.SessionIDKey, sessionID) + ctx := context.WithValue(parent, sessionIDContextKey, sessionID) ctx = withCurrentPageContext(ctx, input.Context) - s := &state{Input: input, Session: sessionID, Usage: &ai.GenerationUsage{}} - return ctx, s, history, nil -} - -// GetMemory returns the agent's chat history store, or nil if it cannot be -// resolved from the agent's memory context. -func (ra *ReActAgent) GetMemory() *memory.HistoryMemory { - h, err := memory.GetHistoryMemory(ra.memoryCtx, memory.ChatHistoryKey) - if err != nil { - return nil + s := &state{ + Input: input, + Session: sessionID, + persistCtx: context.WithoutCancel(parent), Review Comment: `context.WithoutCancel` removes both cancellation and the request deadline. A stalled database write can therefore retain this interaction goroutine indefinitely after the client disconnects. Use a cancellation-independent but bounded persistence context, and cancel it when the interaction finishes. ########## ai/store/gorm/store.go: ########## @@ -0,0 +1,512 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package gormstore + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + conversationstore "dubbo-admin-ai/store" + + "github.com/firebase/genkit/go/ai" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +const sessionExpiration = 24 * time.Hour + +// DefaultMaxTurns matches MemorySpec's default conversation limit. +const DefaultMaxTurns = 100 + +// GormStore persists sessions, turns, and messages in a relational database. +// It deliberately does not use database foreign keys: relationship checks and +// deletion ordering are handled explicitly by the store transaction. +type GormStore struct { + db *gorm.DB + limit int +} + +var _ conversationstore.Store = (*GormStore)(nil) + +// NewGormStore creates a store around an already opened Gorm database. The +// optional limit exists for runtime configuration and tests; the default +// matches MemorySpec. +func NewGormStore(db *gorm.DB, limits ...int) (*GormStore, error) { + if db == nil { + return nil, fmt.Errorf("gorm database is nil") + } + // The Store owns relationship validation and deletion ordering. Keep Gorm + // from creating database foreign-key constraints if models gain fields in + // the future. + if db.Config == nil { + db.Config = &gorm.Config{} + } + db.Config.DisableForeignKeyConstraintWhenMigrating = true + limit := DefaultMaxTurns + if len(limits) > 0 && limits[0] > 0 { + limit = limits[0] + } + return &GormStore{db: db, limit: limit}, nil +} + +// Migrate creates or updates the Store tables. Gorm's model associations are +// not declared, so this migration does not create foreign-key constraints. +func (s *GormStore) Migrate(ctx context.Context) error { + if err := s.checkContext(ctx); err != nil { + return err + } + return s.db.WithContext(normalizeContext(ctx)).AutoMigrate( + &SessionModel{}, &TurnModel{}, &MessageModel{}, + ) +} + +// DB returns the underlying database for connection-pool configuration and +// test inspection. Callers must not replace the database instance. +func (s *GormStore) DB() *gorm.DB { return s.db } + +// Close closes the underlying SQL database connection. +func (s *GormStore) Close() error { + if s == nil || s.db == nil { + return nil + } + sqlDB, err := s.db.DB() + if err != nil { + return err + } + return sqlDB.Close() +} + +func (s *GormStore) Create(ctx context.Context, session *conversationstore.Session) error { + if err := s.checkContext(ctx); err != nil { + return err + } + if err := validateSession(session); err != nil { + return err + } + model := sessionModelFromDomain(session) + return s.db.WithContext(normalizeContext(ctx)).Create(&model).Error +} + +func (s *GormStore) Get(ctx context.Context, sessionID string) (*conversationstore.Session, error) { + if err := s.checkContext(ctx); err != nil { + return nil, err + } + var model SessionModel + err := s.db.WithContext(normalizeContext(ctx)).Where("id = ?", sessionID).First(&model).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, conversationstore.ErrSessionNotFound + } + if err != nil { + return nil, err + } + if isExpired(model.UpdatedAt, time.Now()) { + return nil, conversationstore.ErrSessionExpired + } + return sessionDomainFromModel(&model), nil +} + +func (s *GormStore) List(ctx context.Context) ([]*conversationstore.Session, error) { + if err := s.checkContext(ctx); err != nil { + return nil, err + } + var models []SessionModel + cutoff := time.Now().Add(-sessionExpiration) + err := s.db.WithContext(normalizeContext(ctx)). + Where("status = ? AND updated_at >= ?", "active", cutoff). + Find(&models).Error + if err != nil { + return nil, err + } + result := make([]*conversationstore.Session, 0, len(models)) + for i := range models { + result = append(result, sessionDomainFromModel(&models[i])) + } + return result, nil +} + +func (s *GormStore) Touch(ctx context.Context, sessionID string, updatedAt time.Time) error { + if err := s.checkContext(ctx); err != nil { + return err + } + return s.db.WithContext(normalizeContext(ctx)).Transaction(func(tx *gorm.DB) error { + model, err := s.findSession(tx, sessionID, true) + if err != nil { + return err + } + if model.Status != "active" { + return fmt.Errorf("session %q is not active", sessionID) + } + if isExpired(model.UpdatedAt, time.Now()) { + return conversationstore.ErrSessionExpired + } + return tx.Model(&SessionModel{}).Where("id = ?", sessionID).Update("updated_at", updatedAt).Error + }) +} + +func (s *GormStore) Delete(ctx context.Context, sessionID string) error { + if err := s.checkContext(ctx); err != nil { + return err + } + return s.db.WithContext(normalizeContext(ctx)).Transaction(func(tx *gorm.DB) error { + if _, err := s.findSession(tx, sessionID, true); err != nil { + return err + } + return deleteSessionData(tx, sessionID) + }) +} + +func (s *GormStore) DeleteExpired(ctx context.Context, now time.Time) (int, error) { + if err := s.checkContext(ctx); err != nil { + return 0, err + } + cutoff := now.Add(-sessionExpiration) + deleted := 0 + err := s.db.WithContext(normalizeContext(ctx)).Transaction(func(tx *gorm.DB) error { + query := tx.Where("updated_at < ?", cutoff) + if supportsRowLock(tx) { + query = query.Clauses(clause.Locking{Strength: "UPDATE"}) + } + var sessions []SessionModel + if err := query.Find(&sessions).Error; err != nil { + return err + } + for i := range sessions { + // Re-read the row after acquiring the transaction lock. This prevents + // cleanup from deleting a session refreshed by another instance. + current, err := s.findSession(tx, sessions[i].ID, true) + if err != nil { + if errors.Is(err, conversationstore.ErrSessionNotFound) { + continue + } + return err + } + if !isExpired(current.UpdatedAt, now) { + continue + } + if err := deleteSessionData(tx, sessions[i].ID); err != nil { + return err + } + deleted++ + } + return nil + }) + if err != nil { + return 0, err + } + return deleted, nil +} + +func (s *GormStore) AddHistory(ctx context.Context, sessionID string, messages ...*ai.Message) error { + if err := s.checkContext(ctx); err != nil { + return err + } + return s.db.WithContext(normalizeContext(ctx)).Transaction(func(tx *gorm.DB) error { + model, err := s.findSession(tx, sessionID, true) + if err != nil { + return err + } + if model.Status != "active" { + return fmt.Errorf("session %q is not active", sessionID) + } + if isExpired(model.UpdatedAt, time.Now()) { + return conversationstore.ErrSessionExpired + } + + var turn TurnModel + err = tx.Where("session_id = ? AND completed_at IS NULL", sessionID). + Order("id DESC").First(&turn).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + var completedCount int64 + if err := tx.Model(&TurnModel{}). + Where("session_id = ? AND completed_at IS NOT NULL", sessionID). + Count(&completedCount).Error; err != nil { + return err + } + if completedCount >= int64(s.limit) { + return fmt.Errorf("%w: current session's context is full, please create a new session", conversationstore.ErrTurnLimitReached) + } + turn = TurnModel{SessionID: sessionID, CreatedAt: time.Now()} + if err := tx.Create(&turn).Error; err != nil { + return err + } + } else if err != nil { + return err + } + + var last MessageModel + err = tx.Where("turn_id = ?", turn.ID).Order("sequence DESC").First(&last).Error + nextSequence := uint64(0) + if err == nil { + nextSequence = last.Sequence + 1 + } else if !errors.Is(err, gorm.ErrRecordNotFound) { + return err + } + for _, message := range messages { + if message == nil || !supportedRole(message.Role) { + continue + } + payload, err := encodeMessage(message) + if err != nil { + return fmt.Errorf("failed to encode message: %w", err) + } + stored := MessageModel{ + TurnID: turn.ID, + Sequence: nextSequence, + Payload: payload, + CreatedAt: time.Now(), + } + if err := tx.Create(&stored).Error; err != nil { + return err + } + nextSequence++ + } + return nil + }) +} + +func (s *GormStore) IsEmpty(ctx context.Context, sessionID string) (bool, error) { + if err := s.checkContext(ctx); err != nil { + return false, err + } + var turn TurnModel + err := s.db.WithContext(normalizeContext(ctx)). + Where("session_id = ? AND completed_at IS NULL", sessionID). + First(&turn).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return true, nil + } + return false, err +} + +func (s *GormStore) WindowMemory(ctx context.Context, sessionID string) ([]*ai.Message, error) { + if err := s.checkContext(ctx); err != nil { + return nil, err + } + var turn TurnModel + err := s.db.WithContext(normalizeContext(ctx)). + Where("session_id = ? AND completed_at IS NULL", sessionID). + First(&turn).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return s.readTurnMessages(ctx, turn.ID) +} + +func (s *GormStore) AllMemory(ctx context.Context, sessionID string) ([]*ai.Message, error) { + if err := s.checkContext(ctx); err != nil { + return nil, err + } + var active *TurnModel + var activeModel TurnModel + err := s.db.WithContext(normalizeContext(ctx)). + Where("session_id = ? AND completed_at IS NULL", sessionID). + First(&activeModel).Error Review Comment: This active-turn read and the completed-turn query below are separate autocommit snapshots. If `NextTurn` commits between them, the same turn is returned once as active and again as completed, so `AllMemory` duplicates all of its messages. Read the turn set and messages from one consistent snapshot (or one ordered query). ########## ai/store/gorm/store.go: ########## @@ -0,0 +1,512 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package gormstore + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + conversationstore "dubbo-admin-ai/store" + + "github.com/firebase/genkit/go/ai" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +const sessionExpiration = 24 * time.Hour + +// DefaultMaxTurns matches MemorySpec's default conversation limit. +const DefaultMaxTurns = 100 + +// GormStore persists sessions, turns, and messages in a relational database. +// It deliberately does not use database foreign keys: relationship checks and +// deletion ordering are handled explicitly by the store transaction. +type GormStore struct { + db *gorm.DB + limit int +} + +var _ conversationstore.Store = (*GormStore)(nil) + +// NewGormStore creates a store around an already opened Gorm database. The +// optional limit exists for runtime configuration and tests; the default +// matches MemorySpec. +func NewGormStore(db *gorm.DB, limits ...int) (*GormStore, error) { + if db == nil { + return nil, fmt.Errorf("gorm database is nil") + } + // The Store owns relationship validation and deletion ordering. Keep Gorm + // from creating database foreign-key constraints if models gain fields in + // the future. + if db.Config == nil { + db.Config = &gorm.Config{} + } + db.Config.DisableForeignKeyConstraintWhenMigrating = true + limit := DefaultMaxTurns + if len(limits) > 0 && limits[0] > 0 { + limit = limits[0] + } + return &GormStore{db: db, limit: limit}, nil +} + +// Migrate creates or updates the Store tables. Gorm's model associations are +// not declared, so this migration does not create foreign-key constraints. +func (s *GormStore) Migrate(ctx context.Context) error { + if err := s.checkContext(ctx); err != nil { + return err + } + return s.db.WithContext(normalizeContext(ctx)).AutoMigrate( + &SessionModel{}, &TurnModel{}, &MessageModel{}, + ) +} + +// DB returns the underlying database for connection-pool configuration and +// test inspection. Callers must not replace the database instance. +func (s *GormStore) DB() *gorm.DB { return s.db } + +// Close closes the underlying SQL database connection. +func (s *GormStore) Close() error { + if s == nil || s.db == nil { + return nil + } + sqlDB, err := s.db.DB() + if err != nil { + return err + } + return sqlDB.Close() +} + +func (s *GormStore) Create(ctx context.Context, session *conversationstore.Session) error { + if err := s.checkContext(ctx); err != nil { + return err + } + if err := validateSession(session); err != nil { + return err + } + model := sessionModelFromDomain(session) + return s.db.WithContext(normalizeContext(ctx)).Create(&model).Error +} + +func (s *GormStore) Get(ctx context.Context, sessionID string) (*conversationstore.Session, error) { + if err := s.checkContext(ctx); err != nil { + return nil, err + } + var model SessionModel + err := s.db.WithContext(normalizeContext(ctx)).Where("id = ?", sessionID).First(&model).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, conversationstore.ErrSessionNotFound + } + if err != nil { + return nil, err + } + if isExpired(model.UpdatedAt, time.Now()) { + return nil, conversationstore.ErrSessionExpired + } + return sessionDomainFromModel(&model), nil +} + +func (s *GormStore) List(ctx context.Context) ([]*conversationstore.Session, error) { + if err := s.checkContext(ctx); err != nil { + return nil, err + } + var models []SessionModel + cutoff := time.Now().Add(-sessionExpiration) + err := s.db.WithContext(normalizeContext(ctx)). + Where("status = ? AND updated_at >= ?", "active", cutoff). + Find(&models).Error + if err != nil { + return nil, err + } + result := make([]*conversationstore.Session, 0, len(models)) + for i := range models { + result = append(result, sessionDomainFromModel(&models[i])) + } + return result, nil +} + +func (s *GormStore) Touch(ctx context.Context, sessionID string, updatedAt time.Time) error { + if err := s.checkContext(ctx); err != nil { + return err + } + return s.db.WithContext(normalizeContext(ctx)).Transaction(func(tx *gorm.DB) error { + model, err := s.findSession(tx, sessionID, true) + if err != nil { + return err + } + if model.Status != "active" { + return fmt.Errorf("session %q is not active", sessionID) + } + if isExpired(model.UpdatedAt, time.Now()) { + return conversationstore.ErrSessionExpired + } + return tx.Model(&SessionModel{}).Where("id = ?", sessionID).Update("updated_at", updatedAt).Error + }) +} + +func (s *GormStore) Delete(ctx context.Context, sessionID string) error { + if err := s.checkContext(ctx); err != nil { + return err + } + return s.db.WithContext(normalizeContext(ctx)).Transaction(func(tx *gorm.DB) error { + if _, err := s.findSession(tx, sessionID, true); err != nil { + return err + } + return deleteSessionData(tx, sessionID) + }) +} + +func (s *GormStore) DeleteExpired(ctx context.Context, now time.Time) (int, error) { + if err := s.checkContext(ctx); err != nil { + return 0, err + } + cutoff := now.Add(-sessionExpiration) + deleted := 0 + err := s.db.WithContext(normalizeContext(ctx)).Transaction(func(tx *gorm.DB) error { + query := tx.Where("updated_at < ?", cutoff) + if supportsRowLock(tx) { + query = query.Clauses(clause.Locking{Strength: "UPDATE"}) + } + var sessions []SessionModel + if err := query.Find(&sessions).Error; err != nil { + return err + } Review Comment: This query loads and locks every expired session, then keeps all locks until every session's messages and turns are deleted in the same transaction. On a long-lived persistent deployment, hourly cleanup can create an unbounded transaction and block session touches for a prolonged period. Delete in bounded batches and commit between batches. ########## ai/component/server/engine/handlers.go: ########## @@ -116,59 +130,58 @@ func (h *AgentHandler) StreamChat(c *gin.Context) { } } - case <-c.Request.Context().Done(): + case <-requestCtx.Done(): rt.GetLogger().Info("Client disconnected from stream") return default: - if channels.Closed() { - // Drain remaining messages before finishing - rt.GetLogger().Info("Channels closed, draining remaining messages", "session_id", sessionID) - drainLoop: - for { - select { - case feedback, ok = <-channels.UserRespChan: - if !ok { - channels.UserRespChan = nil - break drainLoop - } - if feedback.IsFinal() { - h.MessageDelta(sseHandler, feedback.Final()) - } else if feedback.IsDone() { - if err := sseHandler.HandleContentBlockStop(feedback.Index()); err != nil { - rt.GetLogger().Error("Failed to handle content block stop", "error", err) - } - } else { - if err := sseHandler.HandleText(feedback.Text(), feedback.Index()); err != nil { - rt.GetLogger().Error("Failed to handle text", "error", err) - } - } - case err, ok = <-channels.ErrorChan: - if !ok { - channels.ErrorChan = nil - break drainLoop - } - if err != nil { - sseHandler.HandleError("agent_error", fmt.Sprintf("agent error: %v", err)) - } - default: + if !channels.Closed() { + continue Review Comment: This polls `Channels.closed` from the handler while the agent goroutine writes it in `Close` without synchronization, which is a data race; the `default`/`continue` path also busy-spins whenever no event is ready. Close a dedicated done channel (or the output channels) and select on it so completion is synchronized and the loop blocks while idle. ########## ai/store/memory/store.go: ########## @@ -0,0 +1,431 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package memory + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + "time" + + conversationstore "dubbo-admin-ai/store" + + "github.com/firebase/genkit/go/ai" +) + +// DefaultMaxTurns matches MemorySpec's default conversation limit. +const DefaultMaxTurns = 100 + +const sessionExpiration = 24 * time.Hour + +type turn struct { + id uint64 + createdAt time.Time + userMessages []*ai.Message + modelMessages []*ai.Message + systemMessages []*ai.Message +} + +func (t *turn) messages() []*ai.Message { + messages := make([]*ai.Message, 0, + len(t.systemMessages)+len(t.userMessages)+len(t.modelMessages)) + messages = append(messages, t.systemMessages...) + messages = append(messages, t.userMessages...) + messages = append(messages, t.modelMessages...) + return messages +} + +// turnWindow mirrors the current HistoryMemory window's bounded behavior. The +// end index is intentionally monotonic, so a full window remains full after +// repeated Pop/Push cycles just as it does in the existing implementation. +type turnWindow struct { + limit int + begin int + end int + data []*turn +} + +func newTurnWindow(limit int) *turnWindow { + return &turnWindow{limit: limit, data: make([]*turn, limit+1)} +} + +func (w *turnWindow) empty() bool { + return w.begin == w.end +} + +func (w *turnWindow) full() bool { + return w.end == w.limit +} + +func (w *turnWindow) push(value *turn) bool { + if w.full() { + return false + } + w.data[w.end] = value + w.end++ + return true +} + +func (w *turnWindow) pop() *turn { + if w.empty() { + return nil + } + value := w.data[w.begin] + w.data[w.begin] = nil + w.begin++ + return value +} + +func (w *turnWindow) values() []*turn { + return w.data[w.begin:w.end] +} + +type sessionHistory struct { + window *turnWindow + history []*turn + nextID uint64 +} + +// MemoryStore is the in-process implementation of the conversation Store. +type MemoryStore struct { + mu sync.RWMutex + limit int + sessions map[string]conversationstore.Session + history map[string]*sessionHistory +} + +var _ conversationstore.Store = (*MemoryStore)(nil) + +// NewMemoryStore creates a MemoryStore. The optional limit exists for runtime +// configuration and tests; the default matches MemorySpec. +func NewMemoryStore(limits ...int) *MemoryStore { + limit := DefaultMaxTurns + if len(limits) > 0 && limits[0] > 0 { + limit = limits[0] + } + return &MemoryStore{ + limit: limit, + sessions: make(map[string]conversationstore.Session), + history: make(map[string]*sessionHistory), + } +} + +func (m *MemoryStore) Create(ctx context.Context, session *conversationstore.Session) error { + if err := checkContext(ctx); err != nil { + return err + } + if session == nil { + return errors.New("session is nil") + } + if session.ID == "" { + return errors.New("session id is required") + } + if session.CreatedAt.IsZero() { + return errors.New("session created_at is required") + } + if session.UpdatedAt.IsZero() { + return errors.New("session updated_at is required") + } + if session.Status == "" { + return errors.New("session status is required") + } + + m.mu.Lock() + defer m.mu.Unlock() + if _, exists := m.sessions[session.ID]; exists { + return fmt.Errorf("session %q already exists", session.ID) + } + m.sessions[session.ID] = *session + return nil +} + +func (m *MemoryStore) Get(ctx context.Context, sessionID string) (*conversationstore.Session, error) { + if err := checkContext(ctx); err != nil { + return nil, err + } + + m.mu.RLock() + defer m.mu.RUnlock() + session, exists := m.sessions[sessionID] + if !exists { + return nil, conversationstore.ErrSessionNotFound + } + if isExpired(session.UpdatedAt, time.Now()) { + return nil, conversationstore.ErrSessionExpired + } + copy := session + return ©, nil +} + +func (m *MemoryStore) List(ctx context.Context) ([]*conversationstore.Session, error) { + if err := checkContext(ctx); err != nil { + return nil, err + } + + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]*conversationstore.Session, 0, len(m.sessions)) + now := time.Now() + for _, session := range m.sessions { + if session.Status != "active" || isExpired(session.UpdatedAt, now) { + continue + } + copy := session + result = append(result, ©) + } + return result, nil +} + +func (m *MemoryStore) Touch(ctx context.Context, sessionID string, updatedAt time.Time) error { + if err := checkContext(ctx); err != nil { + return err + } + + m.mu.Lock() + defer m.mu.Unlock() + session, exists := m.sessions[sessionID] + if !exists { + return conversationstore.ErrSessionNotFound + } + if session.Status != "active" { + return fmt.Errorf("session %q is not active", sessionID) + } + if isExpired(session.UpdatedAt, time.Now()) { + return conversationstore.ErrSessionExpired + } + session.UpdatedAt = updatedAt + m.sessions[sessionID] = session + return nil +} + +func (m *MemoryStore) Delete(ctx context.Context, sessionID string) error { + if err := checkContext(ctx); err != nil { + return err + } + + m.mu.Lock() + defer m.mu.Unlock() + if _, exists := m.sessions[sessionID]; !exists { + return conversationstore.ErrSessionNotFound + } + delete(m.sessions, sessionID) + delete(m.history, sessionID) + return nil +} + +func (m *MemoryStore) DeleteExpired(ctx context.Context, now time.Time) (int, error) { + if err := checkContext(ctx); err != nil { + return 0, err + } + + m.mu.Lock() + defer m.mu.Unlock() + deleted := 0 + for sessionID, session := range m.sessions { + if !isExpired(session.UpdatedAt, now) { + continue + } + delete(m.sessions, sessionID) + delete(m.history, sessionID) + deleted++ + } + return deleted, nil +} + +func (m *MemoryStore) AddHistory(ctx context.Context, sessionID string, messages ...*ai.Message) error { + if err := checkContext(ctx); err != nil { + return err + } + + m.mu.Lock() + defer m.mu.Unlock() + if err := m.validateSessionLocked(sessionID); err != nil { + return err + } + + history := m.ensureHistoryLocked(sessionID) + if history.window.empty() { + if len(history.history) >= m.limit { + return fmt.Errorf("%w: current session's context is full, please create a new session", conversationstore.ErrTurnLimitReached) + } + history.nextID++ + if !history.window.push(&turn{id: history.nextID, createdAt: time.Now()}) { + return fmt.Errorf("failed to create active turn: %w", conversationstore.ErrTurnLimitReached) + } + } + active := history.window.values()[len(history.window.values())-1] + for _, message := range messages { + if message == nil || !supportedRole(message.Role) { + continue + } + copy, err := cloneMessage(message) + if err != nil { + return fmt.Errorf("failed to copy message: %w", err) Review Comment: If cloning a later message fails, messages cloned earlier in this loop have already been appended, and a newly created active turn also remains, even though `AddHistory` returns an error. This differs from the transactional GORM backend and can expose partial writes. Clone and classify the entire accepted batch before mutating `history` or creating the active turn. ########## ai/component/server/engine/handlers.go: ########## @@ -53,28 +66,29 @@ func (h *AgentHandler) StreamChat(c *gin.Context) { } sessionID = req.SessionID - // Validate session exists and update activity time - session, err = h.sessionMgr.GetSession(sessionID) - if err != nil { + requestCtx := c.Request.Context() + if _, err = h.sessionMgr.GetSession(requestCtx, sessionID); err != nil { + c.JSON(http.StatusBadRequest, NewErrorResponse("Invalid session ID: "+err.Error())) Review Comment: With the new persistent backend, `GetSession` or the following `TouchSession` can fail because of a database outage, but both branches report every error as a client-side 400 "Invalid session ID". Distinguish `ErrSessionNotFound`/`ErrSessionExpired` from backend failures and return a 5xx response without exposing the raw database error. This issue also appears in the following locations of the same file: - line 209 - line 243 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
