From d5a8a3fb90291f53ffde953fa9962a8f1f811f7a Mon Sep 17 00:00:00 2001
From: filip <44206832+filipkujawa@users.noreply.github.com>
Date: Mon, 20 Jul 2026 14:13:07 -0700
Subject: [PATCH] fix(session): create inventory tables atomically with schema
 version (#10586)

---
 crates/goose/src/providers/inventory/mod.rs | 49 +--------------------
 crates/goose/src/session/session_manager.rs | 15 ++-----
 2 files changed, 5 insertions(+), 59 deletions(-)

diff --git a/crates/goose/src/providers/inventory/mod.rs b/crates/goose/src/providers/inventory/mod.rs
index 25ef89d3d..121a76775 100644
--- a/crates/goose/src/providers/inventory/mod.rs
+++ b/crates/goose/src/providers/inventory/mod.rs
@@ -18,7 +18,7 @@ use chrono::{DateTime, Duration, Utc};
 use futures::FutureExt;
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
-use sqlx::{Pool, Row, Sqlite, Transaction};
+use sqlx::{Row, Sqlite, Transaction};
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::panic::AssertUnwindSafe;
 use std::sync::{Arc, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard};
@@ -1142,52 +1142,7 @@ fn enriched_model(
     }
 }
 
-pub async fn create_tables(pool: &Pool<Sqlite>) -> Result<()> {
-    sqlx::query(
-        r#"
-        CREATE TABLE IF NOT EXISTS provider_inventory_entries (
-            inventory_key TEXT PRIMARY KEY,
-            provider_id TEXT NOT NULL,
-            provider_family TEXT NOT NULL,
-            last_updated_at TEXT,
-            last_refresh_attempt_at TEXT,
-            last_refresh_error TEXT,
-            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-        )
-        "#,
-    )
-    .execute(pool)
-    .await?;
-
-    sqlx::query(
-        r#"
-        CREATE TABLE IF NOT EXISTS provider_inventory_models (
-            inventory_key TEXT NOT NULL REFERENCES provider_inventory_entries(inventory_key) ON DELETE CASCADE,
-            ordinal INTEGER NOT NULL,
-            model_id TEXT NOT NULL,
-            name TEXT NOT NULL,
-            family TEXT,
-            context_limit INTEGER,
-            reasoning BOOLEAN,
-            recommended BOOLEAN,
-            PRIMARY KEY (inventory_key, ordinal)
-        )
-        "#,
-    )
-    .execute(pool)
-    .await?;
-
-    sqlx::query(
-        "CREATE INDEX IF NOT EXISTS idx_provider_inventory_provider_id ON provider_inventory_entries(provider_id)",
-    )
-    .execute(pool)
-    .await?;
-
-    Ok(())
-}
-
-pub async fn create_tables_in_tx(tx: &mut Transaction<'_, Sqlite>) -> Result<()> {
+pub async fn create_tables(tx: &mut Transaction<'_, Sqlite>) -> Result<()> {
     sqlx::query(
         r#"
         CREATE TABLE IF NOT EXISTS provider_inventory_entries (
diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs
index 0452519e2..cc7c9c865 100644
--- a/crates/goose/src/session/session_manager.rs
+++ b/crates/goose/src/session/session_manager.rs
@@ -891,12 +891,6 @@ impl SessionStorage {
         Ok(&self.pool)
     }
 
-    pub async fn create(session_dir: &Path) -> Result<Self> {
-        let storage = Self::new(session_dir.to_path_buf());
-        Self::create_schema(&storage.pool).await?;
-        Ok(storage)
-    }
-
     async fn create_schema(pool: &Pool<Sqlite>) -> Result<()> {
         // Run schema creation under `BEGIN IMMEDIATE` so SQLite serializes
         // writers across processes. Combined with `IF NOT EXISTS` on every
@@ -1031,12 +1025,9 @@ impl SessionStorage {
         .execute(&mut *tx)
         .await?;
 
-        tx.commit().await?;
+        crate::providers::inventory::create_tables(&mut tx).await?;
 
-        // The inventory tables already use `CREATE TABLE IF NOT EXISTS`
-        // and run on the shared pool, so they don't need to be inside
-        // the same transaction.
-        crate::providers::inventory::create_tables(pool).await?;
+        tx.commit().await?;
 
         Ok(())
     }
@@ -1390,7 +1381,7 @@ impl SessionStorage {
                     .await?;
             }
             11 => {
-                crate::providers::inventory::create_tables_in_tx(tx).await?;
+                crate::providers::inventory::create_tables(tx).await?;
             }
             12 => {
                 // Add archived_at, project_id columns to sessions.