diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 1c063249bb..9b7bafc9c6 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -12551a748d759c7c45a9da39b63117a059eba845 +1a8f29cd14fce3b7c6eb20d27f5b58b8ae56d9f5 diff --git a/backend/migrations/20260915091518_ai_sessions_backup_generation.down.sql b/backend/migrations/20260915091518_ai_sessions_backup_generation.down.sql new file mode 100644 index 0000000000..2662a9f709 --- /dev/null +++ b/backend/migrations/20260915091518_ai_sessions_backup_generation.down.sql @@ -0,0 +1 @@ +ALTER TABLE workspace_settings DROP COLUMN IF EXISTS ai_sessions_backup_generation; diff --git a/backend/migrations/20260915091518_ai_sessions_backup_generation.up.sql b/backend/migrations/20260915091518_ai_sessions_backup_generation.up.sql new file mode 100644 index 0000000000..a62caefa60 --- /dev/null +++ b/backend/migrations/20260915091518_ai_sessions_backup_generation.up.sql @@ -0,0 +1,4 @@ +-- Bumped by every workspace key rotation: the AI session backups in the workspace storage +-- live under a prefix named by it, so a rotation moves to a fresh prefix and the previous +-- ones can be deleted at leisure without ever touching live objects. +ALTER TABLE workspace_settings ADD COLUMN ai_sessions_backup_generation BIGINT NOT NULL DEFAULT 0; diff --git a/backend/summarized_schema.txt b/backend/summarized_schema.txt index 463afd0d33..b93a4bf6fc 100644 --- a/backend/summarized_schema.txt +++ b/backend/summarized_schema.txt @@ -234,7 +234,7 @@ workspace_protection_rule: workspace_id(char), name(char), rules(int), bypass_gr FK: (workspace_id) -> workspace(id) workspace_runnable_dependencies: flow_path(char), runnable_path(char), script_hash(bigint), runnable_is_flow(bool), workspace_id(char), app_path(char), id(bigint) FK: (app_path, workspace_id) -> app(path, workspace_id) | (flow_path, workspace_id) -> flow(path, workspace_id) -workspace_settings: workspace_id(char), slack_team_id(char), slack_name(char), slack_command_script(char), slack_email(char), customer_id(char), plan(char), webhook(text), ai_config(jsonb), large_file_storage(jsonb), git_sync(jsonb), default_app(char), default_scripts(jsonb), deploy_ui(jsonb), mute_critical_alerts(bool), color(char), operator_settings(jsonb), teams_command_script(text), teams_team_id(text), teams_team_name(text), git_app_installations(jsonb), ducklake(jsonb), slack_oauth_client_id(char), slack_oauth_client_secret(char), datatable(jsonb), teams_team_guid(text), auto_invite(jsonb), error_handler(jsonb), success_handler(jsonb), public_app_execution_limit_per_minute(int), dbt_warehouses(jsonb), guest_access_enabled(bool), guest_jwt_public_key(text), guest_jwt_jwks_url(text) +workspace_settings: workspace_id(char), slack_team_id(char), slack_name(char), slack_command_script(char), slack_email(char), customer_id(char), plan(char), webhook(text), ai_config(jsonb), large_file_storage(jsonb), git_sync(jsonb), default_app(char), default_scripts(jsonb), deploy_ui(jsonb), mute_critical_alerts(bool), color(char), operator_settings(jsonb), teams_command_script(text), teams_team_id(text), teams_team_name(text), git_app_installations(jsonb), ducklake(jsonb), slack_oauth_client_id(char), slack_oauth_client_secret(char), datatable(jsonb), teams_team_guid(text), auto_invite(jsonb), error_handler(jsonb), success_handler(jsonb), public_app_execution_limit_per_minute(int), dbt_warehouses(jsonb), guest_access_enabled(bool), guest_jwt_public_key(text), guest_jwt_jwks_url(text), ai_sessions_backup_generation(int) FK: (workspace_id) -> workspace(id) zombie_job_counter: job_id(uuid), counter(int) FK: (job_id) -> v2_job(id) diff --git a/backend/tests/ai_sessions.rs b/backend/tests/ai_sessions.rs new file mode 100644 index 0000000000..f76385d4e1 --- /dev/null +++ b/backend/tests/ai_sessions.rs @@ -0,0 +1,1085 @@ +//! The AI session backup routes (`/w/{w}/ai/sessions/*`): a browser pushes pieces of its +//! sessions into the workspace's object storage and pulls them back whole. Pinned against a +//! FilesystemStorage LFS so the test needs no object store, which also lets it read what +//! landed on disk: the objects must be ciphertext, since bucket credentials are shared far +//! more widely than a user's transcripts. +#![cfg(all(feature = "private", feature = "parquet"))] + +use serde_json::{json, Value}; +use sqlx::{Pool, Postgres}; +use windmill_common::utils::calculate_hash; +use windmill_test_utils::*; + +fn client() -> reqwest::Client { + reqwest::Client::new() +} + +fn authed(builder: reqwest::RequestBuilder, token: &str) -> reqwest::RequestBuilder { + builder.header("Authorization", format!("Bearer {}", token)) +} + +async fn configure_primary_lfs(db: &Pool, root_path: &str) -> anyhow::Result<()> { + sqlx::query!( + "UPDATE workspace_settings SET large_file_storage = $1 WHERE workspace_id = $2", + json!({ + "type": "FilesystemStorage", + "root_path": root_path, + "public_resource": null, + "advanced_permissions": null + }), + "test-workspace" + ) + .execute(db) + .await?; + Ok(()) +} + +async fn list(base: &str, token: &str) -> anyhow::Result { + let resp = authed(client().get(format!("{base}/ai/sessions/list")), token) + .send() + .await?; + assert_eq!(resp.status(), 200, "{}", resp.text().await?); + Ok(resp.json().await?) +} + +async fn pull(base: &str, token: &str, ids: &[&str]) -> anyhow::Result { + let resp = authed(client().post(format!("{base}/ai/sessions/pull")), token) + .json(&json!({ "ids": ids })) + .send() + .await?; + assert_eq!(resp.status(), 200, "{}", resp.text().await?); + Ok(resp.json().await?) +} + +async fn push(base: &str, token: &str, body: Value) -> anyhow::Result { + Ok( + authed(client().post(format!("{base}/ai/sessions/push")), token) + .json(&body) + .send() + .await?, + ) +} + +fn copy_dir(from: &std::path::Path, to: &std::path::Path) -> std::io::Result<()> { + std::fs::create_dir_all(to)?; + for entry in std::fs::read_dir(from)? { + let entry = entry?; + let target = to.join(entry.file_name()); + if entry.path().is_dir() { + copy_dir(&entry.path(), &target)?; + } else { + std::fs::copy(entry.path(), target)?; + } + } + Ok(()) +} + +async fn rotate(base: &str, key: &str) -> anyhow::Result<()> { + let resp = authed( + client().post(format!("{base}/workspaces/encryption_key")), + "SECRET_TOKEN", + ) + .json(&json!({ "new_key": key })) + .send() + .await?; + assert_eq!(resp.status(), 200, "{}", resp.text().await?); + Ok(()) +} + +/// The user's prefix on disk, `windmill_ai_sessions/{w_id}/g{generation}/{email hash}`, +/// under whichever generation is current. +fn user_root(storage_dir: &std::path::Path, email: &str) -> std::path::PathBuf { + let workspace = storage_dir.join("windmill_ai_sessions/test-workspace"); + let hash = calculate_hash(email); + std::fs::read_dir(&workspace) + .ok() + .into_iter() + .flatten() + .flatten() + .map(|entry| entry.path().join(&hash)) + .find(|path| path.exists()) + .expect("the user has backups under the current key") +} + +/// Every file under the storage root, as bytes. +fn files_under(root: &std::path::Path) -> Vec<(std::path::PathBuf, Vec)> { + let mut out = vec![]; + let mut stack = vec![root.to_path_buf()]; + while let Some(dir) = stack.pop() { + let Ok(entries) = std::fs::read_dir(&dir) else { + continue; + }; + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + stack.push(path); + } else { + out.push((path.clone(), std::fs::read(&path).unwrap_or_default())); + } + } + } + out +} + +#[sqlx::test(fixtures("base"))] +async fn test_backups_round_trip_encrypted_and_scoped_to_the_user( + db: Pool, +) -> anyhow::Result<()> { + initialize_tracing().await; + let server = ApiServer::start(db.clone()).await?; + let base = format!( + "http://localhost:{}/api/w/test-workspace", + server.addr.port() + ); + + // No storage configured: the browser is told to stop trying. + let listing = list(&base, "SECRET_TOKEN").await?; + assert_eq!(listing["enabled"], false); + assert_eq!(listing["sessions"], json!([])); + + let storage_dir = tempfile::tempdir()?; + configure_primary_lfs(&db, &storage_dir.path().to_string_lossy()).await?; + + let head = + json!({ "id": "s1", "workspace_id": "test-workspace", "createdAt": 1, "chatId": "c1" }); + let chat = json!({ "id": "c1", "sessionId": "s1", "title": "MARKER_PLAINTEXT_TITLE", "lastModified": 2, + "actualMessages": [], "displayMessages": [{"role": "user", "content": "hello"}] }); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ + "id": "s1", + "whole": true, + "head": head, + "chats": [{ "id": "c1", "record": chat }, { "id": "c2", "record": { "id": "c2" } }], + "images": [{ "chat_id": "c1", "id": "img1", "data_url": "data:image/png;base64,AAAA" }], + "artifacts": { "items": [], "versions": [] } + }] + }), + ) + .await?; + assert_eq!(resp.status(), 200, "{}", resp.text().await?); + let pushed: Value = resp.json().await?; + assert_eq!(pushed["enabled"], true); + assert_eq!(pushed["results"], json!([{ "id": "s1" }])); + + let listing = list(&base, "SECRET_TOKEN").await?; + assert_eq!(listing["enabled"], true); + assert_eq!(listing["sessions"][0]["id"], "s1"); + + // A part more parts follow names its push, or the session would stay listed between + // the parts: one that does not is refused before anything of it lands. + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s1", "partial": true, "delete_chats": ["c1"] }] + }), + ) + .await?; + assert_eq!(resp.status(), 400, "{}", resp.text().await?); + assert_eq!( + list(&base, "SECRET_TOKEN").await?["sessions"][0]["id"], + "s1" + ); + assert_eq!( + pull(&base, "SECRET_TOKEN", &["s1"]).await?["sessions"][0]["chats"] + .as_array() + .unwrap() + .len(), + 2 + ); + + // A part with more of the session to follow lists nothing; the part that completes + // the push does, newest first. + let ids = |listing: &Value| -> Vec { + listing["sessions"] + .as_array() + .unwrap() + .iter() + .map(|s| s["id"].as_str().unwrap().to_string()) + .collect() + }; + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s2", "whole": true, "push": "p2", "opens": true, "head": { "id": "s2", "workspace_id": "test-workspace", "createdAt": 2, "chatId": "c" }, "partial": true }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + assert_eq!(ids(&list(&base, "SECRET_TOKEN").await?), vec!["s1"]); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s2", "whole": true, "push": "p2", "chats": [{ "id": "c", "record": { "id": "c" } }] }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + assert_eq!(ids(&list(&base, "SECRET_TOKEN").await?), vec!["s2", "s1"]); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "removed": ["s2"] }), + ) + .await?; + assert_eq!(resp.status(), 200); + + // A session that outgrew one answer (three chats of 12 MB against the 32 MB budget) + // comes in pages, each naming where the next picks up, and nothing is left out. + let big = "y".repeat(12 * 1024 * 1024); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s4", "whole": true, "push": "p4", "opens": true, "head": { "id": "s4", "workspace_id": "test-workspace", "createdAt": 4, "chatId": "c1" }, "partial": true }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + for cid in ["c1", "c2", "c3"] { + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s4", "whole": true, "push": "p4", "chats": [{ "id": cid, "record": { "id": cid, "big": big } }], "partial": cid != "c3" }] + }), + ) + .await?; + assert_eq!(resp.status(), 200, "{}", resp.text().await?); + } + let mut pages = vec![]; + let mut resume = json!(null); + loop { + let body = if resume.is_null() { + json!({ "ids": ["s4"] }) + } else { + json!({ "ids": ["s4"], "resume": resume }) + }; + let resp = authed( + client().post(format!("{base}/ai/sessions/pull")), + "SECRET_TOKEN", + ) + .json(&body) + .send() + .await?; + assert_eq!(resp.status(), 200, "{}", resp.text().await?); + let pulled: Value = resp.json().await?; + let page = pulled["sessions"][0].clone(); + assert_eq!(page["id"], "s4"); + resume = page["next"].clone(); + pages.push(page); + if resume.is_null() { + break; + } + assert!(pages.len() < 5, "a paged pull must end"); + } + assert!(pages.len() >= 2, "36 MB must not fit one answer"); + // Every page of an unchanged backup carries the same listing fingerprint; a chat added + // to the session changes it, which is what tells a browser its pages do not belong + // together any more. + let listing = pages[0]["listing"].clone(); + assert!(listing.is_string()); + assert!(pages.iter().all(|p| p["listing"] == listing)); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s4", "chats": [{ "id": "c0", "record": { "id": "c0", "n": 1 } }] }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let pulled = pull(&base, "SECRET_TOKEN", &["s4"]).await?; + assert_ne!(pulled["sessions"][0]["listing"], listing); + // So does a chat rewritten at the same size: the fingerprint takes in the entity tag, + // not only the size and a modification time the store may report coarsely. + let listing = pulled["sessions"][0]["listing"].clone(); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s4", "chats": [{ "id": "c0", "record": { "id": "c0", "n": 2 } }] }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let pulled = pull(&base, "SECRET_TOKEN", &["s4"]).await?; + assert_ne!(pulled["sessions"][0]["listing"], listing); + let mut chat_ids: Vec = pages + .iter() + .flat_map(|p| p["chats"].as_array().unwrap().iter()) + .map(|c| c["id"].as_str().unwrap().to_string()) + .collect(); + chat_ids.sort(); + assert_eq!(chat_ids, vec!["c1", "c2", "c3"]); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "removed": ["s4"] }), + ) + .await?; + assert_eq!(resp.status(), 200); + + // More objects than a page keeps listing metadata for, from a store that lists in no + // order: the pages still carry every one of them, each once. + let many = 5001; + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s5", "whole": true, "push": "p5", "opens": true, "head": { "id": "s5", "workspace_id": "test-workspace", "createdAt": 5, "chatId": "c00000" }, "partial": true }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + for start in (0..many).step_by(100) { + let chats: Vec = (start..(start + 100).min(many)) + .map(|i| json!({ "id": format!("c{i:05}"), "record": { "id": format!("c{i:05}") } })) + .collect(); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s5", "whole": true, "push": "p5", "chats": chats, "partial": start + 100 < many }] + }), + ) + .await?; + assert_eq!(resp.status(), 200, "{}", resp.text().await?); + } + let mut seen = std::collections::HashSet::new(); + let mut resume = json!(null); + let mut pages = 0; + loop { + let body = if resume.is_null() { + json!({ "ids": ["s5"] }) + } else { + json!({ "ids": ["s5"], "resume": resume }) + }; + let resp = authed( + client().post(format!("{base}/ai/sessions/pull")), + "SECRET_TOKEN", + ) + .json(&body) + .send() + .await?; + assert_eq!(resp.status(), 200, "{}", resp.text().await?); + let pulled: Value = resp.json().await?; + let page = &pulled["sessions"][0]; + for c in page["chats"].as_array().unwrap() { + assert!( + seen.insert(c["id"].as_str().unwrap().to_string()), + "a chat came twice" + ); + } + pages += 1; + resume = page["next"].clone(); + if resume.is_null() { + break; + } + assert!(pages < 5, "a paged pull must end"); + } + assert_eq!(pages, 2); + assert_eq!(seen.len(), many); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "removed": ["s5"] }), + ) + .await?; + assert_eq!(resp.status(), 200); + + // A head at exactly its cap round-trips: the ciphertext read back is a block larger. + let mut big_head = json!({ "id": "s3", "workspace_id": "test-workspace", "createdAt": 3, "chatId": "c", "pad": "" }); + let pad = 1024 * 1024 - serde_json::to_string(&big_head)?.len(); + big_head["pad"] = json!("x".repeat(pad)); + assert_eq!(serde_json::to_string(&big_head)?.len(), 1024 * 1024); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "sessions": [{ "id": "s3", "whole": true, "head": big_head }] }), + ) + .await?; + assert_eq!(resp.status(), 200, "{}", resp.text().await?); + let pulled = pull(&base, "SECRET_TOKEN", &["s3"]).await?; + assert_eq!(pulled["sessions"][0]["head"], big_head); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "removed": ["s3"] }), + ) + .await?; + assert_eq!(resp.status(), 200); + + // An object larger than any push writes, planted with the bucket's credentials at a + // predictable key, is not read. + let planted = + user_root(storage_dir.path(), "test@windmill.dev").join("sessions/planted/head.json"); + std::fs::create_dir_all(planted.parent().unwrap())?; + std::fs::File::create(&planted)?.set_len(32 * 1024 * 1024 + 1)?; + let pulled = pull(&base, "SECRET_TOKEN", &["planted"]).await?; + assert_eq!(pulled["sessions"], json!([])); + std::fs::remove_dir_all(planted.parent().unwrap())?; + // Under a session that exists, a planted chat is skipped without buffering and without + // the page ending before it, so the pull neither balloons nor loops. + let planted_chat = + user_root(storage_dir.path(), "test@windmill.dev").join("sessions/s1/chats/planted.json"); + std::fs::File::create(&planted_chat)?.set_len(32 * 1024 * 1024 + 1)?; + let mut resume = json!(null); + let mut pages = 0; + let mut chat_ids = vec![]; + loop { + let body = if resume.is_null() { + json!({ "ids": ["s1"] }) + } else { + json!({ "ids": ["s1"], "resume": resume }) + }; + let resp = authed( + client().post(format!("{base}/ai/sessions/pull")), + "SECRET_TOKEN", + ) + .json(&body) + .send() + .await?; + assert_eq!(resp.status(), 200, "{}", resp.text().await?); + let pulled: Value = resp.json().await?; + assert_eq!(pulled["sessions"][0]["head"], head); + for c in pulled["sessions"][0]["chats"].as_array().unwrap() { + chat_ids.push(c["id"].as_str().unwrap().to_string()); + } + pages += 1; + resume = pulled["sessions"][0]["next"].clone(); + if resume.is_null() { + break; + } + assert!(pages < 5, "a planted object must not keep the pull going"); + } + chat_ids.sort(); + assert_eq!(chat_ids, vec!["c1", "c2"]); + std::fs::remove_file(&planted_chat)?; + + let pulled = pull(&base, "SECRET_TOKEN", &["s1", "never-pushed"]).await?; + assert_eq!(pulled["deferred"], json!([])); + let sessions = pulled["sessions"].as_array().unwrap(); + assert_eq!(sessions.len(), 1, "an id with no backup is simply absent"); + let s1 = &sessions[0]; + assert_eq!(s1["head"], head); + let mut chats = s1["chats"].as_array().unwrap().clone(); + chats.sort_by_key(|c| c["id"].as_str().unwrap().to_string()); + assert_eq!(chats[0]["record"], chat); + assert_eq!(chats[1]["id"], "c2"); + assert_eq!( + s1["images"], + json!([{ "chat_id": "c1", "id": "img1", "data_url": "data:image/png;base64,AAAA" }]) + ); + assert_eq!(s1["artifacts"], json!({ "items": [], "versions": [] })); + + // Nothing on disk carries the transcript in the clear. + let files = files_under(storage_dir.path()); + assert!( + files.len() >= 4, + "expected the pushed objects on disk, got {files:?}" + ); + for (path, bytes) in &files { + let text = String::from_utf8_lossy(bytes); + assert!( + !text.contains("MARKER_PLAINTEXT_TITLE") && !text.contains("base64,AAAA"), + "{} holds plaintext", + path.display() + ); + } + let key_paths: Vec = files + .iter() + .map(|(p, _)| { + p.strip_prefix(storage_dir.path()) + .unwrap() + .to_string_lossy() + .to_string() + }) + .collect(); + assert!( + key_paths + .iter() + .all(|p| p.starts_with("windmill_ai_sessions/test-workspace/") + && !p.contains("test@windmill.dev")), + "keys carry the workspace and never the email: {key_paths:?}" + ); + + // Another member of the workspace sees none of it. + let other = list(&base, "SECRET_TOKEN_2").await?; + assert_eq!(other["enabled"], true); + assert_eq!(other["sessions"], json!([])); + let other = pull(&base, "SECRET_TOKEN_2", &["s1"]).await?; + assert_eq!(other["sessions"], json!([])); + + // Nor after copying the first user's ciphertext under their own prefix, which anyone + // holding the bucket credentials can do: the key is bound to the user, not the workspace. + let first = user_root(storage_dir.path(), "test@windmill.dev"); + let second = first + .parent() + .unwrap() + .join(calculate_hash("test2@windmill.dev")); + copy_dir(&first, &second)?; + let other = pull(&base, "SECRET_TOKEN_2", &["s1"]).await?; + assert_eq!( + other["sessions"], + json!([]), + "relocated ciphertext must not decrypt for another user" + ); + std::fs::remove_dir_all(&second)?; + + // Deleting a chat takes its images along; a head-only push leaves the rest in place. + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s1", "delete_chats": ["c1"] }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let pulled = pull(&base, "SECRET_TOKEN", &["s1"]).await?; + let s1 = &pulled["sessions"][0]; + assert_eq!(s1["head"], head); + assert_eq!(s1["chats"].as_array().unwrap().len(), 1); + assert_eq!(s1["chats"][0]["id"], "c2"); + assert_eq!(s1["images"], json!([])); + + // Rotating the workspace key moves the routes to a fresh generation's prefix and deletes + // the older ones off the request rather than re-key anything; the answers name the new + // generation (`backup_generation`), which is what makes every browser push its sessions + // whole again, while `storage_id` names the storage and stays. + let before = list(&base, "SECRET_TOKEN").await?; + rotate(&base, &"b".repeat(64)).await?; + for _ in 0..100 { + if files_under(storage_dir.path()).is_empty() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + assert!( + files_under(storage_dir.path()).is_empty(), + "a rotation must leave no backup object behind" + ); + let listing = list(&base, "SECRET_TOKEN").await?; + assert_eq!(listing["sessions"], json!([])); + assert_eq!(listing["storage_id"], before["storage_id"]); + assert_ne!( + listing["backup_generation"], before["backup_generation"], + "a rotation must bump the backup generation" + ); + assert_eq!( + pull(&base, "SECRET_TOKEN", &["s1"]).await?["sessions"], + json!([]) + ); + // The browser's next push fills the storage back under the new key. + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s1", "whole": true, "head": head, "chats": [{ "id": "c2", "record": { "id": "c2" } }] }] + }), + ) + .await?; + assert_eq!(resp.status(), 200, "{}", resp.text().await?); + let pulled = pull(&base, "SECRET_TOKEN", &["s1"]).await?; + assert_eq!(pulled["sessions"][0]["head"], head); + assert_eq!(pulled["sessions"][0]["chats"][0]["id"], "c2"); + // Setting the key already in place is not a rotation the browsers would notice, so it + // keeps the backups. + let same = list(&base, "SECRET_TOKEN").await?; + rotate(&base, &"b".repeat(64)).await?; + let listing = list(&base, "SECRET_TOKEN").await?; + assert_eq!(listing["storage_id"], same["storage_id"]); + assert_eq!(listing["backup_generation"], same["backup_generation"]); + assert_eq!(listing["sessions"][0]["id"], "s1"); + assert_eq!( + pull(&base, "SECRET_TOKEN", &["s1"]).await?["sessions"][0]["head"], + head + ); + + // A push that does not open the session whole rides on the head in the storage; with + // none there (another device removed the backup, or nothing was ever pushed) it is + // refused and lists nothing, head or no head on it, until the session goes whole. + let s6_head = + json!({ "id": "s6", "workspace_id": "test-workspace", "createdAt": 6, "chatId": "c" }); + let not_listed = |listing: Value| { + listing["sessions"] + .as_array() + .unwrap() + .iter() + .all(|s| s["id"] != "s6") + }; + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s6", "chats": [{ "id": "c", "record": { "id": "c" } }] }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let answer: Value = resp.json().await?; + assert_eq!(answer["results"][0]["needs_whole"], true); + assert!(not_listed(list(&base, "SECRET_TOKEN").await?)); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s6", "whole": true, "head": s6_head, "chats": [{ "id": "c", "record": { "id": "c" } }] }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let answer: Value = resp.json().await?; + assert!(answer["results"][0]["needs_whole"].is_null()); + assert_eq!( + list(&base, "SECRET_TOKEN").await?["sessions"][0]["id"], + "s6" + ); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "removed": ["s6"] }), + ) + .await?; + assert_eq!(resp.status(), 200); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s6", "head": s6_head, "chats": [{ "id": "c2", "record": { "id": "c2" } }] }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let answer: Value = resp.json().await?; + assert_eq!(answer["results"][0]["needs_whole"], true); + assert!(not_listed(list(&base, "SECRET_TOKEN").await?)); + assert!( + files_under(&user_root(storage_dir.path(), "test@windmill.dev").join("sessions/s6")) + .is_empty() + ); + + // Between the parts of a whole push (head landed, marker not yet), an incremental push + // from another device is refused too: it rides on a listed session, and there is none + // until the last part, which lists it. + let s8_head = + json!({ "id": "s8", "workspace_id": "test-workspace", "createdAt": 8, "chatId": "c1" }); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s8", "whole": true, "push": "p8", "opens": true, "head": s8_head, "partial": true }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s8", "chats": [{ "id": "c9", "record": { "id": "c9" } }] }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let answer: Value = resp.json().await?; + assert_eq!(answer["results"][0]["needs_whole"], true); + let listed = |listing: Value| { + listing["sessions"] + .as_array() + .unwrap() + .iter() + .any(|s| s["id"] == "s8") + }; + assert!(!listed(list(&base, "SECRET_TOKEN").await?)); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s8", "whole": true, "push": "p8", "chats": [{ "id": "c1", "record": { "id": "c1" } }] }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + assert!(listed(list(&base, "SECRET_TOKEN").await?)); + assert_eq!( + pull(&base, "SECRET_TOKEN", &["s8"]).await?["sessions"][0]["chats"] + .as_array() + .unwrap() + .iter() + .map(|c| c["id"].as_str().unwrap().to_string()) + .collect::>(), + vec!["c1"] + ); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "removed": ["s8"] }), + ) + .await?; + assert_eq!(resp.status(), 200); + + // A whole push replaces the backup: what the storage held of the session and the new + // push does not carry (a chat deleted while the workspace was on another storage) goes. + let s9_head = + json!({ "id": "s9", "workspace_id": "test-workspace", "createdAt": 9, "chatId": "c1" }); + let s9_chats = |ids: &[&str]| -> Vec { + ids.iter() + .map(|c| json!({ "id": c, "record": { "id": c } })) + .collect() + }; + let pulled_chats = |pulled: Value| -> Vec { + pulled["sessions"][0]["chats"] + .as_array() + .unwrap() + .iter() + .map(|c| c["id"].as_str().unwrap().to_string()) + .collect() + }; + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s9", "whole": true, "head": s9_head, "chats": s9_chats(&["c1", "c2"]) }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + assert_eq!( + pulled_chats(pull(&base, "SECRET_TOKEN", &["s9"]).await?), + vec!["c1", "c2"] + ); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s9", "whole": true, "epoch": 1, "head": s9_head, "chats": s9_chats(&["c1"]) }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + assert_eq!( + pulled_chats(pull(&base, "SECRET_TOKEN", &["s9"]).await?), + vec!["c1"] + ); + // The marker carries the move count the push named, once; an incremental push at + // another count rides on nothing, one at the same count lands. + let s9_epochs = |listing: Value| -> Vec { + listing["sessions"] + .as_array() + .unwrap() + .iter() + .filter(|s| s["id"] == "s9") + .map(|s| s["epoch"].clone()) + .collect() + }; + assert_eq!( + s9_epochs(list(&base, "SECRET_TOKEN").await?), + vec![json!(1)] + ); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s9", "chats": s9_chats(&["c7"]) }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let answer: Value = resp.json().await?; + assert_eq!(answer["results"][0]["needs_whole"], true); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s9", "epoch": 1, "chats": s9_chats(&["c7"]) }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + assert_eq!( + pulled_chats(pull(&base, "SECRET_TOKEN", &["s9"]).await?), + vec!["c1", "c7"] + ); + + // An incremental push split over parts unlists the session while it is in progress (a + // pull between two parts would take a mix of old and new pieces for the backup) and + // lists it again with the last part; while one is in progress or abandoned, a push that + // is not part of it is refused, so the browser's next push of the session goes whole. + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s9", "epoch": 1, "push": "i1", "opens": true, "chats": s9_chats(&["c8"]), "partial": true }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + assert_eq!( + s9_epochs(list(&base, "SECRET_TOKEN").await?), + Vec::::new() + ); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s9", "epoch": 1, "chats": s9_chats(&["c11"]) }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let answer: Value = resp.json().await?; + assert_eq!(answer["results"][0]["needs_whole"], true); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s9", "epoch": 1, "push": "i1", "chats": s9_chats(&["c9"]) }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + assert_eq!( + s9_epochs(list(&base, "SECRET_TOKEN").await?), + vec![json!(1)] + ); + assert_eq!( + pulled_chats(pull(&base, "SECRET_TOKEN", &["s9"]).await?), + vec!["c1", "c7", "c8", "c9"] + ); + assert_eq!( + s9_epochs(list(&base, "SECRET_TOKEN").await?), + vec![json!(1)] + ); + + // Two devices pushing the session whole at once: the push that opened later replaced + // the earlier one's pieces, so the earlier one's last part is refused and lists nothing, + // and the session is listed with the later push's pieces alone. + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s9", "whole": true, "push": "t3", "opens": true, "head": s9_head, "chats": s9_chats(&["c3"]), "partial": true }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s9", "whole": true, "push": "t4", "opens": true, "head": s9_head, "chats": s9_chats(&["c4"]), "partial": true }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s9", "whole": true, "push": "t3", "chats": s9_chats(&["c5"]) }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + let answer: Value = resp.json().await?; + assert_eq!(answer["results"][0]["needs_whole"], true); + assert!(list(&base, "SECRET_TOKEN").await?["sessions"] + .as_array() + .unwrap() + .iter() + .all(|s| s["id"] != "s9")); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s9", "whole": true, "push": "t4", "chats": s9_chats(&["c6"]) }] + }), + ) + .await?; + assert_eq!(resp.status(), 200); + assert_eq!( + pulled_chats(pull(&base, "SECRET_TOKEN", &["s9"]).await?), + vec!["c4", "c6"] + ); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "removed": ["s9"] }), + ) + .await?; + assert_eq!(resp.status(), 200); + + // Removal empties both prefixes. + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "removed": ["s1"] }), + ) + .await?; + assert_eq!(resp.status(), 200); + let listing = list(&base, "SECRET_TOKEN").await?; + assert_eq!(listing["sessions"], json!([])); + assert!( + listing["storage_id"].is_string(), + "an answer names its storage: {listing}" + ); + assert!( + files_under(storage_dir.path()).is_empty(), + "removal must leave no object behind" + ); + + Ok(()) +} + +#[sqlx::test(fixtures("base", "jobs_read_auth"))] +async fn test_backup_writes_are_refused_for_the_wrong_owner_token_or_id( + db: Pool, +) -> anyhow::Result<()> { + initialize_tracing().await; + let server = ApiServer::start(db.clone()).await?; + let base = format!( + "http://localhost:{}/api/w/test-workspace", + server.addr.port() + ); + let storage_dir = tempfile::tempdir()?; + configure_primary_lfs(&db, &storage_dir.path().to_string_lossy()).await?; + + // A push prepared for another user must not land under the caller's prefix. + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test2@windmill.dev", "sessions": [{ "id": "s1", "head": { "id": "s1" } }] }), + ) + .await?; + assert_eq!(resp.status(), 409, "{}", resp.text().await?); + + // Ids are what the server builds keys from. + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "sessions": [{ "id": "../s1", "head": { "id": "../s1" } }] }), + ) + .await?; + assert_eq!(resp.status(), 400, "{}", resp.text().await?); + + // A whole push opens with its head; one without is refused before anything of it lands, + // and nothing lists the session. + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "sessions": [{ "id": "s7", "whole": true, "chats": [{ "id": "c", "record": { "id": "c" } }] }] }), + ) + .await?; + assert_eq!(resp.status(), 400, "{}", resp.text().await?); + assert!(list(&base, "SECRET_TOKEN").await?["sessions"] + .as_array() + .unwrap() + .iter() + .all(|s| s["id"] != "s7")); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "removed": ["a/b"] }), + ) + .await?; + assert_eq!(resp.status(), 400, "{}", resp.text().await?); + + // An image is a base64 data URL, stored and served verbatim; anything JSON would have + // to escape (and so inflate past the pull budget) is refused. + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s1", "images": [{ "chat_id": "c1", "id": "i1", "data_url": "data:image/png;base64,\u{0001}\u{0001}\"" }] }] + }), + ) + .await?; + assert_eq!(resp.status(), 400, "{}", resp.text().await?); + + // Nested lists are bounded too: each entry is an object-store call. + let resp = push( + &base, + "SECRET_TOKEN", + json!({ + "owner": "test@windmill.dev", + "sessions": [{ "id": "s1", "delete_chats": (0..1001).map(|i| format!("c{i}")).collect::>() }] + }), + ) + .await?; + assert_eq!(resp.status(), 400, "{}", resp.text().await?); + + // ...and across the whole request, not only per entry. + let sessions: Vec = (0..100) + .map(|i| { + json!({ "id": format!("s{i}"), "delete_chats": (0..50).map(|j| format!("c{j}")).collect::>() }) + }) + .collect(); + let resp = push( + &base, + "SECRET_TOKEN", + json!({ "owner": "test@windmill.dev", "sessions": sessions }), + ) + .await?; + assert_eq!(resp.status(), 400, "{}", resp.text().await?); + + // A pull body is a handful of ids; a large one is refused before it is parsed. + let resp = authed( + client().post(format!("{base}/ai/sessions/pull")), + "SECRET_TOKEN", + ) + .header("Content-Type", "application/json") + .body(format!("{{\"ids\":[\"{}\"]}}", "a".repeat(100_000))) + .send() + .await?; + assert_eq!(resp.status(), 413, "{}", resp.text().await?); + + // A scoped token (here `jobs:read`) is minted for something narrower than the user's + // whole assistant history. + let resp = authed( + client().get(format!("{base}/ai/sessions/list")), + "SCOPED_DENO_TOKEN", + ) + .send() + .await?; + assert_eq!(resp.status(), 403, "{}", resp.text().await?); + + assert!(files_under(storage_dir.path()).is_empty()); + Ok(()) +} diff --git a/backend/windmill-api-workspaces/src/ai_session_backups.rs b/backend/windmill-api-workspaces/src/ai_session_backups.rs new file mode 100644 index 0000000000..bbdbfbc8c9 --- /dev/null +++ b/backend/windmill-api-workspaces/src/ai_session_backups.rs @@ -0,0 +1,150 @@ +//! What the workspace key rotation and the AI session backup routes +//! (`windmill-api/src/ai_sessions.rs`) share about the backups in the workspace storage. +//! +//! The backups are ciphertext under the workspace key and live under a prefix named by a +//! generation the rotation bumps (`workspace_settings.ai_sessions_backup_generation`) in the +//! transaction that commits the new key. A rotation does not re-key them: once committed, +//! the routes read and write under the new generation's prefix and answer with its number +//! (`backup_generation`; the storage identity, `storage_id`, names the storage and does not +//! change), so every browser marks its sync state stale and pushes its sessions whole again +//! there, and every older generation, which nothing writes to any +//! more, is deleted off the request at leisure. Sessions no browser holds any more are lost, +//! which a rotation (a rare operation) accepts in exchange for having no key but the current +//! one to read with and nothing to rewrite in place. A generation is never reused, so no +//! deletion, however late, can touch live objects; a rotation that fails before its commit +//! bumps nothing and deletes nothing; two rotations racing serialize on the key row. + +use std::sync::Arc; + +use futures::TryStreamExt; +use windmill_common::error::{Error, Result}; +use windmill_common::utils::calculate_hash; +use windmill_common::DB; +use windmill_object_store::object_store_reexports::{ + ObjectStore, ObjectStoreError, Path as ObjectPath, +}; +use windmill_object_store::{object_store_error_to_error, ObjectStoreResource}; +use windmill_types::s3::LargeFileStorage; + +/// The root of every AI session backup key in a workspace's storage. +pub const ROOT: &str = "windmill_ai_sessions"; +/// The push body cap: no object written through the routes is larger. One that is was +/// planted by whoever holds the bucket's credentials, and is left unread. +pub const MAX_OBJECT_BYTES: usize = 32 * 1024 * 1024; + +const IO_CONCURRENCY: usize = 8; + +/// The prefix of one generation's objects: `windmill_ai_sessions/{w_id}/g{generation}/`. +pub fn generation_prefix(w_id: &str, generation: i64) -> String { + format!("{ROOT}/{w_id}/g{generation}") +} + +/// Names the storage the backups are in, by what locates its objects (endpoint, region, +/// bucket; never the credentials, which rotate), so a browser tells that its sync state was +/// recorded against another storage; the generation, answered alongside, tells it a +/// rotation happened in this one. +pub fn storage_id(resource: &ObjectStoreResource) -> String { + let location = match resource { + ObjectStoreResource::S3(s) => format!( + "s3:{}:{}:{}:{}", + s.endpoint, + s.port.unwrap_or_default(), + s.region, + s.bucket + ), + ObjectStoreResource::Azure(a) => format!( + "azure:{}:{}:{}", + a.endpoint.as_deref().unwrap_or_default(), + a.account_name, + a.container_name + ), + ObjectStoreResource::Gcs(g) => format!("gcs:{}", g.bucket), + ObjectStoreResource::Filesystem(f) => format!("fs:{}", f.root_path), + }; + calculate_hash(&location)[..16].to_string() +} + +/// The workspace's primary storage, resolved without a caller: a rotation runs the +/// deletion off its own request. +async fn primary_store(db: &DB, w_id: &str) -> Result>> { + let Some(lfs_json) = sqlx::query_scalar!( + "SELECT large_file_storage FROM workspace_settings WHERE workspace_id = $1", + w_id + ) + .fetch_optional(db) + .await? + .flatten() else { + return Ok(None); + }; + let lfs: LargeFileStorage = serde_json::from_value(lfs_json) + .map_err(|e| Error::internal_err(format!("parsing large_file_storage: {e}")))?; + let resource_value = if matches!(lfs, LargeFileStorage::FilesystemStorage(_)) { + serde_json::Value::Null + } else { + let path = lfs.get_s3_resource_path(); + let path = path.strip_prefix("$res:").unwrap_or(path); + windmill_common::workspaces::transform_json_value_unchecked( + &serde_json::Value::String(format!("$res:{path}")), + w_id, + db, + ) + .await? + }; + let resource = windmill_object_store::lfs_to_object_store_resource(&lfs, resource_value)?; + Ok(Some( + windmill_object_store::build_object_store_client(&resource).await?, + )) +} + +/// The generation an object key sits under, `None` for a key of no generation (an older +/// layout), which counts as older than any. +fn generation_of(w_id: &str, key: &ObjectPath) -> Option { + key.as_ref() + .strip_prefix(&format!("{ROOT}/{w_id}/g"))? + .split('/') + .next()? + .parse() + .ok() +} + +/// Deletes, off the request and as the listing streams, every object of the workspace's +/// backups from a generation older than `current`, once the rotation that made `current` +/// the generation has committed: nothing writes there any more but a push that resolved its +/// prefix before the commit, junk the browser's next push of that session rewrites under the +/// current prefix, as is anything a deletion cut short left behind. For the rotation route, +/// which authorized its caller as a superadmin. +pub(crate) fn spawn_delete_older(db: DB, w_id: String, current: i64) { + tokio::spawn(async move { + let store = match primary_store(&db, &w_id).await { + Ok(Some(store)) => store, + Ok(None) => return, + Err(e) => { + tracing::warn!("older AI session backups of {w_id} left in place: {e:#}"); + return; + } + }; + let prefix = ObjectPath::from(format!("{ROOT}/{w_id}")); + let deleted = store + .list(Some(&prefix)) + .map_err(object_store_error_to_error) + .try_for_each_concurrent(IO_CONCURRENCY, |meta| { + let (store, w_id) = (&store, &w_id); + async move { + if generation_of(w_id, &meta.location).is_some_and(|g| g >= current) { + return Ok(()); + } + match store.delete(&meta.location).await { + Ok(()) | Err(ObjectStoreError::NotFound { .. }) => Ok(()), + Err(e) => Err(object_store_error_to_error(e)), + } + } + }) + .await; + match deleted { + Ok(()) => { + tracing::info!("deleted the AI session backups of {w_id} older than g{current}") + } + Err(e) => tracing::warn!("deleting the older AI session backups of {w_id}: {e:#}"), + } + }); +} diff --git a/backend/windmill-api-workspaces/src/lib.rs b/backend/windmill-api-workspaces/src/lib.rs index 017c9702a5..22f2a2c5bb 100644 --- a/backend/windmill-api-workspaces/src/lib.rs +++ b/backend/windmill-api-workspaces/src/lib.rs @@ -1,6 +1,8 @@ +#[cfg(feature = "parquet")] +pub mod ai_session_backups; +pub mod data_metrics; pub mod datatable_migrations; pub mod deployment_requests; -pub mod data_metrics; pub mod workspaces; pub mod workspaces_extra; pub mod workspaces_oss; diff --git a/backend/windmill-api-workspaces/src/workspaces.rs b/backend/windmill-api-workspaces/src/workspaces.rs index 6048772ff6..1a3c24f096 100644 --- a/backend/windmill-api-workspaces/src/workspaces.rs +++ b/backend/windmill-api-workspaces/src/workspaces.rs @@ -5339,6 +5339,17 @@ async fn set_encryption_key( let mut tx = db.begin().await?; + // Under the row's lock, so two rotations racing serialize and each sees the key the + // other committed. The AI session backups in the workspace storage live under a prefix + // named by a generation this bumps (with the key, in this transaction) rather than + // being re-keyed; the older generations are deleted once this one has committed (see + // `ai_session_backups`). The same key set again is no rotation to them. + let previous_key: String = sqlx::query_scalar( + "SELECT key FROM workspace_key WHERE workspace_id = $1 AND kind = 'cloud' FOR UPDATE", + ) + .bind(&w_id) + .fetch_one(&mut *tx) + .await?; sqlx::query!( "UPDATE workspace_key SET key = $1 WHERE workspace_id = $2", request.new_key.clone(), @@ -5346,6 +5357,18 @@ async fn set_encryption_key( ) .execute(&mut *tx) .await?; + let backups_generation: Option = if previous_key != request.new_key { + sqlx::query_scalar( + "UPDATE workspace_settings SET ai_sessions_backup_generation = \ + ai_sessions_backup_generation + 1 WHERE workspace_id = $1 \ + RETURNING ai_sessions_backup_generation", + ) + .bind(&w_id) + .fetch_optional(&mut *tx) + .await? + } else { + None + }; let mut reencrypted_secret_paths: Vec = Vec::new(); if !request.skip_reencrypt.unwrap_or(false) { @@ -5402,6 +5425,15 @@ async fn set_encryption_key( // Invalidate the cache only after the transaction has committed WORKSPACE_CRYPT_CACHE.remove(w_id.as_str()); + // Nothing writes under the older generations any more; the browsers push their + // sessions again under the new one. + #[cfg(feature = "parquet")] + if let Some(generation) = backups_generation { + crate::ai_session_backups::spawn_delete_older(db.clone(), w_id.clone(), generation); + } + #[cfg(not(feature = "parquet"))] + let _ = backups_generation; + // Build the batch: one event for the encryption key itself plus one per // re-encrypted secret variable. The batch entrypoint dispatches a single // git-sync job per repo carrying all items, so repos with Secrets sync diff --git a/backend/windmill-api/openapi.yaml b/backend/windmill-api/openapi.yaml index 874653b6f9..f3f311f697 100644 --- a/backend/windmill-api/openapi.yaml +++ b/backend/windmill-api/openapi.yaml @@ -13188,6 +13188,153 @@ paths: type: boolean description: more buckets matched than were returned, so summing them under-reports + /w/{workspace}/ai/sessions/list: + get: + summary: list the calling user's AI session backups in the workspace object storage + operationId: listAiSessionBackups + tags: + - ai + parameters: + - $ref: "#/components/parameters/WorkspaceId" + responses: + "200": + description: backups, newest first; `enabled` is false when the workspace has no storage for them + content: + application/json: + schema: + type: object + required: + - enabled + - sessions + properties: + enabled: + type: boolean + storage_id: + type: string + description: names the storage answered from; sync state recorded against another one is void + backup_generation: + type: integer + description: bumped by every workspace key rotation; sync state recorded under another one is void + sessions: + type: array + description: the newest 500 at most + items: + $ref: "#/components/schemas/AISessionBackupListing" + truncated: + type: boolean + description: the user has more sessions than the answer names + + /w/{workspace}/ai/sessions/pull: + post: + summary: fetch whole AI session backups + operationId: pullAiSessionBackups + tags: + - ai + parameters: + - $ref: "#/components/parameters/WorkspaceId" + requestBody: + required: true + content: + application/json: + schema: + type: object + required: + - ids + properties: + ids: + type: array + maxItems: 20 + items: + type: string + resume: + $ref: "#/components/schemas/AISessionBackupCursor" + responses: + "200": + description: the backups found; `deferred` lists ids that did not fit the response budget + content: + application/json: + schema: + type: object + required: + - enabled + - sessions + - deferred + properties: + enabled: + type: boolean + storage_id: + type: string + backup_generation: + type: integer + sessions: + type: array + items: + $ref: "#/components/schemas/AISessionBackup" + deferred: + type: array + items: + type: string + + /w/{workspace}/ai/sessions/push: + post: + summary: write changed pieces of AI sessions to their backups, and remove deleted ones + operationId: pushAiSessionBackups + tags: + - ai + parameters: + - $ref: "#/components/parameters/WorkspaceId" + requestBody: + required: true + content: + application/json: + schema: + type: object + required: + - owner + properties: + owner: + type: string + description: the email the push was prepared for; refused with a 409 when it is not the caller's + sessions: + type: array + items: + $ref: "#/components/schemas/AISessionBackupPush" + removed: + type: array + items: + type: string + responses: + "200": + description: one result per session written or removed, in request order + content: + application/json: + schema: + type: object + required: + - enabled + - results + properties: + enabled: + type: boolean + storage_id: + type: string + backup_generation: + type: integer + results: + type: array + items: + type: object + required: + - id + properties: + id: + type: string + error: + type: string + needs_whole: + type: boolean + description: nothing was written and the session must be pushed whole again; an incremental part found no listed session to ride on (the backup was removed, or a push split over parts is in progress or was abandoned), or a later part of a push split over parts found another push had superseded it + /w/{workspace}/ai/shared_artifacts/share: post: summary: share an AI session artifact with the workspace @@ -28278,6 +28425,155 @@ components: fixes) from the workspace UI. Read from the workspace's own settings even when the providers served fall back to the instance config. AI agent steps and the AI sandbox in flows are unaffected. + sessions_storage_disabled: + type: boolean + description: >- + Stops browsers from backing their AI sessions up to the workspace's object + storage. Read from the workspace's own settings like `copilot_disabled`. + + AISessionBackupListing: + type: object + required: + - id + - updated_at + - epoch + properties: + id: + type: string + updated_at: + type: string + format: date-time + epoch: + type: integer + description: the session's move count when this copy was pushed; of a session two workspaces list, the copy with the higher one is the later + + AISessionBackupImage: + type: object + required: + - chat_id + - id + - data_url + properties: + chat_id: + type: string + id: + type: string + data_url: + type: string + + AISessionBackupChat: + type: object + required: + - id + - record + properties: + id: + type: string + record: + type: object + additionalProperties: true + + AISessionBackup: + type: object + required: + - id + - head + - chats + - images + - listing + properties: + id: + type: string + head: + type: object + additionalProperties: true + chats: + type: array + items: + $ref: "#/components/schemas/AISessionBackupChat" + images: + type: array + items: + $ref: "#/components/schemas/AISessionBackupImage" + artifacts: + type: object + additionalProperties: true + next: + $ref: "#/components/schemas/AISessionBackupCursor" + listing: + type: string + description: a fingerprint of the session's listing; pages of one session whose fingerprints differ do not belong together + moved: + type: boolean + description: the backup kept changing while this page was read, so it may mix two versions; the browser starts the session over + + AISessionBackupCursor: + type: object + description: where a pull of a session that did not fit one answer whole picks up; the rest of the session follows a pull naming that session alone with this as `resume` + required: + - id + - images + - after + properties: + id: + type: string + images: + type: boolean + after: + type: string + + AISessionBackupPush: + type: object + required: + - id + properties: + id: + type: string + head: + type: object + additionalProperties: true + chats: + type: array + items: + $ref: "#/components/schemas/AISessionBackupChat" + images: + type: array + items: + $ref: "#/components/schemas/AISessionBackupImage" + artifacts: + type: object + additionalProperties: true + delete_chats: + type: array + items: + type: string + delete_images: + type: array + items: + type: object + required: + - chat_id + - id + properties: + chat_id: + type: string + id: + type: string + partial: + type: boolean + description: more parts of this session follow, in this push or a later one; the session is not listed on this one. Such a part names its push (`push`), or it is refused + whole: + type: boolean + description: a part of a push of the session whole; the head is on the part that opens it, which replaces whatever the storage holds of the session, and every piece the browser has is on one of them. An incremental part instead rides on a session the storage lists and is refused with needs_whole when it lists none + push: + type: string + description: a push split over several parts names itself on each with a token the browser draws; the part that opens it unlists the session and the last part lists it again, and a later part is written only while that token is the one there (refused with needs_whole otherwise) + opens: + type: boolean + description: this part opens the push named by `push` + epoch: + type: integer + description: the session's move count (its record's `moves`), kept with the marker that lists the session; an incremental part rides on the marker of the same count FreeTierInfo: type: object diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index 8101064c3d..8b9c7929ab 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -451,6 +451,10 @@ pub struct AIConfig { /// and the AI sandbox are unaffected, so the providers stay in force. #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub copilot_disabled: bool, + /// Stops browsers from backing their AI sessions up to the workspace's object storage + /// (`ai_sessions.rs`). Read from the workspace's own row like `copilot_disabled`. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub sessions_storage_disabled: bool, } /// Negotiated rates in USD per million tokens. An unset cache rate is read as the @@ -527,6 +531,9 @@ pub fn workspaced_service() -> Router { #[cfg(feature = "bedrock")] let router = router.route("/check_bedrock_credentials", get(check_bedrock_credentials)); + #[cfg(feature = "parquet")] + let router = router.nest("/sessions", crate::ai_sessions::workspaced_service()); + router } diff --git a/backend/windmill-api/src/ai_sessions.rs b/backend/windmill-api/src/ai_sessions.rs new file mode 100644 index 0000000000..fca9c35fca --- /dev/null +++ b/backend/windmill-api/src/ai_sessions.rs @@ -0,0 +1,1316 @@ +//! Lazily replicated backups of the browser's AI sessions in the workspace's object storage. +//! +//! The browser keeps the sessions in IndexedDB and pushes changed pieces here in batches; an +//! empty browser restores from what was pushed. The server owns the key layout, keeps the +//! caller's own prefix the only one it can reach, and encrypts every object with the +//! workspace key so bucket credentials do not read transcripts: +//! +//! ```text +//! windmill_ai_sessions/{w_id}/g{generation}/{sha256(email)}/sessions/{sid}/head.json +//! windmill_ai_sessions/{w_id}/g{generation}/{sha256(email)}/sessions/{sid}/chats/{cid}.json +//! windmill_ai_sessions/{w_id}/g{generation}/{sha256(email)}/sessions/{sid}/artifacts.json +//! windmill_ai_sessions/{w_id}/g{generation}/{sha256(email)}/images/{sid}/{cid}/{iid} +//! windmill_ai_sessions/{w_id}/g{generation}/{sha256(email)}/index/{sid}/{epoch} +//! ``` +//! +//! The index marker is empty, written last by every push of the session, and is what a +//! listing reads: one object per session, whatever the session holds, its `last_modified` +//! the session's `updated_at`. + +use crate::db::{ApiAuthed, DB}; +use axum::{ + extract::{DefaultBodyLimit, Path}, + routing::{get, post}, + Extension, Json, Router, +}; +use futures::{StreamExt, TryStreamExt}; +use magic_crypt::{MagicCrypt256, MagicCryptTrait}; +use serde::{Deserialize, Serialize}; +use serde_json::value::RawValue; +use std::sync::Arc; +use windmill_api_auth::is_effectively_unscoped; +use windmill_api_workspaces::ai_session_backups::{ + generation_prefix, storage_id, MAX_OBJECT_BYTES, +}; +use windmill_common::error::{Error, JsonResult, Result}; +use windmill_common::utils::calculate_hash; +use windmill_common::variables::{crypt_from_key_with_suffix, get_workspace_key}; +use windmill_object_store::object_store_reexports::{ + ObjectStore, ObjectStoreError, Path as ObjectPath, PutPayload, +}; +use windmill_object_store::{build_object_store_client, object_store_error_to_error}; + +const PUSH_BODY_LIMIT: usize = MAX_OBJECT_BYTES; +/// A pull names at most MAX_PULL_IDS ids of 64 bytes; anything larger is not a pull. +const PULL_BODY_LIMIT: usize = 64 * 1024; +/// A pull answer larger than this hands the remaining ids back as `deferred`. +const PULL_RESPONSE_BUDGET: usize = 32 * 1024 * 1024; +const MAX_HEAD_BYTES: usize = 1024 * 1024; +/// What the cipher adds to a plaintext at most (a block of padding): an object stored at a +/// cap is that much larger than the cap when read back. +const CIPHER_PADDING: usize = 16; +/// The browser bounds an image to a 1568 px edge and re-encodes past 700 KB; this is +/// well above what that produces. +const MAX_IMAGE_BYTES: usize = 4 * 1024 * 1024; +const MAX_PULL_IDS: usize = 20; +const MAX_PUSH_SESSIONS: usize = 100; +const MAX_REMOVED: usize = 200; +const MAX_CHATS_PER_ENTRY: usize = 100; +const MAX_IMAGES_PER_ENTRY: usize = 500; +const MAX_DELETES_PER_ENTRY: usize = 1000; +const MAX_OPERATIONS_PER_PUSH: usize = 4000; +/// Entries of listing metadata a pull holds per page of a session. +const MAX_LISTED_OBJECTS: usize = 5000; +/// Session markers a listing scans, and the newest sessions it answers with. +const MAX_LIST_SCAN: usize = 50_000; +const LIST_MAX: usize = 500; +const IO_CONCURRENCY: usize = 8; + +pub fn workspaced_service() -> Router { + Router::new() + .route("/list", get(list)) + .route( + "/pull", + post(pull).layer(DefaultBodyLimit::max(PULL_BODY_LIMIT)), + ) + .route( + "/push", + post(push).layer(DefaultBodyLimit::max(PUSH_BODY_LIMIT)), + ) +} + +/// What reading an object yields. `Gone`: not there (deleted since the listing, or never +/// pushed). `Grown`: larger than expected, so replaced since the listing (or planted), and +/// left unread; a pull answers with a page ending before it rather than without it. +/// `Foreign`: it does not decrypt for this user (written under another user's or +/// workspace's key), and must not take the rest of the session down with it. +enum Read { + Text(String), + Gone, + Grown, + Foreign, +} + +/// The user's prefix in the workspace storage, plus what reads and writes it. +struct Backend { + store: Arc, + mc: MagicCrypt256, + prefix: String, + /// Name the storage and the generation the objects are under, for the browser's sync + /// state: a row recorded against another storage or generation is stale, a removal is + /// owed to the storage alone (a rotation deleted the older generation's copy anyway). + storage_id: String, + generation: i64, +} + +impl Backend { + fn index_prefix(&self) -> ObjectPath { + ObjectPath::from(format!("{}/index/", self.prefix)) + } + + /// The marker that lists the session, named by the session's move count so that of a + /// session two workspaces list, the copy moved last is told from the listing alone. + fn index_key(&self, sid: &str, epoch: u32) -> ObjectPath { + ObjectPath::from(format!("{}/index/{sid}/{epoch}", self.prefix)) + } + + fn index_session_prefix(&self, sid: &str) -> ObjectPath { + ObjectPath::from(format!("{}/index/{sid}/", self.prefix)) + } + + /// The token of the push split over parts in progress, under the session so a removal + /// or the next whole push clears it with the rest. + fn push_key(&self, sid: &str) -> ObjectPath { + ObjectPath::from(format!("{}/sessions/{sid}/push", self.prefix)) + } + + fn session_prefix(&self, sid: &str) -> ObjectPath { + ObjectPath::from(format!("{}/sessions/{sid}/", self.prefix)) + } + + fn head_key(&self, sid: &str) -> ObjectPath { + ObjectPath::from(format!("{}/sessions/{sid}/head.json", self.prefix)) + } + + fn chat_key(&self, sid: &str, cid: &str) -> ObjectPath { + ObjectPath::from(format!("{}/sessions/{sid}/chats/{cid}.json", self.prefix)) + } + + fn artifacts_key(&self, sid: &str) -> ObjectPath { + ObjectPath::from(format!("{}/sessions/{sid}/artifacts.json", self.prefix)) + } + + fn images_prefix(&self, sid: &str) -> ObjectPath { + ObjectPath::from(format!("{}/images/{sid}/", self.prefix)) + } + + fn chat_images_prefix(&self, sid: &str, cid: &str) -> ObjectPath { + ObjectPath::from(format!("{}/images/{sid}/{cid}/", self.prefix)) + } + + fn image_key(&self, sid: &str, cid: &str, iid: &str) -> ObjectPath { + ObjectPath::from(format!("{}/images/{sid}/{cid}/{iid}", self.prefix)) + } + + fn seal(&self, plaintext: &[u8]) -> Vec { + self.mc.encrypt_bytes_to_bytes(plaintext) + } + + /// Bytes written. + async fn put_sealed(&self, key: &ObjectPath, ciphertext: Vec) -> Result { + let written = ciphertext.len(); + self.store + .put(key, PutPayload::from(ciphertext)) + .await + .map_err(object_store_error_to_error)?; + Ok(written) + } + + async fn put(&self, key: &ObjectPath, plaintext: &[u8]) -> Result { + self.put_sealed(key, self.seal(plaintext)).await + } + + /// `max` is what the listing said the object holds, or the cap of its kind for one read + /// without a listing: checked before buffering, since whoever holds the bucket's + /// credentials can put anything at a predictable key. + async fn get(&self, key: &ObjectPath, max: usize) -> Result { + let result = match self.store.get(key).await { + Ok(result) => result, + Err(ObjectStoreError::NotFound { .. }) => return Ok(Read::Gone), + Err(e) => return Err(object_store_error_to_error(e)), + }; + let size = result.meta.size as usize; + // Larger than any push writes: planted, whatever the listing said, and skipped like + // an object of another key rather than retried like one that grew. + if size > MAX_OBJECT_BYTES { + tracing::warn!("AI session backup object {key} is larger than any push writes"); + return Ok(Read::Foreign); + } + if size > max { + return Ok(Read::Grown); + } + let bytes = result.bytes().await.map_err(object_store_error_to_error)?; + // The objects are JSON and data URLs: a wrong key's output failing UTF-8 tells it + // apart beyond the cipher's padding check, which a wrong key passes now and then. + match self + .mc + .decrypt_bytes_to_bytes(&bytes) + .ok() + .and_then(|plaintext| String::from_utf8(plaintext).ok()) + { + Some(text) => Ok(Read::Text(text)), + None => { + tracing::warn!("AI session backup object {key} does not decrypt for its reader"); + Ok(Read::Foreign) + } + } + } + + async fn delete(&self, key: &ObjectPath) -> Result<()> { + match self.store.delete(key).await { + Ok(()) | Err(ObjectStoreError::NotFound { .. }) => Ok(()), + Err(e) => Err(object_store_error_to_error(e)), + } + } + + /// The entries under `prefix` past `after` in key order, as many as fit `budget` bytes; + /// `true` when more follow. Every key past `after` is seen and the MAX_LISTED_OBJECTS + /// smallest kept (a max-heap dropping its largest), since a page is defined by key + /// order and the store promises none; that cap is what bounds a pull's memory, a + /// session growing by valid pushes without limit. With `at_least_one`, the first entry + /// is taken whatever its size, so an answer owed the session makes progress on it (no + /// object exceeds the push body cap). + async fn list_within( + &self, + prefix: &ObjectPath, + after: Option<&ObjectPath>, + budget: usize, + at_least_one: bool, + ) -> Result<(Vec<(ObjectPath, usize)>, bool)> { + let mut kept: std::collections::BinaryHeap<(ObjectPath, usize)> = Default::default(); + let mut dropped = false; + let mut stream = match after { + Some(after) => self.store.list_with_offset(Some(prefix), after), + None => self.store.list(Some(prefix)), + }; + while let Some(meta) = stream.next().await { + let meta = meta.map_err(object_store_error_to_error)?; + kept.push((meta.location, meta.size as usize)); + if kept.len() > MAX_LISTED_OBJECTS { + kept.pop(); + dropped = true; + } + } + let mut entries = vec![]; + let mut total = 0; + for (key, size) in kept.into_sorted_vec() { + if total + size > budget && !(at_least_one && entries.is_empty()) { + return Ok((entries, true)); + } + total += size; + entries.push((key, size)); + } + Ok((entries, dropped)) + } + + /// Bytes written. Sealed up front so every stream item is owned: an item borrowing + /// from the request makes the future higher-ranked over that lifetime, which the + /// handler's `Send` bound cannot prove. + async fn put_all(&self, puts: Vec<(ObjectPath, Vec)>) -> Result { + futures::stream::iter(puts) + .map(|(key, ciphertext)| async move { self.put_sealed(&key, ciphertext).await }) + .buffer_unordered(IO_CONCURRENCY) + .try_fold(0, |acc, n| async move { Ok::<_, Error>(acc + n) }) + .await + } + + async fn delete_all(&self, keys: Vec) -> Result<()> { + futures::stream::iter(keys) + .map(|key| async move { self.delete(&key).await }) + .buffer_unordered(IO_CONCURRENCY) + .try_collect::>() + .await?; + Ok(()) + } + + /// A fingerprint of the session's marker and of everything listed under its two + /// prefixes (key, size, modification time, entity tag and version), combined as the + /// listing streams and in no particular order, so a session of any size costs bounded + /// memory. `None` for a session + /// the storage does not list. Taken before and after a page is read, so a page a push + /// changed under is read again; pages of one pull carry it, and the browser starts the + /// session over when it moved between two of them. + async fn listing_fingerprint(&self, sid: &str) -> Result> { + use std::hash::{DefaultHasher, Hash, Hasher}; + // The entity tag and version go in with the key, size and time: a store reports + // modification times coarsely, and an object rewritten at the same size within that + // grain would otherwise fingerprint the same. + fn fold( + acc: u64, + location: &str, + size: S, + modified: i64, + e_tag: Option<&str>, + version: Option<&str>, + ) -> u64 { + let mut hasher = DefaultHasher::new(); + (location, size, modified, e_tag, version).hash(&mut hasher); + acc.wrapping_add(hasher.finish()) + } + let mut acc = 0u64; + let mut listed = false; + for (marker, prefix) in [ + (true, self.index_session_prefix(sid)), + (false, self.session_prefix(sid)), + (false, self.images_prefix(sid)), + ] { + let mut stream = self.store.list(Some(&prefix)); + while let Some(meta) = stream.next().await { + let meta = meta.map_err(object_store_error_to_error)?; + listed |= marker; + acc = fold( + acc, + meta.location.as_ref(), + meta.size, + meta.last_modified.timestamp_millis(), + meta.e_tag.as_deref(), + meta.version.as_deref(), + ); + } + } + if !listed { + return Ok(None); + } + Ok(Some(format!("{acc:016x}"))) + } + + async fn exists(&self, key: &ObjectPath) -> Result { + match self.store.head(key).await { + Ok(_) => Ok(true), + Err(ObjectStoreError::NotFound { .. }) => Ok(false), + Err(e) => Err(object_store_error_to_error(e)), + } + } + + /// Deletes as the listing streams, so a prefix of any size costs bounded memory. + async fn delete_prefix(&self, prefix: &ObjectPath) -> Result<()> { + self.store + .list(Some(prefix)) + .map_err(object_store_error_to_error) + .try_for_each_concurrent(IO_CONCURRENCY, |meta| async move { + self.delete(&meta.location).await + }) + .await + } +} + +/// Backups are the user's own browser state and nothing else may reach them: a job token +/// may carry an `on_behalf_of` identity, and every scoped token (guest, embed, app policy, +/// MCP) is minted for something narrower than the user's whole assistant history. +fn require_plain_user_token(authed: &ApiAuthed) -> Result<()> { + if authed.job_id.is_some() || !is_effectively_unscoped(authed.scopes.as_deref()) { + return Err(Error::PermissionDenied( + "AI session backups are only reachable with an unscoped user token".to_string(), + )); + } + Ok(()) +} + +/// Every key is assembled server-side from ids the browser mints (`createLongHash` and +/// `randomUUID` forms), so anything outside this alphabet is a forged id, not a real one. +fn require_valid_id(kind: &str, id: &str) -> Result<()> { + let ok = !id.is_empty() + && id.len() <= 64 + && id + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-'); + if !ok { + return Err(Error::BadRequest(format!("invalid {kind} id: {id:?}"))); + } + Ok(()) +} + +/// Images travel as base64 data URLs and are stored verbatim, so they serialize back into +/// a pull answer at exactly their stored size; anything else (control characters, +/// quotes) could grow several times under JSON escaping and defeat the pull budget. +fn require_data_url(data_url: &str) -> Result<()> { + let ok = data_url.len() <= MAX_IMAGE_BYTES + && data_url + .strip_prefix("data:") + .and_then(|rest| rest.split_once(";base64,")) + .is_some_and(|(mime, payload)| { + !mime.is_empty() + && mime.bytes().all(|b| { + b.is_ascii_alphanumeric() || matches!(b, b'/' | b'.' | b'+' | b'-') + }) + && payload + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'+' | b'/' | b'=')) + }); + if !ok { + return Err(Error::BadRequest( + "an image must be a base64 data URL within the size cap".to_string(), + )); + } + Ok(()) +} + +fn require_json_object(kind: &str, raw: &RawValue, max_bytes: usize) -> Result<()> { + let text = raw.get(); + if !text.trim_start().starts_with('{') { + return Err(Error::BadRequest(format!("{kind} must be a JSON object"))); + } + if text.len() > max_bytes { + return Err(Error::BadRequest(format!( + "{kind} exceeds {max_bytes} bytes" + ))); + } + Ok(()) +} + +/// `None` when the workspace has nowhere to keep backups: no primary storage configured, or +/// the admin switched them off. Both read as `enabled: false` so the browser stops trying. +async fn backend(authed: &ApiAuthed, db: &DB, w_id: &str) -> Result> { + let (disabled, generation) = sqlx::query_as::<_, (Option, i64)>( + "SELECT (ai_config->>'sessions_storage_disabled')::bool, ai_sessions_backup_generation \ + FROM workspace_settings WHERE workspace_id = $1", + ) + .bind(w_id) + .fetch_optional(db) + .await? + .unwrap_or((None, 0)); + if disabled.unwrap_or(false) { + return Ok(None); + } + let (_, resource) = + crate::job_helpers_oss::get_workspace_s3_resource(authed, db, None, w_id, None).await?; + let Some(resource) = resource else { + return Ok(None); + }; + let store = build_object_store_client(&resource).await?; + let user = calculate_hash(&authed.email); + // Keyed per user, not per workspace: anyone who can write the bucket could otherwise copy + // another member's ciphertext under their own prefix and have `pull` decrypt it for them. + let key = get_workspace_key(w_id, db).await?; + let mc = crypt_from_key_with_suffix(&key, &user); + let storage_id = storage_id(&resource); + let prefix = format!("{}/{user}", generation_prefix(w_id, generation)); + Ok(Some(Backend { store, mc, prefix, storage_id, generation })) +} + +#[derive(Serialize)] +struct SessionListing { + id: String, + updated_at: chrono::DateTime, + /// The session's move count when this copy was pushed (see `PushedSession::epoch`). + epoch: u32, +} + +#[derive(Serialize)] +struct ListResponse { + enabled: bool, + /// The storage answered from, and the generation a key rotation bumps; a browser whose + /// sync state names another storage or generation starts over. + #[serde(skip_serializing_if = "Option::is_none")] + storage_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + backup_generation: Option, + sessions: Vec, + /// The user has more sessions than the answer names. + #[serde(skip_serializing_if = "std::ops::Not::not")] + truncated: bool, +} + +/// A session is listed once a push entry of it landed whole (its marker is written last); +/// a push that failed before that left objects the listing does not name. +async fn list( + authed: ApiAuthed, + Extension(db): Extension, + Path(w_id): Path, +) -> JsonResult { + require_plain_user_token(&authed)?; + let Some(backend) = backend(&authed, &db, &w_id).await? else { + return Ok(Json(ListResponse { + enabled: false, + storage_id: None, + backup_generation: None, + sessions: vec![], + truncated: false, + })); + }; + let prefix = backend.index_prefix(); + let mut stream = backend.store.list(Some(&prefix)); + // One marker per session, whatever the session holds: the newest LIST_MAX are kept as + // the scan goes (a min-heap drops the oldest), and the scan itself is bounded. + let mut newest: std::collections::BinaryHeap< + std::cmp::Reverse<(chrono::DateTime, u32, String)>, + > = Default::default(); + let mut scanned = 0; + let mut truncated = false; + while let Some(meta) = stream.next().await { + let meta = meta.map_err(object_store_error_to_error)?; + scanned += 1; + if scanned > MAX_LIST_SCAN { + truncated = true; + break; + } + // `Path` drops the trailing delimiter, so the remainder starts with one. + let Some(rel) = meta.location.as_ref().strip_prefix(prefix.as_ref()) else { + continue; + }; + let Some((sid, epoch)) = rel.trim_start_matches('/').split_once('/') else { + continue; + }; + let Ok(epoch) = epoch.parse::() else { + continue; + }; + if sid.is_empty() || sid.contains('/') { + continue; + } + newest.push(std::cmp::Reverse(( + meta.last_modified, + epoch, + sid.to_string(), + ))); + if newest.len() > LIST_MAX { + newest.pop(); + truncated = true; + } + } + let mut sessions: Vec = newest + .into_iter() + .map(|std::cmp::Reverse((updated_at, epoch, id))| SessionListing { id, updated_at, epoch }) + .collect(); + sessions.sort_by(|a, b| b.updated_at.cmp(&a.updated_at)); + Ok(Json(ListResponse { + enabled: true, + storage_id: Some(backend.storage_id.clone()), + backup_generation: Some(backend.generation), + sessions, + truncated, + })) +} + +#[derive(Deserialize)] +struct PullRequest { + ids: Vec, + /// Picks the session an earlier answer cut up from where it stopped; `ids` then names + /// that session alone. + #[serde(default)] + resume: Option, +} + +/// Where a pull of a session that outgrew one answer picks up: the last key the earlier +/// answer carried, in the session's prefix or, once that one is done, in its images prefix. +#[derive(Serialize, Deserialize, Clone)] +struct PullCursor { + id: String, + images: bool, + after: String, +} + +#[derive(Serialize)] +struct PulledChat { + id: String, + record: Box, +} + +#[derive(Serialize, Deserialize)] +struct ImageObject { + chat_id: String, + id: String, + data_url: String, +} + +#[derive(Serialize)] +struct PulledSession { + id: String, + head: Box, + chats: Vec, + images: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + artifacts: Option>, + /// The session did not fit this answer whole: the rest follows a pull with this cursor. + #[serde(skip_serializing_if = "Option::is_none")] + next: Option, + /// A fingerprint of the session's listing (marker, and every key, size, modification + /// time, entity tag and version), so the browser tells that the backup changed between + /// the pages it assembled. + listing: String, + /// The backup kept changing while this page was read (a push landing object by object), + /// so the page may mix two versions: the browser starts the session over. + #[serde(skip_serializing_if = "std::ops::Not::not")] + moved: bool, +} + +/// How many times a page whose listing moved while it was read is read again before it is +/// handed over as `moved`. +const PULL_REREADS: usize = 3; + +#[derive(Serialize)] +struct PullResponse { + enabled: bool, + #[serde(skip_serializing_if = "Option::is_none")] + storage_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + backup_generation: Option, + sessions: Vec, + /// Ids that did not fit the response budget; ask for them again. + deferred: Vec, +} + +fn raw(kind: &str, text: String) -> Result> { + RawValue::from_string(text) + .map_err(|e| Error::internal_err(format!("stored {kind} is not JSON: {e}"))) +} + +enum PullStep { + Absent, + Deferred, + Fetched(PulledSession, usize), +} + +/// Fetch one session: its head, then every chat and artifact object under its prefix, and +/// as many of its images as the budget allows (a missing image hydrates to a placeholder +/// in the browser). Sizes come from the listings, so a session that would not fit is +/// deferred before anything of it is read, unless it is the first of the response, which +/// must carry something. `Absent` when it has no head. +async fn pull_session( + backend: &Backend, + sid: &str, + budget: usize, + first: bool, + resume: Option<&PullCursor>, +) -> Result { + // A page read while a push lands object by object may mix two versions of the session: + // the listing is taken again once the page is read, and a page it moved under is read + // again, a few times, then handed over as such for the browser to start over. + for reread in 0..PULL_REREADS { + let step = pull_page(backend, sid, budget, first, resume).await?; + let PullStep::Fetched(mut page, size) = step else { + return Ok(step); + }; + if backend.listing_fingerprint(sid).await?.as_deref() == Some(page.listing.as_str()) { + return Ok(PullStep::Fetched(page, size)); + } + if reread + 1 == PULL_REREADS { + page.moved = true; + return Ok(PullStep::Fetched(page, size)); + } + } + unreachable!("a page is answered on the last reread") +} + +async fn pull_page( + backend: &Backend, + sid: &str, + budget: usize, + first: bool, + resume: Option<&PullCursor>, +) -> Result { + // Taken before anything of the page is listed or read: an object landing after it is + // in the next page's fingerprint, whereas one landing after the reads but before a + // fingerprint taken then would have certified a page without it. A session the storage + // does not list (removed, or a whole push in progress) is absent. + let Some(listing) = backend.listing_fingerprint(sid).await? else { + return Ok(PullStep::Absent); + }; + let Read::Text(head) = backend + .get(&backend.head_key(sid), MAX_HEAD_BYTES + CIPHER_PADDING) + .await? + else { + return Ok(PullStep::Absent); + }; + let session_prefix = backend.session_prefix(sid); + let images_prefix = backend.images_prefix(sid); + let mut size = head.len(); + let mut chats = vec![]; + let mut artifacts = None; + let mut images = vec![]; + let mut next = None; + let cursor = |images: bool, after: String| PullCursor { id: sid.to_string(), images, after }; + // Sizes come from the listings, and the listings stop at the budget, so nothing is read + // past it even for the first session of the answer. One that outgrew it (chats + // accumulate over pushes) comes back in pages, in key order, each answer naming where + // the next picks up; the browser imports nothing before the last page. An object that + // grew since the listing (a push replaced it) ends the page just before it, and the + // answer names that spot: a new listing sizes it, whereas dropping it would import the + // session without it for good. + let in_images = resume.is_some_and(|c| c.images); + if !in_images { + let after = resume.map(|c| ObjectPath::from(c.after.as_str())); + let (entries, cut) = backend + .list_within( + &session_prefix, + after.as_ref(), + budget.saturating_sub(size), + first, + ) + .await?; + if cut && !first { + return Ok(PullStep::Deferred); + } + if cut { + next = entries + .last() + .map(|(key, _)| cursor(false, key.to_string())); + } + let to_read: Vec<(ObjectPath, usize, Option)> = entries + .into_iter() + .filter_map(|(key, bytes)| { + let rel = key + .as_ref() + .strip_prefix(session_prefix.as_ref()) + .unwrap_or_default() + .trim_start_matches('/'); + if rel == "artifacts.json" { + Some((key, bytes, None)) + } else { + let cid = rel + .strip_prefix("chats/")? + .strip_suffix(".json")? + .to_string(); + Some((key, bytes, Some(cid))) + } + }) + .collect(); + let reads: Vec<(ObjectPath, usize, Option, Read)> = futures::stream::iter(to_read) + .map(|(key, bytes, cid)| async move { + let read = backend.get(&key, bytes).await?; + Ok::<_, Error>((key, bytes, cid, read)) + }) + .buffered(IO_CONCURRENCY) + .try_collect() + .await?; + let mut before = resume.map(|c| c.after.clone()).unwrap_or_default(); + for (key, bytes, cid, read) in reads { + match (cid, read) { + (_, Read::Grown) => { + next = Some(cursor(false, before)); + break; + } + (None, Read::Text(text)) => { + artifacts = Some(raw("artifacts", text)?); + size += bytes; + } + (Some(cid), Read::Text(text)) => { + chats.push(PulledChat { id: cid, record: raw("chat", text)? }); + size += bytes; + } + _ => {} + } + before = key.to_string(); + } + } + if next.is_none() { + let after = resume + .filter(|c| c.images) + .map(|c| ObjectPath::from(c.after.as_str())); + let (entries, cut) = backend + .list_within( + &images_prefix, + after.as_ref(), + budget.saturating_sub(size), + first, + ) + .await?; + let mut before = after.map(|a| a.to_string()).unwrap_or_default(); + if cut { + // An answer with no room for a single image names where it stood, so the pull + // owed the session alone picks it up there. + next = Some(cursor( + true, + entries + .last() + .map(|(key, _)| key.to_string()) + .unwrap_or_else(|| before.clone()), + )); + } + let to_read: Vec<(ObjectPath, usize, String, String)> = entries + .into_iter() + .filter_map(|(key, bytes)| { + let rel = key.as_ref().strip_prefix(images_prefix.as_ref())?; + let (cid, iid) = rel.trim_start_matches('/').split_once('/')?; + Some((key.clone(), bytes, cid.to_string(), iid.to_string())) + }) + .collect(); + let reads: Vec<(ObjectPath, usize, String, String, Read)> = futures::stream::iter(to_read) + .map(|(key, bytes, cid, iid)| async move { + let read = backend.get(&key, bytes).await?; + Ok::<_, Error>((key, bytes, cid, iid, read)) + }) + .buffered(IO_CONCURRENCY) + .try_collect() + .await?; + for (key, bytes, chat_id, id, read) in reads { + match read { + Read::Grown => { + next = Some(cursor(true, before)); + break; + } + Read::Text(data_url) => { + images.push(ImageObject { chat_id, id, data_url }); + size += bytes; + } + _ => {} + } + before = key.to_string(); + } + } + Ok(PullStep::Fetched( + PulledSession { + id: sid.to_string(), + head: raw("head", head)?, + chats, + images, + artifacts, + next, + listing, + moved: false, + }, + size, + )) +} + +async fn pull( + authed: ApiAuthed, + Extension(db): Extension, + Path(w_id): Path, + Json(req): Json, +) -> JsonResult { + require_plain_user_token(&authed)?; + if req.ids.len() > MAX_PULL_IDS { + return Err(Error::BadRequest(format!( + "at most {MAX_PULL_IDS} sessions per pull" + ))); + } + for id in &req.ids { + require_valid_id("session", id)?; + } + if let Some(cursor) = &req.resume { + if req.ids.len() != 1 || req.ids[0] != cursor.id || cursor.after.len() > 1024 { + return Err(Error::BadRequest( + "a resumed pull names the resumed session alone".to_string(), + )); + } + } + let Some(backend) = backend(&authed, &db, &w_id).await? else { + return Ok(Json(PullResponse { + enabled: false, + storage_id: None, + backup_generation: None, + sessions: vec![], + deferred: vec![], + })); + }; + let mut sessions = vec![]; + let mut deferred = vec![]; + let mut budget = PULL_RESPONSE_BUDGET; + for sid in req.ids { + let resume = req.resume.as_ref().filter(|c| c.id == sid); + match pull_session(&backend, &sid, budget, sessions.is_empty(), resume).await? { + PullStep::Absent => {} + PullStep::Deferred => deferred.push(sid), + PullStep::Fetched(session, size) => { + budget = budget.saturating_sub(size); + sessions.push(session); + } + } + } + Ok(Json(PullResponse { + enabled: true, + storage_id: Some(backend.storage_id), + backup_generation: Some(backend.generation), + sessions, + deferred, + })) +} + +#[derive(Deserialize)] +struct PushedChat { + id: String, + record: Box, +} + +#[derive(Deserialize)] +struct ImageRef { + chat_id: String, + id: String, +} + +#[derive(Deserialize)] +struct PushedSession { + id: String, + #[serde(default)] + head: Option>, + #[serde(default)] + chats: Vec, + #[serde(default)] + images: Vec, + #[serde(default)] + artifacts: Option>, + #[serde(default)] + delete_chats: Vec, + #[serde(default)] + delete_images: Vec, + /// More parts of the session follow, in this push or a later one: the session is not + /// listed on this one. + #[serde(default)] + partial: bool, + /// A part of a push of the session whole: the head is on the part that opens it, which + /// replaces whatever the storage holds of the session, and every piece the browser has + /// is on one of them. An incremental part instead rides on a session the storage lists, + /// and is refused with `needs_whole` when it lists none. + #[serde(default)] + whole: bool, + /// A push split over several parts names itself on each of them with a token the + /// browser draws; the part that `opens` it unlists the session (a pull between two parts + /// would otherwise take a mix of old and new pieces for the backup) and the last part + /// lists it again. A later part is written only while that token is the one there, so a + /// part of a push another one superseded is refused with `needs_whole`. + #[serde(default)] + push: Option, + #[serde(default)] + opens: bool, + /// The session's move count (its record's `moves`), the marker that lists the session + /// is named by: a session moved to another workspace is listed by both until the old + /// copy's removal lands, and the copy with the higher count is the later one. An + /// incremental part rides on the marker of the same count. + #[serde(default)] + epoch: u32, +} + +#[derive(Deserialize)] +struct PushRequest { + /// The email the browser believes it is acting for. An in-place account switch can + /// leave a flush prepared for the previous user; the server refuses it rather than + /// filing that user's sessions under the caller's prefix. + owner: String, + #[serde(default)] + sessions: Vec, + #[serde(default)] + removed: Vec, +} + +#[derive(Serialize)] +struct PushResult { + id: String, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, + /// Nothing was written; the session must be pushed whole again: an incremental part + /// found no listed session to ride on (another device removed the backup, or a push + /// split over parts is in progress or was abandoned), or a later part of a push split + /// over parts found another push had superseded it. + #[serde(skip_serializing_if = "std::ops::Not::not")] + needs_whole: bool, +} + +#[derive(Serialize)] +struct PushResponse { + enabled: bool, + #[serde(skip_serializing_if = "Option::is_none")] + storage_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + backup_generation: Option, + results: Vec, +} + +fn validate_push(req: &PushRequest) -> Result<()> { + if req.sessions.len() > MAX_PUSH_SESSIONS || req.removed.len() > MAX_REMOVED { + return Err(Error::BadRequest( + "too many sessions in one push".to_string(), + )); + } + // Every nested entry costs an object-store call (a deleted chat two), so the lists are + // bounded per entry and across the request; the browser sends far fewer. + let mut operations = req.removed.len(); + for s in &req.sessions { + if s.chats.len() > MAX_CHATS_PER_ENTRY + || s.images.len() > MAX_IMAGES_PER_ENTRY + || s.delete_chats.len() > MAX_DELETES_PER_ENTRY + || s.delete_images.len() > MAX_DELETES_PER_ENTRY + { + return Err(Error::BadRequest(format!( + "too many pieces for session {} in one push", + s.id + ))); + } + operations += s.chats.len() + s.images.len() + s.delete_chats.len() + s.delete_images.len(); + } + if operations > MAX_OPERATIONS_PER_PUSH { + return Err(Error::BadRequest("too many pieces in one push".to_string())); + } + for sid in &req.removed { + require_valid_id("session", sid)?; + } + for s in &req.sessions { + require_valid_id("session", &s.id)?; + if let Some(token) = &s.push { + require_valid_id("push", token)?; + } else if s.opens || s.partial { + // A part more parts follow belongs to a push split over parts, which names + // itself: without the token the session would stay listed between the parts. + return Err(Error::BadRequest(format!( + "session {} is pushed in parts with no push token", + s.id + ))); + } + if s.whole && (s.push.is_none() || s.opens) && s.head.is_none() { + return Err(Error::BadRequest(format!( + "session {} is pushed whole without its head", + s.id + ))); + } + if let Some(head) = &s.head { + require_json_object("head", head, MAX_HEAD_BYTES)?; + } + for c in &s.chats { + require_valid_id("chat", &c.id)?; + require_json_object("chat record", &c.record, PUSH_BODY_LIMIT)?; + } + if let Some(a) = &s.artifacts { + require_json_object("artifacts", a, PUSH_BODY_LIMIT)?; + } + for i in &s.images { + require_valid_id("chat", &i.chat_id)?; + require_valid_id("image", &i.id)?; + require_data_url(&i.data_url)?; + } + for c in &s.delete_chats { + require_valid_id("chat", c)?; + } + for i in &s.delete_images { + require_valid_id("chat", &i.chat_id)?; + require_valid_id("image", &i.id)?; + } + } + Ok(()) +} + +fn push_payload_bytes(req: &PushRequest) -> usize { + req.sessions + .iter() + .map(|s| { + s.head.as_ref().map_or(0, |h| h.get().len()) + + s.artifacts.as_ref().map_or(0, |a| a.get().len()) + + s.chats.iter().map(|c| c.record.get().len()).sum::() + + s.images.iter().map(|i| i.data_url.len()).sum::() + }) + .sum() +} + +/// Runs under the session's lock (see `lock_session`). The part that opens a whole push +/// (the one with the head) replaces the backup: the marker goes first, so nothing lists the +/// session until the last part, then everything else. An incremental part assumes the rest +/// of the session is in the storage, which a removal since would have taken, or a push +/// split over parts may still be bringing: it is refused unless the session is listed, and +/// unlists the session itself while it changes more than one object (a pull between two +/// writes would otherwise take a mix of old and new pieces for the backup). A push split +/// over parts names itself with a token: the part that opens it unlists the session and +/// writes the token, and a later part is written only while that token is +/// the one there, so two devices pushing the session at once cannot list a mix of their +/// pieces: the push that opened later wins, the other is refused and goes again. Every +/// refusal comes before anything of the part lands. Deletes run last, and the marker only +/// by the last part, so a push cut short never leaves a listed session pointing at chats +/// that are not there. +/// Bytes written, and whether the part was refused for the session to go whole. +async fn push_session(backend: &Backend, s: &PushedSession) -> Result<(usize, bool)> { + match &s.push { + Some(token) if !s.opens => { + match backend + .get(&backend.push_key(&s.id), token.len() + CIPHER_PADDING) + .await? + { + Read::Text(current) if current == *token => {} + _ => return Ok((0, true)), + } + } + _ => { + if s.whole { + backend + .delete_prefix(&backend.index_session_prefix(&s.id)) + .await?; + backend + .delete_prefix(&backend.session_prefix(&s.id)) + .await?; + backend.delete_prefix(&backend.images_prefix(&s.id)).await?; + } else { + if !backend.exists(&backend.index_key(&s.id, s.epoch)).await? { + return Ok((0, true)); + } + // Unlisted while more than one object changes (a push split over parts, or + // one part touching several pieces): a pull between two of the writes, or + // after one of them failed, would otherwise take a mix of old and new + // pieces for the backup. One object changing is one write. + let pieces = s.chats.len() + + s.images.len() + + usize::from(s.artifacts.is_some()) + + usize::from(s.head.is_some()) + + s.delete_chats.len() + + s.delete_images.len(); + if s.push.is_some() || pieces > 1 { + backend + .delete_prefix(&backend.index_session_prefix(&s.id)) + .await?; + } + } + if let Some(token) = &s.push { + backend + .put(&backend.push_key(&s.id), token.as_bytes()) + .await?; + } + } + } + let mut written = 0; + written += backend + .put_all( + s.images + .iter() + .map(|img| { + ( + backend.image_key(&s.id, &img.chat_id, &img.id), + backend.seal(img.data_url.as_bytes()), + ) + }) + .collect(), + ) + .await?; + written += backend + .put_all( + s.chats + .iter() + .map(|c| { + ( + backend.chat_key(&s.id, &c.id), + backend.seal(c.record.get().as_bytes()), + ) + }) + .collect(), + ) + .await?; + if let Some(a) = &s.artifacts { + written += backend + .put(&backend.artifacts_key(&s.id), a.get().as_bytes()) + .await?; + } + if let Some(h) = &s.head { + written += backend + .put(&backend.head_key(&s.id), h.get().as_bytes()) + .await?; + } + for cid in &s.delete_chats { + backend.delete(&backend.chat_key(&s.id, cid)).await?; + backend + .delete_prefix(&backend.chat_images_prefix(&s.id, cid)) + .await?; + } + backend + .delete_all( + s.delete_images + .iter() + .map(|i| backend.image_key(&s.id, &i.chat_id, &i.id)) + .collect(), + ) + .await?; + if s.partial { + return Ok((written, false)); + } + // Last, and by the last part only, so a session is listed once its whole entry landed. + backend + .store + .put(&backend.index_key(&s.id, s.epoch), PutPayload::new()) + .await + .map_err(object_store_error_to_error)?; + if s.push.is_some() { + backend.delete(&backend.push_key(&s.id)).await?; + } + Ok((written, false)) +} + +/// The marker goes first so a removal cut short leaves nothing listed, then the head so +/// nothing pulls either, and no push takes it for a session still there (see `push_session`). +async fn remove_session(backend: &Backend, sid: &str) -> Result<()> { + backend + .delete_prefix(&backend.index_session_prefix(sid)) + .await?; + backend.delete(&backend.head_key(sid)).await?; + backend.delete_prefix(&backend.session_prefix(sid)).await?; + backend.delete_prefix(&backend.images_prefix(sid)).await +} + +/// One writer per session at a time, across servers: a push and a removal of the same +/// session interleaving object by object could leave a listed session missing pieces, or a +/// marker over nothing. The lock lives in a transaction that writes no rows; it is released +/// when the transaction ends. A wait past the timeout fails that entry only, and the +/// browser retries it with backoff. +async fn lock_session( + db: &DB, + backend: &Backend, + sid: &str, +) -> Result> { + let mut tx = db.begin().await?; + sqlx::query("SET LOCAL lock_timeout = '30s'") + .execute(&mut *tx) + .await?; + sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0::int8))") + .bind(format!("ai_session_backup:{}/{sid}", backend.prefix)) + .execute(&mut *tx) + .await + .map_err(|e| { + Error::internal_err(format!( + "another device is writing the backup of session {sid}; retried later ({e})" + )) + })?; + Ok(tx) +} + +async fn push_session_locked( + db: &DB, + backend: &Backend, + s: &PushedSession, +) -> Result<(usize, bool)> { + let tx = lock_session(db, backend, &s.id).await?; + let result = push_session(backend, s).await; + tx.commit().await?; + result +} + +async fn remove_session_locked(db: &DB, backend: &Backend, sid: &str) -> Result<()> { + let tx = lock_session(db, backend, sid).await?; + let result = remove_session(backend, sid).await; + tx.commit().await?; + result +} + +async fn push( + authed: ApiAuthed, + Extension(db): Extension, + Path(w_id): Path, + Json(req): Json, +) -> JsonResult { + require_plain_user_token(&authed)?; + if req.owner != authed.email { + return Err(Error::Generic( + http::StatusCode::CONFLICT, + "this push was prepared for another user".to_string(), + )); + } + validate_push(&req)?; + let Some(backend) = backend(&authed, &db, &w_id).await? else { + return Ok(Json(PushResponse { + enabled: false, + storage_id: None, + backup_generation: None, + results: vec![], + })); + }; + #[cfg(not(feature = "enterprise"))] + { + let remaining = + crate::job_helpers_oss::ce_storage_quota_remaining(&db, &w_id, None).await?; + if push_payload_bytes(&req) as i64 > remaining { + return Err(Error::QuotaExceeded( + "the workspace storage quota leaves no room for this AI session backup".to_string(), + )); + } + } + #[cfg(feature = "enterprise")] + let _ = push_payload_bytes(&req); + + let mut results = Vec::with_capacity(req.sessions.len() + req.removed.len()); + let mut written: usize = 0; + // A session split into several entries is listed by the last: once one part failed, the + // later ones are not written, or the marker would list a session missing a part. + let mut failed: std::collections::HashSet<&str> = Default::default(); + for s in &req.sessions { + let (error, needs_whole) = if failed.contains(s.id.as_str()) { + ( + Some("an earlier part of this session in the push failed".to_string()), + false, + ) + } else { + match push_session_locked(&db, &backend, s).await { + Ok((n, needs_whole)) => { + written += n; + (None, needs_whole) + } + Err(e) => { + tracing::warn!("AI session backup push failed for {} in {w_id}: {e}", s.id); + failed.insert(&s.id); + (Some(e.to_string()), false) + } + } + }; + results.push(PushResult { id: s.id.clone(), error, needs_whole }); + } + for sid in &req.removed { + let error = remove_session_locked(&db, &backend, sid) + .await + .err() + .map(|e| { + tracing::warn!("AI session backup removal failed for {sid} in {w_id}: {e}"); + e.to_string() + }); + results.push(PushResult { id: sid.clone(), error, needs_whole: false }); + } + // Overwrites and deletes make this an over-count; the periodic recount the quota check + // schedules once usage is stale settles it. + #[cfg(not(feature = "enterprise"))] + if written > 0 { + crate::job_helpers_oss::bump_storage_usage( + &db, + &w_id, + windmill_object_store::DEFAULT_STORAGE, + written as i64, + ) + .await; + } + #[cfg(feature = "enterprise")] + let _ = written; + Ok(Json(PushResponse { + enabled: true, + storage_id: Some(backend.storage_id), + backup_generation: Some(backend.generation), + results, + })) +} diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index fc90852ae6..715f0b6556 100644 --- a/backend/windmill-api/src/lib.rs +++ b/backend/windmill-api/src/lib.rs @@ -69,6 +69,8 @@ mod ai; #[cfg(feature = "private")] mod ai_free_tier_ee; mod ai_free_tier_oss; +#[cfg(feature = "parquet")] +mod ai_sessions; mod ai_shared_artifacts; mod apps; mod apps_raw_bundle; diff --git a/backend/windmill-api/src/workspaces.rs b/backend/windmill-api/src/workspaces.rs index da8532e9d1..55f244c143 100644 --- a/backend/windmill-api/src/workspaces.rs +++ b/backend/windmill-api/src/workspaces.rs @@ -147,6 +147,7 @@ async fn edit_copilot_config( let workspace_has_config = ai_config.has_providers(); let copilot_disabled = ai_config.copilot_disabled; + let sessions_storage_disabled = ai_config.sessions_storage_disabled; let instance_ai_config = sqlx::query_scalar!("SELECT value FROM global_settings WHERE name = 'ai_config'") .fetch_optional(&db) @@ -174,6 +175,7 @@ async fn edit_copilot_config( AIConfig::default() }; effective_ai_config.copilot_disabled = copilot_disabled; + effective_ai_config.sessions_storage_disabled = sessions_storage_disabled; Ok(Json(EditCopilotConfigResponse { effective_ai_config, @@ -212,6 +214,9 @@ async fn get_copilot_info( let copilot_disabled = workspace_ai_config .as_ref() .is_some_and(|c| c.0.copilot_disabled); + let sessions_storage_disabled = workspace_ai_config + .as_ref() + .is_some_and(|c| c.0.sessions_storage_disabled); let instance_config = sqlx::query_scalar!("SELECT value FROM global_settings WHERE name = 'ai_config'") .fetch_optional(&db) @@ -236,6 +241,7 @@ async fn get_copilot_info( AIConfig::default() }; effective.copilot_disabled = copilot_disabled; + effective.sessions_storage_disabled = sessions_storage_disabled; Ok(Json(effective)) } diff --git a/backend/windmill-common/src/variables.rs b/backend/windmill-common/src/variables.rs index 045ddcd95b..1feb4182ab 100644 --- a/backend/windmill-common/src/variables.rs +++ b/backend/windmill-common/src/variables.rs @@ -252,12 +252,17 @@ pub async fn build_crypt_with_key_suffix( key_suffix: &str, ) -> crate::error::Result { let key = get_workspace_key(w_id, db).await?; + Ok(crypt_from_key_with_suffix(&key, key_suffix)) +} + +/// The cipher `build_crypt_with_key_suffix` builds, from a key string in hand. +pub fn crypt_from_key_with_suffix(key: &str, key_suffix: &str) -> MagicCrypt256 { let crypt_key = if let Some(ref salt) = SECRET_SALT.as_ref() { format!("{}{}{}", key, salt, key_suffix) } else { format!("{}{}", key, key_suffix) }; - Ok(magic_crypt::new_magic_crypt!(crypt_key, 256)) + magic_crypt::new_magic_crypt!(crypt_key, 256) } pub async fn get_workspace_key(w_id: &str, db: &DB) -> crate::error::Result { diff --git a/backend/windmill-object-store/src/lib.rs b/backend/windmill-object-store/src/lib.rs index c403625e00..25dc6ea19c 100644 --- a/backend/windmill-object-store/src/lib.rs +++ b/backend/windmill-object-store/src/lib.rs @@ -68,8 +68,8 @@ pub mod object_store_reexports { pub use object_store::path::Path; pub use object_store::{ Attribute, Attributes, Error as ObjectStoreError, GetOptions, GetRange, GetResult, - ObjectStore, PutMultipartOpts, PutPayload, PutResult, Result as ObjectStoreResult, - WriteMultipart, + ObjectMeta, ObjectStore, PutMode, PutMultipartOpts, PutOptions, PutPayload, PutResult, + Result as ObjectStoreResult, UpdateVersion, WriteMultipart, }; } diff --git a/docs/ai-session-backups.md b/docs/ai-session-backups.md new file mode 100644 index 0000000000..0b4c550456 --- /dev/null +++ b/docs/ai-session-backups.md @@ -0,0 +1,266 @@ +# AI session backups + +AI sessions live in the browser: the session list (`windmill-sessions`), chat transcripts and +image blobs (`copilot-chat-history`) and artifacts (`copilot-artifacts`), all per-user IndexedDB +stores. This is the design of their backup in the workspace's object storage, and the +constraints future work on either side must keep. + +Backend: `backend/windmill-api/src/ai_sessions.rs` (`/w/{w}/ai/sessions/{list,pull,push}`). +Frontend: `frontend/src/lib/components/sessions/sessionMirror*.ts`. + +## Why it is lazy + +A session changes at the local write rate: a transcript write every 2 s while streaming, a +session-record write per new message on screen. An object in S3 is replaced whole and every PUT +is billed, so the backup deliberately does not follow that rate. Local writes only mark a session +dirty (`sessionMirrorSignal.ts`, import-free so the stores never depend on the backup). A flush +runs 15 s after the marks go quiet, at most 2 min after the first unflushed mark, when the tab is +hidden, and 10 s after load for marks a crash left behind. Marks are persisted in localStorage +(shared by the user's tabs) for that reason, one key per mark: a shared blob would let two tabs +marking different sessions at once rewrite each other's mark away. A dirty mark is a counter +bumped on every write; a push retires it by recording the counter it covered on the session's +sync row rather than deleting the mark, since two localStorage calls cannot compare-and-delete +and a bump landing between them would be lost; retired marks are not reclaimed (one small key per +session ever backed up), and the marks of unsent drafts and of workspaces that are off stay too, +each costing one lookup per flush. Only a session gone from the store has its mark deleted. Losing the last +seconds of a device that never comes back is accepted; a tab that closes normally keeps its marks. +A signal names the user whose store the write landed in (read off the store's scoped name), so +a write that completes after the logged-in user changed marks that user's session, for their +next load, rather than the current user's. + +A flush plans and sends one session at a time, filling requests of about 8 MB as it goes, so a +first backfill of a large history never holds more than one request's worth of records and +images in memory. + +## What a push carries + +The pure planner (`sessionMirrorPlan.ts`) compares each piece against the marker of what was +last pushed, kept per session in the `windmill-sessions-mirror` store: + +| Piece | Object | Sent when | +|---|---|---| +| session record | `sessions/{sid}/head.json` | its signature changed | +| chat | `sessions/{sid}/chats/{cid}.json` | its `lastModified` moved | +| artifacts | `sessions/{sid}/artifacts.json` | their fingerprint changed | +| image | `images/{sid}/{cid}/{iid}` | never pushed before (write-once) | +| index marker | `index/{sid}/{epoch}` | last, by the part that completes a push of the session (empty; named by the record's move count) | + +All under `windmill_ai_sessions/{w_id}/g{generation}/{sha256(email)}/` in the workspace's +primary storage (the generation is what a key rotation moves, see below). +The listing reads only `index/`: one object per session whatever the session holds, so a +session with many chats cannot crowd newer ones out of a bounded scan, and its +`last_modified` is the session's `updated_at`. Written last, and only by an entry no unsent +part follows (a session split over several entries says `partial` on all but the last, and +names the push on each, or the part is refused), it +lists a session only once a whole push landed; the parts of a session after a failed one are +not written either, on the server within one push and on the client across pushes, so the +marker on the last part never lists a session missing a chat, and a new session whose last part +never lands is not listed at all. A push of the session whole (no sync row, or a stale one, +`whole` on every part) opens with the head on the first: that part replaces the backup (the +marker goes first, then everything under the session), so what an old storage still held of +the session and the push does not carry is gone. An incremental part rides on a listed +session, and the server refuses it with `needs_whole`, writing nothing, when none is listed +(a removal deletes the marker first), rather than write a marker over a session missing what +earlier parts or earlier pushes carried; one that changes more than one object unlists the +session before its writes and lists it again after them, so a pull between two of the writes, +or after one failed, finds it absent rather than a mix of old and new pieces (one object +changing is one write, and stays listed). A push split over several parts, whole or +incremental, names itself on each with a token the browser draws (`push`, `opens` on the +first): the opening part unlists the session, so a pull between two parts finds it absent +rather than a mix of old and new pieces, the last part lists it again, and a later part is +written only while that token is the one there, so two devices pushing the session at once +cannot list a mix of their pieces (the push that opened later wins; the other is refused with +`needs_whole` and goes again), and one abandoned leaves the session unlisted, so the next +push of it goes whole. A push and a removal of one session +are serialized on the server by a Postgres advisory lock keyed on the session's prefix, so +the two never interleave object by object. + +The head signature leaves out `name` (a per-browser counter the sessions page routes by), +the unsent-draft fields, `workspace_root_id` (recomputed on import), and the two fields reading +a session bumps (`lastSeenCount`, `lastActivityAt`). Reading a session must never cost a push; +keep that property when adding fields to `Session`. + +Unsent drafts (no `workspace_id`) and attached files (Blobs, directory handles) are not backed up. + +## Encryption and access + +Every object is encrypted with a key derived from the workspace key and the user +(`build_crypt_with_key_suffix` with the email hash), because workspace storage credentials are +shared far more widely than a user's transcripts: `public_resource` storages and legacy-mode +READ/WRITE hand any member the bucket. The key is per user rather than per workspace so that a +member who copies another user's ciphertext under their own prefix gets nothing from `pull`; an +object that does not decrypt for its reader is treated as absent. Rotating the workspace key +(`set_encryption_key`) does not re-key the backups the way it re-encrypts the workspace's +secrets. The objects live under a prefix named by a generation +(`workspace_settings.ai_sessions_backup_generation`) that the rotation bumps in the +transaction committing the new key; once committed, the routes read and write under the new +generation's prefix, the answers name it (`backup_generation`, below; `storage_id` names the +storage and does not change), so every browser marks its sync rows stale and pushes its +sessions whole again there, and every older generation, which nothing writes to any more, is deleted off the +request at leisure (`windmill-api-workspaces/src/ai_session_backups.rs`). Sessions no browser +holds any more are lost. A generation is never reused, so no deletion, however late, can +touch live objects; a rotation that fails before its commit bumps nothing and deletes +nothing; two rotations racing serialize on the key row; the same key set again bumps +nothing. A rotation is rare, and the alternative, +rewriting every object in place while pushes, restarts, storage switches and further +rotations race the rewrite, is where the complexity would be; with this, nothing but the +current key ever reads an object. The +server builds every key from ids it validated +(`[A-Za-z0-9_-]{1,64}`) and the caller's own email; the client never names a key, and the +workspace storage permission rules are not consulted (the same stance as volumes). Only an +unscoped user token may reach the routes: a job token can carry an `on_behalf_of` identity and +every scoped token (guest, embed, app policy, MCP) was minted for something narrower. + +The backup is keyed by the email like the browser's own stores are (`userScopedDb` scopes +IndexedDB by it): a user whose email changes starts from an empty history on both sides, and +the objects under the old hash stay in the bucket unread. Carrying them over would need a +server-side re-key (decrypt with the old suffix, encrypt with the new, move every object) in the +email-change flow, which this design leaves out. + +An image is accepted only as a base64 data URL of at most 4 MB and stored verbatim, so it +serializes back into a pull answer at its stored size; anything JSON would escape could grow +several times and defeat the pull budget. + +`push` carries `owner`, the email the browser prepared the batch for, and the server refuses a +mismatch with 409: an in-place account switch must not file one user's sessions under another's +prefix. The client captures its user at flush start and checks every store handle's name +against it for the same reason. + +The feature is on wherever the workspace has primary storage, and off with +`ai_config.sessions_storage_disabled` (the `copilot_disabled` pattern: no migration, carried by +settings export and the CLI). A build without `parquet` has no routes (404), a workspace without +storage answers `enabled: false`; either turns the backup off for ten minutes, after which the +page asks again on its own (a flush for whatever is pending, and a restore), and the AI +settings page tells the mirror at once when the switch is saved there (the off state is +forgotten, the rows that went stale are marked again, a restore runs). + +## Conflicts and deletion + +Last write wins across devices. The head carries no manifest; `pull` lists the session's prefix +instead, so a stale device that renames or archives a session rewrites only the head and cannot +hide chats a newer device wrote. Two devices continuing the same chat still collide. + +Restore brings back only sessions the browser does not have (`importSessions` is write-if-absent, +and skips ids the user deleted in this page) and never overwrites or deletes a local one from +remote state. It covers the workspace and its forks together, and only once every one of them +that keeps backups has listed (a listing that failed leaves the family for the next page load +or workspace switch, or the copy that did list could be the stale one): a session listed by +two of them (moved between them, the old copy not yet removed, since that mark is the moving +browser's, which may never come back) is brought back from the copy that moved last (`epoch`, +the record's move count, which names the marker), the storage's own modification time +deciding between two of the same count, and not from the other, which would otherwise take +the id first and keep the later copy out for good. A workspace's records land together once +its pulls are done, and just before they do the whole family is listed again (members whose +backups were off included, since a move from another device can land in a workspace between +the first listings and the pulls; a family of one, with nowhere else for a copy to show up, is +not): a session a later copy of which showed up elsewhere is left, with the family, for the +next time. Only a user-initiated `deleteSession` removes the backup; the next push from +another device that still has the session is refused with `needs_whole` (nothing of it is +written), its row goes stale without a backoff, and that device's next flush sends the session +whole; the workspace-lifecycle +removals (`reconcileSessionsLifecycle`, `deleteSessionsForWorkspace`) leave it, so a session +dropped by a wrong reconcile comes back on the next restore. Objects of deleted workspaces stay +in the bucket. A session moved to another workspace is pushed whole into the new one, and once +that push has landed the copy in the old one gets a removal mark of its own, naming the storages holding +that copy (the row that knew is the new workspace's by then), retried independently until +each of them has answered, even when the old workspace's backups are off at the time (they +may hold the copy still). Filing the removal only after the new copy is acknowledged keeps the +session backed up somewhere at every point. + +A restore writes a session's artifacts and chats before its record, and records nothing for a +session whose pieces could not be written: recording it would let the next flush push the +half-empty local state over the backup. + +Every answer names the storage it came from (`storage_id`, a hash of what locates the objects, +endpoint, region and bucket, not the credentials, which rotate) and the backup generation a +key rotation bumps (`backup_generation`). A sync row records both, and a row naming another +storage or generation goes stale and its session is marked again: a workspace pointed at a +new bucket, or whose key was rotated, holds nothing, and the server looks nowhere else, so +the next flush carries the session whole. A switch leaves the old copy where it was, so the +row rewritten under the new storage records the old one (`alsoIn`, one entry per storage +the workspace was on), and a removal is done only once every storage holding a copy answered +it, whatever the generation (a rotation deleted the older generation's copy anyway): each +answer narrows the row to the storages still holding one, and the mark waits for them to +answer, so a switch back never brings a deleted session back. That includes the rows a flush has just written, when a later answer of the +same flush names another storage or the session was pushed in part on top of a row from the +old one; a session whose own parts were answered from different storages is not settled at +all. The listing a restore starts with runs the same check, so a storage switch is noticed at +the first push after it or on the next page load, whichever comes first. + +## Limits + +Push bodies are packed to about 8 MB (UTF-8 bytes as sent), at most 100 entries, 200 removals and +4000 pieces each (the server's caps, with 32 MB on the body, and 100 chats, 500 images or 1000 +deletes per entry, since every piece is an object-store call); an entry that +outgrows the target is split into chat-only parts (the artifacts and deletes on the last, the +head on the last too for an incremental push and on the first part, whatever it carries, for a +push of the session whole), and deletes +past the per-entry cap are carried over to the next push, which the session stays marked for. A chat above +16 MB or a session's artifacts above 8 MB are left out with a console warning; a chat that grew +past the cap after it was backed up has its copy deleted, so a restore never presents the old +transcript as the current one. A 413 fails only the sessions of that request. A +request the server refuses (any other 4xx but 404/403/409) stops the backup for the page but keeps +the marks and the sync state, so the next load tries again; a session the server reports it could +not store stays marked and is retried with backoff. A move files the old workspace's removal +mark before recording the new copy's row, so a mark that could not be written leaves the move +to be planned again. A workspace that answers `enabled: false` +marks its sync rows stale (the next push after storage returns carries every session whole, +since a new storage may be a new bucket) and leaves its dirty marks where they are (a move into +it must still remember the old copy); it keeps the removal marks of sessions that had been backed +up, so one deleted while backups are off does not come back once they are on, and drops the +removals of sessions never backed up from this browser, so a storage-less instance does not +collect one mark per deleted session forever. A user delete whose removal mark cannot be written +(localStorage full) is carried by the session's sync row instead (`removed`), which the flush +and the restore read like a mark; a session without a row yet (its first push may be in +flight) gets a row saying only that, and every row write keeps a removal filed meanwhile, so +the push's own row cannot erase it. Pull bodies are +capped at 64 KB. Pull answers up to 20 ids within a 32 MB +budget: a session's size is known from the listings before anything of it is read, one that +would not fit is deferred unless it is the first of the answer, in which case it comes in +pages: the answer carries what fits in key order (at least one object, so every page makes +progress) and names where the next picks up (`next`, a cursor the browser sends back as +`resume` with that session alone). Every page carries a fingerprint of the session's listing +(marker, keys, sizes, modification times, entity tags and versions, since a store reports +modification times coarsely and an object rewritten at the same size within that grain would +otherwise fingerprint the same) taken before anything of it is read, and the server +takes it again once the page is read: a page the backup moved under (a push landing object by +object) is read again, a few times, then answered as `moved`, and the browser starts the +session over on that or on two pages whose fingerprints differ. The browser writes each page's pieces as it arrives, over whatever +an earlier restore cut short had staged (the session is absent locally, so its pieces have no +local edits to keep, and the backup may have moved on), and the record, which is what makes +the session visible, only with the last page. A restore in progress keeps a staging row for +the session (the ids of every chat, image, artifact and version it wrote), which outlives it +if it is cut short; the next restore deletes the staged pieces the backup no longer has, by id +and never by clock, before the record lands (once the record is there no restore looks at the +session again, and a flush would push them back), and a prune that could not run leaves the +session, its pieces and its staging row for the restore after. A restore holds the user's tab +lock while it runs, so two tabs cannot each write the same absent session's pieces over the +other's, and runs only where Web Locks exist (a secure context: https, or localhost); on a plain +http origin the browser still backs up, and its sessions come back on a secure one. A +restore never writes an older record over a newer +one; between pages it holds nothing but the sync +row being assembled, whose chats also admit the images of a later page. Every page carries a +fingerprint of the session's listing (`listing`), taken before anything of the page is +listed or read, so an object landing after it is in the next page's; a session whose +fingerprint moved between two of its pages (a chat added by another device could sort before +the cursor and be missed) starts over, up to three times, then waits for the next restore. An object that grew +since the listing (a push replaced it) ends its page just before it and the answer names that +spot, so the next page sizes it anew rather than the session being imported without it. A pull sees every key of a session's listing but keeps the 5000 smallest +past its cursor (a page is defined by key order, and the store promises none), so a session +grown without bound by valid pushes cannot grow the answer's memory through its metadata +either; removing a prefix and a rotation's deletion stream their listings. `list` scans at most 50 000 index markers, keeps the newest 500 as it goes and answers with +them (`truncated` says when there were more); the restore takes 50 of them. Every read checks the object's size before buffering it: one larger than any push writes +(32 MB) is planted, whatever its listing said, and skipped, since whoever holds the bucket's +credentials can put anything at a predictable key; one larger than its listing said grew +since (a push replaced it) and ends its page, for the next pull to size anew. A dirty mark that cannot be written +(localStorage full) records its bump on the session's sync row instead (`extraV`, counted +with the mark's counter and kept by every row write, so a push in flight cannot retire it); +a session without a row yet (its first push in flight) keeps the bump in the page, and the +next row write takes it onto the row in the same transaction, so the row the push writes +cannot retire the mark with the bump unseen; every load's backfill marks again a session +without a row, with a stale one, or with one carrying bumps no push has covered. +A mark or removal for another user (a write that landed after a switch) reaches that user's +rows through a connection of its own, since the shared handle follows the current user. Nothing is read past the budget, whatever a session holds. A +restore takes the newest 50 sessions per workspace: every visible session gets a runtime, and +each runtime's history load reads the whole chat store. On CE the push checks the storage quota +and bumps usage by bytes written (an over-count on overwrites; the periodic recount settles it). diff --git a/frontend/src/lib/components/InstanceSettings.svelte b/frontend/src/lib/components/InstanceSettings.svelte index 2543865889..f1c94874fb 100644 --- a/frontend/src/lib/components/InstanceSettings.svelte +++ b/frontend/src/lib/components/InstanceSettings.svelte @@ -1084,12 +1084,13 @@ it with none, whether SSO logins evaluate an IdP groups claim (SAML or OIDC) and change a membership, the plan tier and quota shown when the execution meter is opened, whether app sandbox isolation is turned on, whether a step's workspace script is - edited from the flow editor, which skin approval steps are given, whether AI agent - steps set their own memory id as a fixed id or an expression, how data tables and - their migrations are set up and used, how often an empty workspace home is seen, how - often the home page’s create menu and hub-project picker are opened and from which - entry point, the name of any public hub project imported from the home page and how - far that import got, and whether a pre-approved trial offer was opened, last 30 days)
  • feature adoption (counts of which flow, script, trigger, worker and data table @@ -1150,12 +1151,13 @@ it with none, whether SSO logins evaluate an IdP groups claim (SAML or OIDC) and change a membership, the plan tier and quota shown when the execution meter is opened, whether app sandbox isolation is turned on, whether a step's workspace script is - edited from the flow editor, which skin approval steps are given, whether AI agent - steps set their own memory id as a fixed id or an expression, how data tables and - their migrations are set up and used, how often an empty workspace home is seen, how - often the home page’s create menu and hub-project picker are opened and from which - entry point, the name of any public hub project imported from the home page and how - far that import got, and whether a pre-approved trial offer was opened, last 30 days)
  • feature adoption (counts of which flow, script, trigger, worker and data table diff --git a/frontend/src/lib/components/copilot/chat/HistoryManager.svelte.ts b/frontend/src/lib/components/copilot/chat/HistoryManager.svelte.ts index 592b6217c8..4edac6dc0a 100644 --- a/frontend/src/lib/components/copilot/chat/HistoryManager.svelte.ts +++ b/frontend/src/lib/components/copilot/chat/HistoryManager.svelte.ts @@ -1,9 +1,15 @@ -import { type DBSchema as IDBSchema, type IDBPDatabase } from 'idb' +import { + type DBSchema as IDBSchema, + type IDBPDatabase, + type IDBPTransaction, + type StoreNames +} from 'idb' import type { ChatJob, DisplayMessage } from './shared' import { expanded, messageDraft } from './chatDraft' import { createLongHash } from '$lib/editorLangUtils' import { userScopedDb, type UserScopedDbMigrateDeps } from '$lib/userScopedDb' -import { scopedKey } from '$lib/userScopedStorage' +import { emailOfScopedKey, scopedKey, scopedKeyFor } from '$lib/userScopedStorage' +import { markSessionDirty } from '$lib/components/sessions/sessionMirrorSignal' import type { ChatCompletionMessageParam } from 'openai/resources/index.mjs' import type { PersistedContextUsage } from './tokenUsage' import { IMAGE_OMITTED_PLACEHOLDER, type AttachedImage } from './imageUtils' @@ -14,8 +20,9 @@ import { randomUUID } from '$lib/utils/uuid' // shared browser. The bare name is also the legacy (pre-namespacing) DB, claimed // once on first login. const DB_NAME = 'copilot-chat-history' -// v3 adds the images blob store (replacing v2's short-lived toolImages store). -const DB_VERSION = 3 +// v3 adds the images blob store (replacing v2's short-lived toolImages store); v4 indexes +// chats by session. +const DB_VERSION = 4 /** Newest image blobs kept per chat; each is a bounded (≤1568px) data URL. */ const MAX_IMAGES_PER_CHAT = 30 /** Marks a persisted image whose bytes live in the `images` store. */ @@ -45,6 +52,7 @@ interface ChatSchema extends IDBSchema { // chats predating this feature. Persisted out-of-band like modifiedItems. backgroundJobs?: ChatJob[] } + indexes: { 'by-session': string } } // Image bytes, out-of-band from the chat record on purpose: the record is // re-cloned into IndexedDB on every saveChat, while a blob is written once @@ -64,9 +72,19 @@ interface ChatSchema extends IDBSchema { } } -function createChatStore(db: IDBPDatabase): void { - if (!db.objectStoreNames.contains('chats')) { - db.createObjectStore('chats', { keyPath: 'id' }) +/** A persisted chat, exactly as the store holds it (image refs, not bytes). */ +export type StoredChat = ChatSchema['chats']['value'] + +function createChatStore( + db: IDBPDatabase, + tx: IDBPTransaction[], 'versionchange'> +): void { + const chats = db.objectStoreNames.contains('chats') + ? tx.objectStore('chats') + : db.createObjectStore('chats', { keyPath: 'id' }) + // Lets the session backup find a session's chats without reading every record. + if (!chats.indexNames.contains('by-session')) { + chats.createIndex('by-session', 'sessionId') } // v2 briefly kept full-resolution tool screenshots in their own store; the // general blob store below covers them now. @@ -117,7 +135,9 @@ async function claimLegacyChatDb( { openDB, deleteDB }: UserScopedDbMigrateDeps ): Promise { if ((await scopedDb.count('chats')) > 0) return - const legacy = await openDB(DB_NAME, 1, { upgrade: createChatStore }) + const legacy = await openDB(DB_NAME, 1, { + upgrade: (db, _oldVersion, _newVersion, tx) => createChatStore(db, tx) + }) const legacyChats = await legacy.getAll('chats') if (legacyChats.length > 0) { const tx = scopedDb.transaction('chats', 'readwrite') @@ -156,6 +176,129 @@ export async function readChatModifiedItems(chatId: string): Promise(DB_NAME, { + version: DB_VERSION, + upgrade: createChatStore, + migrate: migrateLegacyChatDb +}) + +async function backupDb(email: string): Promise | undefined> { + const db = await backupDbh.whenReady() + return db && db.name === scopedKeyFor(DB_NAME, email) ? db : undefined +} + +/** Test-only: let go of the backup's handle so the next call opens the test's fresh + * IndexedDB rather than the connection a previous test left. */ +export function __resetBackupStoreForTesting(): void { + backupDbh.close() +} + +/** Ids of the chats tagged with this session, or undefined when the store is unavailable. */ +export async function listSessionChatIds( + sessionId: string, + email: string +): Promise { + const db = await backupDb(email) + if (!db) return undefined + return (await db.getAllKeysFromIndex('chats', 'by-session', sessionId)).map(String) +} + +export async function readStoredChat(id: string, email: string): Promise { + const db = await backupDb(email) + return db?.get('chats', id) +} + +/** Ids of the image blobs a chat owns, or undefined when the store is unavailable. */ +export async function listChatImageIds( + chatId: string, + email: string +): Promise { + const db = await backupDb(email) + if (!db) return undefined + return (await imageKeysForChat(db, chatId)).map(String) +} + +export async function readImageDataUrl(id: string, email: string): Promise { + const db = await backupDb(email) + return (await db?.get('images', id))?.dataUrl +} + +export interface RestoredImage { + id: string + chatId: string + dataUrl: string +} + +/** + * Write restored chats and image blobs, leaving any that already exist alone: a record + * this browser wrote since is newer than the backup it came from. False when the store + * could not be reached, which the caller must not record as a restore. + */ +export async function importStoredChats( + chats: StoredChat[], + images: RestoredImage[], + email: string, + overwrite = false +): Promise { + const db = await backupDb(email) + if (!db) return false + const tx = db.transaction(['chats', 'images'], 'readwrite') + const chatStore = tx.objectStore('chats') + const imageStore = tx.objectStore('images') + const savedAt = Date.now() + for (const image of images) { + if (overwrite || (await imageStore.getKey(image.id)) === undefined) { + await imageStore.put({ id: image.id, chatId: image.chatId, dataUrl: image.dataUrl, savedAt }) + } + } + // An overwrite never puts an older record over a newer one: without a cross-tab lock, + // another restore may have landed a newer backup's copy meanwhile. + for (const chat of chats) { + const existing = await chatStore.get(chat.id) + if (existing === undefined || (overwrite && existing.lastModified <= chat.lastModified)) { + await chatStore.put(chat) + } + } + await tx.done + return true +} + +/** Deletes these chats of the session (with their images) and these images: what an earlier + * restore staged for it and the backup no longer has. False when nothing could be deleted. */ +export async function pruneSessionChats( + sessionId: string, + chats: Set, + images: Set, + email: string +): Promise { + if (chats.size === 0 && images.size === 0) return true + const db = await backupDb(email) + if (!db) return false + try { + const tx = db.transaction(['chats', 'images'], 'readwrite') + const chatStore = tx.objectStore('chats') + const imageStore = tx.objectStore('images') + for (const chatId of await chatStore.index('by-session').getAllKeys(sessionId)) { + if (!chats.has(String(chatId))) continue + await chatStore.delete(chatId) + const keys = await imageStore + .index('by-chat') + .getAllKeys(IDBKeyRange.bound([chatId, -Infinity], [chatId, Infinity])) + for (const key of keys) await imageStore.delete(key) + } + for (const id of images) await imageStore.delete(id) + await tx.done + return true + } catch (err) { + console.error('Could not prune chats for session', err) + return false + } +} + export default class HistoryManager { // Per-instance handle to the shared per-user DB lifecycle. There is one // HistoryManager per AIChatManager (the singleton + one per session runtime), @@ -289,7 +432,10 @@ export default class HistoryManager { const snapshot = $state.snapshot(existing) const updated = { ...snapshot, sessionId } this.savedChats = { ...this.savedChats, [chatId]: updated } - await this.enqueueDbWrite((db) => db.put('chats', updated)) + await this.enqueueDbWrite(async (db) => { + await db.put('chats', updated) + markSessionDirty(sessionId, chatId, emailOfScopedKey(DB_NAME, db.name)) + }) } getPastChats() { @@ -565,6 +711,13 @@ export default class HistoryManager { const keep = this.keptImageIds(refs) await this.writeKeptImageBlobs(db, updatedChat.id, blobs, keep) await db.put('chats', updatedChat) + if (updatedChat.sessionId) { + markSessionDirty( + updatedChat.sessionId, + updatedChat.id, + emailOfScopedKey(DB_NAME, db.name) + ) + } // Best-effort: the record is already committed, so a failed cleanup // (e.g. a user switch closed this handle mid-op) must not turn a // successful save into a rejection — the orphans are reclaimed by @@ -589,6 +742,7 @@ export default class HistoryManager { } deletePastChat(id: string) { + const sessionId = this.savedChats[id]?.sessionId this.savedChats = Object.fromEntries( Object.entries(this.savedChats).filter(([key]) => key !== id) ) @@ -596,6 +750,7 @@ export default class HistoryManager { await db.delete('chats', id) const keys = await imageKeysForChat(db, id) await Promise.all(keys.map((key) => db.delete('images', key))) + if (sessionId) markSessionDirty(sessionId, id, emailOfScopedKey(DB_NAME, db.name)) }).catch((err) => console.error('Could not delete chat', err)) } diff --git a/frontend/src/lib/components/copilot/chat/artifacts/artifactsDB.ts b/frontend/src/lib/components/copilot/chat/artifacts/artifactsDB.ts index 94afac404e..e7ff4b1f60 100644 --- a/frontend/src/lib/components/copilot/chat/artifacts/artifactsDB.ts +++ b/frontend/src/lib/components/copilot/chat/artifacts/artifactsDB.ts @@ -2,6 +2,8 @@ // active chat's rotation, so chatId-keying would drop artifacts on each new conversation. import { type DBSchema as IDBSchema, type IDBPObjectStore, type IDBPTransaction } from 'idb' import { userScopedDb } from '$lib/userScopedDb' +import { emailOfScopedKey, scopedKeyFor } from '$lib/userScopedStorage' +import { markSessionDirty } from '$lib/components/sessions/sessionMirrorSignal' export type ArtifactKind = 'md' | 'html' @@ -89,9 +91,11 @@ interface ArtifactsSchema extends IDBSchema { } } +const ARTIFACTS_DB = 'copilot-artifacts' + // User-scoped like the chat-history store these are keyed against: no cross-user // co-residency on a shared browser. -const dbh = userScopedDb('copilot-artifacts', { +const dbh = userScopedDb(ARTIFACTS_DB, { version: 2, // Runs for a fresh database and for the v1 upgrade alike, so create each store only // when it is missing. @@ -118,11 +122,69 @@ export async function putArtifact(artifact: PersistedArtifact): Promise { // A rejected write (most likely QuotaExceededError) leaves the artifact usable for the // session but unpersisted — degrade like the reads rather than throwing at the caller. await db.put('items', artifact) + markSessionDirty(artifact.sessionId, undefined, emailOfScopedKey(ARTIFACTS_DB, db.name)) } catch (err) { console.error('Could not persist artifact', err) } } +/** A session's artifacts with their history, or undefined when the store is unavailable + * or no longer the named user's (see `readStoredSessions`). */ +export async function readSessionArtifacts( + sessionId: string, + email: string +): Promise<{ items: PersistedArtifact[]; versions: ArtifactVersion[] } | undefined> { + const db = await getDB() + if (!db || db.name !== scopedKeyFor(ARTIFACTS_DB, email)) return undefined + try { + const items = await db.getAllFromIndex('items', 'by-session', sessionId) + const versions = ( + await Promise.all(items.map((i) => db.getAllFromIndex('versions', 'by-artifact', i.id))) + ).flat() + return { items, versions } + } catch (err) { + console.error('Could not read artifacts', err) + return undefined + } +} + +/** Write restored artifacts and snapshots, leaving any that already exist alone. Reports + * whether they are all in the store now: unlike the other writes here, a caller records the + * restore as done on the strength of this answer. */ +export async function importArtifacts( + items: PersistedArtifact[], + versions: ArtifactVersion[], + email: string, + overwrite = false +): Promise { + if (items.length === 0 && versions.length === 0) return true + const db = await getDB() + if (!db || db.name !== scopedKeyFor(ARTIFACTS_DB, email)) return false + try { + const tx = db.transaction(['items', 'versions'], 'readwrite') + const itemStore = tx.objectStore('items') + const versionStore = tx.objectStore('versions') + // An overwrite never puts an older record over a newer one: without a cross-tab lock, + // another restore may have landed a newer backup's copy meanwhile. + for (const item of items) { + const existing = await itemStore.get(item.id) + if (existing === undefined || (overwrite && existing.updatedAt <= item.updatedAt)) { + await itemStore.put(item) + } + } + for (const version of versions) { + if (overwrite || (await versionStore.getKey(version.key)) === undefined) { + await versionStore.put(version) + } + } + await tx.done + return true + } catch (err) { + console.error('Could not import artifacts', err) + return false + } +} + export async function getArtifact(id: string): Promise { const db = await getDB() if (!db) return undefined @@ -273,7 +335,11 @@ export async function mutateArtifact( reportFailure = false abort() } - return { outcome: await settled, artifact: edit.artifact } + const outcome = await settled + if (outcome === 'saved' && db) { + markSessionDirty(edit.artifact.sessionId, undefined, emailOfScopedKey(ARTIFACTS_DB, db.name)) + } + return { outcome, artifact: edit.artifact } } /** @@ -342,14 +408,47 @@ export async function deleteArtifact(id: string): Promise { if (!db) return try { const tx = db.transaction(['items', 'versions'], 'readwrite') - await tx.objectStore('items').delete(id) + const items = tx.objectStore('items') + const sessionId = (await items.get(id))?.sessionId + await items.delete(id) await deleteVersionsIn(tx.objectStore('versions'), id) await tx.done + if (sessionId) markSessionDirty(sessionId, undefined, emailOfScopedKey(ARTIFACTS_DB, db.name)) } catch (err) { console.error('Could not delete artifact', err) } } +/** Deletes these artifacts of the session (with their versions) and these versions: what + * an earlier restore staged for it and the backup no longer has. False when nothing could + * be deleted. */ +export async function pruneSessionArtifacts( + sessionId: string, + itemIds: Set, + versionKeys: Set, + email: string +): Promise { + if (itemIds.size === 0 && versionKeys.size === 0) return true + const db = await getDB() + if (!db || db.name !== scopedKeyFor(ARTIFACTS_DB, email)) return false + try { + const tx = db.transaction(['items', 'versions'], 'readwrite') + const items = tx.objectStore('items') + const versions = tx.objectStore('versions') + for (const id of await items.index('by-session').getAllKeys(sessionId)) { + if (!itemIds.has(String(id))) continue + await items.delete(id) + await deleteVersionsIn(versions, String(id)) + } + for (const key of versionKeys) await versions.delete(key) + await tx.done + return true + } catch (err) { + console.error('Could not prune artifacts for session', err) + return false + } +} + export async function deleteArtifactsForSession(sessionId: string): Promise { const db = await getDB() if (!db) return diff --git a/frontend/src/lib/components/sessions/sessionMirror.svelte.ts b/frontend/src/lib/components/sessions/sessionMirror.svelte.ts new file mode 100644 index 0000000000..b1ac4bcd06 --- /dev/null +++ b/frontend/src/lib/components/sessions/sessionMirror.svelte.ts @@ -0,0 +1,1644 @@ +// Lazily backs the browser's AI sessions up to their workspace's object storage, and +// restores the ones this browser does not have. +// +// IndexedDB stays the store every write lands in; the funnels there only mark a session +// dirty (sessionMirrorSignal). A flush runs once the marks have been quiet for a while, +// bounded by a maximum delay so a long turn still gets backed up part-way, and sends +// batched requests per workspace carrying only the pieces whose marker moved +// (sessionMirrorPlan). Marks are persisted, so a crash leaves them for the next load. +// +// What a flush is for is decided per workspace: `enabled: false` (no storage, or the +// admin switch) turns it off for ten minutes, after which the page asks again on its own, +// and at once when the switch is saved from this page (`backupSettingsChanged`). +import { BROWSER } from 'esm-env' +import { get } from 'svelte/store' +import { openDB, type DBSchema, type IDBPDatabase } from 'idb' +import { + AiService, + ApiError, + type AISessionBackup, + type AISessionBackupCursor, + type AISessionBackupImage, + type AISessionBackupPush +} from '$lib/gen' +import { userWorkspaces } from '$lib/stores' +import { userScopedDb } from '$lib/userScopedDb' +import { getCurrentUserEmail, onUserChange, scopedKey, scopedKeyFor } from '$lib/userScopedStorage' +import { logFeatureUsage } from '$lib/utils/featureUsage' +import { randomUUID } from '$lib/utils/uuid' +import { workspaceRootId } from './sessionScope.svelte' +import { onMirrorSignal } from './sessionMirrorSignal' +import { + importSessions, + isSessionTombstoned, + readStoredSessions, + sessionState, + type Session +} from './sessionState.svelte' +import { + importStoredChats, + pruneSessionChats, + listChatImageIds, + listSessionChatIds, + readImageDataUrl, + readStoredChat, + type RestoredImage, + type StoredChat +} from '../copilot/chat/HistoryManager.svelte' +import { + importArtifacts, + pruneSessionArtifacts, + readSessionArtifacts, + type ArtifactVersion, + type PersistedArtifact +} from '../copilot/chat/artifacts/artifactsDB' +import { + artifactsFingerprint, + headSig, + jsonBytes, + operationsOf, + planSessionPush, + splitEntry, + type ChatSnapshot, + type MirrorSyncState, + type PlannedPush, + type PushBody +} from './sessionMirrorPlan' + +/** A flush waits for the marks to go quiet this long. */ +const QUIET_MS = 15_000 +/** ...but never longer than this after the first unflushed mark. */ +const MAX_DELAY_MS = 120_000 +const STARTUP_DELAY_MS = 10_000 +const RETRY_MIN_MS = 30_000 +const RETRY_MAX_MS = 600_000 +/** Requests are packed up to about this many bytes; the server accepts four times that. */ +const REQUEST_TARGET_BYTES = 8 * 1024 * 1024 +/** The server's caps per request and per entry. */ +const MAX_ENTRIES_PER_REQUEST = 100 +const MAX_REMOVED_PER_REQUEST = 200 +const MAX_IMAGES_PER_ENTRY = 500 +const MAX_OPERATIONS_PER_REQUEST = 4000 +/** A chat or a session's artifacts beyond this are left out of the backup rather than + * sent: with the record and the deletes riding along, the largest entry stays well under + * the server's 32 MB body cap. */ +const MAX_CHAT_BYTES = 16 * 1024 * 1024 +const MAX_ARTIFACTS_BYTES = 8 * 1024 * 1024 +const PULL_BATCH = 5 +/** Newest sessions restored per workspace: every visible session gets a runtime, and each + * runtime's history load reads the whole chat store. */ +const RESTORE_MAX = 50 +const PENDING_PREFIX = 'windmill_sessions_mirror_pending' +const SYNC_DB = 'windmill-sessions-mirror' + +interface MirrorSchema extends DBSchema { + sync: { key: string; value: MirrorSyncState } +} + +function createSyncStore(db: IDBPDatabase): void { + if (!db.objectStoreNames.contains('sync')) db.createObjectStore('sync', { keyPath: 'id' }) +} + +const syncDbh = userScopedDb(SYNC_DB, { version: 1, upgrade: createSyncStore }) + +async function syncDb(email: string) { + const db = await syncDbh.whenReady() + return db && db.name === scopedKeyFor(SYNC_DB, email) ? db : undefined +} + +// --- Pending marks --- +// +// One localStorage key per mark, shared by every tab of the user: a dirty mark holds a +// counter bumped on every write, a removal mark the workspace to remove from and, for a +// session that moved to another workspace, the storages holding the old copy (the row that +// knew is the new workspace's by then). Keying each mark on its own is what lets two tabs +// mark different sessions at the same time without one rewriting the other's mark away, as +// a single JSON blob would. + +interface PendingMarks { + dirty: { id: string; v: number }[] + removed: { id: string; ws?: string; key: string; storages?: string[] }[] +} + +function pendingBase(): string | undefined { + return scopedKey(PENDING_PREFIX) +} + +/** The marks of the user a write landed for: the current user unless the signal says + * otherwise (its store's user, when the user changed while the write was pending). */ +function pendingBaseFor(email: string | undefined): string | undefined { + return email ? scopedKeyFor(PENDING_PREFIX, email) : pendingBase() +} + +function dirtyKey(base: string, sessionId: string): string { + return `${base}::d::${sessionId}` +} + +function removedKey(base: string, sessionId: string, ws: string | undefined): string { + return `${base}::r::${sessionId}::${ws ?? ''}` +} + +function readPending(): PendingMarks { + const marks: PendingMarks = { dirty: [], removed: [] } + const base = pendingBase() + if (!base) return marks + try { + for (let i = 0; i < localStorage.length; i++) { + const key = localStorage.key(i) + if (!key || !key.startsWith(`${base}::`)) continue + const rest = key.slice(base.length + 2) + if (rest.startsWith('d::')) { + const v = Number(localStorage.getItem(key)) + marks.dirty.push({ id: rest.slice(3), v: Number.isFinite(v) ? v : 0 }) + } else if (rest.startsWith('r::')) { + const [id, ws] = rest.slice(3).split('::') + marks.removed.push({ id, ws: ws || undefined, key, storages: storagesOf(key) }) + } + } + } catch (e) { + console.error('Could not read session backup marks', e) + } + // Marks this page could not write: their rows carry the bumps, or this page does until + // a row exists (see `bumpViaSyncRow`); the counter counts with those. + for (const [id, bumps] of unwritableMarks) { + const mark = marks.dirty.find((d) => d.id === id) + if (mark) mark.v += bumps + else marks.dirty.push({ id, v: bumps }) + } + return marks +} + +/** False when the mark could not be written (storage full): the caller must carry the + * change some other way. */ +function bumpDirty(sessionId: string, email?: string): boolean { + const base = pendingBaseFor(email) + if (!base) return false + try { + const key = dirtyKey(base, sessionId) + const v = Number(localStorage.getItem(key) ?? '0') + localStorage.setItem(key, String((Number.isFinite(v) ? v : 0) + 1)) + return true + } catch (e) { + console.error('Could not persist session backup mark', e) + return false + } +} + +/** The durable fallback for a dirty mark that could not be written: the bump goes on the + * sync row, where a flush in flight cannot lose it (`writeSync` keeps it). A session with + * no row yet (its first push in flight, say) keeps it in this page until a row is written, + * which takes it over (see `writeSync`): the row the push writes would otherwise retire the + * mark with the bump unseen. Another user's session with no row is left to that user's next + * load, whose backfill marks every session without a row. */ +async function bumpViaSyncRow(id: string, email = getCurrentUserEmail()): Promise { + if (!email) return + let carried = false + await updateSyncRow(id, email, (row) => { + if (!row) return undefined + carried = true + return { ...row, extraV: (row.extraV ?? 0) + 1 } + }) + if (!carried && email === getCurrentUserEmail()) { + unwritableMarks.set(id, (unwritableMarks.get(id) ?? 0) + 1) + } +} + +/** Reads a row and writes what `update` makes of it, in the store of `email`: through the + * shared handle when that is the current user, and through a connection of its own + * otherwise, since the shared handle follows the current user and a write that landed + * after a switch must still reach the store it belongs to. */ +async function updateSyncRow( + id: string, + email: string, + update: (row: MirrorSyncState | undefined) => MirrorSyncState | undefined +): Promise { + if (email === getCurrentUserEmail()) { + const next = update(await readSync(id, email)) + if (next) await writeSync([next], email) + return + } + let db: IDBPDatabase | undefined + try { + db = await openDB(scopedKeyFor(SYNC_DB, email), 1, { upgrade: createSyncStore }) + const next = update(await db.get('sync', id)) + if (next) await db.put('sync', next) + } catch (e) { + console.error('Could not update the session backup state of another user', e) + } finally { + db?.close() + } +} + +/** Drop the dirty mark of a session gone from the store, the one case nothing can bump + * again. Every other mark stays: one whose push landed is retired through `flushedV` on + * the sync row (two localStorage calls cannot compare-and-delete, and a bump landing + * between them would be lost), and a draft's or an off workspace's waits its turn. */ +function dropDirty(sessionId: string): void { + const base = pendingBase() + if (base) removeKey(dirtyKey(base, sessionId)) +} + +/** The storages a removal mark names, when it does. */ +function storagesOf(key: string): string[] | undefined { + try { + const parsed: unknown = JSON.parse(localStorage.getItem(key) ?? '') + if (Array.isArray(parsed) && parsed.every((s) => typeof s === 'string')) return parsed + } catch {} + return undefined +} + +/** False when the mark could not be written (storage full): the caller must not act as if + * the removal were scheduled. */ +function addRemoved( + sessionId: string, + ws: string | undefined, + dropDirty: boolean, + email?: string, + storages?: string[] +): boolean { + const base = pendingBaseFor(email) + if (!base) return false + try { + if (dropDirty) localStorage.removeItem(dirtyKey(base, sessionId)) + localStorage.setItem( + removedKey(base, sessionId, ws), + storages && storages.length > 0 ? JSON.stringify(storages) : '1' + ) + return true + } catch (e) { + console.error('Could not persist session backup mark', e) + return false + } +} + +function removeKey(key: string): void { + try { + localStorage.removeItem(key) + } catch {} +} + +// --- Scheduling --- + +let quietTimer: ReturnType | undefined +let maxTimer: ReturnType | undefined +let retryTimer: ReturnType | undefined +let retryAt = 0 +let retryMs = RETRY_MIN_MS +/** + * Workspaces whose storage answered this page load. `off`: nowhere to keep backups. + * `refused`: the server rejected what this page sends; nothing more is sent for the page, + * but the marks and the sync state stay, so the next load tries again. + */ +const wsState = new Map() +/** Backups an admin turns on again elsewhere are noticed `OFF_RETRY_MS` after the workspace + * was found off: a timer per off workspace forgets the state then and asks again, so a + * pending mark left for the workspace does not wait for something else to flush. The + * admin's own page hears of the switch at once (see `backupSettingsChanged`). */ +const offTimers = new Map>() +const OFF_RETRY_MS = 10 * 60_000 +let offRetryMs = OFF_RETRY_MS + +function markOff(ws: string): void { + wsState.set(ws, 'off') + clearTimeout(offTimers.get(ws)) + offTimers.set( + ws, + setTimeout(() => { + offTimers.delete(ws) + if (wsState.get(ws) === 'off') wsState.delete(ws) + restoredWorkspaces.delete(ws) + void enqueue(flush) + restoreSessionBackups(ws) + }, offRetryMs) + ) +} + +function forgetOff(ws: string): void { + clearTimeout(offTimers.get(ws)) + offTimers.delete(ws) + if (wsState.get(ws) === 'off') wsState.delete(ws) +} + +function clearOffTimers(): void { + for (const timer of offTimers.values()) clearTimeout(timer) + offTimers.clear() +} + +function isOff(ws: string): boolean { + return wsState.get(ws) === 'off' +} +let backfilled = false +/** Sessions of the current user whose dirty mark could not be written this page. */ +/** Sessions whose dirty mark localStorage refused this page, with the bumps of theirs no + * sync row could take yet. */ +let unwritableMarks = new Map() +const restoredWorkspaces = new Set() + +// One flush or restore at a time in this tab; each reads the marks fresh. +let chain: Promise = Promise.resolve() +function enqueue(fn: () => Promise): Promise { + const run = chain.then(fn, fn) + chain = run.catch((e) => console.error('Session backup failed', e)) + return run +} + +function clearTimers(): void { + clearTimeout(quietTimer) + clearTimeout(maxTimer) + quietTimer = undefined + maxTimer = undefined +} + +function scheduleFlush(): void { + clearTimeout(quietTimer) + quietTimer = setTimeout(runFlush, QUIET_MS) + maxTimer ??= setTimeout(runFlush, MAX_DELAY_MS) +} + +function runFlush(): void { + clearTimers() + void enqueue(flush) +} + +function backOff(): void { + retryAt = Date.now() + retryMs + retryMs = Math.min(retryMs * 2, RETRY_MAX_MS) + clearTimeout(retryTimer) + retryTimer = setTimeout(runFlush, retryAt - Date.now()) +} + +/** Serialize with the other tabs of the same user where the browser lets us; on plain + * http there is no lock, and two tabs at worst upload the same bytes twice. */ +/** One tab of the user at a time in the flush and the restore. A flush finding the lock + * taken reschedules itself; a restore waits its turn, since two tabs restoring the same + * absent session would each write its pieces over the other's. */ +function webLocks(): LockManager | undefined { + return typeof navigator === 'undefined' ? undefined : (navigator as { locks?: LockManager }).locks +} + +function hasWebLocks(): boolean { + return webLocks() !== undefined +} + +async function withUserLock(email: string, fn: () => Promise, wait = false): Promise { + const locks = webLocks() + if (!locks) return fn() + await locks.request(`wm-ai-sessions-mirror::${email}`, { ifAvailable: !wait }, async (lock) => { + if (lock) await fn() + // The other tab's flush read the marks before this one's were written: try again + // once it is done, rather than wait for the next write or load. + else scheduleFlush() + }) +} + +// --- Flush --- + +function statusOf(e: unknown): number | undefined { + return e instanceof ApiError ? e.status : undefined +} + +async function readSync(id: string, email: string): Promise { + return (await syncDb(email))?.get('sync', id) +} + +async function allSyncRows(email: string): Promise { + return (await (await syncDb(email))?.getAll('sync')) ?? [] +} + +/** The durable fallback for a user delete whose localStorage mark could not be written. + * A session with a row was backed up; one without may have its first push in flight, so + * it gets a row saying only that, which the push's own row write keeps (`writeSync`). An + * unsent draft gets nothing: it was never pushed. */ +async function removeViaSyncRow( + id: string, + ws: string | undefined, + email = getCurrentUserEmail() +): Promise { + if (!email) return + await updateSyncRow(id, email, (row) => + row + ? { ...row, removed: true } + : ws + ? { id, ws, head: '', chats: {}, images: {}, removed: true } + : undefined + ) +} + +/** Writes rows whole, except that a removal filed on a row meanwhile survives: the flush + * writes a session's row from state it read before the push, and the user may have deleted + * the session in between. */ +/** A row that says where else the copy is (`alsoIn`) is believed; one that does not + * inherits it, plus the storage the row it replaces was on when that is another one of this + * workspace's: the copy there stays where it was. */ +async function writeSync(states: MirrorSyncState[], email: string): Promise { + const db = await syncDb(email) + if (!db || states.length === 0) return + // Bumps this page kept for want of a row (see `bumpViaSyncRow`) move onto the row now; + // what arrives while the write is in flight stays counted here. + const taken = new Map() + const tx = db.transaction('sync', 'readwrite') + for (const s of states) { + const cur = await tx.store.get(s.id) + const removed = s.removed || cur?.removed + const kept = email === getCurrentUserEmail() ? (unwritableMarks.get(s.id) ?? 0) : 0 + if (kept > 0) taken.set(s.id, kept) + const extraV = Math.max(s.extraV ?? 0, cur?.extraV ?? 0) + kept + const alsoIn = new Set(s.alsoIn ?? []) + if (s.alsoIn === undefined && cur?.ws === s.ws) { + for (const id of cur.alsoIn ?? []) alsoIn.add(id) + if (cur.storageId !== undefined) alsoIn.add(cur.storageId) + } + if (s.storageId !== undefined) alsoIn.delete(s.storageId) + const { alsoIn: _, ...rest } = s + await tx.store.put({ + ...rest, + ...(removed ? { removed: true } : {}), + ...(extraV > 0 ? { extraV } : {}), + ...(alsoIn.size > 0 ? { alsoIn: [...alsoIn] } : {}) + }) + } + await tx.done + for (const [id, n] of taken) unwritableMarks.set(id, (unwritableMarks.get(id) ?? 0) - n) +} + +async function deleteSync(ids: string[], email: string): Promise { + const db = await syncDb(email) + if (!db || ids.length === 0) return + const tx = db.transaction('sync', 'readwrite') + for (const id of ids) await tx.store.delete(id) + await tx.done +} + +/** A workspace's storage went away: what was pushed there can no longer be trusted to be + * where a later storage looks, so every session goes whole on the next push. The rows stay + * (stale) so a removal still knows a backup existed. */ +async function staleWorkspaceSync(ws: string, email: string): Promise { + const db = await syncDb(email) + if (!db) return + const rows = (await db.getAll('sync')).filter((s) => s.ws === ws && !s.stale) + await writeSync( + rows.map((s) => ({ ...s, stale: true })), + email + ) +} + +/** The workspace's live rows recorded against another storage than the one the server + * answers from: they describe objects it no longer looks at (a new bucket starts empty). */ +function foreignRows( + ws: string, + storageId: string, + generation: number, + rows: Iterable +): MirrorSyncState[] { + return [...rows].filter( + (row) => + row.ws === ws && + !row.stale && + !row.removed && + !row.staging && + (row.storageId !== storageId || (row.generation ?? 0) !== generation) + ) +} + +/** Stale rows plan like no row at all, and the mark makes the next flush pick them up. */ +async function markStale(rows: MirrorSyncState[], email: string): Promise { + await writeSync( + rows.map((row) => ({ ...row, stale: true })), + email + ) + for (const row of rows) bumpDirty(row.id) +} + +/** Sessions this browser has backed up nothing of yet, whose backup went stale, or whose + * row carries bumps (localStorage refused their marks) no push has covered: everything + * committed to a workspace gets a mark, once per page load. */ +async function backfillMarks(marks: PendingMarks, email: string): Promise { + if (backfilled) return true + const sessions = await readStoredSessions(email) + const db = await syncDb(email) + if (!sessions || !db) return false + backfilled = true + const rows = new Map((await db.getAll('sync')).map((s) => [s.id, s])) + const marked = new Set(marks.dirty.map((d) => d.id)) + for (const s of sessions) { + const row = rows.get(s.id) + // A restored row has neither counter: only bumps it carries make it owed. + const owed = + !row || row.stale || ((row.extraV ?? 0) > 0 && (row.flushedV ?? -1) < (row.extraV ?? 0)) + if (s.workspace_id && owed && !marked.has(s.id)) { + // The counter as the next `readPending` will see it: the mark's, or this page's + // when none could be written, so what this flush retires is what a later bump + // counts from. + let v = 1 + if (!bumpDirty(s.id)) { + if (!unwritableMarks.has(s.id)) unwritableMarks.set(s.id, 0) + v = unwritableMarks.get(s.id) ?? 0 + } + marks.dirty.push({ id: s.id, v }) + } + } + return true +} + +/** Read what the stores hold for a dirty session and plan its push. `undefined` when the + * session has nowhere to go (unsent), `unavailable` when a store could not be read. */ +async function planFor( + session: Session, + sync: MirrorSyncState | undefined, + email: string +): Promise { + const chatIds = await listSessionChatIds(session.id, email) + if (!chatIds) return 'unavailable' + const chats: ChatSnapshot[] = [] + for (const id of chatIds) { + const record = await readStoredChat(id, email) + if (!record) continue + const imageIds = (await listChatImageIds(id, email)) ?? [] + if (jsonBytes(record) > MAX_CHAT_BYTES) { + console.warn(`AI session chat ${id} is too large to back up; leaving it out`) + chats.push({ id, lastModified: record.lastModified, imageIds, omitted: true }) + continue + } + chats.push({ id, lastModified: record.lastModified, record, imageIds }) + } + let artifacts = await readSessionArtifacts(session.id, email) + if (!artifacts) return 'unavailable' + if (jsonBytes(artifacts) > MAX_ARTIFACTS_BYTES) { + console.warn(`AI session ${session.id} artifacts are too large to back up; leaving them out`) + artifacts = { items: [], versions: [] } + } + return planSessionPush({ session, chats, artifacts, sync }) +} + +interface WorkspaceWork { + items: { session: Session; v: number; sync?: MirrorSyncState }[] + /** `key` is the localStorage mark; absent when the removal rides on the sync row. + * `storageId` and `alsoIn` are the storages holding a copy, when a sync row says: a + * removal is done only once each of them answered it. */ + removed: { id: string; key?: string; storageId?: string; alsoIn?: string[] }[] +} + +type SendStatus = 'ok' | 'off' | 'refused' | 'abort' | 'transient' + +interface WorkspaceOutcome { + status: SendStatus + /** The storage and backup generation the server answered from, once it answered. */ + storageId?: string + generation?: number + /** Sessions every part of which the server stored. */ + settled: { + id: string + v: number + next: MirrorSyncState + removeFrom?: string + carried: boolean + storageId?: string + generation?: number + }[] + /** Marks with nothing behind them (an unsent draft, a session gone from the store). */ + dropped: { id: string; v: number }[] + /** Sessions the server holds no head of any more (another device removed the backup): + * their rows go stale, so the next flush sends them whole. */ + needsWhole: string[] + /** Removal marks the server carried out. */ + removedDone: { id: string; key?: string }[] + /** Removals one storage carried out while others still hold a copy: the row narrows to + * those, and the mark waits for them to answer. */ + removedFrom: { id: string; key?: string; remaining: string[] }[] + /** Some session the server could not store. */ + anyFailed: boolean + /** Some store could not be read; its marks stay for a later flush. */ + unavailable: boolean +} + +/** + * Plan and send a workspace's sessions one at a time, filling requests of about + * REQUEST_TARGET_BYTES as it goes, so a first backfill never holds more than one + * request's worth of records and images at once. Removals go first, in their own + * request(s). + */ +async function pushWorkspace( + ws: string, + work: WorkspaceWork, + email: string +): Promise { + const out: WorkspaceOutcome = { + status: 'ok', + settled: [], + dropped: [], + needsWhole: [], + removedDone: [], + removedFrom: [], + anyFailed: false, + unavailable: false + } + // Parts of a session still to be acknowledged, and the sessions the server refused. + // `complete` once every part of the session has been appended: a session whose first + // part is still to come when a request fails has nothing behind it yet. + const attempted = new Map< + string, + { + v: number + next: MirrorSyncState + parts: number + complete: boolean + removeFrom?: string + carried: boolean + storageId?: string + generation?: number + /** The storage and generation the last answer for this session came from. */ + answered?: string + } + >() + const failed = new Set() + const headless = new Set() + let current: PushBody | undefined + let size = 0 + let ops = 0 + + const send = async (body: PushBody): Promise => { + if (getCurrentUserEmail() !== email) return 'abort' + let res + try { + res = await AiService.pushAiSessionBackups({ workspace: ws, requestBody: body }) + } catch (e) { + const status = statusOf(e) + // 404: a build without object storage. 403: nothing this token may back up. + if (status === 404 || status === 403) return 'off' + if (status === 409) return 'abort' + // Too large is a fact about the sessions in this body, not the workspace: they + // stay marked (and retried with backoff), the others go on. + if (status === 413) { + console.error('Session backup push too large', e) + for (const entry of body.sessions) failed.add(entry.id) + for (const id of body.removed ?? []) failed.add(id) + return 'ok' + } + if (status !== undefined && status < 500 && status !== 429) { + console.error('Session backup push refused', e) + return 'refused' + } + console.warn('Session backup push failed, retrying later', e) + return 'transient' + } + if (!res.enabled) return 'off' + out.storageId = res.storage_id + out.generation = res.backup_generation + const answered = `${res.storage_id}:${res.backup_generation}` + const errors = new Set() + for (const r of res.results) { + if (r.error) { + console.warn(`Session backup of ${r.id} failed: ${r.error}`) + errors.add(r.id) + } else if (r.needs_whole) { + // The part rode on a head the storage no longer has: not a failure to back + // off from, but nothing to settle either. + headless.add(r.id) + out.needsWhole.push(r.id) + } + } + for (const entry of body.sessions) { + const a = attempted.get(entry.id) + if (!a) continue + a.parts -= 1 + // Parts answered from different storages sit in different buckets: nothing to + // settle, the session goes again whole. + if (a.answered !== undefined && a.answered !== answered) failed.add(entry.id) + a.answered = answered + a.storageId = res.storage_id + a.generation = res.backup_generation + if (errors.has(entry.id)) failed.add(entry.id) + } + for (const id of body.removed ?? []) { + if (errors.has(id)) failed.add(id) + else { + const mark = work.removed.find((r) => r.id === id) + if (!mark) continue + const holding = new Set( + [mark.storageId, ...(mark.alsoIn ?? [])].filter((s): s is string => s !== undefined) + ) + // Answered from a storage holding no copy: the copies are still where they + // were, and the mark waits for those storages to answer. + if (holding.size === 0 || res.storage_id === undefined) out.removedDone.push(mark) + else if (holding.has(res.storage_id)) { + holding.delete(res.storage_id) + if (holding.size === 0) out.removedDone.push(mark) + else out.removedFrom.push({ id, key: mark.key, remaining: [...holding] }) + } + } + } + return 'ok' + } + const flushCurrent = async (): Promise => { + if (!current) return 'ok' + const body = current + current = undefined + size = 0 + ops = 0 + return send(body) + } + const append = async (entry: AISessionBackupPush): Promise => { + const bytes = jsonBytes(entry) + const entryOps = operationsOf(entry) + if ( + current && + (current.sessions.length >= MAX_ENTRIES_PER_REQUEST || + ops + entryOps > MAX_OPERATIONS_PER_REQUEST || + size + bytes > REQUEST_TARGET_BYTES) + ) { + const status = await flushCurrent() + if (status !== 'ok') return status + } + // A part the server refused holds the rest of its session back: the head rides on + // the last part, and would list a session missing a piece. + if (failed.has(entry.id)) return 'ok' + current ??= { owner: email, sessions: [] } + current.sessions.push(entry) + size += bytes + ops += entryOps + const a = attempted.get(entry.id) + if (a) a.parts += 1 + return 'ok' + } + const finish = (status: SendStatus): WorkspaceOutcome => { + out.status = status + for (const [id, a] of attempted) { + if (a.complete && a.parts === 0 && !failed.has(id) && !headless.has(id)) { + out.settled.push({ + id, + v: a.v, + next: a.next, + removeFrom: a.removeFrom, + carried: a.carried, + storageId: a.storageId, + generation: a.generation + }) + } + } + out.anyFailed = failed.size > 0 + return out + } + + for (let i = 0; i < work.removed.length; i += MAX_REMOVED_PER_REQUEST) { + const chunk = work.removed.slice(i, i + MAX_REMOVED_PER_REQUEST) + const status = await send({ owner: email, sessions: [], removed: chunk.map((r) => r.id) }) + if (status !== 'ok') return finish(status) + } + + for (const item of work.items) { + if (getCurrentUserEmail() !== email) return finish('abort') + const plan = await planFor(item.session, item.sync, email) + if (plan === 'unavailable') { + out.unavailable = true + continue + } + if (!plan) { + out.dropped.push({ id: item.session.id, v: item.v }) + continue + } + const nothingToSend = !plan.entry && plan.images.length === 0 + // A move: the copy in the old workspace is filed for removal once this push has + // landed (see `settled`), never before, so the session is backed up somewhere at + // every point. + // A session with nothing to send gets no answer to name its storage: it keeps the + // one its row has, which the storage check below then judges like any other. + attempted.set(item.session.id, { + v: item.v, + next: plan.next, + parts: 0, + complete: nothingToSend, + removeFrom: plan.removeFrom, + carried: plan.carried, + storageId: item.sync?.storageId, + generation: item.sync?.generation + }) + if (nothingToSend) continue + // A push of the session whole names itself on every part and opens with its head on + // the first, whichever that is: the server replaces the backup on that part, lists + // the session by the last, and meanwhile refuses any other push of it. + const epoch = item.session.moves ?? 0 + let opened = false + let token: string | undefined + const open = (part: AISessionBackupPush): AISessionBackupPush => { + const first = !opened + opened = true + // A push split over parts (the first says more follow) names itself on each, so + // the server keeps the session unlisted between them; a whole one opens with the + // head on whichever part goes first. + if (first && part.partial) token = randomUUID() + return { + ...part, + epoch, + ...(plan.whole ? { whole: true } : {}), + ...(plan.whole && first ? { head: plan.entry?.head } : {}), + ...(token ? { push: token } : {}), + ...(token && first ? { opens: true } : {}) + } + } + let images: AISessionBackupImage[] = [] + let imagesBytes = 0 + for (const { chat_id, id } of plan.images) { + const data_url = await readImageDataUrl(id, email) + // Evicted since the plan was made; the next save of that chat drops the id. + if (!data_url) continue + if ( + images.length > 0 && + (images.length >= MAX_IMAGES_PER_ENTRY || + imagesBytes + data_url.length > REQUEST_TARGET_BYTES) + ) { + const status = await append(open({ id: item.session.id, images, partial: true })) + if (status !== 'ok') return finish(status) + images = [] + imagesBytes = 0 + } + images.push({ chat_id, id, data_url }) + imagesBytes += data_url.length + } + // Every part but the last says so: the server lists a session on the part that + // completes its entry, never on one an unsent part still follows. + const parts = plan.entry + ? splitEntry( + plan.whole ? { ...plan.entry, head: undefined } : plan.entry, + REQUEST_TARGET_BYTES + ) + : [] + if (images.length > 0) { + const status = await append( + open({ + id: item.session.id, + images, + partial: parts.length > 0 || undefined + }) + ) + if (status !== 'ok') return finish(status) + } + for (const [i, part] of parts.entries()) { + if (failed.has(item.session.id)) break + const status = await append(open(i < parts.length - 1 ? { ...part, partial: true } : part)) + if (status !== 'ok') return finish(status) + } + attempted.get(item.session.id)!.complete = true + } + return finish(await flushCurrent()) +} + +async function flush(): Promise { + const email = getCurrentUserEmail() + if (!email) return + if (Date.now() < retryAt) { + clearTimeout(retryTimer) + retryTimer = setTimeout(runFlush, retryAt - Date.now()) + return + } + await withUserLock(email, async () => { + const marks = readPending() + // A store that cannot be opened right now (another tab's upgrade in progress) is + // retried with backoff, as any other unavailable store below. + if (!(await backfillMarks(marks, email))) { + backOff() + return + } + // Read once: retired marks are not reclaimed (a mark cannot be deleted without a + // window in which a bump is lost), so there is one per session ever backed up, and + // telling them apart from live ones is what lets a flush with nothing to do stop + // here, before the sessions store. + const syncRows = new Map((await allSyncRows(email)).map((row) => [row.id, row])) + // A mark's counter counts with the bumps its row carries (the ones localStorage refused). + const live = marks.dirty + .map((d) => ({ id: d.id, v: d.v + (syncRows.get(d.id)?.extraV ?? 0) })) + .filter((d) => { + const sync = syncRows.get(d.id) + return !(sync && !sync.stale && (sync.flushedV ?? -1) >= d.v) + }) + // A removal whose mark could not be written rides on the sync row instead. + const removals: { id: string; ws?: string; key?: string; storages?: string[] }[] = [ + ...marks.removed + ] + for (const row of syncRows.values()) { + if (row.removed && !removals.some((r) => r.id === row.id)) { + removals.push({ id: row.id, ws: row.ws }) + } + } + if (live.length === 0 && removals.length === 0) return + const stored = await readStoredSessions(email) + if (!stored) { + backOff() + return + } + const byId = new Map(stored.map((s) => [s.id, s])) + const work = new Map() + const workFor = (ws: string) => { + let w = work.get(ws) + if (!w) work.set(ws, (w = { items: [], removed: [] })) + return w + } + const droppedDirty: string[] = [] + const consumedRemoved: string[] = [] + + for (const r of removals) { + const sync = syncRows.get(r.id) + const ws = r.ws ?? sync?.ws + // Nowhere to remove it from (an unsent draft): done. A workspace whose backups are + // off keeps the removal of a session that was backed up, for when they are on + // again, or the session would come back; one never backed up from here has + // nothing there, so its mark goes, or a storage-less instance would collect one + // per deleted session forever. + if (!ws) { + if (r.key) consumedRemoved.push(r.key) + } else if (!isOff(ws) && wsState.get(ws) !== 'refused') { + // The row says where the copy is only for its own workspace: a session that + // moved on has the row its new workspace wrote, and its mark says instead. + const own = sync?.ws === ws ? sync : undefined + const holding = (own ? [own.storageId, ...(own.alsoIn ?? [])] : (r.storages ?? [])).filter( + (x): x is string => x !== undefined + ) + workFor(ws).removed.push({ + id: r.id, + key: r.key, + storageId: holding[0], + alsoIn: holding.slice(1) + }) + } else if (isOff(ws) && !sync && r.key) consumedRemoved.push(r.key) + } + for (const d of live) { + const session = byId.get(d.id) + // Gone from the store, so nothing can bump it again; an unsent draft keeps its + // mark, since it may commit to a workspace while this flush runs. + if (!session) { + droppedDirty.push(d.id) + continue + } + if (session.transient || !session.workspace_id) continue + // Left in place for a workspace that is off or refused: a move into it must still + // remember the old copy, and a mark costs one lookup per flush. + if (isOff(session.workspace_id) || wsState.get(session.workspace_id) === 'refused') { + continue + } + const sync = syncRows.get(d.id) + // A stale row plans like no row at all: the whole session goes again. One naming + // another workspace still says where the old copy is, stale or not. + workFor(session.workspace_id).items.push({ + session, + v: d.v, + sync: sync?.stale && sync.ws === session.workspace_id ? undefined : sync + }) + } + + let settledAny = false + let leftForNext = false + for (const [ws, w] of work) { + const out = await pushWorkspace(ws, w, email) + if (out.status === 'abort') return + if (out.status === 'off') { + markOff(ws) + await staleWorkspaceSync(ws, email) + // Same rule as the loop above: a removal is worth keeping only for a session + // that was backed up from here. + for (const r of w.removed) { + if (!syncRows.has(r.id) && r.key) consumedRemoved.push(r.key) + } + continue + } + // Refused stops the workspace for the page, but what the earlier requests of this + // flush stored is recorded like any other. + if (out.status === 'refused') wsState.set(ws, 'refused') + if (out.status === 'transient' || out.anyFailed || out.unavailable) backOff() + else settledAny = true + // The new workspace holds a moved session now, so the old copy can go: its removal + // mark is written before the row that forgets where the old copy was, and a + // session whose mark could not be written keeps its old row, so the next flush + // plans the move again rather than orphan the copy. + const recorded = out.settled.filter((s) => { + if (s.removeFrom) { + const old = syncRows.get(s.id) + const storages = + old?.ws === s.removeFrom + ? [old.storageId, ...(old.alsoIn ?? [])].filter((x): x is string => x !== undefined) + : undefined + if (!addRemoved(s.id, s.removeFrom, false, undefined, storages)) return false + } + if (s.removeFrom || s.carried) leftForNext = true + return true + }) + // A session with deletes carried over stays one bump short of retired, so the + // next flush sends the rest. + const written = recorded.map((s) => ({ + ...s.next, + flushedV: s.carried ? s.v - 1 : s.v, + storageId: s.storageId, + generation: s.generation + })) + await writeSync(written, email) + // The server names the storage and generation every answer comes from. Rows + // naming another go stale, the ones written just now included: a session answered + // from a storage the later answers left behind, or pushed in part on top of a row + // from another one, has its backup split across buckets the server no longer + // looks at as a whole. + if (out.storageId !== undefined) { + const storageId = out.storageId + const generation = out.generation ?? 0 + const elsewhere = (row: MirrorSyncState) => + row.storageId !== storageId || (row.generation ?? 0) !== generation + const removed = new Set(out.removedDone.map((r) => r.id)) + const writtenIds = new Set(written.map((row) => row.id)) + const untouched = [...syncRows.values()].filter( + (row) => !removed.has(row.id) && !writtenIds.has(row.id) + ) + const foreign = [ + ...foreignRows(ws, storageId, generation, untouched), + ...written.filter((row) => { + const prior = syncRows.get(row.id) + const partial = prior && !prior.stale && prior.ws === ws && elsewhere(prior) + return elsewhere(row) || partial + }) + ] + if (foreign.length > 0) { + await markStale(foreign, email) + leftForNext = true + } + } + droppedDirty.push(...out.dropped.map((d) => d.id)) + // The server holds no head of these any more (another device removed the backup): + // stale rows send them whole next. + const headless = out.needsWhole + .map((id) => syncRows.get(id)) + .filter((row): row is MirrorSyncState => row !== undefined && !row.stale) + if (headless.length > 0) { + await markStale(headless, email) + leftForNext = true + } + for (const r of out.removedDone) { + // The row describes this workspace's copy only; a session that moved on keeps + // the row its new workspace wrote. + if ((await readSync(r.id, email))?.ws === ws) await deleteSync([r.id], email) + if (r.key) consumedRemoved.push(r.key) + } + for (const r of out.removedFrom) { + const row = await readSync(r.id, email) + const [storageId, ...alsoIn] = r.remaining + if (row?.ws === ws) await writeSync([{ ...row, storageId, alsoIn }], email) + else if (r.key) { + try { + localStorage.setItem(r.key, JSON.stringify(r.remaining)) + } catch {} + } + } + } + if (settledAny && Date.now() >= retryAt) retryMs = RETRY_MIN_MS + if (getCurrentUserEmail() !== email) return + for (const id of droppedDirty) dropDirty(id) + for (const key of consumedRemoved) removeKey(key) + // Whatever is still marked is either waiting on the backoff timer, on a write that + // scheduled its own flush, or on a workspace that is off for a while or refused for the page; + // none of it wants another flush in 15 s. What this flush left for the next one + // does: a moved session's removal from its old workspace, carried-over deletes, or + // the sessions of a storage the server no longer answers from. + if (leftForNext) scheduleFlush() + }) +} + +// --- Restore --- + +function asRecord(value: unknown): Record | undefined { + return value && typeof value === 'object' && !Array.isArray(value) + ? (value as Record) + : undefined +} + +/** Turn one pulled backup into store rows, dropping anything that does not name this + * session: the bucket is written by the server from validated ids, but a record is + * still data from outside this browser. */ +function unpackBackup( + ws: string, + backup: AISessionBackup, + updatedAt: number, + storageId: string | undefined, + generation: number | undefined, + earlierChats: Iterable = [] +): + | { + session: Session + chats: StoredChat[] + images: RestoredImage[] + artifacts: { items: PersistedArtifact[]; versions: ArtifactVersion[] } + sync: MirrorSyncState + } + | undefined { + const head = asRecord(backup.head) + if (!head || head.id !== backup.id || head.workspace_id !== ws) return undefined + if (typeof head.createdAt !== 'number') return undefined + const chats: StoredChat[] = [] + for (const c of backup.chats) { + const record = asRecord(c.record) + if ( + !record || + record.id !== c.id || + record.sessionId !== backup.id || + !Array.isArray(record.actualMessages) || + !Array.isArray(record.displayMessages) || + typeof record.lastModified !== 'number' + ) { + continue + } + chats.push(record as unknown as StoredChat) + } + // An image's chat may have come on an earlier page of the session. + const chatIds = new Set([...chats.map((c) => c.id), ...earlierChats]) + const images: RestoredImage[] = backup.images + .filter((i) => chatIds.has(i.chat_id) && typeof i.data_url === 'string') + .map((i) => ({ id: i.id, chatId: i.chat_id, dataUrl: i.data_url })) + const artifactsRecord = asRecord(backup.artifacts) + const items = ( + Array.isArray(artifactsRecord?.items) ? (artifactsRecord.items as PersistedArtifact[]) : [] + ).filter((i) => asRecord(i)?.sessionId === backup.id && typeof i.id === 'string') + const itemIds = new Set(items.map((i) => i.id)) + const versions = ( + Array.isArray(artifactsRecord?.versions) ? (artifactsRecord.versions as ArtifactVersion[]) : [] + ).filter((v) => asRecord(v) && itemIds.has(v.artifactId) && typeof v.key === 'string') + const active = chats.find((c) => c.id === head.chatId) + const session: Session = { + ...(head as unknown as Session), + name: '', + // Everything in the backup has been read here: no unread badge, and the last + // activity is the backup's own time. + lastSeenCount: active?.displayMessages.length ?? 0, + lastActivityAt: updatedAt + } + const sync: MirrorSyncState = { + id: backup.id, + ws, + head: headSig(session), + chats: Object.fromEntries(chats.map((c) => [c.id, c.lastModified])), + images: Object.fromEntries(images.map((i) => [i.id, i.chatId])), + artifacts: artifactsFingerprint({ items, versions }), + storageId, + generation + } + return { session, chats, images, artifacts: { items, versions }, sync } +} + +type BackupListing = Awaited> + +/** The workspace's listing; `off` when it keeps no backups, `failed` when it could not be + * listed this time. */ +async function listWorkspace(ws: string, email: string): Promise { + let listing + try { + listing = await AiService.listAiSessionBackups({ workspace: ws }) + } catch (e) { + const status = statusOf(e) + if (status === 404 || status === 403) { + markOff(ws) + return 'off' + } + console.warn('Could not list session backups', e) + return 'failed' + } + if (!listing.enabled) { + markOff(ws) + return 'off' + } + if (wsState.get(ws) !== 'refused') wsState.set(ws, 'on') + const foreign = + listing.storage_id === undefined + ? [] + : foreignRows( + ws, + listing.storage_id, + listing.backup_generation ?? 0, + await allSyncRows(email) + ) + if (foreign.length > 0) { + await markStale(foreign, email) + scheduleFlush() + } + return listing +} + +/** The workspaces of one family are restored together: a session moved between two of them + * is listed by both until the old copy's removal lands (that mark is the moving browser's, + * which may never come back), and whichever were restored first would take the id and keep + * the other out. The copy that moved last (`epoch`, the record's move count, kept with the + * marker) is the one brought back, the storage's own modification time deciding between two + * of the same count; the other is left where it is. A family one of whose workspaces could + * not be listed is not restored at all this time, or the copy that lists could be the stale + * one; the next page load or workspace switch tries again. */ +async function restoreFamily(family: string[], todo: string[], email: string): Promise { + const listings = new Map() + for (const ws of todo) { + if (getCurrentUserEmail() !== email) return + const listing = await listWorkspace(ws, email) + if (listing === 'failed') { + for (const w of todo) restoredWorkspaces.delete(w) + return + } + if (listing !== 'off') listings.set(ws, listing) + // Probed again once its backups may be on again (see `isOff`). + else restoredWorkspaces.delete(ws) + } + const newest = new Map() + for (const [ws, listing] of listings) { + for (const s of listing.sessions) { + const at = Date.parse(s.updated_at) + const cur = newest.get(s.id) + if (!cur || s.epoch > cur.epoch || (s.epoch === cur.epoch && at > cur.at)) { + newest.set(s.id, { ws, epoch: s.epoch, at }) + } + } + } + // A move landing between the listings and the pulls (another device pushing the session + // into a workspace listed before it held it, or one whose backups were off then) would + // make a copy about to be imported the stale one: the whole family is listed again once + // a workspace's pulls are done and its records are about to land, and a session a later + // copy of which showed up elsewhere is left, with the family, for next time. A family of + // one has nowhere else for a copy to show up. + const verify = async (ids: string[]): Promise> => { + const superseded = new Set() + if (family.length < 2) return superseded + for (const w of family) { + const listing = await listWorkspace(w, email) + if (listing === 'failed') { + for (const id of ids) superseded.add(id) + break + } + if (listing === 'off') continue + for (const s of listing.sessions) { + const cur = newest.get(s.id) + if (!cur || !ids.includes(s.id) || w === cur.ws) continue + const at = Date.parse(s.updated_at) + if (s.epoch > cur.epoch || (s.epoch === cur.epoch && at > cur.at)) superseded.add(s.id) + } + } + if (superseded.size > 0) for (const w of todo) restoredWorkspaces.delete(w) + return superseded + } + for (const [ws, listing] of listings) { + if (getCurrentUserEmail() !== email) return + const elsewhere = listing.sessions.filter((s) => newest.get(s.id)?.ws !== ws).map((s) => s.id) + await restoreWorkspace(ws, email, listing, new Set(elsewhere), verify) + } +} + +async function restoreWorkspace( + ws: string, + email: string, + listing: BackupListing, + elsewhere: Set, + verify: (ids: string[]) => Promise> +): Promise { + const rows = await allSyncRows(email) + const local = new Set() + for (const s of (await readStoredSessions(email)) ?? []) local.add(s.id) + for (const s of sessionState.sessions) local.add(s.id) + // A removal names the workspace it is for: a session that moved here from another one + // still has that one's removal pending, and is ours to restore. + for (const r of readPending().removed) if (r.ws === ws) local.add(r.id) + for (const row of rows) if (row.removed && row.ws === ws) local.add(row.id) + const candidates = listing.sessions + .filter((s) => !local.has(s.id) && !elsewhere.has(s.id) && !isSessionTombstoned(s.id)) + .slice(0, RESTORE_MAX) + const updatedAt = new Map(candidates.map((s) => [s.id, Date.parse(s.updated_at)])) + let ids = candidates.map((s) => s.id) + const toImport: Staged[] = [] + // A session that did not fit one answer whole comes in pages, kept here until the last + // one: importing a page alone would leave a session the next restore takes for whole. + type Pieces = { + chats: Set + images: Set + items: Set + versions: Set + } + type Staged = { + session: Session + sync: MirrorSyncState + /** Everything the pages so far wrote for the session. */ + pieces: Pieces + /** The listing fingerprint the pages so far were answered with. */ + listing?: string + } + const staged = new Map() + const noPieces = (): Pieces => ({ + chats: new Set(), + images: new Set(), + items: new Set(), + versions: new Set() + }) + const union = (a: Pieces, b: Pieces): Pieces => ({ + chats: new Set([...a.chats, ...b.chats]), + images: new Set([...a.images, ...b.images]), + items: new Set([...a.items, ...b.items]), + versions: new Set([...a.versions, ...b.versions]) + }) + // What earlier attempts (a restore cut short, a start over) wrote for a session that has + // no record yet, from their staging rows: the pieces of it the backup no longer has go + // before the record lands, or a later flush would push them back. Ids, never clocks. + const earlierStaging = new Map() + for (const row of rows) { + if (row.staging) { + earlierStaging.set(row.id, { + chats: new Set(row.staging.chats), + images: new Set(row.staging.images), + items: new Set(row.staging.items), + versions: new Set(row.staging.versions) + }) + } + } + // A session whose backup moved between two of its pages starts over, a few times. + const restarts = new Map() + const MAX_RESTARTS = 3 + const resumes: AISessionBackupCursor[] = [] + while (ids.length > 0 || resumes.length > 0) { + if (getCurrentUserEmail() !== email) return + const resume = resumes.shift() + const batch = resume ? [resume.id] : ids.slice(0, PULL_BATCH) + if (!resume) ids = ids.slice(PULL_BATCH) + let pulled + try { + pulled = await AiService.pullAiSessionBackups({ + workspace: ws, + requestBody: { ids: batch, resume } + }) + } catch (e) { + console.warn('Could not pull session backups', e) + return + } + if (!pulled.enabled) { + markOff(ws) + return + } + // Ask again for what did not fit, one at a time so each answer is as small as can be. + for (const id of pulled.deferred) if (!ids.includes(id)) ids.unshift(id) + // A session's pieces land before its record, page by page (the writes are absent-only, + // so a restore cut short leaves nothing a later one cannot finish), and the record, + // which is what makes the session visible, only with the last page; a session whose + // pieces could not be written is left for the next restore, since recording it now + // would let the next flush push its half-empty local state over the backup. What a + // page leaves for the next is the sync row being assembled, never its pieces. + const ready: Staged[] = [] + for (const b of pulled.sessions) { + const earlier = staged.get(b.id) + staged.delete(b.id) + // Brought back meanwhile by this tab itself (a session moved here from another + // workspace, say; the lock keeps other tabs out): its pieces are not ours to write + // over any more. + if ((await readStoredSessions(email))?.some((s) => s.id === b.id)) continue + // The backup moved between two pages (a chat sorting before the cursor would be + // missed) or under this one (the page may mix two versions): the pages so far do + // not belong together, the session starts over. + if (b.moved || (earlier && b.listing !== earlier.listing)) { + const n = (restarts.get(b.id) ?? 0) + 1 + restarts.set(b.id, n) + if (earlier) { + earlierStaging.set(b.id, union(earlierStaging.get(b.id) ?? noPieces(), earlier.pieces)) + } + if (n < MAX_RESTARTS) ids.unshift(b.id) + else console.warn(`Session backup ${b.id} kept changing while restoring; left for later`) + continue + } + const u = unpackBackup( + ws, + b, + updatedAt.get(b.id) ?? Date.now(), + pulled.storage_id, + pulled.backup_generation, + Object.keys(earlier?.sync.chats ?? {}) + ) + if (!u) continue + const written: Pieces = { + chats: new Set(u.chats.map((c) => c.id)), + images: new Set(u.images.map((i) => i.id)), + items: new Set(u.artifacts.items.map((i) => i.id)), + versions: new Set(u.artifacts.versions.map((v) => v.key)) + } + const merged: Staged = earlier + ? { + session: { + ...u.session, + lastSeenCount: Math.max( + earlier.session.lastSeenCount ?? 0, + u.session.lastSeenCount ?? 0 + ) + }, + sync: { + ...u.sync, + chats: { ...earlier.sync.chats, ...u.sync.chats }, + images: { ...earlier.sync.images, ...u.sync.images }, + artifacts: b.artifacts !== undefined ? u.sync.artifacts : earlier.sync.artifacts + }, + pieces: union(earlier.pieces, written) + } + : { session: u.session, sync: u.sync, pieces: written } + merged.listing = b.listing + // The staging row goes before the pieces, and outlives a restore cut short: the + // next one reads it to know what to delete. The record replaces it. + const stagingPieces = union(earlierStaging.get(b.id) ?? noPieces(), merged.pieces) + await writeSync( + [ + { + id: b.id, + ws, + head: '', + chats: {}, + images: {}, + staging: { + chats: [...stagingPieces.chats], + images: [...stagingPieces.images], + items: [...stagingPieces.items], + versions: [...stagingPieces.versions] + } + } + ], + email + ) + // Written over whatever is there: the session is absent locally, so its pieces + // can only be what an earlier restore staged before it was cut short, and the + // backup may have moved on since. + try { + if (!(await importArtifacts(u.artifacts.items, u.artifacts.versions, email, true))) continue + if (!(await importStoredChats(u.chats, u.images, email, true))) continue + } catch (e) { + console.error(`Could not restore session ${u.session.id}`, e) + continue + } + if (b.next) { + staged.set(b.id, merged) + resumes.push(b.next) + continue + } + ready.push(merged) + } + if (ready.length === 0) continue + // The staged pieces the backup no longer has go before the record: once the record + // is there, no restore looks at the session again, and a flush would push them back. + // A prune that could not run leaves the session, its pieces and its staging row for + // the next restore. + const pruned: Staged[] = [] + for (const r of ready) { + const prior = earlierStaging.get(r.session.id) + if (prior) { + const gone = (was: Set, now: Set) => + new Set([...was].filter((id) => !now.has(id))) + const ok = + (await pruneSessionChats( + r.session.id, + gone(prior.chats, r.pieces.chats), + gone(prior.images, r.pieces.images), + email + )) && + (await pruneSessionArtifacts( + r.session.id, + gone(prior.items, r.pieces.items), + gone(prior.versions, r.pieces.versions), + email + )) + if (!ok) { + console.warn(`Session backup ${r.session.id} could not be tidied; left for later`) + continue + } + earlierStaging.delete(r.session.id) + } + pruned.push(r) + } + toImport.push(...pruned) + } + // The records land together once every pull is done, so the family is listed again + // once per workspace rather than per answer; the pieces are in place either way, and a + // restore cut short before this leaves them staged for the next. + if (toImport.length === 0) return + const superseded = await verify(toImport.map((r) => r.session.id)) + const current = toImport.filter((r) => !superseded.has(r.session.id)) + if (current.length === 0) return + const imported = new Set( + await importSessions( + current.map((r) => r.session), + email + ) + ) + await writeSync( + current.filter((r) => imported.has(r.session.id)).map((r) => r.sync), + email + ) + if (imported.size > 0) logFeatureUsage('ai_session', 'restored', { value: imported.size }) +} + +/** + * Restore the sessions of the workspace family the user is looking at (the workspace and + * its forks), once per workspace per page load. Sessions that exist here are never + * touched; only ones this browser lacks are brought back. + */ +export function restoreSessionBackups(currentWorkspace: string): void { + if (!BROWSER) return + // A restore writes an absent session's pieces page by page and prunes what an earlier + // one staged: two tabs doing that at once would write over each other, so it runs only + // under the tab lock. Where Web Locks do not exist (a plain http origin), the browser + // still backs up; its sessions come back on a secure one. + if (!hasWebLocks()) return + const email = getCurrentUserEmail() + if (!email) return + const all = get(userWorkspaces) + const root = workspaceRootId(currentWorkspace, all) ?? currentWorkspace + const family = new Set([currentWorkspace]) + for (const w of all) if ((workspaceRootId(w.id, all) ?? w.id) === root) family.add(w.id) + const todo = [...family].filter((ws) => !restoredWorkspaces.has(ws) && !isOff(ws)) + for (const ws of todo) restoredWorkspaces.add(ws) + if (todo.length > 0) { + void enqueue(() => withUserLock(email, () => restoreFamily([...family], todo, email), true)) + } +} + +/** The workspace's backups were just turned on or off from this page: what was learnt of + * them is forgotten, and the next flush and a restore find out afresh. */ +export function backupSettingsChanged(ws: string): void { + forgetOff(ws) + wsState.delete(ws) + restoredWorkspaces.delete(ws) + // The rows went stale when the backups went off: the sessions are marked again. + backfilled = false + scheduleFlush() + restoreSessionBackups(ws) +} + +// --- Wiring --- + +if (BROWSER) { + onMirrorSignal((signal) => { + // A mark for another user waits for that user's next load. + const mine = !signal.email || signal.email === getCurrentUserEmail() + if (signal.kind === 'dirty') { + if (bumpDirty(signal.sessionId, signal.email)) { + if (mine) scheduleFlush() + } else { + if (mine && !unwritableMarks.has(signal.sessionId)) { + unwritableMarks.set(signal.sessionId, 0) + } + void bumpViaSyncRow(signal.sessionId, signal.email).finally(() => { + if (mine) scheduleFlush() + }) + } + return + } + if (!addRemoved(signal.sessionId, signal.workspaceId, true, signal.email)) { + void removeViaSyncRow(signal.sessionId, signal.workspaceId, signal.email) + } + if (mine) scheduleFlush() + }) + onUserChange((email) => { + clearTimers() + clearTimeout(retryTimer) + retryAt = 0 + retryMs = RETRY_MIN_MS + wsState.clear() + clearOffTimers() + restoredWorkspaces.clear() + backfilled = false + unwritableMarks.clear() + if (email) setTimeout(runFlush, STARTUP_DELAY_MS) + }) + // A tab going to the background may not come back: carry what it has now. + if (typeof document !== 'undefined') { + document.addEventListener('visibilitychange', () => { + if (document.visibilityState === 'hidden') runFlush() + }) + } +} + +/** Test-only: run a flush now, outside the timers. */ +export function __flushForTesting(): Promise { + return enqueue(flush) +} + +/** Test-only: what the sync table holds for the user. */ +export async function __syncRowsForTesting(email: string): Promise { + return allSyncRows(email) +} + +/** Test-only: how long a workspace found off is left alone before the page asks again. */ +export function __setOffRetryForTesting(ms: number): void { + offRetryMs = ms +} + +/** Test-only: plant sync rows, as an earlier page load would have left them. */ +export function __writeSyncForTesting(rows: MirrorSyncState[], email: string): Promise { + return writeSync(rows, email) +} + +/** Test-only: wait for whatever flush or restore is queued. */ +export function __settleForTesting(): Promise { + return enqueue(async () => {}) +} + +/** Test-only: forget every page-lifetime decision, and let go of the sync store so the + * next open lands in the test's fresh IndexedDB rather than the cached connection. */ +export function __resetMirrorForTesting(): void { + syncDbh.close() + clearTimers() + clearTimeout(retryTimer) + retryAt = 0 + retryMs = RETRY_MIN_MS + wsState.clear() + clearOffTimers() + offRetryMs = OFF_RETRY_MS + restoredWorkspaces.clear() + backfilled = false + unwritableMarks.clear() +} diff --git a/frontend/src/lib/components/sessions/sessionMirror.test.ts b/frontend/src/lib/components/sessions/sessionMirror.test.ts new file mode 100644 index 0000000000..d7260c0878 --- /dev/null +++ b/frontend/src/lib/components/sessions/sessionMirror.test.ts @@ -0,0 +1,1449 @@ +import { describe, it, expect, beforeEach, vi } from 'vitest' +import { IDBFactory } from 'fake-indexeddb' + +// The stores are BROWSER-gated; the vitest "server" env reports false. +vi.mock('esm-env', async (importOriginal) => ({ + ...(await importOriginal()), + BROWSER: true +})) + +const { pushMock, listMock, pullMock } = vi.hoisted(() => ({ + pushMock: vi.fn(), + listMock: vi.fn(), + pullMock: vi.fn() +})) +vi.mock('$lib/gen', async (orig) => { + const actual = await orig() + return { + ...actual, + AiService: { + ...actual.AiService, + pushAiSessionBackups: pushMock, + listAiSessionBackups: listMock, + pullAiSessionBackups: pullMock + }, + WorkspaceService: { + ...actual.WorkspaceService, + listUserWorkspaces: vi.fn().mockResolvedValue([]), + getSessionWorkspaceStatus: vi.fn().mockResolvedValue({}) + } + } +}) + +// The chat store during a restore: real, or unreachable for one session. +const { chatImport } = vi.hoisted(() => ({ + chatImport: { unavailable: false, pruneFails: false } +})) +vi.mock('../copilot/chat/HistoryManager.svelte', async (orig) => { + const actual = await orig() + return { + ...actual, + importStoredChats: (...args: Parameters) => + chatImport.unavailable ? Promise.resolve(false) : actual.importStoredChats(...args), + pruneSessionChats: (...args: Parameters) => + chatImport.pruneFails ? Promise.resolve(false) : actual.pruneSessionChats(...args) + } +}) + +/** The Web Locks API, which the node test environment lacks: one holder per name at a + * time, `ifAvailable` answering null while a holder is there. */ +function fakeLockManager(): LockManager { + const tails = new Map>() + return { + request: async ( + name: string, + options: LockOptions | undefined, + cb: (lock: Lock | null) => unknown + ) => { + const prev = tails.get(name) + if (options?.ifAvailable && prev) return cb(null) + const run = (prev ?? Promise.resolve()).then(() => cb({ name, mode: 'exclusive' })) + const tail = run.catch(() => {}) + tails.set(name, tail) + try { + return await run + } finally { + if (tails.get(name) === tail) tails.delete(name) + } + } + } as unknown as LockManager +} + +function setWebLocks(locks: LockManager | undefined): void { + if (typeof navigator === 'undefined') { + Object.defineProperty(globalThis, 'navigator', { + value: {}, + configurable: true, + writable: true + }) + } + Object.defineProperty(navigator, 'locks', { value: locks, configurable: true }) +} + +import { superadmin, userStore, usersWorkspaceStore, type UserExt } from '$lib/stores' +import HistoryManager, { + __resetBackupStoreForTesting, + __resetLegacyChatClaimForTesting, + readStoredChat +} from '../copilot/chat/HistoryManager.svelte' +import { + deleteSession, + importSessions, + putSession, + sessionState, + type Session +} from './sessionState.svelte' +import { markSessionDirty } from './sessionMirrorSignal' +import { + __flushForTesting, + __resetMirrorForTesting, + __settleForTesting, + __setOffRetryForTesting, + __syncRowsForTesting, + __writeSyncForTesting, + backupSettingsChanged, + restoreSessionBackups +} from './sessionMirror.svelte' + +const EMAIL = 'mirror@x.com' +const IMAGE = 'data:image/png;base64,AAAA' +const PENDING_PREFIX = `windmill_sessions_mirror_pending::${EMAIL}::` + +function asUser(email: string): UserExt { + return { email, username: email.split('@')[0] } as unknown as UserExt +} +const flush = () => new Promise((r) => setTimeout(r, 0)) + +function pendingKeys(): string[] { + const keys: string[] = [] + for (let i = 0; i < localStorage.length; i++) { + const key = localStorage.key(i) + if (key?.startsWith(PENDING_PREFIX)) keys.push(key.slice(PENDING_PREFIX.length)) + } + return keys.sort() +} + +function removalKeys(): string[] { + return pendingKeys().filter((k) => k.startsWith('r::')) +} + +/** Sessions whose dirty counter no push has covered yet. */ +async function pendingDirty(): Promise { + const flushed = new Map() + for (const r of await __syncRowsForTesting(EMAIL)) { + flushed.set(r.id, r.stale ? -1 : (r.flushedV ?? -1) - (r.extraV ?? 0)) + } + const out: string[] = [] + for (const key of pendingKeys()) { + if (!key.startsWith('d::')) continue + const id = key.slice(3) + if (Number(localStorage.getItem(PENDING_PREFIX + key)) > (flushed.get(id) ?? -1)) out.push(id) + } + return out.sort() +} + +beforeEach(async () => { + ;(globalThis as any).indexedDB = new IDBFactory() + localStorage.clear() + __resetLegacyChatClaimForTesting() + __resetBackupStoreForTesting() + __resetMirrorForTesting() + pushMock.mockReset() + listMock.mockReset() + pullMock.mockReset() + chatImport.unavailable = false + chatImport.pruneFails = false + setWebLocks(fakeLockManager()) + superadmin.set(false) + usersWorkspaceStore.set(undefined) + userStore.set(undefined) + await flush() + sessionState.sessions = [] + userStore.set(asUser(EMAIL)) + await vi.waitFor(() => expect(sessionState.hydrated).toBe(true)) +}) + +// The whole loop, end to end against the real stores: local writes mark, a flush sends +// one batch carrying the record, the chat and its image, reading the session sends +// nothing, and a user delete removes the backup. +describe('sessionMirror flush', () => { + it('pushes what changed, once, and removes a deleted session', async () => { + pushMock.mockResolvedValue({ enabled: true, results: [{ id: 's1' }] }) + const s: Session = { + id: 's1', + name: 'session-1', + createdAt: 1, + workspace_id: 'ws', + chatId: 'c1' + } + sessionState.sessions = [s] + await putSession(s) + + const hm = new HistoryManager() + await hm.init() + hm.setSessionId('s1') + hm.setCurrentChatId('c1') + await hm.saveChat( + [{ role: 'user', content: 'hello', images: [{ dataUrl: IMAGE, name: 'a.png' }] } as never], + [{ role: 'user', content: [{ type: 'image_url', image_url: { url: IMAGE } }] } as never] + ) + + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + const body = pushMock.mock.calls[0][0].requestBody + expect(pushMock.mock.calls[0][0].workspace).toBe('ws') + expect(body.owner).toBe(EMAIL) + expect(body.removed).toBeUndefined() + const [imageEntry, entry] = body.sessions + expect(imageEntry.images).toEqual([{ chat_id: 'c1', id: expect.any(String), data_url: IMAGE }]) + // A first push goes whole, under one token on every part: the head opens it on the + // first part, whichever that is. + expect(imageEntry.whole).toBe(true) + expect(imageEntry.partial).toBe(true) + expect(typeof imageEntry.push).toBe('string') + expect(imageEntry.opens).toBe(true) + expect(imageEntry.head).toEqual({ id: 's1', createdAt: 1, workspace_id: 'ws', chatId: 'c1' }) + expect(entry.head).toBeUndefined() + expect(entry.whole).toBe(true) + expect(entry.push).toBe(imageEntry.push) + expect(entry.opens).toBeUndefined() + expect(entry.chats.map((c: { id: string }) => c.id)).toEqual(['c1']) + // The record keeps its blob ref; bytes travel as the image object only. + expect(JSON.stringify(entry.chats[0].record)).not.toContain(IMAGE) + expect(entry.artifacts).toBeUndefined() + expect(await pendingDirty()).toEqual([]) + + // Reading the session is not a change the backup keeps. + await putSession({ ...s, lastSeenCount: 2, lastActivityAt: 99 }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + + // A user delete takes the backup with it. + deleteSession('s1') + await flush() + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + expect(pushMock.mock.calls[1][0].requestBody).toEqual({ + owner: EMAIL, + sessions: [], + removed: ['s1'] + }) + expect(pendingKeys()).toEqual([]) + hm.close() + }) + + it('carries a user delete on the sync row when its localStorage mark cannot be written', async () => { + pushMock.mockResolvedValue({ enabled: true, results: [{ id: 'sd' }] }) + const s: Session = { id: 'sd', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + + // Storage full at the moment of the delete. + const setItem = localStorage.setItem.bind(localStorage) + localStorage.setItem = (key: string, value: string) => { + if (key.includes('::r::')) throw new Error('QuotaExceededError') + setItem(key, value) + } + try { + deleteSession('sd') + } finally { + localStorage.setItem = setItem + } + await flush() + expect(removalKeys()).toEqual([]) + await vi.waitFor(async () => + expect((await __syncRowsForTesting(EMAIL)).find((r) => r.id === 'sd')?.removed).toBe(true) + ) + + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + expect(pushMock.mock.calls[1][0].requestBody).toEqual({ + owner: EMAIL, + sessions: [], + removed: ['sd'] + }) + expect((await __syncRowsForTesting(EMAIL)).some((r) => r.id === 'sd')).toBe(false) + await __settleForTesting() + }) + + it('keeps a delete filed on the sync row while the first push is still in flight', async () => { + const s: Session = { id: 'sr', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + let release!: (value: unknown) => void + pushMock.mockImplementationOnce(() => new Promise((r) => (release = r))) + const inFlight = __flushForTesting() + await vi.waitFor(() => expect(pushMock).toHaveBeenCalledTimes(1)) + + // Storage full at the moment of the delete, the push not yet answered. + const setItem = localStorage.setItem.bind(localStorage) + localStorage.setItem = (key: string, value: string) => { + if (key.includes('::r::')) throw new Error('QuotaExceededError') + setItem(key, value) + } + try { + deleteSession('sr') + } finally { + localStorage.setItem = setItem + } + await vi.waitFor(async () => + expect((await __syncRowsForTesting(EMAIL)).find((r) => r.id === 'sr')?.removed).toBe(true) + ) + release({ enabled: true, results: [{ id: 'sr' }] }) + await inFlight + // The push's own row write did not lose the removal. + expect((await __syncRowsForTesting(EMAIL)).find((r) => r.id === 'sr')?.removed).toBe(true) + + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'sr' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + expect(pushMock.mock.calls[1][0].requestBody.removed).toEqual(['sr']) + expect((await __syncRowsForTesting(EMAIL)).some((r) => r.id === 'sr')).toBe(false) + }) + + it('keeps the marks when the push fails, and stops for a workspace without storage', async () => { + const s: Session = { id: 's2', name: 'session-2', createdAt: 1, workspace_id: 'ws' } + const never: Session = { id: 's2b', name: 'session-3', createdAt: 2, workspace_id: 'ws' } + sessionState.sessions = [s, never] + await putSession(s) + + pushMock.mockRejectedValueOnce(new TypeError('network')) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + expect(await pendingDirty()).toEqual(['s2']) + // Still marked: the retry carries it again once the backoff lapses. + __resetMirrorForTesting() + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 's2' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + expect(await pendingDirty()).toEqual([]) + + // The storage goes away: no further request for the page. A delete keeps its + // removal mark only for a session that was backed up (s2), for when backups are on + // again, or it would come back from the bucket; one never backed up has nothing + // there, so its mark goes rather than piling up on a storage-less instance. + await putSession({ ...s, summary: 'changed' }) + await putSession(never) + __setOffRetryForTesting(50) + pushMock.mockResolvedValueOnce({ enabled: false, results: [] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(3) + // The reconcile that runs after login may have re-read the list from the store + // before `never` was in it; deleteSession only acts on sessions it can see. + sessionState.sessions = [s, never] + deleteSession('s2') + deleteSession('s2b') + await flush() + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(3) + expect(pendingKeys()).toEqual(['r::s2::ws']) + + // Backups an admin turns on again elsewhere are noticed once the page's memory of + // them being off expires, with nothing else prompting it: the removal goes then. + listMock.mockResolvedValue({ enabled: true, sessions: [] }) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 's2' }] }) + await vi.waitFor(() => expect(pushMock).toHaveBeenCalledTimes(4), { timeout: 3000 }) + expect(pushMock.mock.calls[3][0].requestBody.removed).toEqual(['s2']) + await vi.waitFor(() => expect(listMock).toHaveBeenCalledTimes(1)) + await __settleForTesting() + expect(pendingKeys()).toEqual([]) + }) + + it('backs up after a reload an edit whose bump only the sync row carries', async () => { + // Storage full from the session's first mark on: the first push carries the mark this + // page kept, and the edit landing while it is in flight goes onto the row it writes. + const s: Session = { id: 'sx', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + const setItem = localStorage.setItem.bind(localStorage) + localStorage.setItem = (key: string, value: string) => { + if (key.includes('::d::')) throw new Error('QuotaExceededError') + setItem(key, value) + } + try { + sessionState.sessions = [s] + await putSession(s) + let release!: (value: unknown) => void + pushMock.mockImplementationOnce(() => new Promise((r) => (release = r))) + const inFlight = __flushForTesting() + await vi.waitFor(() => expect(pushMock).toHaveBeenCalledTimes(1)) + await putSession({ ...s, summary: 'second' }) + release({ enabled: true, results: [{ id: 'sx' }] }) + await inFlight + expect(pendingKeys()).toEqual([]) + const row = (await __syncRowsForTesting(EMAIL)).find((r) => r.id === 'sx') + expect((row?.extraV ?? 0) > (row?.flushedV ?? -1)).toBe(true) + + // A reload, storage still full: nothing in localStorage or memory names the + // session, the row does. + __resetMirrorForTesting() + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'sx' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + expect(pushMock.mock.calls[1][0].requestBody.sessions[0].head.summary).toBe('second') + expect(await pendingDirty()).toEqual([]) + + // The next refused bump counts from what that push retired (the row carries it; + // `pendingDirty` reads localStorage marks only, so the push is the check). + await putSession({ ...s, summary: 'third' }) + await vi.waitFor(async () => + expect((await __syncRowsForTesting(EMAIL)).find((r) => r.id === 'sx')?.extraV).toBe(3) + ) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'sx' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(3) + expect(pushMock.mock.calls[2][0].requestBody.sessions[0].head.summary).toBe('third') + expect(await pendingDirty()).toEqual([]) + } finally { + localStorage.setItem = setItem + } + }) + + it('backs up and restores again at once when the backups are turned on from this page', async () => { + const s: Session = { id: 'so', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'so' }] }) + await __flushForTesting() + await putSession({ ...s, summary: 'changed' }) + pushMock.mockResolvedValueOnce({ enabled: false, results: [] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + await putSession({ ...s, summary: 'changed again' }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + + listMock.mockResolvedValue({ enabled: true, sessions: [] }) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'so' }] }) + usersWorkspaceStore.set({ email: EMAIL, workspaces: [] } as never) + backupSettingsChanged('ws') + await __settleForTesting() + await __flushForTesting() + expect(listMock).toHaveBeenCalledTimes(1) + expect(pushMock).toHaveBeenCalledTimes(3) + // The rows went stale while the backups were off: the session goes whole. + expect(pushMock.mock.calls[2][0].requestBody.sessions[0].whole).toBe(true) + expect(pushMock.mock.calls[2][0].requestBody.sessions[0].head.summary).toBe('changed again') + expect(await pendingDirty()).toEqual([]) + }) + + it('settles nothing of a request that failed, even a session whose parts were still to come', async () => { + // Enough sessions for two requests: the first fails, the second is never sent. + const ids = Array.from({ length: 101 }, (_, i) => `m${i}`) + for (const id of ids) { + await putSession({ id, name: id, createdAt: 1, workspace_id: 'ws' }) + } + pushMock.mockRejectedValueOnce(new TypeError('network')) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + expect(pushMock.mock.calls[0][0].requestBody.sessions).toHaveLength(100) + expect(await pendingDirty()).toHaveLength(101) + + // Once the backoff lapses, every one of them is carried again. + __resetMirrorForTesting() + pushMock.mockImplementation( + async ({ requestBody }: { requestBody: { sessions: { id: string }[] } }) => ({ + enabled: true, + results: requestBody.sessions.map((s) => ({ id: s.id })) + }) + ) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(3) + expect( + pushMock.mock.calls + .slice(1) + .flatMap((c) => c[0].requestBody.sessions.map((s: { id: string }) => s.id)) + .sort() + ).toEqual([...ids].sort()) + expect(await pendingDirty()).toEqual([]) + }) + + it('backs off when the server could not store a session, keeping its mark', async () => { + const s: Session = { id: 's3', name: 'session-3', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + + pushMock.mockResolvedValue({ enabled: true, results: [{ id: 's3', error: 'bucket refused' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + // Still marked, but not re-sent until the backoff lapses. + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + expect(await pendingDirty()).toEqual(['s3']) + }) + + it('fails only the sessions of a request the server found too large', async () => { + const ids = Array.from({ length: 101 }, (_, i) => `t${i}`) + for (const id of ids) { + await putSession({ id, name: id, createdAt: 1, workspace_id: 'ws' }) + } + const { ApiError } = await import('$lib/gen') + pushMock + .mockRejectedValueOnce( + new ApiError({ method: 'POST', url: '' } as never, { status: 413 } as never, 'too large') + ) + .mockImplementation( + async ({ requestBody }: { requestBody: { sessions: { id: string }[] } }) => ({ + enabled: true, + results: requestBody.sessions.map((s) => ({ id: s.id })) + }) + ) + await __flushForTesting() + // The second request still went out and settled its session; the first's stay marked. + expect(pushMock).toHaveBeenCalledTimes(2) + expect(await pendingDirty()).toHaveLength(100) + const settled = pushMock.mock.calls[1][0].requestBody.sessions[0].id + expect(await pendingDirty()).not.toContain(settled) + }) + + it('moves a session whole into its new workspace and files the old copy for removal', async () => { + const s: Session = { id: 'mv', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'mv' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + + // The old workspace's backups go off before the move: its rows go stale, but they + // still say where the copy is. + await putSession({ ...s, summary: 'changed' }) + pushMock.mockResolvedValueOnce({ enabled: false, results: [] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + + await putSession({ ...s, summary: 'changed', workspace_id: 'ws2' }) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'mv' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(3) + expect(pushMock.mock.calls[2][0].workspace).toBe('ws2') + expect(pushMock.mock.calls[2][0].requestBody.sessions[0].head.workspace_id).toBe('ws2') + // The removal waits for the old workspace's backups to be on again. + expect(removalKeys()).toEqual(['r::mv::ws']) + expect(await pendingDirty()).toEqual([]) + }) + + it('pushes every session whole again once the server answers from another storage', async () => { + const a: Session = { + id: 'sa', + name: 'session-1', + createdAt: 1, + workspace_id: 'ws', + chatId: 'ca' + } + const b: Session = { id: 'sb', name: 'session-2', createdAt: 2, workspace_id: 'ws' } + sessionState.sessions = [a, b] + await putSession(a) + await putSession(b) + const hm = new HistoryManager() + await hm.init() + hm.setSessionId('sa') + hm.setCurrentChatId('ca') + await hm.saveChat( + [{ role: 'user', content: 'hello' } as never], + [{ role: 'user', content: 'hello' } as never] + ) + pushMock.mockResolvedValueOnce({ + enabled: true, + storage_id: 'bucket-1', + results: [{ id: 'sa' }, { id: 'sb' }] + }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + + // Only `sa`'s record changes, and the answer names a new storage: it holds `sb` + // nowhere and `sa` only in the part that went, so both are backed up whole again. + await putSession({ ...a, summary: 'changed' }) + pushMock.mockResolvedValueOnce({ + enabled: true, + storage_id: 'bucket-2', + results: [{ id: 'sa' }] + }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + expect(pushMock.mock.calls[1][0].requestBody.sessions[0].chats).toBeUndefined() + expect(await pendingDirty()).toEqual(['sa', 'sb']) + + pushMock.mockResolvedValueOnce({ + enabled: true, + storage_id: 'bucket-2', + results: [{ id: 'sa' }, { id: 'sb' }] + }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(3) + const entries = pushMock.mock.calls[2][0].requestBody.sessions + expect(entries.map((s: { id: string }) => s.id).sort()).toEqual(['sa', 'sb']) + expect(entries.every((s: { head?: unknown }) => s.head !== undefined)).toBe(true) + expect(entries.find((s: { id: string }) => s.id === 'sa').chats).toHaveLength(1) + expect(await pendingDirty()).toEqual([]) + }) + + /** A session with three chats of 1.5 MB (record ~3 MB each): its entry splits past the + * 8 MB request target, so it spans two requests. */ + async function splitSession(id: string): Promise { + const s: Session = { id, name: 'session-1', createdAt: 1, workspace_id: 'ws', chatId: 'c1' } + sessionState.sessions = [s] + await putSession(s) + const hm = new HistoryManager() + await hm.init() + hm.setSessionId(id) + const big = 'x'.repeat(1.5 * 1024 * 1024) + for (const cid of ['c1', 'c2', 'c3']) { + hm.setCurrentChatId(cid) + await hm.saveChat( + [{ role: 'user', content: big } as never], + [{ role: 'user', content: big } as never] + ) + } + } + + it('settles nothing of a session whose parts were answered from different storages', async () => { + await splitSession('sp') + pushMock + .mockResolvedValueOnce({ enabled: true, storage_id: 'bucket-1', results: [{ id: 'sp' }] }) + .mockResolvedValueOnce({ enabled: true, storage_id: 'bucket-2', results: [{ id: 'sp' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + // The first part opens the session whole with its head; only the last completes + // the entry server-side. + const [first, last] = pushMock.mock.calls.map((c) => c[0].requestBody.sessions[0]) + expect(first.partial).toBe(true) + expect(first.whole).toBe(true) + expect(typeof first.push).toBe('string') + expect(first.opens).toBe(true) + expect(first.head).toBeDefined() + expect(last.partial).toBeUndefined() + expect(last.whole).toBe(true) + expect(last.push).toBe(first.push) + expect(last.opens).toBeUndefined() + expect(last.head).toBeUndefined() + expect(await pendingDirty()).toEqual(['sp']) + expect(await __syncRowsForTesting(EMAIL)).toEqual([]) + }) + + it('names an incremental push split over requests on each of its parts', async () => { + await splitSession('si') + pushMock.mockResolvedValue({ enabled: true, results: [{ id: 'si' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + expect(await pendingDirty()).toEqual([]) + // Every chat changes: the update spans two requests again, incremental this time. + const hm = new HistoryManager() + await hm.init() + hm.setSessionId('si') + const big = 'y'.repeat(1.5 * 1024 * 1024) + for (const cid of ['c1', 'c2', 'c3']) { + hm.setCurrentChatId(cid) + await hm.saveChat( + [{ role: 'user', content: big } as never], + [{ role: 'user', content: big } as never] + ) + } + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(4) + const [first, last] = pushMock.mock.calls.slice(2).map((c) => c[0].requestBody.sessions[0]) + expect(first.whole).toBeUndefined() + expect(first.partial).toBe(true) + expect(typeof first.push).toBe('string') + expect(first.opens).toBe(true) + expect(last.partial).toBeUndefined() + expect(last.push).toBe(first.push) + expect(last.opens).toBeUndefined() + expect(await pendingDirty()).toEqual([]) + hm.close() + }) + + it('retires a removal only once the storage holding the backup answered it', async () => { + const s: Session = { id: 'sr2', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + pushMock.mockResolvedValueOnce({ enabled: true, storage_id: 'A', results: [{ id: 'sr2' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + + // The workspace moved to another storage before the delete: the copy in A stays. + deleteSession('sr2') + await flush() + pushMock.mockResolvedValueOnce({ enabled: true, storage_id: 'B', results: [{ id: 'sr2' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + expect(removalKeys()).toEqual(['r::sr2::ws']) + expect((await __syncRowsForTesting(EMAIL)).some((r) => r.id === 'sr2')).toBe(true) + + // Back on A, the removal lands. + pushMock.mockResolvedValueOnce({ enabled: true, storage_id: 'A', results: [{ id: 'sr2' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(3) + expect(pushMock.mock.calls[2][0].requestBody.removed).toEqual(['sr2']) + expect(removalKeys()).toEqual([]) + expect((await __syncRowsForTesting(EMAIL)).some((r) => r.id === 'sr2')).toBe(false) + }) + + it('removes a deleted session from every storage that holds a copy of it', async () => { + const s: Session = { id: 'sr3', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + pushMock.mockResolvedValueOnce({ enabled: true, storage_id: 'A', results: [{ id: 'sr3' }] }) + await __flushForTesting() + // The workspace moves to B: the row goes stale, the session goes whole to B and + // settles there, and the row remembers the copy A keeps. + await putSession({ ...s, summary: 'changed' }) + pushMock.mockResolvedValue({ enabled: true, storage_id: 'B', results: [{ id: 'sr3' }] }) + await __flushForTesting() + await __flushForTesting() + expect(await pendingDirty()).toEqual([]) + const row = (await __syncRowsForTesting(EMAIL)).find((r) => r.id === 'sr3') + expect(row?.storageId).toBe('B') + expect(row?.alsoIn).toEqual(['A']) + + // Deleted while on B: B's copy goes, and the mark waits for A to answer. + deleteSession('sr3') + await flush() + await __flushForTesting() + expect(removalKeys()).toEqual(['r::sr3::ws']) + const narrowed = (await __syncRowsForTesting(EMAIL)).find((r) => r.id === 'sr3') + expect(narrowed?.storageId).toBe('A') + expect(narrowed?.alsoIn).toBeUndefined() + + // Back on A, the copy there goes too, and only then is the removal done. + pushMock.mockResolvedValue({ enabled: true, storage_id: 'A', results: [{ id: 'sr3' }] }) + await __flushForTesting() + expect(pushMock.mock.lastCall?.[0].requestBody.removed).toEqual(['sr3']) + expect(removalKeys()).toEqual([]) + expect((await __syncRowsForTesting(EMAIL)).some((r) => r.id === 'sr3')).toBe(false) + }) + + it("removes a moved session's old copy from the storage that held it, whatever its old workspace is on now", async () => { + const s: Session = { id: 'mv2', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + pushMock.mockResolvedValueOnce({ enabled: true, storage_id: 'A', results: [{ id: 'mv2' }] }) + await __flushForTesting() + // The old workspace's backups go off, then the session moves: the new workspace's + // row replaces the old one, so the mark carries where the old copy is. + await putSession({ ...s, summary: 'changed' }) + pushMock.mockResolvedValueOnce({ enabled: false, results: [] }) + await __flushForTesting() + await putSession({ ...s, summary: 'changed', workspace_id: 'ws2' }) + pushMock.mockResolvedValueOnce({ enabled: true, storage_id: 'B', results: [{ id: 'mv2' }] }) + await __flushForTesting() + expect(removalKeys()).toEqual(['r::mv2::ws']) + + // On the next page load the old workspace is on again, but on another storage: the + // removal there deletes nothing, and the mark waits for the storage holding the copy. + __resetMirrorForTesting() + pushMock.mockResolvedValueOnce({ enabled: true, storage_id: 'C', results: [{ id: 'mv2' }] }) + await __flushForTesting() + expect(pushMock.mock.lastCall?.[0].requestBody.removed).toEqual(['mv2']) + expect(removalKeys()).toEqual(['r::mv2::ws']) + pushMock.mockResolvedValueOnce({ enabled: true, storage_id: 'A', results: [{ id: 'mv2' }] }) + await __flushForTesting() + expect(pushMock.mock.lastCall?.[0].requestBody.removed).toEqual(['mv2']) + expect(removalKeys()).toEqual([]) + expect((await __syncRowsForTesting(EMAIL)).find((r) => r.id === 'mv2')?.ws).toBe('ws2') + }) + + it('takes a bumped backup generation as a new storage for the rows, not for a removal', async () => { + const a: Session = { id: 'ga', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + const b: Session = { id: 'gb', name: 'session-2', createdAt: 2, workspace_id: 'ws' } + sessionState.sessions = [a, b] + await putSession(a) + await putSession(b) + pushMock.mockResolvedValueOnce({ + enabled: true, + storage_id: 'A', + backup_generation: 0, + results: [{ id: 'ga' }, { id: 'gb' }] + }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + + // The key was rotated before the delete: the removal is done (the old generation is + // gone with the rotation), and `gb` goes whole again under the new one. + deleteSession('ga') + await flush() + pushMock.mockResolvedValueOnce({ + enabled: true, + storage_id: 'A', + backup_generation: 1, + results: [{ id: 'ga' }] + }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + expect(pushMock.mock.calls[1][0].requestBody.removed).toEqual(['ga']) + expect(removalKeys()).toEqual([]) + expect(await pendingDirty()).toEqual(['gb']) + }) + + it('sends a session whole again once the server says its head is gone', async () => { + const s: Session = { + id: 'sh', + name: 'session-1', + createdAt: 1, + workspace_id: 'ws', + chatId: 'c1' + } + sessionState.sessions = [s] + await putSession(s) + const hm = new HistoryManager() + await hm.init() + hm.setSessionId('sh') + hm.setCurrentChatId('c1') + await hm.saveChat( + [{ role: 'user', content: 'a' } as never], + [{ role: 'user', content: 'a' } as never] + ) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'sh' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + + // Another device removed the backup; the next chat-only push finds no head there. + hm.setCurrentChatId('c2') + await hm.saveChat( + [{ role: 'user', content: 'ab' } as never], + [{ role: 'user', content: 'ab' } as never] + ) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'sh', needs_whole: true }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + expect(pushMock.mock.calls[1][0].requestBody.sessions[0].whole).toBeUndefined() + expect(await pendingDirty()).toEqual(['sh']) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'sh' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(3) + const whole = pushMock.mock.calls[2][0].requestBody.sessions[0] + expect(whole.whole).toBe(true) + expect(whole.push).toBeUndefined() + expect(whole.head).toBeDefined() + expect(whole.chats.map((c: { id: string }) => c.id).sort()).toEqual(['c1', 'c2']) + expect(await pendingDirty()).toEqual([]) + hm.close() + }) + + it('holds the rest of a session back once the server refused a part of it', async () => { + await splitSession('sf') + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'sf', error: 'boom' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + expect(await pendingDirty()).toEqual(['sf']) + expect(await __syncRowsForTesting(EMAIL)).toEqual([]) + }) + + it('keeps an edit whose dirty mark cannot be written while the session has no row yet', async () => { + const s: Session = { id: 'sn', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + // The first push is held open; storage is full for the edit that lands meanwhile. + let release!: (value: unknown) => void + pushMock.mockImplementationOnce(() => new Promise((r) => (release = r))) + const inFlight = __flushForTesting() + await vi.waitFor(() => expect(pushMock).toHaveBeenCalledTimes(1)) + const setItem = localStorage.setItem.bind(localStorage) + localStorage.setItem = (key: string, value: string) => { + if (key.includes('::d::')) throw new Error('QuotaExceededError') + setItem(key, value) + } + try { + await putSession({ ...s, summary: 'second' }) + } finally { + localStorage.setItem = setItem + } + release({ enabled: true, results: [{ id: 'sn' }] }) + await inFlight + // The row the push wrote took the bump over, so the edit is still pending. + expect((await __syncRowsForTesting(EMAIL)).find((r) => r.id === 'sn')?.extraV).toBe(1) + expect(await pendingDirty()).toEqual(['sn']) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'sn' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(2) + expect(pushMock.mock.calls[1][0].requestBody.sessions[0].head.summary).toBe('second') + expect(await pendingDirty()).toEqual([]) + }) + + it('carries an edit on the sync row when its dirty mark cannot be written, even during a push', async () => { + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'sm' }] }) + const s: Session = { id: 'sm', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + + // An edit's push is held open; storage is full for the edit that follows. + let release!: (value: unknown) => void + pushMock.mockImplementationOnce(() => new Promise((r) => (release = r))) + await putSession({ ...s, summary: 'first' }) + const inFlight = __flushForTesting() + await vi.waitFor(() => expect(pushMock).toHaveBeenCalledTimes(2)) + const setItem = localStorage.setItem.bind(localStorage) + localStorage.setItem = (key: string, value: string) => { + if (key.includes('::d::')) throw new Error('QuotaExceededError') + setItem(key, value) + } + try { + await putSession({ ...s, summary: 'second' }) + } finally { + localStorage.setItem = setItem + } + await vi.waitFor(async () => + expect((await __syncRowsForTesting(EMAIL)).find((r) => r.id === 'sm')?.extraV).toBe(1) + ) + release({ enabled: true, results: [{ id: 'sm' }] }) + await inFlight + // The push's own row write kept the bump, so the later edit is still pending. + expect(await pendingDirty()).toEqual(['sm']) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'sm' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(3) + expect(pushMock.mock.calls[2][0].requestBody.sessions[0].head.summary).toBe('second') + expect(await pendingDirty()).toEqual([]) + }) + + it('files a mark under the user whose store the write landed in', async () => { + markSessionDirty('sw', undefined, 'other@x.com') + expect(localStorage.getItem('windmill_sessions_mirror_pending::other@x.com::d::sw')).toBe('1') + expect(pendingKeys()).toEqual([]) + + // The mark of another user that cannot be written goes to that user's own rows. + const { openDB } = await import('idb') + const theirs = await openDB('windmill-sessions-mirror::other@x.com', 1, { + upgrade: (db) => db.createObjectStore('sync', { keyPath: 'id' }) + }) + await theirs.put('sync', { id: 'so', ws: 'ws', head: 'h', chats: {}, images: {}, flushedV: 1 }) + theirs.close() + const setItem = localStorage.setItem.bind(localStorage) + localStorage.setItem = (key: string, value: string) => { + if (key.includes('::d::')) throw new Error('QuotaExceededError') + setItem(key, value) + } + try { + markSessionDirty('so', undefined, 'other@x.com') + } finally { + localStorage.setItem = setItem + } + await vi.waitFor(async () => { + const db = await openDB('windmill-sessions-mirror::other@x.com', 1) + try { + expect((await db.get('sync', 'so'))?.extraV).toBe(1) + } finally { + db.close() + } + }) + }) + + it('keeps the marks when the server refuses a request, for the next page load', async () => { + const s: Session = { id: 's4', name: 'session-4', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + + const { ApiError } = await import('$lib/gen') + pushMock.mockRejectedValue( + new ApiError({ method: 'POST', url: '' } as never, { status: 400 } as never, 'quota') + ) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + // Nothing more for this page, marks untouched by the follow-up flush. + await putSession({ ...s, summary: 'changed' }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + expect(await pendingDirty()).toEqual(['s4']) + }) +}) + +describe('sessionMirror restore', () => { + const backup = { + id: 's9', + head: { id: 's9', workspace_id: 'ws', createdAt: 5, chatId: 'c9', summary: 'remote' }, + chats: [ + { + id: 'c9', + record: { + id: 'c9', + sessionId: 's9', + title: 't', + lastModified: 7, + actualMessages: [], + displayMessages: [{ role: 'user', content: 'hi' }] + } + } + ], + images: [] + } + + it('brings back a session the browser lacks, and records nothing for one whose chats could not be written', async () => { + listMock.mockResolvedValue({ + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z' }] + }) + pullMock.mockResolvedValue({ enabled: true, sessions: [backup], deferred: [] }) + usersWorkspaceStore.set({ email: EMAIL, workspaces: [] } as never) + + // A removal pending for another workspace (the session moved here from it) does + // not stand in the way of restoring this workspace's copy. + localStorage.setItem(`${PENDING_PREFIX}r::s9::elsewhere`, '1') + chatImport.unavailable = true + restoreSessionBackups('ws') + await __settleForTesting() + expect(pullMock).toHaveBeenCalledTimes(1) + expect(sessionState.sessions.map((s) => s.id)).toEqual([]) + // Not recorded as restored: the next restore tries again, and no flush can push a + // transcript-less copy over the backup. + __resetMirrorForTesting() + chatImport.unavailable = false + restoreSessionBackups('ws') + await __settleForTesting() + expect(pullMock).toHaveBeenCalledTimes(2) + await vi.waitFor(() => expect(sessionState.sessions.map((s) => s.id)).toEqual(['s9'])) + const restored = sessionState.sessions[0] + expect(restored.name).toBe('session-1') + expect(restored.summary).toBe('remote') + expect(restored.lastSeenCount).toBe(1) + expect((await readStoredChat('c9', EMAIL))?.displayMessages).toHaveLength(1) + + // Restored state is what the backup holds: nothing to push (the other workspace's + // removal is its own request, not part of this check). + localStorage.removeItem(`${PENDING_PREFIX}r::s9::elsewhere`) + await __flushForTesting() + expect(pushMock).not.toHaveBeenCalled() + }) + + it('imports a session that came in pages only once the last page arrived', async () => { + listMock.mockResolvedValue({ + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z' }] + }) + const cursor = { id: 's9', images: false, after: 'sessions/s9/chats/c9.json' } + const c9b = { ...backup.chats[0], id: 'c9b', record: { ...backup.chats[0].record, id: 'c9b' } } + pullMock + .mockResolvedValueOnce({ + enabled: true, + sessions: [{ ...backup, next: cursor }], + deferred: [] + }) + .mockResolvedValueOnce({ + enabled: true, + sessions: [{ ...backup, chats: [c9b] }], + deferred: [] + }) + usersWorkspaceStore.set({ email: EMAIL, workspaces: [] } as never) + restoreSessionBackups('ws') + await __settleForTesting() + expect(pullMock).toHaveBeenCalledTimes(2) + expect(pullMock.mock.calls[1][0].requestBody).toEqual({ ids: ['s9'], resume: cursor }) + await vi.waitFor(() => expect(sessionState.sessions.map((s) => s.id)).toEqual(['s9'])) + expect((await readStoredChat('c9', EMAIL))?.displayMessages).toHaveLength(1) + expect((await readStoredChat('c9b', EMAIL))?.id).toBe('c9b') + }) + + it('starts a session over when its only page was read while the backup moved', async () => { + listMock.mockResolvedValue({ + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z' }] + }) + const c9b = { ...backup.chats[0], id: 'c9b', record: { ...backup.chats[0].record, id: 'c9b' } } + pullMock + .mockResolvedValueOnce({ + enabled: true, + sessions: [{ ...backup, chats: [backup.chats[0]], moved: true }], + deferred: [] + }) + .mockResolvedValueOnce({ + enabled: true, + sessions: [{ ...backup, chats: [backup.chats[0], c9b] }], + deferred: [] + }) + usersWorkspaceStore.set({ email: EMAIL, workspaces: [] } as never) + restoreSessionBackups('ws') + await __settleForTesting() + expect(pullMock).toHaveBeenCalledTimes(2) + await vi.waitFor(() => expect(sessionState.sessions.map((s) => s.id)).toEqual(['s9'])) + expect((await readStoredChat('c9b', EMAIL))?.id).toBe('c9b') + }) + + it('brings a session two workspaces of the family list back from the copy that moved last', async () => { + // Both copies carry the same modification time (a store reports them coarsely): the + // move count tells them apart. + listMock.mockImplementation(async ({ workspace }: { workspace: string }) => ({ + enabled: true, + sessions: [ + { id: 's9', updated_at: '2026-09-14T00:00:00Z', epoch: workspace === 'ws' ? 0 : 1 } + ] + })) + pullMock.mockImplementation(async ({ workspace }: { workspace: string }) => ({ + enabled: true, + sessions: [{ ...backup, head: { ...backup.head, workspace_id: workspace, moves: 1 } }], + deferred: [] + })) + usersWorkspaceStore.set({ + email: EMAIL, + workspaces: [ + { id: 'ws', name: 'ws', username: 'u' }, + { id: 'ws2', name: 'ws2', username: 'u', parent_workspace_id: 'ws' } + ] + } as never) + restoreSessionBackups('ws') + await __settleForTesting() + // The stale copy in the old workspace is left alone, and does not take the id first. + expect(pullMock.mock.calls.map((c) => c[0].workspace)).toEqual(['ws2']) + await vi.waitFor(() => expect(sessionState.sessions.map((s) => s.id)).toEqual(['s9'])) + expect(sessionState.sessions[0].workspace_id).toBe('ws2') + }) + + it('leaves a session whose later copy showed up elsewhere after the listings, and tries again', async () => { + // The move lands in the other workspace between the family's listings and the pull: + // the listing taken again before the record lands shows it, and the copy about to be + // imported is the stale one. + let ws2Listings = 0 + listMock.mockImplementation(async ({ workspace }: { workspace: string }) => { + if (workspace === 'ws') { + return { + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z', epoch: 0 }] + } + } + ws2Listings += 1 + return { + enabled: true, + sessions: + ws2Listings === 1 ? [] : [{ id: 's9', updated_at: '2026-09-14T00:00:00Z', epoch: 1 }] + } + }) + pullMock.mockImplementation(async ({ workspace }: { workspace: string }) => ({ + enabled: true, + sessions: [{ ...backup, head: { ...backup.head, workspace_id: workspace, moves: 1 } }], + deferred: [] + })) + usersWorkspaceStore.set({ + email: EMAIL, + workspaces: [ + { id: 'ws', name: 'ws', username: 'u' }, + { id: 'ws2', name: 'ws2', username: 'u', parent_workspace_id: 'ws' } + ] + } as never) + restoreSessionBackups('ws') + await __settleForTesting() + expect(pullMock.mock.calls.map((c) => c[0].workspace)).toEqual(['ws']) + expect(sessionState.sessions).toEqual([]) + restoreSessionBackups('ws') + await __settleForTesting() + await vi.waitFor(() => expect(sessionState.sessions.map((s) => s.id)).toEqual(['s9'])) + expect(sessionState.sessions[0].workspace_id).toBe('ws2') + }) + + it('checks a family member whose backups were off when the restore started', async () => { + // The other workspace comes on (and gets the moved session) between the family's + // listings and the pull: the listing taken again before the records land covers it. + let ws2Listings = 0 + listMock.mockImplementation(async ({ workspace }: { workspace: string }) => { + if (workspace === 'ws') { + return { + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z', epoch: 0 }] + } + } + ws2Listings += 1 + return ws2Listings === 1 + ? { enabled: false, sessions: [] } + : { + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z', epoch: 1 }] + } + }) + pullMock.mockImplementation(async ({ workspace }: { workspace: string }) => ({ + enabled: true, + sessions: [{ ...backup, head: { ...backup.head, workspace_id: workspace, moves: 1 } }], + deferred: [] + })) + usersWorkspaceStore.set({ + email: EMAIL, + workspaces: [ + { id: 'ws', name: 'ws', username: 'u' }, + { id: 'ws2', name: 'ws2', username: 'u', parent_workspace_id: 'ws' } + ] + } as never) + restoreSessionBackups('ws') + await __settleForTesting() + expect(sessionState.sessions).toEqual([]) + restoreSessionBackups('ws') + await __settleForTesting() + await vi.waitFor(() => expect(sessionState.sessions.map((s) => s.id)).toEqual(['s9'])) + expect(sessionState.sessions[0].workspace_id).toBe('ws2') + }) + + it('marks nothing at load for a session a restore brought back', async () => { + listMock.mockResolvedValue({ + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z', epoch: 0 }] + }) + pullMock.mockResolvedValue({ enabled: true, sessions: [backup], deferred: [] }) + usersWorkspaceStore.set({ email: EMAIL, workspaces: [] } as never) + restoreSessionBackups('ws') + await __settleForTesting() + await vi.waitFor(() => expect(sessionState.sessions.map((s) => s.id)).toEqual(['s9'])) + // The next load's backfill leaves the clean row alone: no mark, no read, no push. + __resetMirrorForTesting() + await __flushForTesting() + expect(pushMock).not.toHaveBeenCalled() + expect(pendingKeys()).toEqual([]) + }) + + it('lists a family of one once, whatever it restores', async () => { + listMock.mockResolvedValue({ + enabled: true, + sessions: [ + { id: 's9', updated_at: '2026-09-14T00:00:00Z', epoch: 0 }, + { id: 's8', updated_at: '2026-09-13T00:00:00Z', epoch: 0 } + ] + }) + pullMock.mockImplementation(async ({ requestBody }: { requestBody: { ids: string[] } }) => ({ + enabled: true, + sessions: requestBody.ids.map((id) => ({ + ...backup, + id, + head: { ...backup.head, id }, + chats: backup.chats.map((c) => ({ + ...c, + id: `${c.id}-${id}`, + record: { ...c.record, id: `${c.id}-${id}`, sessionId: id } + })) + })), + deferred: [] + })) + usersWorkspaceStore.set({ email: EMAIL, workspaces: [] } as never) + restoreSessionBackups('ws') + await __settleForTesting() + await vi.waitFor(() => + expect(sessionState.sessions.map((s) => s.id).sort()).toEqual(['s8', 's9']) + ) + expect(listMock).toHaveBeenCalledTimes(1) + }) + + it('restores nothing of a family one of whose workspaces could not be listed, and tries again', async () => { + listMock + .mockImplementationOnce(async () => ({ + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z', epoch: 0 }] + })) + .mockRejectedValueOnce(new Error('offline')) + .mockImplementation(async ({ workspace }: { workspace: string }) => ({ + enabled: true, + sessions: [ + { id: 's9', updated_at: '2026-09-14T00:00:00Z', epoch: workspace === 'ws' ? 0 : 1 } + ] + })) + pullMock.mockImplementation(async ({ workspace }: { workspace: string }) => ({ + enabled: true, + sessions: [{ ...backup, head: { ...backup.head, workspace_id: workspace, moves: 1 } }], + deferred: [] + })) + usersWorkspaceStore.set({ + email: EMAIL, + workspaces: [ + { id: 'ws', name: 'ws', username: 'u' }, + { id: 'ws2', name: 'ws2', username: 'u', parent_workspace_id: 'ws' } + ] + } as never) + restoreSessionBackups('ws') + await __settleForTesting() + // The copy that listed could be the stale one: nothing is imported this time. + expect(pullMock).not.toHaveBeenCalled() + expect(sessionState.sessions).toEqual([]) + restoreSessionBackups('ws') + await __settleForTesting() + expect(pullMock.mock.calls.map((c) => c[0].workspace)).toEqual(['ws2']) + await vi.waitFor(() => expect(sessionState.sessions.map((s) => s.id)).toEqual(['s9'])) + }) + + it('starts a session over when its backup moved between two pages', async () => { + listMock.mockResolvedValue({ + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z' }] + }) + const cursor = { id: 's9', images: false, after: 'sessions/s9/chats/c9.json' } + const c9b = { ...backup.chats[0], id: 'c9b', record: { ...backup.chats[0].record, id: 'c9b' } } + const c9a = { ...backup.chats[0], id: 'c9a', record: { ...backup.chats[0].record, id: 'c9a' } } + pullMock + // The first attempt: a chat sorting before the cursor lands between the pages. + .mockResolvedValueOnce({ + enabled: true, + sessions: [{ ...backup, next: cursor, listing: 'L1' }], + deferred: [] + }) + .mockResolvedValueOnce({ + enabled: true, + sessions: [{ ...backup, chats: [c9b], listing: 'L2' }], + deferred: [] + }) + // The second attempt sees the whole of it. + .mockResolvedValueOnce({ + enabled: true, + sessions: [{ ...backup, chats: [c9a, backup.chats[0]], next: cursor, listing: 'L2' }], + deferred: [] + }) + .mockResolvedValueOnce({ + enabled: true, + sessions: [{ ...backup, chats: [c9b], listing: 'L2' }], + deferred: [] + }) + usersWorkspaceStore.set({ email: EMAIL, workspaces: [] } as never) + restoreSessionBackups('ws') + await __settleForTesting() + expect(pullMock).toHaveBeenCalledTimes(4) + expect(pullMock.mock.calls[2][0].requestBody).toEqual({ ids: ['s9'], resume: undefined }) + await vi.waitFor(() => expect(sessionState.sessions.map((s) => s.id)).toEqual(['s9'])) + expect((await readStoredChat('c9a', EMAIL))?.id).toBe('c9a') + expect((await readStoredChat('c9b', EMAIL))?.id).toBe('c9b') + }) + + it('restores nothing without Web Locks, and still backs up', async () => { + setWebLocks(undefined) + listMock.mockResolvedValue({ + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z' }] + }) + usersWorkspaceStore.set({ email: EMAIL, workspaces: [] } as never) + restoreSessionBackups('ws') + await __settleForTesting() + expect(listMock).not.toHaveBeenCalled() + expect(sessionState.sessions).toEqual([]) + const s: Session = { id: 'sl', name: 'session-1', createdAt: 1, workspace_id: 'ws' } + sessionState.sessions = [s] + await putSession(s) + pushMock.mockResolvedValueOnce({ enabled: true, results: [{ id: 'sl' }] }) + await __flushForTesting() + expect(pushMock).toHaveBeenCalledTimes(1) + }) + + it('leaves a session for the next restore when what an earlier one staged cannot be pruned', async () => { + listMock.mockResolvedValue({ + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z' }] + }) + // An earlier restore was cut short after staging a chat the backup has since dropped. + const { importStoredChats } = await import('../copilot/chat/HistoryManager.svelte') + await importStoredChats([{ ...backup.chats[0].record, id: 'cx' } as never], [], EMAIL, true) + const staging = { + id: 's9', + ws: 'ws', + head: '', + chats: {}, + images: {}, + staging: { chats: ['cx'], images: [], items: [], versions: [] } + } + await __writeSyncForTesting([staging], EMAIL) + pullMock.mockResolvedValue({ enabled: true, sessions: [backup], deferred: [] }) + chatImport.pruneFails = true + usersWorkspaceStore.set({ email: EMAIL, workspaces: [] } as never) + restoreSessionBackups('ws') + await __settleForTesting() + // No record, so no flush can push the stale chat back, and the staging row stays, + // now naming what this page wrote too. + expect(sessionState.sessions).toEqual([]) + expect((await readStoredChat('cx', EMAIL))?.id).toBe('cx') + expect( + (await __syncRowsForTesting(EMAIL)).find((r) => r.id === 's9')?.staging?.chats?.sort() + ).toEqual(['c9', 'cx']) + + // The restore after prunes and brings the session back. + chatImport.pruneFails = false + __resetMirrorForTesting() + restoreSessionBackups('ws') + await __settleForTesting() + await vi.waitFor(() => expect(sessionState.sessions.map((s) => s.id)).toEqual(['s9'])) + expect(await readStoredChat('cx', EMAIL)).toBeUndefined() + expect((await __syncRowsForTesting(EMAIL)).find((r) => r.id === 's9')?.staging).toBeUndefined() + }) + + it('never writes an older record over a newer one, and prunes only what it is told', async () => { + const { importStoredChats, pruneSessionChats } = await import( + '../copilot/chat/HistoryManager.svelte' + ) + const chat = (lastModified: number, title: string) => + ({ ...backup.chats[0].record, lastModified, title }) as never + await importStoredChats([chat(20, 'newer')], [], EMAIL, true) + await importStoredChats([chat(10, 'older')], [], EMAIL, true) + expect((await readStoredChat('c9', EMAIL))?.title).toBe('newer') + await importStoredChats( + [{ ...backup.chats[0].record, id: 'c9b', lastModified: 30 } as never], + [], + EMAIL, + true + ) + // A prune names what goes; nothing else of the session is touched. + await pruneSessionChats('s9', new Set(['c9']), new Set(), EMAIL) + expect(await readStoredChat('c9', EMAIL)).toBeUndefined() + expect((await readStoredChat('c9b', EMAIL))?.id).toBe('c9b') + }) + + it('keeps an image whose chat came on an earlier page, and restages after a page failed', async () => { + listMock.mockResolvedValue({ + enabled: true, + sessions: [{ id: 's9', updated_at: '2026-09-14T00:00:00Z' }] + }) + const cursor = { id: 's9', images: true, after: '' } + const imagePage = { + ...backup, + chats: [], + images: [{ chat_id: 'c9', id: 'i9', data_url: IMAGE }] + } + // The first restore is cut short after staging the chats and an artifact with two + // versions. + const gone = { ...backup.chats[0], id: 'cx', record: { ...backup.chats[0].record, id: 'cx' } } + const item = { id: 'a1', sessionId: 's9', kind: 'markdown', name: 'a', content: 'x' } + const version = (n: number) => ({ + key: `a1:${n}`, + artifactId: 'a1', + version: n, + name: 'a', + content: 'x', + savedAt: n + }) + pullMock + .mockResolvedValueOnce({ + enabled: true, + sessions: [ + { + ...backup, + chats: [...backup.chats, gone], + artifacts: { items: [item], versions: [version(1), version(2)] }, + next: cursor + } + ], + deferred: [] + }) + .mockRejectedValueOnce(new Error('offline')) + usersWorkspaceStore.set({ email: EMAIL, workspaces: [] } as never) + restoreSessionBackups('ws') + await __settleForTesting() + expect(pullMock).toHaveBeenCalledTimes(2) + expect(sessionState.sessions.map((s) => s.id)).toEqual([]) + expect((await readStoredChat('c9', EMAIL))?.title).toBe('t') + expect((await readStoredChat('cx', EMAIL))?.id).toBe('cx') + + // The backup moved on meanwhile: the retry takes the newer chat over the staged one + // and drops the chat the backup no longer has. + __resetMirrorForTesting() + const newer = { + ...backup, + chats: [{ ...backup.chats[0], record: { ...backup.chats[0].record, title: 'newer' } }], + artifacts: { items: [item], versions: [version(1)] } + } + pullMock + .mockResolvedValueOnce({ + enabled: true, + sessions: [{ ...newer, next: cursor }], + deferred: [] + }) + .mockResolvedValueOnce({ enabled: true, sessions: [imagePage], deferred: [] }) + restoreSessionBackups('ws') + await __settleForTesting() + await vi.waitFor(() => expect(sessionState.sessions.map((s) => s.id)).toEqual(['s9'])) + expect((await readStoredChat('c9', EMAIL))?.title).toBe('newer') + expect(await readStoredChat('cx', EMAIL)).toBeUndefined() + const { readImageDataUrl } = await import('../copilot/chat/HistoryManager.svelte') + expect(await readImageDataUrl('i9', EMAIL)).toBe(IMAGE) + // The version the backup no longer has went with the chat it no longer has. + const { readSessionArtifacts } = await import('../copilot/chat/artifacts/artifactsDB') + const artifacts = await readSessionArtifacts('s9', EMAIL) + expect(artifacts?.items.map((i) => i.id)).toEqual(['a1']) + expect(artifacts?.versions.map((v) => v.key)).toEqual(['a1:1']) + }) +}) diff --git a/frontend/src/lib/components/sessions/sessionMirrorPlan.test.ts b/frontend/src/lib/components/sessions/sessionMirrorPlan.test.ts new file mode 100644 index 0000000000..aa3c8bf735 --- /dev/null +++ b/frontend/src/lib/components/sessions/sessionMirrorPlan.test.ts @@ -0,0 +1,204 @@ +import { describe, expect, it } from 'vitest' +import { + artifactsFingerprint, + headSig, + jsonBytes, + planSessionPush, + splitEntry, + type ChatSnapshot, + type MirrorSyncState +} from './sessionMirrorPlan' +import type { Session } from './sessionState.svelte' + +function session(over: Partial = {}): Session { + return { id: 's1', name: 'session-1', createdAt: 1, workspace_id: 'ws', chatId: 'c1', ...over } +} + +function chat(id: string, lastModified: number, imageIds: string[] = []): ChatSnapshot { + return { id, lastModified, record: { id, lastModified }, imageIds } +} + +const noArtifacts = { items: [], versions: [] } + +function synced(over: Partial = {}): MirrorSyncState { + return { + id: 's1', + ws: 'ws', + head: headSig(session()), + chats: { c1: 10 }, + images: { i1: 'c1' }, + artifacts: artifactsFingerprint(noArtifacts), + ...over + } +} + +describe('planSessionPush', () => { + it('pushes everything for a session never backed up, and nothing for an unsent draft', () => { + const plan = planSessionPush({ + session: session(), + chats: [chat('c1', 10, ['i1'])], + artifacts: noArtifacts + }) + expect(plan?.entry?.head?.id).toBe('s1') + expect(plan?.entry?.chats?.map((c) => c.id)).toEqual(['c1']) + expect(plan?.images).toEqual([{ chat_id: 'c1', id: 'i1' }]) + // Nothing to store yet, so no artifacts object either. + expect(plan?.entry?.artifacts).toBeUndefined() + expect(plan?.next).toEqual(synced()) + + expect( + planSessionPush({ + session: session({ workspace_id: undefined, pending_workspace_id: 'ws' }), + chats: [], + artifacts: noArtifacts + }) + ).toBeUndefined() + }) + + it('sends nothing when only the fields reading a session bumps changed', () => { + const plan = planSessionPush({ + session: session({ lastSeenCount: 7, lastActivityAt: 99, name: 'session-9' }), + chats: [chat('c1', 10, ['i1'])], + artifacts: noArtifacts, + sync: synced() + }) + expect(plan?.entry).toBeUndefined() + expect(plan?.images).toEqual([]) + }) + + it('carries only the chat whose lastModified moved, plus its new images', () => { + const plan = planSessionPush({ + session: session(), + chats: [chat('c1', 10, ['i1']), chat('c2', 20, ['i2'])], + artifacts: noArtifacts, + sync: synced() + }) + expect(plan?.entry?.head).toBeUndefined() + expect(plan?.entry?.chats?.map((c) => c.id)).toEqual(['c2']) + expect(plan?.images).toEqual([{ chat_id: 'c2', id: 'i2' }]) + expect(plan?.next.chats).toEqual({ c1: 10, c2: 20 }) + expect(plan?.next.images).toEqual({ i1: 'c1', i2: 'c2' }) + }) + + it('deletes the copy of a chat that grew too large to back up, instead of keeping a stale one', () => { + const plan = planSessionPush({ + session: session(), + chats: [{ id: 'c1', lastModified: 11, imageIds: ['i1'], omitted: true }], + artifacts: noArtifacts, + sync: synced() + }) + expect(plan?.entry?.delete_chats).toEqual(['c1']) + expect(plan?.entry?.chats).toBeUndefined() + expect(plan?.images).toEqual([]) + expect(plan?.next.chats).toEqual({}) + expect(plan?.next.images).toEqual({}) + }) + + it('carries deletes past the per-entry cap over to the next push', () => { + const prevChats = Object.fromEntries(Array.from({ length: 1005 }, (_, i) => [`c${i}`, 10])) + const plan = planSessionPush({ + session: session(), + chats: [], + artifacts: noArtifacts, + sync: synced({ chats: prevChats, images: {} }) + }) + expect(plan?.entry?.delete_chats).toHaveLength(1000) + // Still listed as pushed, so the next plan finds them gone again, and the session + // stays marked for that plan. + expect(Object.keys(plan?.next.chats ?? {})).toHaveLength(5) + expect(plan?.carried).toBe(true) + }) + + it('deletes a chat that is gone and an image its chat evicted', () => { + const plan = planSessionPush({ + session: session(), + chats: [chat('c1', 11, [])], + artifacts: noArtifacts, + sync: synced({ chats: { c1: 10, c2: 20 }, images: { i1: 'c1', i2: 'c2' } }) + }) + expect(plan?.entry?.delete_chats).toEqual(['c2']) + // i2 goes with c2 server-side; only c1's evicted image is deleted on its own. + expect(plan?.entry?.delete_images).toEqual([{ chat_id: 'c1', id: 'i1' }]) + }) + + it('moves a session as a full push to the new workspace and a removal from the old', () => { + const plan = planSessionPush({ + session: session({ workspace_id: 'ws2' }), + chats: [chat('c1', 10, ['i1'])], + artifacts: noArtifacts, + sync: synced() + }) + expect(plan?.workspaceId).toBe('ws2') + expect(plan?.removeFrom).toBe('ws') + expect(plan?.entry?.head?.workspace_id).toBe('ws2') + expect(plan?.entry?.chats?.map((c) => c.id)).toEqual(['c1']) + expect(plan?.images).toEqual([{ chat_id: 'c1', id: 'i1' }]) + expect(plan?.next.ws).toBe('ws2') + }) + + it('pushes artifacts when their fingerprint changes, including emptying them', () => { + const items = [ + { + id: 'a1', + sessionId: 's1', + kind: 'md' as const, + name: 'notes', + content: 'x', + createdAt: 1, + updatedAt: 2, + version: 1 + } + ] + const withArtifact = planSessionPush({ + session: session(), + chats: [chat('c1', 10, ['i1'])], + artifacts: { items, versions: [] }, + sync: synced() + }) + expect(withArtifact?.entry?.artifacts).toEqual({ items, versions: [] }) + + const emptied = planSessionPush({ + session: session(), + chats: [chat('c1', 10, ['i1'])], + artifacts: noArtifacts, + sync: synced({ artifacts: artifactsFingerprint({ items, versions: [] }) }) + }) + expect(emptied?.entry?.artifacts).toEqual(noArtifacts) + }) +}) + +describe('jsonBytes', () => { + it('counts the bytes the request carries, not UTF-16 code units', () => { + expect(jsonBytes('ab')).toBe(4) + expect(jsonBytes('日本')).toBe(8) + expect(jsonBytes('😀')).toBe(6) + }) +}) + +describe('splitEntry', () => { + it('splits an oversized entry into chat-only parts, the head riding on the last', () => { + const big = (id: string) => ({ id, record: { id, text: 'x'.repeat(150) } }) + const entry = { + id: 's', + head: { id: 's' }, + chats: [big('c1'), big('c2'), big('c3')], + delete_chats: ['old'] + } + const parts = splitEntry(entry, 200) + expect(parts.map((p) => p.chats?.map((c) => c.id))).toEqual([['c1'], ['c2'], ['c3']]) + expect( + parts.slice(0, -1).every((p) => p.head === undefined && p.delete_chats === undefined) + ).toBe(true) + expect(parts.at(-1)?.head).toEqual({ id: 's' }) + expect(parts.at(-1)?.delete_chats).toEqual(['old']) + // Within the target, or a single chat: nothing to split. + expect(splitEntry(entry, 10_000)).toEqual([entry]) + expect(splitEntry({ id: 's', chats: [big('c1')] }, 10)).toHaveLength(1) + // The server's per-entry chat cap splits too, however small the chats. + const many = { + id: 's', + chats: Array.from({ length: 250 }, (_, i) => ({ id: `c${i}`, record: {} })) + } + expect(splitEntry(many, 1_000_000).map((p) => p.chats?.length)).toEqual([100, 100, 50]) + }) +}) diff --git a/frontend/src/lib/components/sessions/sessionMirrorPlan.ts b/frontend/src/lib/components/sessions/sessionMirrorPlan.ts new file mode 100644 index 0000000000..c61bcf2763 --- /dev/null +++ b/frontend/src/lib/components/sessions/sessionMirrorPlan.ts @@ -0,0 +1,309 @@ +// The pure half of the session backup: given what the local stores hold for one session +// and what was last pushed, decide what the next push carries. Every piece is compared +// against its own marker (a signature for the record, `lastModified` for a chat, the id +// for a write-once image) so a session that only changed locally in ways the backup +// does not keep sends nothing. +import type { AISessionBackupPush } from '$lib/gen' +import { orderedJsonStringify } from '$lib/utils' +import type { Session } from './sessionState.svelte' +import type { ArtifactVersion, PersistedArtifact } from '../copilot/chat/artifacts/artifactsDB' + +/** What the backup remembers of a session after a successful push. */ +export interface MirrorSyncState { + id: string + /** The workspace whose storage holds the backup. */ + ws: string + /** `headSig` of the record pushed. */ + head: string + /** `lastModified` of each chat pushed, by chat id. */ + chats: Record + /** The chat each pushed image belongs to, by image id. */ + images: Record + artifacts?: string + /** The workspace's storage went away after this push: what it holds is unknown, so the + * next push carries everything again. Kept rather than deleted, so a removal still + * knows a backup existed. */ + stale?: boolean + /** The dirty mark's counter this push covered. A mark is pending while its counter is + * above this; retiring it here rather than deleting the mark means a tab bumping the + * counter while another flushes can never have its bump erased. */ + flushedV?: number + /** The user deleted the session and its removal mark could not be written to + * localStorage (full): the row itself carries the removal, until it lands. */ + removed?: boolean + /** The storage the push landed in, as the server names it, and the backup generation + * (bumped by a workspace key rotation) it landed under. A row recorded against another + * storage or generation describes objects the server no longer looks at. */ + storageId?: string + generation?: number + /** Other storages this workspace was on that still hold a copy of the backup (a switch + * leaves the old copy where it was): a removal is done only once each has answered it, + * or a switch back would bring a deleted session back. */ + alsoIn?: string[] + /** Bumps of the dirty mark that localStorage refused, recorded here instead: the mark's + * counter plus this is what a push retires, and every row write keeps it. */ + extraV?: number + /** A restore in progress (or cut short): the pieces it wrote for a session that has no + * record yet, so a later restore deletes the ones the backup no longer has. */ + staging?: { chats: string[]; images: string[]; items: string[]; versions: string[] } +} + +/** + * The part of a session record the backup keeps. Left out on purpose: `name` (a + * per-browser counter), the unsent-draft fields (`pending_*`, `draftPrompt`, + * `autoSendDraftAt`), `workspace_root_id` (derived on import), `transient`, and the two + * fields reading a session bumps (`lastSeenCount`, `lastActivityAt`) — so opening a + * session and reading its new messages never costs a push. + */ +export type SessionHead = Pick< + Session, + | 'id' + | 'workspace_id' + | 'chatId' + | 'summary' + | 'summarySource' + | 'createdAt' + | 'archived' + | 'archivedByWorkspace' + | 'moves' + | 'previewTabs' + | 'activePreviewTabId' + | 'previewCollapsed' + | 'previewSize' +> + +export function sessionHead(s: Session): SessionHead { + const head: SessionHead = { id: s.id, createdAt: s.createdAt } + if (s.workspace_id !== undefined) head.workspace_id = s.workspace_id + if (s.chatId !== undefined) head.chatId = s.chatId + if (s.summary !== undefined) head.summary = s.summary + if (s.summarySource !== undefined) head.summarySource = s.summarySource + if (s.archived !== undefined) head.archived = s.archived + if (s.archivedByWorkspace !== undefined) head.archivedByWorkspace = s.archivedByWorkspace + if (s.moves !== undefined) head.moves = s.moves + if (s.previewTabs !== undefined) head.previewTabs = s.previewTabs + if (s.activePreviewTabId !== undefined) head.activePreviewTabId = s.activePreviewTabId + if (s.previewCollapsed !== undefined) head.previewCollapsed = s.previewCollapsed + if (s.previewSize !== undefined) head.previewSize = s.previewSize + return head +} + +export function headSig(s: Session): string { + return orderedJsonStringify(sessionHead(s)) +} + +export interface ArtifactsSnapshot { + items: PersistedArtifact[] + versions: ArtifactVersion[] +} + +/** Cheap to compute from the rows alone: every edit bumps `updatedAt`, every snapshot has + * its own key, and approving a plan changes `approvedVersion`. */ +export function artifactsFingerprint(a: ArtifactsSnapshot): string { + const items = a.items + .map((i) => `${i.id}:${i.updatedAt}:${i.version ?? 1}:${i.approvedVersion ?? ''}`) + .sort() + const versions = a.versions.map((v) => v.key).sort() + return JSON.stringify([items, versions]) +} + +export interface ChatSnapshot { + id: string + lastModified: number + /** The stored record; absent for a chat that did not change since the last push, whose + * bytes the caller did not read. */ + record?: unknown + imageIds: string[] + /** Too large to back up: planned as if it did not exist, so a copy pushed while it was + * smaller is deleted rather than restored one day as the current transcript. */ + omitted?: boolean +} + +export interface PlanInput { + session: Session + /** Every chat the session owns right now. */ + chats: ChatSnapshot[] + artifacts: ArtifactsSnapshot + sync?: MirrorSyncState +} + +export interface PlannedPush { + workspaceId: string + /** Absent when nothing changed that the backup keeps. */ + entry?: AISessionBackupPush + /** Images the entry needs uploaded, whose bytes the caller loads. */ + images: { chat_id: string; id: string }[] + /** The workspace the session was backed up in before it moved. */ + removeFrom?: string + next: MirrorSyncState + /** Deletes past the per-entry cap were left in `next` for the following push, so the + * session must stay marked once this one lands. */ + carried: boolean + /** Nothing of the session is taken to be in the storage: every piece goes, and the + * first part opens the push whole (see `whole` on the entry). */ + whole: boolean +} + +/** `undefined` for a session with nowhere to go: an unsent draft has no workspace yet. */ +export function planSessionPush(input: PlanInput): PlannedPush | undefined { + const { session, chats, artifacts } = input + const workspaceId = session.workspace_id + if (!workspaceId) return undefined + // A move is a full push into the new workspace's storage; the copy in the old one goes. + const prev = input.sync?.ws === workspaceId ? input.sync : undefined + const removeFrom = input.sync && input.sync.ws !== workspaceId ? input.sync.ws : undefined + + const entry: AISessionBackupPush = { id: session.id } + let changed = false + let carried = false + const sig = headSig(session) + if (prev?.head !== sig) { + entry.head = sessionHead(session) + changed = true + } + + const next: MirrorSyncState = { + id: session.id, + ws: workspaceId, + head: sig, + chats: {}, + images: {} + } + const images: { chat_id: string; id: string }[] = [] + const pushedChats: { id: string; record: Record }[] = [] + for (const chat of chats) { + if (chat.omitted) continue + next.chats[chat.id] = chat.lastModified + if (prev?.chats[chat.id] !== chat.lastModified && chat.record !== undefined) { + pushedChats.push({ id: chat.id, record: chat.record as Record }) + } + for (const id of chat.imageIds) { + next.images[id] = chat.id + if (prev?.images[id] === undefined) images.push({ chat_id: chat.id, id }) + } + } + if (pushedChats.length > 0) { + entry.chats = pushedChats + changed = true + } + if (prev) { + const gone = Object.keys(prev.chats).filter((id) => next.chats[id] === undefined) + // An image evicted by the per-chat cap, from a chat that is still there (a deleted + // chat takes its images with it server-side). + const evicted = Object.entries(prev.images).filter( + ([id, chatId]) => next.images[id] === undefined && next.chats[chatId] !== undefined + ) + // Past the server's cap per entry, the rest stays in `next` as if still pushed, so + // the following push finds it gone again. + if (gone.length > 0) { + entry.delete_chats = gone.slice(0, MAX_DELETES_PER_ENTRY) + for (const id of gone.slice(MAX_DELETES_PER_ENTRY)) { + next.chats[id] = prev.chats[id] + carried = true + } + changed = true + } + if (evicted.length > 0) { + entry.delete_images = evicted + .slice(0, MAX_DELETES_PER_ENTRY) + .map(([id, chatId]) => ({ chat_id: chatId, id })) + for (const [id, chatId] of evicted.slice(MAX_DELETES_PER_ENTRY)) { + next.images[id] = chatId + carried = true + } + changed = true + } + } + + const fingerprint = artifactsFingerprint(artifacts) + next.artifacts = fingerprint + if (prev?.artifacts !== fingerprint && (artifacts.items.length > 0 || prev?.artifacts)) { + entry.artifacts = { items: artifacts.items, versions: artifacts.versions } + changed = true + } + + return { + workspaceId, + entry: changed ? entry : undefined, + images, + removeFrom, + next, + carried, + whole: prev === undefined + } +} + +/** Bytes a JSON body would carry for this value, as sent: UTF-8, not UTF-16 code units, + * which would under-count a transcript in a non-Latin script by up to three times. Counted + * rather than encoded: the values measured are the multi-megabyte ones. */ +export function jsonBytes(value: unknown): number { + const text = JSON.stringify(value) + let bytes = 0 + for (let i = 0; i < text.length; i++) { + const c = text.charCodeAt(i) + if (c < 0x80) bytes += 1 + else if (c < 0x800) bytes += 2 + else if (c >= 0xd800 && c <= 0xdbff) { + // A surrogate pair is one four-byte code point. + bytes += 4 + i++ + } else bytes += 3 + } + return bytes +} + +/** Object-store calls the server makes for an entry, the unit its per-request cap counts. */ +export function operationsOf(entry: AISessionBackupPush): number { + return ( + (entry.chats?.length ?? 0) + + (entry.images?.length ?? 0) + + (entry.delete_chats?.length ?? 0) + + (entry.delete_images?.length ?? 0) + ) +} + +export interface PushBody { + owner: string + sessions: AISessionBackupPush[] + removed?: string[] +} + +/** The server's caps on chats and on each delete list per entry. */ +export const MAX_CHATS_PER_ENTRY = 100 +export const MAX_DELETES_PER_ENTRY = 1000 + +/** + * Break an entry that outgrows the target, or the server's per-entry chat cap, into + * chat-only entries, each written on its own, with everything else riding on the last one: + * the entries go out in order and the server lists the session by the last, so the marker + * never lists a chat that has not landed. (A push of the session whole moves the head to + * whichever part goes first; see the mirror.) + */ +export function splitEntry(entry: AISessionBackupPush, targetBytes: number): AISessionBackupPush[] { + if ( + !entry.chats || + entry.chats.length <= 1 || + (entry.chats.length <= MAX_CHATS_PER_ENTRY && jsonBytes(entry) <= targetBytes) + ) { + return [entry] + } + const { chats, ...rest } = entry + const parts: AISessionBackupPush[] = [] + let current: typeof chats = [] + let size = 0 + for (const chat of chats) { + const bytes = jsonBytes(chat) + if ( + current.length > 0 && + (current.length >= MAX_CHATS_PER_ENTRY || size + bytes > targetBytes) + ) { + parts.push({ id: entry.id, chats: current }) + current = [] + size = 0 + } + current.push(chat) + size += bytes + } + parts.push({ ...rest, chats: current }) + return parts +} diff --git a/frontend/src/lib/components/sessions/sessionMirrorSignal.ts b/frontend/src/lib/components/sessions/sessionMirrorSignal.ts new file mode 100644 index 0000000000..eb936d2e22 --- /dev/null +++ b/frontend/src/lib/components/sessions/sessionMirrorSignal.ts @@ -0,0 +1,42 @@ +// The one-way channel from the IndexedDB write funnels (session records, chat history, +// artifacts) to the session backup. Import-free on purpose: the stores it is called from +// must not depend on the backup module, which depends on all of them. + +// `email` names the user whose store the write landed in (from the store's scoped name): +// the current user may have changed while the write was pending, and the mark belongs to +// the store's user, not to whoever is logged in when it completes. +export type MirrorSignal = + | { kind: 'dirty'; sessionId: string; chatId?: string; email?: string } + | { kind: 'removed'; sessionId: string; workspaceId?: string; email?: string } + +let handler: ((signal: MirrorSignal) => void) | undefined +// Signals raised before the backup module registered, replayed to it on registration. +let buffered: MirrorSignal[] = [] + +function emit(signal: MirrorSignal): void { + if (handler) handler(signal) + else buffered.push(signal) +} + +/** A durable local write landed for this session (and, when known, this chat) in the + * store of `email`. */ +export function markSessionDirty(sessionId: string, chatId?: string, email?: string): void { + emit({ kind: 'dirty', sessionId, chatId, email }) +} + +/** The user deleted this session; its backup goes with it. */ +export function markSessionRemoved(sessionId: string, workspaceId?: string, email?: string): void { + emit({ kind: 'removed', sessionId, workspaceId, email }) +} + +export function onMirrorSignal(fn: (signal: MirrorSignal) => void): void { + handler = fn + const replay = buffered + buffered = [] + for (const signal of replay) fn(signal) +} + +export function __resetMirrorSignalForTesting(): void { + handler = undefined + buffered = [] +} diff --git a/frontend/src/lib/components/sessions/sessionState.svelte.ts b/frontend/src/lib/components/sessions/sessionState.svelte.ts index d19faeac92..2670b3fa4d 100644 --- a/frontend/src/lib/components/sessions/sessionState.svelte.ts +++ b/frontend/src/lib/components/sessions/sessionState.svelte.ts @@ -24,8 +24,10 @@ import { workspaceRootId } from './sessionScope.svelte' import { clearSessionRecovered } from './sessionRecoveryNotice.svelte' import { type DBSchema, type IDBPDatabase } from 'idb' import { userScopedDb } from '$lib/userScopedDb' +import { emailOfScopedKey, scopedKeyFor } from '$lib/userScopedStorage' import { deleteItemsForSession } from '../copilot/chat/files/attachedFilesDB' import { deleteArtifactsForSession } from '../copilot/chat/artifacts/artifactsDB' +import { markSessionDirty, markSessionRemoved } from './sessionMirrorSignal' // Switch the global workspace iff the target differs from the active one // and is non-empty. Centralises the "session needs its workspace in focus" @@ -102,6 +104,10 @@ export type Session = { // archived (not by the user). Lets reconciliation auto-unarchive the session // when the workspace is unarchived, while leaving user-archived sessions be. archivedByWorkspace?: boolean + // How many times the session moved to another workspace. The backup keeps it + // with the session's marker, so a restore that finds a copy in two workspaces + // (moved, the old copy not yet removed) takes the later one without a clock. + moves?: number // In-memory-only flag: the session exists but hasn't been written to // IndexedDB yet. Set at creation, cleared on the first genuine user touch // (typed prompt, workspace/fork pick, preview tab, rename) which persists @@ -444,6 +450,7 @@ async function deleteSessionRow(db: IDBPDatabase, id: string): Pr async function putSessionRow(db: IDBPDatabase, s: Session): Promise { if (deletedSessionIds.has(s.id)) return await db.put('sessions', s) + markSessionDirty(s.id, undefined, emailOfScopedKey(SESSIONS_DB, db.name)) } // Write-behind a single session record. Transient sessions are in-memory only @@ -775,6 +782,15 @@ export function findEmptyLandingSession(): Session | undefined { ) } +// Session names are a per-browser counter (`session-N`) that the sessions page puts in +// its URL, so a new or restored record takes the number after the highest in use. +function nextSessionNumber(sessions: Session[]): number { + const numbers = sessions + .map((s) => /^session-(\d+)$/.exec(s.name)?.[1]) + .map((n) => (n ? parseInt(n, 10) : 0)) + return (numbers.length ? Math.max(...numbers) : 0) + 1 +} + export function createSession(): Session { // Reuse an existing untouched draft from the active family rather than pile a // blank entry on every `+`, so several pending sessions can still be built up @@ -795,10 +811,7 @@ export function createSession(): Session { return reusable } sessionState.sessions = sessionState.sessions.filter((s) => !isDiscardableDraft(s)) - const existingNumbers = sessionState.sessions - .map((s) => /^session-(\d+)$/.exec(s.name)?.[1]) - .map((n) => (n ? parseInt(n, 10) : 0)) - const next = (existingNumbers.length ? Math.max(...existingNumbers) : 0) + 1 + const next = nextSessionNumber(sessionState.sessions) // Start in the workspace you're in. The one exception: a root you can't // deploy to (locked, no bypass) steers to its dev, since a session there // couldn't edit anything. The picker lets you switch. @@ -1069,6 +1082,7 @@ export async function moveSessionToWorkspace(id: string, newWorkspaceId: string) const s = sessionState.sessions.find((x) => x.id === id) if (!s) return if (s.workspace_id === newWorkspaceId) return + if (s.workspace_id !== undefined) s.moves = (s.moves ?? 0) + 1 s.workspace_id = newWorkspaceId delete s.pending_workspace_id delete s.pending_fork @@ -1144,9 +1158,63 @@ export function deleteSession(id: string) { // GC any linked files and artifacts persisted for this session. void deleteItemsForSession(id) void deleteArtifactsForSession(id) + // Only a delete the user asked for takes the backup with it: the workspace-lifecycle + // removals above keep theirs, so a session dropped by a wrong reconcile can be restored. + markSessionRemoved(id, s.workspace_id) logFeatureUsage('ai_session', 'deleted', { entityId: id, workspace: s.workspace_id }) } +// --- Session backup support (sessionMirror) --- + +export function isSessionTombstoned(id: string): boolean { + return deletedSessionIds.has(id) +} + +// Every stored record of the named user, or undefined when the store is unavailable or +// already serves someone else: the backup captures its user up front and must not follow +// an in-place account switch. +export async function readStoredSessions(email: string): Promise { + if (!BROWSER) return undefined + const db = await sessionsDb.whenReady() + if (!db || db.name !== scopedKeyFor(SESSIONS_DB, email)) return undefined + try { + return await db.getAll('sessions') + } catch (e) { + console.error('Failed to read sessions from IndexedDB', e) + return undefined + } +} + +// Add restored records for sessions this browser does not have, and re-hydrate the list. +// A record that exists, or was deleted here, is left alone: the local copy is the newer +// one. Returns the ids written. +export async function importSessions(records: Session[], email: string): Promise { + if (!BROWSER) return [] + const db = await sessionsDb.whenReady() + if (!db || db.name !== scopedKeyFor(SESSIONS_DB, email)) return [] + const imported: string[] = [] + try { + const tx = db.transaction('sessions', 'readwrite') + const existing = new Set((await tx.store.getAllKeys()).map(String)) + let next = nextSessionNumber([...(await tx.store.getAll()), ...sessionState.sessions]) + for (const r of records) { + if (existing.has(r.id) || deletedSessionIds.has(r.id)) continue + const record: Session = { ...r, name: `session-${next++}` } + delete record.transient + delete record.workspace_root_id + ensureSessionRootId(record) + await tx.store.put(record) + imported.push(record.id) + } + await tx.done + } catch (e) { + console.error('Failed to import sessions', e) + return [] + } + if (imported.length > 0) await hydrateSessions() + return imported +} + export function setSessionChatId(sessionId: string, chatId: string) { const s = sessionState.sessions.find((x) => x.id === sessionId) if (s && s.chatId !== chatId) { @@ -1172,6 +1240,7 @@ async function patchStoredSessionChatId(s: Session, chatId: string): Promise { userStore.set(undefined) await vi.waitFor(() => expect(sessionState.sessions).toEqual([])) }) + + // A restored backup must never replace what this browser has, come back after the + // user deleted it here, or take a name the sessions page already routes by. + it('importSessions adds only unknown, undeleted records under fresh names', async () => { + const user = freshUser() + await login(user) + const local = session({ id: 'local', name: 'session-3', createdAt: 1, summary: 'mine' }) + await putSession(local) + sessionState.sessions.push(local) + deleteSession('local') + await flush() + await putSession(session({ id: 'kept', name: 'session-5', createdAt: 2, summary: 'kept' })) + + const imported = await importSessions( + [ + session({ id: 'local', name: 'session-1', createdAt: 1, summary: 'remote copy' }), + session({ id: 'kept', name: 'session-1', createdAt: 2, summary: 'remote copy' }), + session({ id: 'new', name: 'session-1', createdAt: 3, workspace_id: 'ws' }) + ], + user.email + ) + expect(imported).toEqual(['new']) + await vi.waitFor(() => + expect(sessionState.sessions.map((s) => [s.id, s.name])).toEqual([ + ['new', 'session-6'], + ['kept', 'session-5'] + ]) + ) + expect(sessionState.sessions.find((s) => s.id === 'kept')?.summary).toBe('kept') + + // The wrong user's name gets nothing written. + expect(await importSessions([session({ id: 'other', createdAt: 4 })], 'nobody@x')).toEqual([]) + }) }) diff --git a/frontend/src/lib/components/workspaceSettings/AISettings.svelte b/frontend/src/lib/components/workspaceSettings/AISettings.svelte index ae9d2640b0..7e8b9cc76e 100644 --- a/frontend/src/lib/components/workspaceSettings/AISettings.svelte +++ b/frontend/src/lib/components/workspaceSettings/AISettings.svelte @@ -28,6 +28,7 @@ import ModelPricing from './ModelPricing.svelte' import AiUsagePanel from './AiUsagePanel.svelte' import { setCopilotInfo } from '$lib/aiStore' + import { backupSettingsChanged } from '$lib/components/sessions/sessionMirror.svelte' import AIPromptsModal from '../settings/AIPromptsModal.svelte' import { Settings } from 'lucide-svelte' import { untrack } from 'svelte' @@ -79,6 +80,7 @@ let usingOpenaiClientCredentialsOauth = $state(false) let workspaceOverrideEditorOpened = $state(false) let copilotDisabled = $state(false) + let sessionsStorageDisabled = $state(false) // --- Initial state for dirty tracking --- let initialAiProviders: Exclude = $state({}) @@ -90,6 +92,7 @@ let initialModelPricing: Record = $state({}) let initialPrompts: Record = $state({}) let initialCopilotDisabled = $state(false) + let initialSessionsStorageDisabled = $state(false) let lastLoadedConfigKey = $state(undefined) function clone(v: T): T { @@ -118,6 +121,7 @@ maxTokensPerModel = clone(config?.max_tokens_per_model ?? {}) modelPricing = clone(config?.model_pricing ?? {}) copilotDisabled = config?.copilot_disabled === true + sessionsStorageDisabled = config?.sessions_storage_disabled === true for (const mode of ['edit', 'fix', 'gen']) { if (!(mode in customPrompts)) { customPrompts[mode] = '' @@ -135,6 +139,7 @@ initialModelPricing = clone(modelPricing) initialPrompts = clone(customPrompts) initialCopilotDisabled = copilotDisabled + initialSessionsStorageDisabled = sessionsStorageDisabled } export function loadFromConfig(config: AIConfig | undefined) { @@ -151,6 +156,7 @@ maxTokensPerModel = clone(initialMaxTokensPerModel) modelPricing = clone(initialModelPricing) copilotDisabled = initialCopilotDisabled + sessionsStorageDisabled = initialSessionsStorageDisabled } $effect(() => { @@ -186,7 +192,8 @@ JSON.stringify(customPrompts) !== JSON.stringify(initialCustomPrompts) || JSON.stringify(maxTokensPerModel) !== JSON.stringify(initialMaxTokensPerModel) || JSON.stringify(modelPricing) !== JSON.stringify(initialModelPricing) || - copilotDisabled !== initialCopilotDisabled + copilotDisabled !== initialCopilotDisabled || + sessionsStorageDisabled !== initialSessionsStorageDisabled ) $effect(() => { @@ -291,8 +298,9 @@ .filter(([_, prompt]) => prompt.trim().length > 0) .reduce((acc, [mode, prompt]) => ({ ...acc, [mode]: prompt }), {}) - // The flag is the one thing a workspace on instance defaults still stores of its own. + // The flags are what a workspace on instance defaults still stores of its own. const copilot_disabled = copilotDisabled ? true : undefined + const sessions_storage_disabled = sessionsStorageDisabled ? true : undefined return Object.keys(aiProviders ?? {}).length > 0 ? { providers: aiProviders, @@ -303,9 +311,10 @@ max_tokens_per_model: Object.keys(maxTokensPerModel).length > 0 ? maxTokensPerModel : undefined, model_pricing: Object.keys(modelPricing).length > 0 ? modelPricing : undefined, - copilot_disabled + copilot_disabled, + sessions_storage_disabled } - : { copilot_disabled } + : { copilot_disabled, sessions_storage_disabled } } function isSaveDisabled(): boolean { @@ -332,6 +341,7 @@ async function editCopilotConfig(): Promise { const config = buildConfig() + const backupsToggled = sessionsStorageDisabled !== initialSessionsStorageDisabled let settingsState: GetCopilotSettingsStateResponse | undefined if (customSave) { @@ -348,6 +358,9 @@ instance_ai_summary: response.instance_ai_summary } sendUserToast('AI settings updated') + // This page's session backups follow the switch at once, rather than at the + // next page load. + if (backupsToggled) backupSettingsChanged(effectiveWorkspace) } storeInitialState() // Hand the parent what was persisted: it owns `initialConfig`, and this component is @@ -646,6 +659,18 @@ options={{ right: 'Hide AI sessions in this workspace' }} /> + + { + sessionsStorageDisabled = e.detail + }} + options={{ right: 'Do not back AI sessions up to the workspace storage' }} + /> + {/if} diff --git a/frontend/src/lib/userScopedDb.ts b/frontend/src/lib/userScopedDb.ts index 71d80a369f..d1283b2c73 100644 --- a/frontend/src/lib/userScopedDb.ts +++ b/frontend/src/lib/userScopedDb.ts @@ -1,4 +1,11 @@ -import { openDB as idbOpenDB, deleteDB as idbDeleteDB, type DBSchema, type IDBPDatabase } from 'idb' +import { + openDB as idbOpenDB, + deleteDB as idbDeleteDB, + type DBSchema, + type IDBPDatabase, + type IDBPTransaction, + type StoreNames +} from 'idb' import { scopedKey } from '$lib/userScopedStorage' // Per-user IndexedDB lifecycle, shared by the session list and the copilot @@ -21,7 +28,12 @@ export interface UserScopedDbMigrateDeps { export interface UserScopedDbOptions { version: number - upgrade: (db: IDBPDatabase) => void + // The version-change transaction is the only way to add an index to a store that + // already exists; a store being created gets it from the store handle instead. + upgrade: ( + db: IDBPDatabase, + tx: IDBPTransaction[], 'versionchange'> + ) => void // Invoked once per scoped name right after a successful open. The fn owns its // own "already migrated / not applicable" gate (e.g. checking a store's // count) — claim-then-delete legacy data lives here. @@ -100,11 +112,11 @@ export function userScopedDb( try { let handle: IDBPDatabase | undefined const db = await openDB(name, opts.version, { - upgrade(database) { + upgrade(database, _oldVersion, _newVersion, transaction) { // The version-change transaction is ours: nothing is queued ahead of this // open any more, and what remains is our own upgrade running. stopWaiting() - opts.upgrade(database) + opts.upgrade(database, transaction) }, // Another tab is opening this database at a higher version, which our open // connection would block indefinitely. Let go so their upgrade lands; this diff --git a/frontend/src/lib/userScopedStorage.ts b/frontend/src/lib/userScopedStorage.ts index 905d609b61..e3ee1d842d 100644 --- a/frontend/src/lib/userScopedStorage.ts +++ b/frontend/src/lib/userScopedStorage.ts @@ -52,7 +52,19 @@ export function getCurrentUserEmail(): string | undefined { // treat that as "do not read/write" so we never touch a browser-global key. export function scopedKey(base: string): string | undefined { if (!currentEmail) return undefined - return `${base}::${currentEmail}` + return scopedKeyFor(base, currentEmail) +} + +// The key a base name has for a given user, for work that captured its user up front and +// must not follow an in-place account switch (the session backup flush). +export function scopedKeyFor(base: string, email: string): string { + return `${base}::${email}` +} + +// The email a scoped key or database name was built for, so a write that landed in a +// store can name the user it belongs to even after the current user changed. +export function emailOfScopedKey(base: string, key: string): string | undefined { + return key.startsWith(`${base}::`) ? key.slice(base.length + 2) : undefined } // Register a callback invoked whenever the scoping email changes. Fired once diff --git a/frontend/src/routes/(root)/(logged)/+layout.svelte b/frontend/src/routes/(root)/(logged)/+layout.svelte index 4447a8d4b8..225b9443d3 100644 --- a/frontend/src/routes/(root)/(logged)/+layout.svelte +++ b/frontend/src/routes/(root)/(logged)/+layout.svelte @@ -92,6 +92,7 @@ import { parsePreviewItemRoute } from '$lib/components/sessions/previewPaths' import { rememberNavRoute } from '$lib/components/sessions/sessionSwitch.svelte' import { sessionState } from '$lib/components/sessions/sessionState.svelte' + import { restoreSessionBackups } from '$lib/components/sessions/sessionMirror.svelte' import { currentWorkspaceRootId } from '$lib/components/sessions/sessionScope.svelte' import WorkspaceScopeHeader from '$lib/components/sidebar/WorkspaceScopeHeader.svelte' import { DEFAULT_HUB_BASE_URL } from '$lib/hub' @@ -725,6 +726,16 @@ $workspaceStore untrack(() => updateUserStore($workspaceStore)) }) + // Bring back the AI sessions this browser lacks for the workspace family in view, once + // the local list is known (so nothing it has is fetched again) and the memberships have + // resolved (the family is derived from them). + $effect(() => { + const ws = $workspaceStore + const ready = sessionState.hydrated && $usersWorkspaceStore !== undefined + if (globalAiEnabled && ready && ws && !$userStore?.operator) { + untrack(() => restoreSessionBackups(ws)) + } + }) // While a fork is reachable, mirror its parent linkage to localStorage so a // later reload landing on a now-deleted fork can return to the parent (see // forkParentMemory + the deleted-fork recovery in the root layout).