mirror of
https://github.com/mudler/LocalAI.git
synced 2026-11-05 20:13:23 -05:00
fix(studio): stop background work before shutdown and wait for tool results in specs
The connector spec "keeps asking when the setting is ..." read the tool
results as soon as the bot had an answer. The harness records a result
after it emitted the turn events, so under load the answer could come
first and the list was empty. The spec now waits for the result.
Shutdown left two background writers running: the file text indexer loop
ran on the application context, and the chat search backfill had no stop.
A test that removed its temp directory right after Shutdown could meet a
write into the "textindex" folder ("directory not empty"). The indexer
loop now has its own context and Shutdown waits for it, and the chat
store gets a Close that stops the backfill and waits.
Assisted-by: Claude:claude-sonnet-5-5
This commit is contained in:
5 files changed
+53
-5
No files matched your search
@@ -203,6 +203,7 @@ type Application struct {
|
||||
studioFilesIndex io.Closer // closes the workspace index: ends pending blob removals, frees the JSON folder
|
||||
studioFilesTemp string // blob cache folder of a cluster with no data path; removed at shutdown
|
||||
studioTextIndex *textindex.Indexer // the text of the documents in the files; nil when files are off
|
||||
studioTextIndexStop func() // stops the indexer's loop and waits for it; nil when it does not run
|
||||
authDB *gorm.DB
|
||||
metricsService *monitoring.LocalAIMetricsService
|
||||
statsRecorder *billing.Recorder
|
||||
@@ -727,6 +728,13 @@ func (a *Application) Shutdown() error {
|
||||
if a.studioTempDir != "" {
|
||||
_ = os.RemoveAll(a.studioTempDir)
|
||||
}
|
||||
// Before the folders go: a pass in flight would write into them.
|
||||
if a.studioTextIndexStop != nil {
|
||||
a.studioTextIndexStop()
|
||||
}
|
||||
if c, ok := a.studioStore.(io.Closer); ok {
|
||||
_ = c.Close() // ends the background indexing of earlier chats
|
||||
}
|
||||
if a.studioFilesIndex != nil {
|
||||
_ = a.studioFilesIndex.Close()
|
||||
}
|
||||
|
||||
@@ -160,7 +160,13 @@ func (a *Application) initStudioTextIndex(f *files.Service, root string) {
|
||||
if cfg.Distributed.ChatWorker {
|
||||
return
|
||||
}
|
||||
go ix.Run(a.backgroundContext())
|
||||
ctx, cancel := context.WithCancel(a.backgroundContext())
|
||||
done := make(chan struct{})
|
||||
a.studioTextIndexStop = func() { cancel(); <-done }
|
||||
go func() {
|
||||
defer close(done)
|
||||
ix.Run(ctx)
|
||||
}()
|
||||
}
|
||||
|
||||
// syncAddonSkills lets the add-on owner keep the add-on skill layer in step
|
||||
|
||||
@@ -476,6 +476,9 @@ var _ = Describe("how a memory note is saved", func() {
|
||||
Expect(r.conns.Ingest(ctx, "telegram", contract.ConnectorInbound{BindingID: b.ID, ConnectorMessage: contract.ConnectorMessage{
|
||||
Chat: "11", Direct: true, SenderID: "11", SenderName: "S", Text: "remember"}})).To(Succeed())
|
||||
waitFor("the answer", func() bool { return len(bot.answers()) >= 1 })
|
||||
// The harness records a tool result after the turn's events went
|
||||
// out, so the answer can reach the bot first: wait for the result.
|
||||
waitFor("the tool result", func() bool { return len(r.h.ResultsFor("propose_change")) >= 1 })
|
||||
res := r.h.ResultsFor("propose_change")
|
||||
Expect(res).To(HaveLen(1))
|
||||
Expect(res[0].Err).To(BeNil(), "%+v", res)
|
||||
|
||||
@@ -139,6 +139,22 @@ func (pendingRow) TableName() string { return "studio_pending" }
|
||||
type gormStore struct {
|
||||
db *gorm.DB
|
||||
text textIndex
|
||||
|
||||
// The backfill of the chat search index runs in the background: Close
|
||||
// stops it and waits, so nothing writes once the owner of the database
|
||||
// (or of its folder) is gone.
|
||||
stopBackfill context.CancelFunc
|
||||
backfillDone chan struct{}
|
||||
}
|
||||
|
||||
// Close stops the background indexing of earlier chats and waits for it. It
|
||||
// does not close the database, which the caller owns.
|
||||
func (s *gormStore) Close() error {
|
||||
if s.stopBackfill != nil {
|
||||
s.stopBackfill()
|
||||
<-s.backfillDone
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewGORM returns a Store on db (SQLite or PostgreSQL) and creates its tables.
|
||||
@@ -158,7 +174,13 @@ func NewGORM(db *gorm.DB) (Store, error) {
|
||||
// Indexing is true from here, so a search never reports the index as
|
||||
// complete before the first look at the history.
|
||||
s.text.indexing.Store(true)
|
||||
go s.backfillTexts()
|
||||
var ctx context.Context
|
||||
ctx, s.stopBackfill = context.WithCancel(context.Background())
|
||||
s.backfillDone = make(chan struct{})
|
||||
go func() {
|
||||
defer close(s.backfillDone)
|
||||
s.backfillTexts(ctx)
|
||||
}()
|
||||
return s, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -338,9 +338,8 @@ func (s *gormStore) titleChats(ctx context.Context, q chatsearch.Query, terms []
|
||||
// from the chat it reached. A message written while it runs is indexed by the
|
||||
// write itself, and a row that backfill adds for a message that was just
|
||||
// removed is ignored by every search (it joins the entry).
|
||||
func (s *gormStore) backfillTexts() {
|
||||
func (s *gormStore) backfillTexts(ctx context.Context) {
|
||||
defer s.text.indexing.Store(false)
|
||||
ctx := context.Background()
|
||||
var st textStateRow
|
||||
if err := s.db.Where("id = ?", 1).Limit(1).Find(&st).Error; err != nil {
|
||||
xlog.Warn("Chat search: could not read the backfill state", "error", err)
|
||||
@@ -359,6 +358,9 @@ func (s *gormStore) backfillTexts() {
|
||||
}
|
||||
last, started, logged := st.LastChat, time.Now(), time.Now()
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
var ids []string
|
||||
if err := s.db.Model(&chatRow{}).Where("id > ?", last).Order("id").Limit(backfillStep).Pluck("id", &ids).Error; err != nil {
|
||||
xlog.Debug("Chat search: backfill stopped", "error", err)
|
||||
@@ -368,6 +370,9 @@ func (s *gormStore) backfillTexts() {
|
||||
break
|
||||
}
|
||||
for _, id := range ids {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
if err := s.backfillChat(ctx, id); err != nil {
|
||||
xlog.Debug("Chat search: backfill stopped", "error", err)
|
||||
return
|
||||
@@ -383,7 +388,11 @@ func (s *gormStore) backfillTexts() {
|
||||
logged = time.Now()
|
||||
xlog.Info("Chat search: indexing earlier chats", "done", s.text.done.Load(), "of", total)
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond) // leave the database to the chats being written
|
||||
select { // leave the database to the chats being written
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-time.After(10 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
// Rows of a chat deleted while the backfill read it.
|
||||
_ = s.db.Exec(`DELETE FROM studio_entry_texts WHERE NOT EXISTS (SELECT 1 FROM studio_chats c WHERE c.id = studio_entry_texts.chat_id)`).Error
|
||||
|
||||
Reference in new issue
Block a user