use rusqlite::params;
use crate::db::{
count_retryable_extraction_replay_ranges, list_extraction_replay_ranges,
quarantine_extraction_replay_ranges, record_captured_event, retry_extraction_replay_ranges,
CaptureEventInput,
};
use super::*;
fn setup_conn() -> Connection {
let conn = Connection::open_in_memory().expect("in-memory db should open");
crate::migrate::run_migrations(&conn).expect("migrations should run");
conn
}
fn insert_task(conn: &Connection, session_id: &str, task_kind: ExtractionTaskKind) -> Result<i64> {
let outcome = record_captured_event(
conn,
&CaptureEventInput {
host: "codex-cli",
session_id,
project: "/tmp/remem",
cwd: None,
event_type: "tool_result",
role: None,
tool_name: Some("Task"),
content: session_id,
task_kind: Some(task_kind),
},
)?;
outcome
.extraction_task_id
.ok_or_else(|| anyhow::anyhow!("expected extraction task id"))
}
fn task_status(conn: &Connection, task_id: i64) -> (String, i64, Option<i64>, Option<String>) {
conn.query_row(
"SELECT status, attempts, next_retry_epoch, last_error
FROM extraction_tasks WHERE id = ?1",
params![task_id],
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)),
)
.expect("task state should query")
}
#[test]
fn claim_next_extraction_task_orders_by_priority_and_age() {
let mut conn = setup_conn();
let observation_id = insert_task(
&conn,
"sess-observation",
ExtractionTaskKind::ObservationExtract,
)
.expect("observation task should insert");
let rollup_id = insert_task(&conn, "sess-rollup", ExtractionTaskKind::SessionRollup)
.expect("rollup task should insert");
let claimed = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
assert_eq!(claimed.id, rollup_id);
assert_eq!(claimed.task_kind, ExtractionTaskKind::SessionRollup);
assert_eq!(claimed.host, "codex-cli");
assert_eq!(claimed.session_id.as_deref(), Some("sess-rollup"));
let status = task_status(&conn, observation_id).0;
assert_eq!(status, "pending");
}
#[test]
fn claim_next_extraction_task_preserves_ai_profile_from_capture_payload() -> Result<()> {
let mut conn = setup_conn();
record_captured_event(
&conn,
&CaptureEventInput {
host: "codex-cli",
session_id: "sess-profile",
project: "/tmp/remem",
cwd: None,
event_type: "session_stop",
role: None,
tool_name: None,
content: r#"{"session_id":"sess-profile","remem_ai_profile":"custom"}"#,
task_kind: Some(ExtractionTaskKind::SessionRollup),
},
)?;
let claimed = claim_next_extraction_task(&mut conn, "worker-a", 60)?
.ok_or_else(|| anyhow::anyhow!("task should exist"))?;
assert_eq!(claimed.ai_profile.as_deref(), Some("custom"));
Ok(())
}
#[test]
fn claim_next_extraction_task_reads_ai_profile_from_large_capture_blob() -> Result<()> {
let mut conn = setup_conn();
let content = format!(
r#"{{"session_id":"sess-large-profile","prefix":"{}","remem_ai_profile":"large-custom","suffix":"{}"}}"#,
"a".repeat(10 * 1024),
"b".repeat(10 * 1024)
);
record_captured_event(
&conn,
&CaptureEventInput {
host: "codex-cli",
session_id: "sess-large-profile",
project: "/tmp/remem",
cwd: None,
event_type: "session_stop",
role: None,
tool_name: None,
content: &content,
task_kind: Some(ExtractionTaskKind::SessionRollup),
},
)?;
let claimed = claim_next_extraction_task(&mut conn, "worker-a", 60)?
.ok_or_else(|| anyhow::anyhow!("task should exist"))?;
assert_eq!(claimed.ai_profile.as_deref(), Some("large-custom"));
Ok(())
}
#[test]
fn claim_next_extraction_task_does_not_double_claim_active_task() {
let mut conn = setup_conn();
let task_id = insert_task(&conn, "sess-single", ExtractionTaskKind::SessionRollup)
.expect("task should insert");
let first = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("first claim should succeed")
.expect("first claim should return task");
let second =
claim_next_extraction_task(&mut conn, "worker-b", 60).expect("second claim should succeed");
assert_eq!(first.id, task_id);
assert!(second.is_none());
}
#[test]
fn release_expired_extraction_task_leases_requeues_only_expired_tasks() {
let mut conn = setup_conn();
let expired_id = insert_task(&conn, "sess-expired", ExtractionTaskKind::SessionRollup)
.expect("expired task should insert");
let fresh_id = insert_task(&conn, "sess-fresh", ExtractionTaskKind::ObservationExtract)
.expect("fresh task should insert");
claim_next_extraction_task(&mut conn, "worker-expired", 60)
.expect("expired worker claim should succeed")
.expect("expired task should be claimed");
claim_next_extraction_task(&mut conn, "worker-fresh", 60)
.expect("fresh worker claim should succeed")
.expect("fresh task should be claimed");
let now = chrono::Utc::now().timestamp();
conn.execute(
"UPDATE extraction_tasks
SET lease_expires_epoch = ?1
WHERE id = ?2",
params![now - 1, expired_id],
)
.expect("expired lease should update");
let released = release_expired_extraction_task_leases(&conn).expect("release should succeed");
assert_eq!(released, 1);
assert_eq!(task_status(&conn, expired_id).0, "pending");
assert_eq!(task_status(&conn, fresh_id).0, "processing");
}
#[test]
fn mark_extraction_task_done_clears_lease_and_advances_cursor() {
let mut conn = setup_conn();
let task_id = insert_task(&conn, "sess-done", ExtractionTaskKind::SessionRollup)
.expect("task should insert");
let task = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
mark_extraction_task_done(&conn, task.id, "worker-a", task.high_watermark_event_id)
.expect("done should succeed");
let (status, lease_owner, cursor, high_watermark): (
String,
Option<String>,
Option<i64>,
Option<i64>,
) = conn
.query_row(
"SELECT status, lease_owner, cursor_event_id, high_watermark_event_id
FROM extraction_tasks WHERE id = ?1",
params![task_id],
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)),
)
.expect("task should query");
assert_eq!(status, "done");
assert!(lease_owner.is_none());
assert_eq!(cursor, high_watermark);
}
#[test]
fn mark_extraction_task_done_requeues_when_watermark_advanced_after_claim() {
let mut conn = setup_conn();
let task_id = insert_task(&conn, "sess-coalesce", ExtractionTaskKind::SessionRollup)
.expect("task should insert");
let task = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
let claimed_high_watermark = task.high_watermark_event_id;
insert_task(&conn, "sess-coalesce", ExtractionTaskKind::SessionRollup)
.expect("coalesced task should update high watermark");
mark_extraction_task_done(&conn, task.id, "worker-a", claimed_high_watermark)
.expect("done should succeed");
let (status, lease_owner, cursor, high_watermark): (
String,
Option<String>,
Option<i64>,
Option<i64>,
) = conn
.query_row(
"SELECT status, lease_owner, cursor_event_id, high_watermark_event_id
FROM extraction_tasks WHERE id = ?1",
params![task_id],
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)),
)
.expect("task should query");
assert_eq!(status, "pending");
assert!(lease_owner.is_none());
assert_eq!(cursor, claimed_high_watermark);
assert!(high_watermark > cursor);
}
#[test]
fn mark_extraction_task_done_with_partial_progress_requeues_remaining_range() {
let mut conn = setup_conn();
let task_id = insert_task(&conn, "sess-partial", ExtractionTaskKind::MemoryCandidate)
.expect("task should insert");
insert_task(&conn, "sess-partial", ExtractionTaskKind::MemoryCandidate)
.expect("second event should coalesce and advance the watermark");
let task = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
assert_eq!(task.id, task_id);
let watermark = task
.high_watermark_event_id
.expect("claimed task should have a watermark");
let partial_event_id: i64 = conn
.query_row(
"SELECT MIN(id) FROM captured_events WHERE session_id = 'sess-partial'",
[],
|row| row.get(0),
)
.expect("first captured event id should query");
assert!(
partial_event_id < watermark,
"fixture must produce partial progress: to_event_id={partial_event_id} watermark={watermark}"
);
mark_extraction_task_done(&conn, task.id, "worker-a", Some(partial_event_id))
.expect("done with partial progress should succeed");
let (status, lease_owner, cursor): (String, Option<String>, Option<i64>) = conn
.query_row(
"SELECT status, lease_owner, cursor_event_id FROM extraction_tasks WHERE id = ?1",
params![task_id],
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
)
.expect("task should query");
assert_eq!(status, "pending", "partial progress must requeue the task");
assert!(lease_owner.is_none());
assert_eq!(cursor, Some(partial_event_id));
let reclaimed = claim_next_extraction_task(&mut conn, "worker-b", 60)
.expect("claim should succeed")
.expect("requeued task should be claimable");
assert_eq!(reclaimed.id, task_id);
assert_eq!(
reclaimed.cursor_event_id,
Some(partial_event_id),
"next round must resume from the written cursor"
);
assert_eq!(
reclaimed.high_watermark_event_id,
Some(watermark),
"next round must still cover the remaining (cursor, watermark] range"
);
}
#[test]
fn mark_extraction_task_failed_or_retry_keeps_retryable_task_visible() {
let mut conn = setup_conn();
let task_id = insert_task(&conn, "sess-retry", ExtractionTaskKind::SessionRollup)
.expect("task should insert");
claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
mark_extraction_task_failed_or_retry(&conn, task_id, "worker-a", "temporary", 30)
.expect("retry should succeed");
let (status, attempts, next_retry, last_error) = task_status(&conn, task_id);
assert_eq!(status, "pending");
assert_eq!(attempts, 1);
assert!(next_retry.is_some());
assert_eq!(last_error.as_deref(), Some("temporary"));
}
#[test]
fn mark_extraction_task_failed_or_retry_exhausts_after_max_attempts() {
let mut conn = setup_conn();
let task_id = insert_task(&conn, "sess-failed", ExtractionTaskKind::SessionRollup)
.expect("task should insert");
conn.execute(
"UPDATE extraction_tasks SET attempts = ?1 WHERE id = ?2",
params![EXTRACTION_TASK_MAX_ATTEMPTS - 1, task_id],
)
.expect("attempt count should update");
claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
mark_extraction_task_failed_or_retry(&conn, task_id, "worker-a", "exhausted", 30)
.expect("failure should succeed");
let (status, attempts, next_retry, last_error) = task_status(&conn, task_id);
assert_eq!(status, "failed");
assert_eq!(attempts, EXTRACTION_TASK_MAX_ATTEMPTS);
assert!(next_retry.is_none());
assert_eq!(last_error.as_deref(), Some("exhausted"));
}
#[test]
fn mark_extraction_task_failed_records_permanent_failure_without_retry() {
let mut conn = setup_conn();
let task_id = insert_task(&conn, "sess-permanent", ExtractionTaskKind::SessionRollup)
.expect("task should insert");
claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
mark_extraction_task_failed(&conn, task_id, "worker-a", "not implemented")
.expect("permanent failure should succeed");
let (status, attempts, next_retry, last_error) = task_status(&conn, task_id);
assert_eq!(status, "failed");
assert_eq!(attempts, 1);
assert!(next_retry.is_none());
assert_eq!(last_error.as_deref(), Some("not implemented"));
}
#[test]
fn defer_extraction_task_requeues_and_increments_attempts() {
let mut conn = setup_conn();
let task_id = insert_task(&conn, "sess-defer", ExtractionTaskKind::SessionRollup)
.expect("task should insert");
claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
defer_extraction_task(&conn, task_id, "worker-a", "not implemented", 30)
.expect("defer should succeed");
let (status, attempts, next_retry, last_error) = task_status(&conn, task_id);
assert_eq!(status, "pending");
assert_eq!(attempts, 1);
assert!(next_retry.is_some());
assert_eq!(last_error.as_deref(), Some("not implemented"));
}
#[test]
fn defer_exhaustion_does_not_permanently_stall_session_extraction() {
let mut conn = setup_conn();
let task_id = insert_task(&conn, "sess-defer-stall", ExtractionTaskKind::SessionRollup)
.expect("task should insert");
conn.execute(
"UPDATE extraction_tasks SET attempts = ?1 WHERE id = ?2",
params![EXTRACTION_TASK_MAX_ATTEMPTS - 1, task_id],
)
.expect("attempt count should update");
let stuck = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
let stuck_watermark = stuck.high_watermark_event_id;
defer_extraction_task(&conn, task_id, "worker-a", "still ambiguous", 30)
.expect("defer should exhaust");
let ranges = list_extraction_replay_ranges(&conn, None, 10).expect("ranges should list");
assert_eq!(ranges.len(), 1);
assert_eq!(ranges[0].source_task_id, task_id);
assert_eq!(ranges[0].status, "pending");
assert_eq!(Some(ranges[0].to_event_id), stuck_watermark);
insert_task(&conn, "sess-defer-stall", ExtractionTaskKind::SessionRollup)
.expect("new captured event should coalesce into the task");
let reclaimed = claim_next_extraction_task(&mut conn, "worker-b", 60)
.expect("claim should succeed")
.expect("session should stay extractable after defer exhaustion");
assert_eq!(reclaimed.id, task_id);
assert_eq!(
reclaimed.cursor_event_id, stuck_watermark,
"cursor should advance past the exhausted range so new events are not blocked by the stuck region"
);
assert!(
reclaimed.attempts < EXTRACTION_TASK_MAX_ATTEMPTS,
"resurrected task needs retry budget, attempts={} fails terminally on the next defer",
reclaimed.attempts
);
}
#[test]
fn defer_extraction_task_exhausts_after_max_attempts() {
let mut conn = setup_conn();
let task_id = insert_task(
&conn,
"sess-defer-exhaust",
ExtractionTaskKind::SessionRollup,
)
.expect("task should insert");
conn.execute(
"UPDATE extraction_tasks SET attempts = ?1 WHERE id = ?2",
params![EXTRACTION_TASK_MAX_ATTEMPTS - 1, task_id],
)
.expect("attempt count should update");
claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
defer_extraction_task(&conn, task_id, "worker-a", "still ambiguous", 30)
.expect("defer should exhaust");
let (status, attempts, next_retry, last_error) = task_status(&conn, task_id);
assert_eq!(status, "failed");
assert_eq!(attempts, EXTRACTION_TASK_MAX_ATTEMPTS);
assert!(next_retry.is_none());
assert_eq!(last_error.as_deref(), Some("still ambiguous"));
}
#[test]
fn enqueue_followup_revives_exhausted_task_with_fresh_retry_budget() {
let mut conn = setup_conn();
let task_id = insert_task(
&conn,
"sess-followup-revive",
ExtractionTaskKind::MemoryCandidate,
)
.expect("task should insert");
conn.execute(
"UPDATE extraction_tasks SET attempts = ?1 WHERE id = ?2",
params![EXTRACTION_TASK_MAX_ATTEMPTS - 1, task_id],
)
.expect("attempt count should update");
let claimed = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
let stuck_watermark = claimed
.high_watermark_event_id
.expect("claimed task should have a watermark");
defer_extraction_task(&conn, task_id, "worker-a", "still ambiguous", 30)
.expect("defer should exhaust");
let followup_id = enqueue_followup_extraction_task(
&conn,
&claimed,
ExtractionTaskKind::MemoryCandidate,
stuck_watermark + 1,
)
.expect("followup should coalesce into the exhausted task");
assert_eq!(followup_id, task_id);
let (status, attempts, next_retry, _) = task_status(&conn, task_id);
assert_eq!(status, "pending");
assert_eq!(attempts, 0, "revived task needs a fresh retry budget");
assert!(next_retry.is_none());
}
#[test]
fn exhaustion_records_only_claimed_range_when_watermark_advances() {
let mut conn = setup_conn();
let task_id = insert_task(
&conn,
"sess-claim-watermark",
ExtractionTaskKind::SessionRollup,
)
.expect("task should insert");
conn.execute(
"UPDATE extraction_tasks SET attempts = ?1 WHERE id = ?2",
params![EXTRACTION_TASK_MAX_ATTEMPTS - 1, task_id],
)
.expect("attempt count should update");
let claimed = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
let claimed_watermark = claimed
.high_watermark_event_id
.expect("claimed task should have a watermark");
insert_task(
&conn,
"sess-claim-watermark",
ExtractionTaskKind::SessionRollup,
)
.expect("new captured event should coalesce while the task is leased");
defer_claimed_extraction_task(&conn, &claimed, "worker-a", "still ambiguous", 30)
.expect("defer should exhaust the claimed range");
let ranges = list_extraction_replay_ranges(&conn, None, 10).expect("ranges should list");
assert_eq!(ranges.len(), 1);
assert_eq!(
ranges[0].from_event_id,
claimed.cursor_event_id.unwrap_or(0) + 1
);
assert_eq!(ranges[0].to_event_id, claimed_watermark);
let (status, attempts, next_retry, _) = task_status(&conn, task_id);
assert_eq!(status, "pending");
assert_eq!(attempts, 0);
assert!(next_retry.is_none());
let reclaimed = claim_next_extraction_task(&mut conn, "worker-b", 60)
.expect("claim should succeed")
.expect("later range should be claimable");
assert_eq!(reclaimed.id, task_id);
assert_eq!(reclaimed.cursor_event_id, Some(claimed_watermark));
assert!(
reclaimed.high_watermark_event_id > reclaimed.cursor_event_id,
"later event must remain available to the primary task"
);
}
#[test]
fn retry_extraction_replay_range_creates_bounded_replay_task() {
let mut conn = setup_conn();
let task_id = insert_task(&conn, "sess-replay", ExtractionTaskKind::ObservationExtract)
.expect("task should insert");
conn.execute(
"UPDATE extraction_tasks SET attempts = ?1 WHERE id = ?2",
params![EXTRACTION_TASK_MAX_ATTEMPTS - 1, task_id],
)
.expect("attempt count should update");
let claimed = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
defer_claimed_extraction_task(&conn, &claimed, "worker-a", "bad model output", 30)
.expect("defer should exhaust");
let range = list_extraction_replay_ranges(&conn, None, 10)
.expect("ranges should list")
.pop()
.expect("range should exist");
let retried = retry_extraction_replay_ranges(&conn, None, 10).expect("retry should enqueue");
assert_eq!(retried, 1);
let listed = list_extraction_replay_ranges(&conn, None, 10).expect("ranges should list");
assert_eq!(listed.len(), 1);
assert_eq!(listed[0].status, "requeued");
let replay = claim_next_extraction_task(&mut conn, "worker-b", 60)
.expect("claim should succeed")
.expect("replay task should be claimable");
assert_ne!(replay.id, task_id);
assert_eq!(replay.replay_range_id, Some(range.id));
assert_eq!(replay.cursor_event_id, Some(range.from_event_id - 1));
assert_eq!(replay.high_watermark_event_id, Some(range.to_event_id));
}
#[test]
fn replayed_range_clears_terminal_source_failure() {
let mut conn = setup_conn();
let task_id = insert_task(
&conn,
"sess-replay-source-clear",
ExtractionTaskKind::ObservationExtract,
)
.expect("task should insert");
conn.execute(
"UPDATE extraction_tasks SET attempts = ?1 WHERE id = ?2",
params![EXTRACTION_TASK_MAX_ATTEMPTS - 1, task_id],
)
.expect("attempt count should update");
let claimed = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
defer_claimed_extraction_task(&conn, &claimed, "worker-a", "bad model output", 30)
.expect("defer should exhaust");
assert_eq!(task_status(&conn, task_id).0, "failed");
retry_extraction_replay_ranges(&conn, None, 10).expect("retry should enqueue");
let replay = claim_next_extraction_task(&mut conn, "worker-b", 60)
.expect("claim should succeed")
.expect("replay task should be claimable");
mark_extraction_task_done(&conn, replay.id, "worker-b", replay.high_watermark_event_id)
.expect("replay task should finish");
assert_eq!(
task_status(&conn, task_id).0,
"done",
"successful replay should clear the terminal source task failure"
);
assert!(
list_extraction_replay_ranges(&conn, None, 10)
.expect("ranges should list")
.is_empty(),
"replayed ranges should leave the operational queue"
);
}
#[test]
fn replay_followup_stays_scoped_to_replay_range() {
let mut conn = setup_conn();
let task_id = insert_task(
&conn,
"sess-replay-followup",
ExtractionTaskKind::ObservationExtract,
)
.expect("task should insert");
conn.execute(
"UPDATE extraction_tasks SET attempts = ?1 WHERE id = ?2",
params![EXTRACTION_TASK_MAX_ATTEMPTS - 1, task_id],
)
.expect("attempt count should update");
let claimed = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
defer_claimed_extraction_task(&conn, &claimed, "worker-a", "bad model output", 30)
.expect("defer should exhaust");
retry_extraction_replay_ranges(&conn, None, 10).expect("retry should enqueue");
let replay = claim_next_extraction_task(&mut conn, "worker-b", 60)
.expect("claim should succeed")
.expect("replay task should be claimable");
let followup_id = enqueue_followup_extraction_task(
&conn,
&replay,
ExtractionTaskKind::MemoryCandidate,
replay
.high_watermark_event_id
.expect("replay should have watermark"),
)
.expect("followup should enqueue");
let (cursor, high_watermark, replay_range_id): (Option<i64>, Option<i64>, Option<i64>) = conn
.query_row(
"SELECT cursor_event_id, high_watermark_event_id, replay_range_id
FROM extraction_tasks WHERE id = ?1",
params![followup_id],
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
)
.expect("followup should query");
assert_eq!(cursor, replay.cursor_event_id);
assert_eq!(high_watermark, replay.high_watermark_event_id);
assert_eq!(replay_range_id, replay.replay_range_id);
mark_extraction_task_done(&conn, replay.id, "worker-b", replay.high_watermark_event_id)
.expect("replay task should finish");
let ranges = list_extraction_replay_ranges(&conn, None, 10).expect("ranges should list");
assert_eq!(ranges.len(), 1);
assert_eq!(ranges[0].status, "requeued");
let followup = claim_next_extraction_task(&mut conn, "worker-c", 60)
.expect("claim should succeed")
.expect("followup task should be claimable");
assert_eq!(followup.id, followup_id);
mark_extraction_task_done(
&conn,
followup.id,
"worker-c",
followup.high_watermark_event_id,
)
.expect("followup should finish");
assert!(
list_extraction_replay_ranges(&conn, None, 10)
.expect("ranges should list")
.is_empty(),
"range should disappear from operational list only after the replay chain finishes"
);
}
#[test]
fn replay_range_stays_failed_when_followup_fails_before_parent_done() {
let mut conn = setup_conn();
let task_id = insert_task(
&conn,
"sess-replay-followup-failure",
ExtractionTaskKind::ObservationExtract,
)
.expect("task should insert");
conn.execute(
"UPDATE extraction_tasks SET attempts = ?1 WHERE id = ?2",
params![EXTRACTION_TASK_MAX_ATTEMPTS - 1, task_id],
)
.expect("attempt count should update");
let claimed = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
defer_claimed_extraction_task(&conn, &claimed, "worker-a", "bad model output", 30)
.expect("defer should exhaust");
retry_extraction_replay_ranges(&conn, None, 10).expect("retry should enqueue");
let replay = claim_next_extraction_task(&mut conn, "worker-b", 60)
.expect("claim should succeed")
.expect("replay task should be claimable");
let followup_id = enqueue_followup_extraction_task(
&conn,
&replay,
ExtractionTaskKind::MemoryCandidate,
replay
.high_watermark_event_id
.expect("replay should have watermark"),
)
.expect("followup should enqueue");
let followup = claim_next_extraction_task(&mut conn, "worker-c", 60)
.expect("claim should succeed")
.expect("followup task should be claimable");
assert_eq!(followup.id, followup_id);
mark_extraction_task_failed(&conn, followup.id, "worker-c", "followup failed")
.expect("followup failure should mark replay range failed");
let ranges = list_extraction_replay_ranges(&conn, None, 10).expect("ranges should list");
assert_eq!(ranges.len(), 1);
assert_eq!(ranges[0].status, "failed");
assert_eq!(
count_retryable_extraction_replay_ranges(&conn, None, 10).expect("count should succeed"),
0
);
assert_eq!(
retry_extraction_replay_ranges(&conn, None, 10).expect("retry should skip active siblings"),
0
);
assert_eq!(
quarantine_extraction_replay_ranges(&conn, None, 10)
.expect("quarantine should skip active siblings"),
0
);
mark_extraction_task_done(&conn, replay.id, "worker-b", replay.high_watermark_event_id)
.expect("parent replay task should finish");
assert_eq!(
count_retryable_extraction_replay_ranges(&conn, None, 10).expect("count should succeed"),
1
);
assert_eq!(
quarantine_extraction_replay_ranges(&conn, None, 10).expect("quarantine should succeed"),
1
);
assert_eq!(task_status(&conn, followup_id).0, "done");
assert_eq!(task_status(&conn, task_id).0, "done");
let ranges = list_extraction_replay_ranges(&conn, None, 10).expect("ranges should list");
assert_eq!(ranges.len(), 1);
assert_eq!(ranges[0].status, "quarantined");
}
#[test]
fn quarantine_extraction_replay_ranges_removes_retryable_ranges() {
let mut conn = setup_conn();
let task_id = insert_task(&conn, "sess-quarantine", ExtractionTaskKind::SessionRollup)
.expect("task should insert");
conn.execute(
"UPDATE extraction_tasks SET attempts = ?1 WHERE id = ?2",
params![EXTRACTION_TASK_MAX_ATTEMPTS - 1, task_id],
)
.expect("attempt count should update");
let claimed = claim_next_extraction_task(&mut conn, "worker-a", 60)
.expect("claim should succeed")
.expect("task should be claimed");
defer_claimed_extraction_task(&conn, &claimed, "worker-a", "still ambiguous", 30)
.expect("defer should exhaust");
assert_eq!(task_status(&conn, task_id).0, "failed");
assert_eq!(
count_retryable_extraction_replay_ranges(&conn, None, 10).expect("count should succeed"),
1
);
let quarantined =
quarantine_extraction_replay_ranges(&conn, None, 10).expect("quarantine should succeed");
assert_eq!(quarantined, 1);
assert_eq!(
count_retryable_extraction_replay_ranges(&conn, None, 10).expect("count should succeed"),
0
);
let ranges = list_extraction_replay_ranges(&conn, None, 10).expect("ranges should list");
assert_eq!(ranges.len(), 1);
assert_eq!(ranges[0].status, "quarantined");
assert_eq!(task_status(&conn, task_id).0, "done");
}