chore: isolate thread goal storage behind GoalStore (#23295)

## Why

Thread goal persistence is being prepared for a dedicated storage
boundary. Before that split, goal-specific reads, writes, accounting,
and cleanup were exposed directly on `StateRuntime`, so core and
app-server callsites stayed coupled to the full runtime instead of a
goal-specific store.

This PR introduces that boundary without changing the goal wire API or
current persistence behavior. Callers now go through
`StateRuntime::thread_goals()` and the new `GoalStore`, while
`GoalStore` still uses the existing state DB pool underneath.

## What changed

- Added `GoalStore` in `state/src/runtime/goals.rs` and exposed it from
`StateRuntime` via `thread_goals()`.
- Moved thread-goal reads, writes, status updates, pause, delete, and
usage accounting onto `GoalStore`.
- Updated core session goal handling, app-server goal RPCs, resume
snapshots, and goal tests to use the store boundary.
- Kept thread deletion responsible for cascading goal cleanup by
deleting the goal through the store only after a thread row is removed.

## Testing

- Existing goal persistence, resume, and accounting tests were updated
to exercise the new `GoalStore` access path.
This commit is contained in:
jif-oai authored and GitHub committed 2026-05-18 14:47:05 +02:00
1 parent 6a8173588c
commit 7ee7fe239f
10 files changed
+203 -43

No files matched your search

+127 -33
View File
@@ -1,6 +1,18 @@
use super::*;
use crate::model::ThreadGoalRow;
use uuid::Uuid;
#[derive(Clone)]
pub struct GoalStore {
pool: Arc<SqlitePool>,
}
impl GoalStore {
pub(crate) fn new(pool: Arc<SqlitePool>) -> Self {
Self { pool }
}
}
pub struct ThreadGoalUpdate {
pub objective: Option<String>,
pub status: Option<crate::ThreadGoalStatus>,
@@ -21,7 +33,7 @@ pub enum ThreadGoalAccountingMode {
ActiveOrStopped,
}
impl StateRuntime {
impl GoalStore {
pub async fn get_thread_goal(
&self,
thread_id: ThreadId,
@@ -308,29 +320,6 @@ WHERE thread_id = ?
self.get_thread_goal(thread_id).await
}
async fn set_thread_preview_if_empty(
&self,
thread_id: ThreadId,
preview: &str,
) -> anyhow::Result<()> {
let preview = preview.trim();
if preview.is_empty() {
return Ok(());
}
sqlx::query(
r#"
UPDATE threads
SET preview = ?
WHERE id = ? AND preview = ''
"#,
)
.bind(preview)
.bind(thread_id.to_string())
.execute(self.pool.as_ref())
.await?;
Ok(())
}
pub async fn pause_active_thread_goal(
&self,
thread_id: ThreadId,
@@ -463,6 +452,29 @@ RETURNING
let updated = thread_goal_from_row(&row)?;
Ok(ThreadGoalAccountingOutcome::Updated(updated))
}
async fn set_thread_preview_if_empty(
&self,
thread_id: ThreadId,
preview: &str,
) -> anyhow::Result<()> {
let preview = preview.trim();
if preview.is_empty() {
return Ok(());
}
sqlx::query(
r#"
UPDATE threads
SET preview = ?
WHERE id = ? AND preview = ''
"#,
)
.bind(preview)
.bind(thread_id.to_string())
.execute(self.pool.as_ref())
.await?;
Ok(())
}
}
fn thread_goal_from_row(row: &sqlx::sqlite::SqliteRow) -> anyhow::Result<crate::ThreadGoal> {
@@ -519,6 +531,7 @@ mod tests {
upsert_test_thread(&runtime, thread_id).await;
let goal = runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"optimize the benchmark",
@@ -529,7 +542,11 @@ mod tests {
.expect("goal replacement should succeed");
assert_eq!(
Some(goal.clone()),
runtime.get_thread_goal(thread_id).await.unwrap()
runtime
.thread_goals()
.get_thread_goal(thread_id)
.await
.unwrap()
);
let metadata = runtime
.get_thread(thread_id)
@@ -539,6 +556,7 @@ mod tests {
assert_eq!(metadata.preview.as_deref(), Some("hello"));
let updated = runtime
.thread_goals()
.update_thread_goal(
thread_id,
ThreadGoalUpdate {
@@ -560,6 +578,7 @@ mod tests {
assert_eq!(expected, updated);
let replaced = runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"ship the new result",
@@ -574,9 +593,28 @@ mod tests {
assert_eq!(0, replaced.tokens_used);
assert_eq!(0, replaced.time_used_seconds);
assert!(runtime.delete_thread_goal(thread_id).await.unwrap());
assert_eq!(None, runtime.get_thread_goal(thread_id).await.unwrap());
assert!(!runtime.delete_thread_goal(thread_id).await.unwrap());
assert!(
runtime
.thread_goals()
.delete_thread_goal(thread_id)
.await
.unwrap()
);
assert_eq!(
None,
runtime
.thread_goals()
.get_thread_goal(thread_id)
.await
.unwrap()
);
assert!(
!runtime
.thread_goals()
.delete_thread_goal(thread_id)
.await
.unwrap()
);
}
#[tokio::test]
@@ -596,6 +634,7 @@ mod tests {
.expect("test thread should be upserted");
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"optimize the benchmark",
@@ -621,6 +660,7 @@ mod tests {
upsert_test_thread(&runtime, thread_id).await;
let replaced = runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"stay within budget",
@@ -643,6 +683,7 @@ mod tests {
upsert_test_thread(&runtime, thread_id).await;
let inserted = runtime
.thread_goals()
.insert_thread_goal(
thread_id,
"optimize the benchmark",
@@ -654,6 +695,7 @@ mod tests {
.expect("goal should be inserted");
let duplicate = runtime
.thread_goals()
.insert_thread_goal(
thread_id,
"replace the benchmark",
@@ -666,7 +708,11 @@ mod tests {
assert_eq!(None, duplicate);
assert_eq!(
Some(inserted),
runtime.get_thread_goal(thread_id).await.unwrap()
runtime
.thread_goals()
.get_thread_goal(thread_id)
.await
.unwrap()
);
}
@@ -677,6 +723,7 @@ mod tests {
upsert_test_thread(&runtime, thread_id).await;
let inserted = runtime
.thread_goals()
.insert_thread_goal(
thread_id,
"stay within budget",
@@ -700,6 +747,7 @@ mod tests {
upsert_test_thread(&runtime, thread_id).await;
let original = runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"old objective",
@@ -709,6 +757,7 @@ mod tests {
.await
.expect("goal replacement should succeed");
let replacement = runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"new objective",
@@ -719,6 +768,7 @@ mod tests {
.expect("goal replacement should succeed");
let stale_update = runtime
.thread_goals()
.update_thread_goal(
thread_id,
ThreadGoalUpdate {
@@ -735,12 +785,14 @@ mod tests {
assert_eq!(
Some(replacement.clone()),
runtime
.thread_goals()
.get_thread_goal(thread_id)
.await
.expect("goal read should succeed")
);
let fresh_update = runtime
.thread_goals()
.update_thread_goal(
thread_id,
ThreadGoalUpdate {
@@ -763,6 +815,7 @@ mod tests {
upsert_test_thread(&runtime, thread_id).await;
let original = runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"old objective",
@@ -772,6 +825,7 @@ mod tests {
.await
.expect("goal replacement should succeed");
let replacement = runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"new objective",
@@ -782,6 +836,7 @@ mod tests {
.expect("goal replacement should succeed");
let outcome = runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 5,
@@ -809,6 +864,7 @@ mod tests {
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"draft the report",
@@ -818,6 +874,7 @@ mod tests {
.await
.expect("goal replacement should succeed");
let outcome = runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 12,
@@ -832,6 +889,7 @@ mod tests {
};
let updated = runtime
.thread_goals()
.update_thread_goal(
thread_id,
ThreadGoalUpdate {
@@ -860,6 +918,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"optimize the benchmark",
@@ -869,7 +928,7 @@ mod tests {
.await
.expect("goal replacement should succeed");
let status_update = runtime.update_thread_goal(
let status_update = runtime.thread_goals().update_thread_goal(
thread_id,
ThreadGoalUpdate {
objective: None,
@@ -878,7 +937,7 @@ mod tests {
expected_goal_id: None,
},
);
let budget_update = runtime.update_thread_goal(
let budget_update = runtime.thread_goals().update_thread_goal(
thread_id,
ThreadGoalUpdate {
objective: None,
@@ -892,6 +951,7 @@ mod tests {
budget_update.expect("budget update should succeed");
let goal = runtime
.thread_goals()
.get_thread_goal(thread_id)
.await
.expect("goal read should succeed")
@@ -906,6 +966,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
let goal = runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"optimize the benchmark",
@@ -916,6 +977,7 @@ mod tests {
.expect("goal replacement should succeed");
let paused = runtime
.thread_goals()
.pause_active_thread_goal(thread_id)
.await
.expect("active pause should succeed")
@@ -928,6 +990,7 @@ mod tests {
assert_eq!(expected, paused);
let complete = runtime
.thread_goals()
.update_thread_goal(
thread_id,
ThreadGoalUpdate {
@@ -941,6 +1004,7 @@ mod tests {
.expect("goal update should succeed")
.expect("goal should exist");
let pause_result = runtime
.thread_goals()
.pause_active_thread_goal(thread_id)
.await
.expect("terminal pause attempt should succeed");
@@ -948,6 +1012,7 @@ mod tests {
assert_eq!(
Some(complete),
runtime
.thread_goals()
.get_thread_goal(thread_id)
.await
.expect("goal read should succeed")
@@ -960,6 +1025,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"stay within budget",
@@ -970,6 +1036,7 @@ mod tests {
.expect("goal replacement should succeed");
let outcome = runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 7,
@@ -987,6 +1054,7 @@ mod tests {
assert_eq!(7, goal.time_used_seconds);
let outcome = runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 3,
@@ -1004,6 +1072,7 @@ mod tests {
assert_eq!(10, goal.time_used_seconds);
let outcome = runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 5,
@@ -1027,6 +1096,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"stay stopped",
@@ -1037,6 +1107,7 @@ mod tests {
.expect("goal replacement should succeed");
let outcome = runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 5,
@@ -1060,6 +1131,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"stop before overrun",
@@ -1069,6 +1141,7 @@ mod tests {
.await
.expect("goal replacement should succeed");
runtime
.thread_goals()
.update_thread_goal(
thread_id,
crate::ThreadGoalUpdate {
@@ -1082,6 +1155,7 @@ mod tests {
.expect("goal update should succeed");
let outcome = runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 3,
@@ -1105,6 +1179,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"stay within budget",
@@ -1114,6 +1189,7 @@ mod tests {
.await
.expect("goal replacement should succeed");
runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 1,
@@ -1125,6 +1201,7 @@ mod tests {
.expect("usage accounting should succeed");
let lowered = runtime
.thread_goals()
.update_thread_goal(
thread_id,
ThreadGoalUpdate {
@@ -1149,6 +1226,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"stay within budget",
@@ -1158,6 +1236,7 @@ mod tests {
.await
.expect("goal replacement should succeed");
runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 1,
@@ -1169,6 +1248,7 @@ mod tests {
.expect("usage accounting should succeed");
let reactivated = runtime
.thread_goals()
.update_thread_goal(
thread_id,
ThreadGoalUpdate {
@@ -1197,6 +1277,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"stay within budget",
@@ -1206,6 +1287,7 @@ mod tests {
.await
.expect("goal replacement should succeed");
runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 1,
@@ -1217,6 +1299,7 @@ mod tests {
.expect("usage accounting should succeed");
let paused = runtime
.thread_goals()
.update_thread_goal(
thread_id,
ThreadGoalUpdate {
@@ -1241,6 +1324,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"finish the report",
@@ -1251,6 +1335,7 @@ mod tests {
.expect("goal replacement should succeed");
let active_only = runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 30,
@@ -1268,6 +1353,7 @@ mod tests {
assert_eq!(0, goal.time_used_seconds);
let completing_turn = runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 30,
@@ -1291,6 +1377,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"finish the report",
@@ -1300,6 +1387,7 @@ mod tests {
.await
.expect("goal replacement should succeed");
runtime
.thread_goals()
.update_thread_goal(
thread_id,
ThreadGoalUpdate {
@@ -1314,6 +1402,7 @@ mod tests {
.expect("goal should exist");
let active_only = runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 30,
@@ -1331,6 +1420,7 @@ mod tests {
assert_eq!(0, goal.time_used_seconds);
let in_flight_turn = runtime
.thread_goals()
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 30,
@@ -1354,6 +1444,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"count every token",
@@ -1363,14 +1454,14 @@ mod tests {
.await
.expect("goal replacement should succeed");
let first = runtime.account_thread_goal_usage(
let first = runtime.thread_goals().account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 4,
/*token_delta*/ 40,
ThreadGoalAccountingMode::ActiveOnly,
/*expected_goal_id*/ None,
);
let second = runtime.account_thread_goal_usage(
let second = runtime.thread_goals().account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 6,
/*token_delta*/ 60,
@@ -1382,6 +1473,7 @@ mod tests {
second.expect("second usage accounting should succeed");
let goal = runtime
.thread_goals()
.get_thread_goal(thread_id)
.await
.expect("goal read should succeed")
@@ -1396,6 +1488,7 @@ mod tests {
let thread_id = test_thread_id();
upsert_test_thread(&runtime, thread_id).await;
runtime
.thread_goals()
.replace_thread_goal(
thread_id,
"clean up with the thread",
@@ -1413,6 +1506,7 @@ mod tests {
assert_eq!(
None,
runtime
.thread_goals()
.get_thread_goal(thread_id)
.await
.expect("goal read should succeed")