From d1e63702cdd7d9bd9976bb7fcd74c914f4a6d362 Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Fri, 31 Jul 2026 13:32:29 +0300 Subject: [PATCH] Keep optimizer store's in-memory database alive The optimizer's tool store runs on a shared-cache in-memory SQLite database, which exists only while at least one connection to it is open. When the database/sql pool drops to zero connections SQLite discards the whole database, and nothing recreates it because the schema is executed only in the constructor. Every later session build and find_tool then fails with "no such table: llm_capabilities" until the pod restarts, while the process stays healthy and /readyz stays green. The pool does reach zero. modernc.org/sqlite implements context cancellation with sqlite3_interrupt and reports an interrupted connection as unusable from conn.IsValid, so database/sql discards a connection whose statement was canceled rather than returning it to the pool. Because the store is touched only on session build and find_tool, a single pooled connection is the normal steady state, so one canceled request is enough to destroy the database for the life of the process. An embedding-service outage makes this far more likely, since slow and hanging calls cause clients to give up mid-statement, but the fault is not specific to embeddings. Pin a dedicated connection for the store's lifetime. It is acquired with a background context and never used to run statements, so it can never be interrupted and never discarded. Co-Authored-By: Claude Opus 5 (1M context) --- .../internal/toolstore/sqlite_store.go | 51 +++++++++++++++++-- .../internal/toolstore/sqlite_store_test.go | 41 +++++++++++++++ 2 files changed, 89 insertions(+), 3 deletions(-) diff --git a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go index f400397c53..acb8a154a4 100644 --- a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go +++ b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go @@ -50,7 +50,29 @@ var schemaSQL string // and optional vector embedding-based semantic search. // It satisfies the types.ToolStore interface. type sqliteToolStore struct { - db *sql.DB + db *sql.DB + + // keepAlive is a dedicated connection held open for the store's whole lifetime. + // + // The database is in-memory with a shared cache, so it exists only while at + // least one connection to it is open. If the database/sql pool ever drops to + // zero connections, SQLite discards the entire database — table, FTS5 index + // and triggers — and nothing recreates it, because the schema is executed + // only in the constructor. Every later operation then fails with + // "no such table: llm_capabilities" until the process restarts. + // + // The pool does reach zero in practice. modernc.org/sqlite implements context + // cancellation with sqlite3_interrupt and reports an interrupted connection as + // unusable from conn.IsValid, so database/sql discards a connection whose + // statement was canceled instead of returning it to the pool. A single + // canceled request while the pool holds one connection is enough to destroy + // the database for the life of the process (#5889). + // + // This connection is acquired with a background context and never used to run + // statements, so it can never be interrupted and never discarded. It exists + // solely to keep the database alive. + keepAlive *sql.Conn + embeddingClient types.EmbeddingClient // nil = FTS5-only maxToolsToReturn int hybridSemanticRatio float64 @@ -77,8 +99,17 @@ func newSQLiteToolStore( return sqliteToolStore{}, fmt.Errorf("failed to open sqlite database: %w", err) } + // Pin the keep-alive connection before anything else, so the database cannot + // be discarded from this point on. See the keepAlive field for why. + keepAlive, err := db.Conn(context.Background()) + if err != nil { + _ = db.Close() + return sqliteToolStore{}, fmt.Errorf("failed to acquire keep-alive connection: %w", err) + } + // Execute schema if _, err := db.Exec(schemaSQL); err != nil { + _ = keepAlive.Close() _ = db.Close() return sqliteToolStore{}, fmt.Errorf("failed to initialize schema: %w", err) } @@ -100,6 +131,7 @@ func newSQLiteToolStore( store := sqliteToolStore{ db: db, + keepAlive: keepAlive, embeddingClient: embeddingClient, maxToolsToReturn: maxTools, hybridSemanticRatio: hybridRatio, @@ -254,14 +286,27 @@ func (s sqliteToolStore) Search(ctx context.Context, q types.SearchQuery, allowe return merged, nil } -// Close releases the underlying database connection. +// Close releases the underlying database connections. Releasing the keep-alive +// connection drops the in-memory database, so it happens only here, on the way +// to closing the pool itself. +// +// Close may be called more than once: a keep-alive connection that was already +// released reports sql.ErrConnDone, which is not a failure. func (s sqliteToolStore) Close() error { var embErr error if s.embeddingClient != nil { embErr = s.embeddingClient.Close() } + + var keepAliveErr error + if s.keepAlive != nil { + if err := s.keepAlive.Close(); err != nil && !errors.Is(err, sql.ErrConnDone) { + keepAliveErr = err + } + } + dbErr := s.db.Close() - return errors.Join(embErr, dbErr) + return errors.Join(embErr, keepAliveErr, dbErr) } // searchFTS5 performs a full-text search using FTS5 MATCH with BM25 ranking. diff --git a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go index 13231f166d..e2a8e2caa4 100644 --- a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go +++ b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go @@ -10,6 +10,7 @@ import ( "sync" "sync/atomic" "testing" + "time" "github.com/stretchr/testify/require" @@ -55,6 +56,7 @@ func TestNewSQLiteToolStore(t *testing.T) { t.Parallel() store := newTestStore(t, nil, nil) require.NotNil(t, store.db) + require.NotNil(t, store.keepAlive) require.Nil(t, store.embeddingClient) }) @@ -66,6 +68,45 @@ func TestNewSQLiteToolStore(t *testing.T) { }) } +// TestSQLiteToolStore_SurvivesCanceledStatement is a regression test for #5889. +// +// Canceling a statement mid-flight makes modernc.org/sqlite interrupt the +// connection, which database/sql then discards rather than returning to the +// pool. Without the store's keep-alive connection that empties the pool, and +// emptying the pool destroys the shared in-memory database: every later +// operation fails with "no such table: llm_capabilities" until the process +// restarts. +func TestSQLiteToolStore_SurvivesCanceledStatement(t *testing.T) { + t.Parallel() + + store := newTestStore(t, nil, nil) + tools := makeTools(mcp.NewTool("read_file", mcp.WithDescription("Read a file from disk"))) + require.NoError(t, store.UpsertTools(context.Background(), tools)) + + // Counting to 200 million takes far longer than the deadline below, so the + // statement is always canceled while it is still executing. + const slowQuery = `WITH RECURSIVE counter(x) AS ( + SELECT 1 UNION ALL SELECT x + 1 FROM counter WHERE x < 200000000 + ) SELECT count(*) FROM counter` + + for range 3 { + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + var count int + require.Error(t, store.db.QueryRowContext(ctx, slowQuery).Scan(&count)) + cancel() + } + + // The store still works: schema, FTS5 index and rows all intact. + results, err := store.Search( + context.Background(), + types.SearchQuery{Description: "read a file"}, + toolNames(tools), + ) + require.NoError(t, err) + require.Len(t, results, 1) + require.Equal(t, "read_file", results[0].Name) +} + func TestSQLiteToolStore_UpsertTools(t *testing.T) { t.Parallel()