diff --git a/crates/tinyagents-harness/src/storage/mod.rs b/crates/tinyagents-harness/src/storage/mod.rs index 38c1e149e..52a1f37d0 100644 --- a/crates/tinyagents-harness/src/storage/mod.rs +++ b/crates/tinyagents-harness/src/storage/mod.rs @@ -110,36 +110,52 @@ pub fn collections() -> Vec { pub async fn open_storage( storage: ScopedStorage, ) -> tinyagents_storage_kit::error::Result { - let owned_storage = storage.clone(); - tinyagents_storage_kit::with_migration_lease_once( - &storage, - "harness", - move |lease| async move { - let storage = owned_storage; - let registry = Registry::compose([storage_schema()])?; - for spec in collections() { - storage.documents().ensure_collection(&spec).await?; - } - storage - .documents() - .ensure_collection(&tinyagents_storage_kit::primitives::kv_collection_spec()) - .await?; - registry.migrate(&storage, &MigrateOptions::new()).await?; - lease.renew().await?; - for (collection, kind) in [ - ("harness_store", LegacyKv::Flat), - ("harness_namespaced_store", LegacyKv::Namespaced), - ("harness_response_cache", LegacyKv::ResponseCache), - ] { - import_legacy_kv_with_lease(&storage, "harness", collection, kind, &lease).await?; - } - Ok(()) - }, - ) + tinyagents_storage_kit::with_migration_lease_once(&storage, "harness", { + let owned_storage = storage.clone(); + move |lease| migrate_storage_with_lease(owned_storage, lease) + }) .await?; Ok(storage) } +/// Run the harness's owner-local migrations even if this process has already +/// opened the same storage port and scope. +/// +/// Hosts call this after copying legacy rows into a scope and clearing the +/// durable progress records. Unlike [`open_storage`], this deliberately +/// bypasses the process-local once cache so the newly copied rows are visited. +pub async fn migrate_storage(storage: ScopedStorage) -> tinyagents_storage_kit::error::Result<()> { + let owned_storage = storage.clone(); + tinyagents_storage_kit::with_migration_lease(&storage, "harness", move |lease| { + migrate_storage_with_lease(owned_storage, lease) + }) + .await +} + +async fn migrate_storage_with_lease( + storage: ScopedStorage, + lease: tinyagents_storage_kit::MigrationLease, +) -> tinyagents_storage_kit::error::Result<()> { + let registry = Registry::compose([storage_schema()])?; + for spec in collections() { + storage.documents().ensure_collection(&spec).await?; + } + storage + .documents() + .ensure_collection(&tinyagents_storage_kit::primitives::kv_collection_spec()) + .await?; + registry.migrate(&storage, &MigrateOptions::new()).await?; + lease.renew().await?; + for (collection, kind) in [ + ("harness_store", LegacyKv::Flat), + ("harness_namespaced_store", LegacyKv::Namespaced), + ("harness_response_cache", LegacyKv::ResponseCache), + ] { + import_legacy_kv_with_lease(&storage, "harness", collection, kind, &lease).await?; + } + Ok(()) +} + const BLOB_STEP_IDS: [&str; 5] = [ "attachment_blobs_v2", "tool_result_blobs_v2", diff --git a/crates/tinyagents-harness/src/storage/mod_tests.rs b/crates/tinyagents-harness/src/storage/mod_tests.rs index 4f4a4ff9e..3a1609914 100644 --- a/crates/tinyagents-harness/src/storage/mod_tests.rs +++ b/crates/tinyagents-harness/src/storage/mod_tests.rs @@ -1,8 +1,8 @@ use super::*; use serde_json::{Value, json}; use tinyagents_storage_kit::{ - Filter, MigrateOptions, Precondition, Registry, Scope, ScopedStorage, StorageBackend, - WriteOptions, + DocumentStoreExt, Filter, MIGRATIONS_COLLECTION, MigrateOptions, Precondition, Query, Registry, + Scope, ScopedStorage, StorageBackend, WriteOptions, }; use tinystoragedrivers_testkit::MemoryBackend; @@ -140,6 +140,52 @@ async fn run_status_moves_to_a_bare_snake_case_document_under_its_run_id() { assert_eq!(left, 0, "the v1 rows are removed"); } +#[tokio::test] +async fn remigration_reads_v1_rows_copied_after_the_first_open() { + let storage = scoped(); + open_storage(storage.clone()).await.unwrap(); + put_v1( + &storage, + "harness_run_status", + "late-row", + json!({ + "runId": "run.late", "rootRunId": "run.late", "threadId": "thread-late", + "active": false, + "record": { + "run_id": "run.late", "root_run_id": "run.late", "thread_id": "thread-late", + "component": "model", "status": "completed", "current_phase": "model", + "model_calls": 1, "tool_calls": 0, "active_tool_calls": [], + "usage": {}, "cost": {}, "started_at": 1_700_000_000_u64, + "updated_at": 1_700_000_010_u64, "metadata": {} + } + }), + ) + .await; + + let docs = storage.documents(); + let progress = docs + .query_all(MIGRATIONS_COLLECTION, &Query::all()) + .await + .unwrap(); + for row in progress + .into_iter() + .filter(|row| row.doc.get("owner").and_then(Value::as_str) == Some("harness")) + { + docs.delete(MIGRATIONS_COLLECTION, &row.id, Precondition::None) + .await + .unwrap(); + } + + migrate_storage(storage.clone()).await.unwrap(); + + let migrated = docs + .get("harness.run_status", "run.late") + .await + .unwrap() + .expect("the copied v1 row is visible after remigration"); + assert_eq!(migrated.doc["run_id"], "run.late"); +} + #[tokio::test] async fn embedding_vectors_move_to_a_key_of_signature_and_hex() { let storage = scoped();