mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-09-21 16:02:36 +00:00
Ruben FiszelandClaude Fable 5.1
ee6d317e31
* feat: retention for AI sessions, swept on the object store and in the browser Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * fix: make the retention sweeps retryable and safe against pushes Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix: spare other tabs' sessions, reclaim abandoned split pushes Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix: sweep under an exclusive session lock, keep the captured user Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * docs: say a tab selecting a session mid-sweep is not held back Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix: sweep local sessions only while no other tab has them loaded Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix: one tab sweeps at a time, and keeps the switched user's hold Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * test: push the fallback session again before the rotation assertions Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * refactor: keep retention server-side here, move the browser sweep out Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * docs: the retention setting no longer touches browser-local sessions Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
1583 lines
57 KiB
Rust
1583 lines
57 KiB
Rust
//! 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<Postgres>, 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(())
|
|
}
|
|
|
|
/// Configures the primary storage through the route, which is what sweeps the workspace's
|
|
/// backups out of the instance store.
|
|
async fn configure_primary_lfs_via_route(base: &str, root_path: &str) -> anyhow::Result<()> {
|
|
let resp = authed(
|
|
client().post(format!("{base}/workspaces/edit_large_file_storage_config")),
|
|
"SECRET_TOKEN",
|
|
)
|
|
.json(&json!({ "large_file_storage": {
|
|
"type": "FilesystemStorage",
|
|
"root_path": root_path,
|
|
"public_resource": false,
|
|
"advanced_permissions": null,
|
|
"secondary_storage": {}
|
|
}}))
|
|
.send()
|
|
.await?;
|
|
assert_eq!(resp.status(), 200, "{}", resp.text().await?);
|
|
Ok(())
|
|
}
|
|
|
|
/// The instance setting allowing the instance store to stand in for a workspace without
|
|
/// storage: `None` leaves it unset, which is on.
|
|
async fn set_instance_fallback(db: &Pool<Postgres>, on: Option<bool>) -> anyhow::Result<()> {
|
|
sqlx::query("DELETE FROM global_settings WHERE name = 'ai_sessions_instance_storage_fallback'")
|
|
.execute(db)
|
|
.await?;
|
|
if let Some(on) = on {
|
|
sqlx::query(
|
|
"INSERT INTO global_settings (name, value) VALUES ('ai_sessions_instance_storage_fallback', $1)",
|
|
)
|
|
.bind(json!(on))
|
|
.execute(db)
|
|
.await?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Polls until nothing is under the directory, for a deletion that runs off the request.
|
|
async fn wait_until_empty(dir: &std::path::Path, what: &str) {
|
|
for _ in 0..100 {
|
|
if files_under(dir).is_empty() {
|
|
return;
|
|
}
|
|
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
|
}
|
|
panic!("{what}: objects left under {}", dir.display());
|
|
}
|
|
|
|
async fn list(base: &str, token: &str) -> anyhow::Result<Value> {
|
|
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<Value> {
|
|
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<reqwest::Response> {
|
|
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 the newest generation: deleting an older generation's objects leaves its
|
|
/// directories behind, and `read_dir` order differs across filesystems.
|
|
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()
|
|
.filter_map(|entry| {
|
|
let generation: i64 = entry
|
|
.file_name()
|
|
.to_str()?
|
|
.strip_prefix('g')?
|
|
.parse()
|
|
.ok()?;
|
|
Some((generation, entry.path().join(&hash)))
|
|
})
|
|
.filter(|(_, path)| path.exists())
|
|
.max_by_key(|(generation, _)| *generation)
|
|
.map(|(_, path)| path)
|
|
.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<u8>)> {
|
|
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<Postgres>,
|
|
) -> 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, and the instance store (another test of this process may
|
|
// have loaded one) not allowed to stand in: the browser is told to stop trying.
|
|
set_instance_fallback(&db, Some(false)).await?;
|
|
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<String> {
|
|
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<String> = 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<Value> = (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<String> = 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<_>>(),
|
|
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<Value> {
|
|
ids.iter()
|
|
.map(|c| json!({ "id": c, "record": { "id": c } }))
|
|
.collect()
|
|
};
|
|
let pulled_chats = |pulled: Value| -> Vec<String> {
|
|
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<Value> {
|
|
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::<Value>::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);
|
|
|
|
// An incremental part changing more than one object unlists the session before its
|
|
// writes, so one write failing after another landed leaves it absent rather than listed
|
|
// as a mix of old and new pieces. A directory planted at `artifacts.json` fails that write.
|
|
let s10_head =
|
|
json!({ "id": "s10", "workspace_id": "test-workspace", "createdAt": 10, "chatId": "c1" });
|
|
let resp = push(
|
|
&base,
|
|
"SECRET_TOKEN",
|
|
json!({
|
|
"owner": "test@windmill.dev",
|
|
"sessions": [{ "id": "s10", "whole": true, "head": s10_head, "chats": s9_chats(&["c1"]), "artifacts": { "items": ["a1"] } }]
|
|
}),
|
|
)
|
|
.await?;
|
|
assert_eq!(resp.status(), 200);
|
|
let s10_dir = user_root(storage_dir.path(), "test@windmill.dev").join("sessions/s10");
|
|
let artifacts_path = s10_dir.join("artifacts.json");
|
|
std::fs::remove_file(&artifacts_path)?;
|
|
std::fs::create_dir(&artifacts_path)?;
|
|
let resp = push(
|
|
&base,
|
|
"SECRET_TOKEN",
|
|
json!({
|
|
"owner": "test@windmill.dev",
|
|
"sessions": [{ "id": "s10", "chats": s9_chats(&["c2"]), "artifacts": { "items": ["a2"] } }]
|
|
}),
|
|
)
|
|
.await?;
|
|
assert_eq!(resp.status(), 200);
|
|
let answer: Value = resp.json().await?;
|
|
assert!(
|
|
answer["results"][0]["error"].is_string(),
|
|
"the artifacts write must fail: {answer}"
|
|
);
|
|
assert!(answer["results"][0]["needs_whole"].is_null());
|
|
assert!(
|
|
s10_dir.join("chats/c2.json").is_file(),
|
|
"the chat landed before the artifacts failed"
|
|
);
|
|
let s10_listed = |listing: Value| {
|
|
listing["sessions"]
|
|
.as_array()
|
|
.unwrap()
|
|
.iter()
|
|
.any(|s| s["id"] == "s10")
|
|
};
|
|
assert!(!s10_listed(list(&base, "SECRET_TOKEN").await?));
|
|
assert_eq!(
|
|
pull(&base, "SECRET_TOKEN", &["s10"]).await?["sessions"],
|
|
json!([])
|
|
);
|
|
let resp = push(
|
|
&base,
|
|
"SECRET_TOKEN",
|
|
json!({
|
|
"owner": "test@windmill.dev",
|
|
"sessions": [{ "id": "s10", "chats": s9_chats(&["c3"]) }]
|
|
}),
|
|
)
|
|
.await?;
|
|
assert_eq!(resp.status(), 200);
|
|
let answer: Value = resp.json().await?;
|
|
assert_eq!(answer["results"][0]["needs_whole"], true);
|
|
assert!(!s10_listed(list(&base, "SECRET_TOKEN").await?));
|
|
std::fs::remove_dir(&artifacts_path)?;
|
|
let resp = push(
|
|
&base,
|
|
"SECRET_TOKEN",
|
|
json!({
|
|
"owner": "test@windmill.dev",
|
|
"sessions": [{ "id": "s10", "whole": true, "head": s10_head, "chats": s9_chats(&["c1", "c2", "c3"]), "artifacts": { "items": ["a2"] } }]
|
|
}),
|
|
)
|
|
.await?;
|
|
assert_eq!(resp.status(), 200);
|
|
assert!(s10_listed(list(&base, "SECRET_TOKEN").await?));
|
|
let pulled = pull(&base, "SECRET_TOKEN", &["s10"]).await?;
|
|
assert_eq!(
|
|
pulled["sessions"][0]["artifacts"],
|
|
json!({ "items": ["a2"] })
|
|
);
|
|
assert_eq!(pulled_chats(pulled), vec!["c1", "c2", "c3"]);
|
|
let resp = push(
|
|
&base,
|
|
"SECRET_TOKEN",
|
|
json!({ "owner": "test@windmill.dev", "removed": ["s10"] }),
|
|
)
|
|
.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<Postgres>,
|
|
) -> 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::<Vec<_>>() }]
|
|
}),
|
|
)
|
|
.await?;
|
|
assert_eq!(resp.status(), 400, "{}", resp.text().await?);
|
|
|
|
// ...and across the whole request, not only per entry.
|
|
let sessions: Vec<Value> = (0..100)
|
|
.map(|i| {
|
|
json!({ "id": format!("s{i}"), "delete_chats": (0..50).map(|j| format!("c{j}")).collect::<Vec<_>>() })
|
|
})
|
|
.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(())
|
|
}
|
|
|
|
/// Sets the object's modification time `days` back: the FilesystemStorage answers
|
|
/// `last_modified` from it, so this is a session no push touched since.
|
|
fn age_object(path: &std::path::Path, days: u64) -> std::io::Result<()> {
|
|
let at = std::time::SystemTime::now() - std::time::Duration::from_secs(days * 86_400);
|
|
std::fs::File::options()
|
|
.write(true)
|
|
.open(path)?
|
|
.set_modified(at)
|
|
}
|
|
|
|
#[sqlx::test(fixtures("base"))]
|
|
async fn test_expired_backups_are_swept_by_age_and_left_out_of_the_listing(
|
|
db: Pool<Postgres>,
|
|
) -> 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?;
|
|
|
|
let chat = |sid: &str, cid: &str| {
|
|
json!({ "id": cid, "record": { "id": cid, "sessionId": sid, "lastModified": 2,
|
|
"actualMessages": [], "displayMessages": [] } })
|
|
};
|
|
let whole = |sid: &str| {
|
|
json!({
|
|
"id": sid, "whole": true, "epoch": 0,
|
|
"head": { "id": sid, "workspace_id": "test-workspace", "createdAt": 1, "chatId": "c1" },
|
|
"chats": [chat(sid, "c1")],
|
|
"images": [{ "chat_id": "c1", "id": "img1", "data_url": "data:image/png;base64,AAAA" }],
|
|
"artifacts": { "items": [], "versions": [] }
|
|
})
|
|
};
|
|
// Two pushes split over parts of which only the first part landed: one a browser
|
|
// abandoned long ago (its token aged past the retention), one still in flight.
|
|
let opening = |sid: &str| {
|
|
json!({
|
|
"id": sid, "whole": true, "epoch": 0, "push": format!("t-{sid}"), "opens": true,
|
|
"partial": true, "chats": [chat(sid, "c1")],
|
|
"head": { "id": sid, "workspace_id": "test-workspace", "createdAt": 1, "chatId": "c1" }
|
|
})
|
|
};
|
|
let resp = push(
|
|
&base,
|
|
"SECRET_TOKEN",
|
|
json!({ "owner": "test@windmill.dev", "sessions": [
|
|
whole("old"), whole("live"), opening("abandoned"), opening("inflight")
|
|
] }),
|
|
)
|
|
.await?;
|
|
assert_eq!(resp.status(), 200, "{}", resp.text().await?);
|
|
let root = user_root(storage_dir.path(), "test@windmill.dev");
|
|
age_object(&root.join("index/old/0"), 40)?;
|
|
age_object(&root.join("index/abandoned/push"), 40)?;
|
|
|
|
let listed = |listing: Value| -> Vec<String> {
|
|
let mut ids: Vec<String> = listing["sessions"]
|
|
.as_array()
|
|
.unwrap()
|
|
.iter()
|
|
.map(|s| s["id"].as_str().unwrap().to_string())
|
|
.collect();
|
|
ids.sort();
|
|
ids
|
|
};
|
|
let objects = |root: &std::path::Path| -> Vec<String> {
|
|
files_under(root)
|
|
.into_iter()
|
|
.map(|(p, _)| p.strip_prefix(root).unwrap().to_string_lossy().into_owned())
|
|
.collect()
|
|
};
|
|
|
|
// Without a retention nothing is swept, however old.
|
|
windmill_api::sweep_expired_ai_session_backups(&db).await;
|
|
assert_eq!(listed(list(&base, "SECRET_TOKEN").await?), ["live", "old"]);
|
|
|
|
let set_retention = |days: Value| {
|
|
authed(
|
|
client().post(format!("{base}/workspaces/edit_copilot_config")),
|
|
"SECRET_TOKEN",
|
|
)
|
|
.json(&json!({ "sessions_retention_days": days }))
|
|
.send()
|
|
};
|
|
let resp = set_retention(json!(0)).await?;
|
|
assert_eq!(resp.status(), 400, "{}", resp.text().await?);
|
|
let resp = set_retention(json!(30)).await?;
|
|
assert_eq!(resp.status(), 200, "{}", resp.text().await?);
|
|
|
|
// The listing leaves the expired session out before the sweep reaches it.
|
|
assert_eq!(listed(list(&base, "SECRET_TOKEN").await?), ["live"]);
|
|
assert!(root.join("sessions/old/head.json").exists());
|
|
|
|
// A removal cut short (a directory stands where the head is, so it cannot be unlinked)
|
|
// leaves the sweep's record with the markers gone; the next pass finds it and finishes.
|
|
let head = root.join("sessions/old/head.json");
|
|
std::fs::remove_file(&head)?;
|
|
std::fs::create_dir(&head)?;
|
|
std::fs::write(head.join("planted"), b"")?;
|
|
windmill_api::sweep_expired_ai_session_backups(&db).await;
|
|
assert!(root.join("index/old/sweep").exists());
|
|
assert!(!root.join("index/old/0").exists());
|
|
assert!(root.join("sessions/old/chats/c1.json").exists());
|
|
assert_eq!(listed(list(&base, "SECRET_TOKEN").await?), ["live"]);
|
|
std::fs::remove_dir_all(&head)?;
|
|
|
|
windmill_api::sweep_expired_ai_session_backups(&db).await;
|
|
let remaining = objects(&root);
|
|
assert!(
|
|
remaining
|
|
.iter()
|
|
.all(|p| !p.contains("/old/") && !p.contains("/abandoned/")),
|
|
"{remaining:?}"
|
|
);
|
|
for kept in [
|
|
"index/live/0",
|
|
"sessions/live/head.json",
|
|
"sessions/live/chats/c1.json",
|
|
"images/live/c1/img1",
|
|
"index/inflight/push",
|
|
"sessions/inflight/head.json",
|
|
] {
|
|
assert!(
|
|
remaining.iter().any(|p| p == kept),
|
|
"{kept} in {remaining:?}"
|
|
);
|
|
}
|
|
assert_eq!(listed(list(&base, "SECRET_TOKEN").await?), ["live"]);
|
|
|
|
// A second pass has nothing to do; a session pushed again since its marker aged is
|
|
// renewed by the push, which rewrites the marker.
|
|
windmill_api::sweep_expired_ai_session_backups(&db).await;
|
|
age_object(&root.join("index/live/0"), 40)?;
|
|
let resp = push(
|
|
&base,
|
|
"SECRET_TOKEN",
|
|
json!({ "owner": "test@windmill.dev",
|
|
"sessions": [{ "id": "live", "epoch": 0, "chats": [chat("live", "c2")] }] }),
|
|
)
|
|
.await?;
|
|
assert_eq!(resp.status(), 200, "{}", resp.text().await?);
|
|
windmill_api::sweep_expired_ai_session_backups(&db).await;
|
|
let mut after = objects(&root);
|
|
after.sort();
|
|
let mut expected = remaining.clone();
|
|
expected.push("sessions/live/chats/c2.json".to_string());
|
|
expected.sort();
|
|
assert_eq!(after, expected);
|
|
assert_eq!(listed(list(&base, "SECRET_TOKEN").await?), ["live"]);
|
|
Ok(())
|
|
}
|
|
|
|
/// Puts the process-wide instance store back to none, even when an assertion fails.
|
|
struct ResetInstanceStore;
|
|
impl Drop for ResetInstanceStore {
|
|
fn drop(&mut self) {
|
|
if let Ok(mut store) = windmill_object_store::OBJECT_STORE_SETTINGS.try_write() {
|
|
*store = None;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Puts the process-wide license key id back to none, an Enterprise plan in this build,
|
|
/// even when an assertion fails.
|
|
struct ResetLicensePlan;
|
|
impl Drop for ResetLicensePlan {
|
|
fn drop(&mut self) {
|
|
windmill_common::ee::LICENSE_KEY_ID.store(std::sync::Arc::new(String::new()));
|
|
}
|
|
}
|
|
|
|
/// A workspace without storage of its own backs up to the instance object store, every
|
|
/// answer saying so (`fallback`); a storage of its own, once configured, answers instead,
|
|
/// under a generation past everything the workspace left in the instance store, which the
|
|
/// change deletes; a plan switched to Pro stops the fallback with the store still loaded.
|
|
#[sqlx::test(fixtures("base"))]
|
|
async fn test_backups_fall_back_to_the_instance_storage(db: Pool<Postgres>) -> anyhow::Result<()> {
|
|
initialize_tracing().await;
|
|
let server = ApiServer::start(db.clone()).await?;
|
|
let base = format!(
|
|
"http://localhost:{}/api/w/test-workspace",
|
|
server.addr.port()
|
|
);
|
|
|
|
// The instance store, built from its settings as `reload_object_store_setting` does.
|
|
let instance_dir = tempfile::tempdir()?;
|
|
let instance_root = instance_dir.path().to_string_lossy().to_string();
|
|
*windmill_object_store::OBJECT_STORE_SETTINGS.write().await = Some(
|
|
windmill_object_store::build_object_store_from_settings(
|
|
windmill_object_store::ObjectSettings::Filesystem(
|
|
windmill_object_store::FilesystemSettings { root_path: instance_root.clone() },
|
|
),
|
|
None,
|
|
)
|
|
.await?,
|
|
);
|
|
let _reset = ResetInstanceStore;
|
|
let in_instance = instance_dir
|
|
.path()
|
|
.join("windmill_ai_sessions/test-workspace");
|
|
|
|
// Turned off by the instance setting: the browser is told to stop trying.
|
|
set_instance_fallback(&db, Some(false)).await?;
|
|
assert_eq!(list(&base, "SECRET_TOKEN").await?["enabled"], false);
|
|
set_instance_fallback(&db, None).await?;
|
|
|
|
// On, as it is unless turned off: the backups land in the instance store, under the
|
|
// workspace's prefix, and every answer says which kind of store it came from.
|
|
let head =
|
|
json!({ "id": "s1", "workspace_id": "test-workspace", "createdAt": 1, "chatId": "c1" });
|
|
let entry = json!({ "id": "s1", "whole": true, "head": head, "chats": [{ "id": "c1", "record": { "id": "c1" } }] });
|
|
let push_whole = || {
|
|
push(
|
|
&base,
|
|
"SECRET_TOKEN",
|
|
json!({ "owner": "test@windmill.dev", "sessions": [entry.clone()] }),
|
|
)
|
|
};
|
|
let resp = push_whole().await?;
|
|
assert_eq!(resp.status(), 200, "{}", resp.text().await?);
|
|
let pushed: Value = resp.json().await?;
|
|
assert_eq!(pushed["fallback"], true);
|
|
assert_eq!(pushed["results"], json!([{ "id": "s1" }]));
|
|
let listing = list(&base, "SECRET_TOKEN").await?;
|
|
assert_eq!(listing["enabled"], true);
|
|
assert_eq!(listing["fallback"], true);
|
|
assert_eq!(listing["sessions"][0]["id"], "s1");
|
|
let fallback_storage_id = listing["storage_id"].clone();
|
|
let pulled = pull(&base, "SECRET_TOKEN", &["s1"]).await?;
|
|
assert_eq!(pulled["fallback"], true);
|
|
assert_eq!(pulled["sessions"][0]["head"], head);
|
|
assert!(!files_under(&in_instance).is_empty());
|
|
|
|
// The workspace's storage usage counts them, under a name of their own.
|
|
let resp = authed(
|
|
client().get(format!("{base}/job_helpers/storage_usage")),
|
|
"SECRET_TOKEN",
|
|
)
|
|
.send()
|
|
.await?;
|
|
assert_eq!(resp.status(), 200, "{}", resp.text().await?);
|
|
let usage: Value = resp.json().await?;
|
|
let fallback_usage = usage["storages"]
|
|
.as_array()
|
|
.unwrap()
|
|
.iter()
|
|
.find(|s| s["storage"] == "_ai_sessions_fallback_")
|
|
.unwrap_or_else(|| panic!("no fallback usage in {usage}"));
|
|
assert!(fallback_usage["bytes"].as_i64().unwrap() > 0);
|
|
|
|
// The retention sweep reaches what the instance store keeps for the workspace, choosing
|
|
// that store from the row it reads the generation from.
|
|
let instance_user = user_root(instance_dir.path(), "test@windmill.dev");
|
|
age_object(&instance_user.join("index/s1/0"), 40)?;
|
|
sqlx::query(
|
|
"UPDATE workspace_settings SET ai_config = '{\"sessions_retention_days\": 30}' \
|
|
WHERE workspace_id = 'test-workspace'",
|
|
)
|
|
.execute(&db)
|
|
.await?;
|
|
windmill_api::sweep_expired_ai_session_backups(&db).await;
|
|
assert!(!instance_user.join("index/s1/0").exists());
|
|
assert!(!instance_user.join("sessions/s1/head.json").exists());
|
|
// Pushed again, so the rotation below has a backup to delete.
|
|
assert_eq!(push_whole().await?.status(), 200);
|
|
assert!(!files_under(&in_instance).is_empty());
|
|
|
|
// A key rotation sweeps the older generation out of the instance store too.
|
|
rotate(&base, &"c".repeat(64)).await?;
|
|
wait_until_empty(&in_instance, "a rotation on the instance store").await;
|
|
assert_eq!(list(&base, "SECRET_TOKEN").await?["sessions"], json!([]));
|
|
assert_eq!(push_whole().await?.status(), 200);
|
|
assert!(!files_under(&in_instance).is_empty());
|
|
|
|
// A storage of its own answers instead, under a generation the configuration moved past
|
|
// everything the workspace left in the instance store: nothing there is read again,
|
|
// whichever store a later return to the fallback finds, and it is deleted.
|
|
let before = list(&base, "SECRET_TOKEN").await?;
|
|
let storage_dir = tempfile::tempdir()?;
|
|
configure_primary_lfs_via_route(&base, &storage_dir.path().to_string_lossy()).await?;
|
|
wait_until_empty(&in_instance, "configuring a workspace storage").await;
|
|
let listing = list(&base, "SECRET_TOKEN").await?;
|
|
assert_eq!(listing["enabled"], true);
|
|
assert!(listing.get("fallback").is_none(), "{listing}");
|
|
assert_ne!(listing["storage_id"], fallback_storage_id);
|
|
assert_eq!(
|
|
listing["backup_generation"].as_i64(),
|
|
before["backup_generation"].as_i64().map(|g| g + 1),
|
|
"configuring a storage over the fallback must move the generation on"
|
|
);
|
|
assert_eq!(listing["sessions"], json!([]));
|
|
assert_eq!(push_whole().await?.status(), 200);
|
|
assert!(!files_under(storage_dir.path()).is_empty());
|
|
assert!(files_under(&in_instance).is_empty());
|
|
let resp = authed(
|
|
client().get(format!("{base}/job_helpers/storage_usage?refresh=true")),
|
|
"SECRET_TOKEN",
|
|
)
|
|
.send()
|
|
.await?;
|
|
assert_eq!(resp.status(), 200, "{}", resp.text().await?);
|
|
let usage: Value = resp.json().await?;
|
|
assert!(
|
|
!usage.to_string().contains("_ai_sessions_fallback_"),
|
|
"nothing is counted in the instance store for a workspace with storage: {usage}"
|
|
);
|
|
|
|
// Pointed at the instance store's own bucket, a storage of its own keeps its live
|
|
// backups there under the current generation, which no storage change deletes.
|
|
configure_primary_lfs_via_route(&base, &instance_root).await?;
|
|
assert_eq!(push_whole().await?.status(), 200);
|
|
assert!(!files_under(&in_instance).is_empty());
|
|
let same = list(&base, "SECRET_TOKEN").await?;
|
|
configure_primary_lfs_via_route(&base, &instance_root).await?;
|
|
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
|
|
assert!(!files_under(&in_instance).is_empty());
|
|
let listing = list(&base, "SECRET_TOKEN").await?;
|
|
assert!(listing.get("fallback").is_none(), "{listing}");
|
|
assert_eq!(listing["backup_generation"], same["backup_generation"]);
|
|
assert_eq!(listing["sessions"][0]["id"], "s1");
|
|
|
|
// Back to no storage of its own, the fallback answers; a plan switched to Pro while the
|
|
// instance store stays loaded stops it at once, for the listing and the push alike.
|
|
let resp = authed(
|
|
client().post(format!("{base}/workspaces/edit_large_file_storage_config")),
|
|
"SECRET_TOKEN",
|
|
)
|
|
.json(&json!({ "large_file_storage": null }))
|
|
.send()
|
|
.await?;
|
|
assert_eq!(resp.status(), 200, "{}", resp.text().await?);
|
|
assert_eq!(list(&base, "SECRET_TOKEN").await?["fallback"], true);
|
|
let _enterprise_again = ResetLicensePlan;
|
|
windmill_common::ee::LICENSE_KEY_ID.store(std::sync::Arc::new("test_pro".to_string()));
|
|
assert_eq!(list(&base, "SECRET_TOKEN").await?["enabled"], false);
|
|
let resp = push_whole().await?;
|
|
assert_eq!(resp.status(), 200, "{}", resp.text().await?);
|
|
let pushed: Value = resp.json().await?;
|
|
assert_eq!(pushed["enabled"], false, "{pushed}");
|
|
|
|
Ok(())
|
|
}
|