Merge branch 'goal-mode-1-state' into goal-mode-2-app-server

This commit is contained in:
Eric Traut
2026-04-16 11:07:42 -07:00
3 changed files with 29 additions and 2 deletions

View File

@@ -4,6 +4,7 @@ CREATE TABLE thread_goals (
status TEXT NOT NULL CHECK(status IN ('active', 'paused', 'budget_limited', 'complete')),
token_budget INTEGER,
tokens_used INTEGER NOT NULL DEFAULT 0,
time_used_seconds INTEGER NOT NULL DEFAULT 0,
created_at_ms INTEGER NOT NULL,
updated_at_ms INTEGER NOT NULL
);

View File

@@ -56,6 +56,7 @@ pub struct ThreadGoal {
pub status: ThreadGoalStatus,
pub token_budget: Option<i64>,
pub tokens_used: i64,
pub time_used_seconds: i64,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
@@ -66,6 +67,7 @@ pub(crate) struct ThreadGoalRow {
pub status: String,
pub token_budget: Option<i64>,
pub tokens_used: i64,
pub time_used_seconds: i64,
pub created_at_ms: i64,
pub updated_at_ms: i64,
}
@@ -78,6 +80,7 @@ impl ThreadGoalRow {
status: row.try_get("status")?,
token_budget: row.try_get("token_budget")?,
tokens_used: row.try_get("tokens_used")?,
time_used_seconds: row.try_get("time_used_seconds")?,
created_at_ms: row.try_get("created_at_ms")?,
updated_at_ms: row.try_get("updated_at_ms")?,
})
@@ -94,6 +97,7 @@ impl TryFrom<ThreadGoalRow> for ThreadGoal {
status: ThreadGoalStatus::try_from(row.status.as_str())?,
token_budget: row.token_budget,
tokens_used: row.tokens_used,
time_used_seconds: row.time_used_seconds,
created_at: epoch_millis_to_datetime(row.created_at_ms)?,
updated_at: epoch_millis_to_datetime(row.updated_at_ms)?,
})

View File

@@ -29,6 +29,7 @@ SELECT
status,
token_budget,
tokens_used,
time_used_seconds,
created_at_ms,
updated_at_ms
FROM thread_goals
@@ -59,14 +60,16 @@ INSERT INTO thread_goals (
status,
token_budget,
tokens_used,
time_used_seconds,
created_at_ms,
updated_at_ms
) VALUES (?, ?, ?, ?, 0, ?, ?)
) VALUES (?, ?, ?, ?, 0, 0, ?, ?)
ON CONFLICT(thread_id) DO UPDATE SET
objective = excluded.objective,
status = excluded.status,
token_budget = excluded.token_budget,
tokens_used = 0,
time_used_seconds = 0,
created_at_ms = excluded.created_at_ms,
updated_at_ms = excluded.updated_at_ms
"#,
@@ -155,10 +158,13 @@ WHERE thread_id = ?
pub async fn account_thread_goal_usage(
&self,
thread_id: ThreadId,
time_delta_seconds: i64,
token_delta: i64,
mode: ThreadGoalAccountingMode,
) -> anyhow::Result<ThreadGoalAccountingOutcome> {
if token_delta <= 0 {
let time_delta_seconds = time_delta_seconds.max(0);
let token_delta = token_delta.max(0);
if time_delta_seconds == 0 && token_delta == 0 {
return Ok(ThreadGoalAccountingOutcome::Unchanged(
self.get_thread_goal(thread_id).await?,
));
@@ -173,6 +179,7 @@ WHERE thread_id = ?
r#"
UPDATE thread_goals
SET
time_used_seconds = time_used_seconds + ?,
tokens_used = tokens_used + ?,
status = CASE
WHEN status = 'active' AND token_budget IS NOT NULL AND tokens_used + ? >= token_budget
@@ -186,6 +193,7 @@ WHERE thread_id = ?
);
let result = sqlx::query(&query)
.bind(time_delta_seconds)
.bind(token_delta)
.bind(token_delta)
.bind(crate::ThreadGoalStatus::BudgetLimited.as_str())
@@ -275,6 +283,7 @@ mod tests {
assert_eq!(crate::ThreadGoalStatus::Active, replaced.status);
assert_eq!(None, replaced.token_budget);
assert_eq!(0, replaced.tokens_used);
assert_eq!(0, replaced.time_used_seconds);
}
#[tokio::test]
@@ -335,6 +344,7 @@ mod tests {
let outcome = runtime
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 7,
/*token_delta*/ 5,
ThreadGoalAccountingMode::ActiveOnly,
)
@@ -345,10 +355,12 @@ mod tests {
};
assert_eq!(crate::ThreadGoalStatus::Active, goal.status);
assert_eq!(5, goal.tokens_used);
assert_eq!(7, goal.time_used_seconds);
let outcome = runtime
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 3,
/*token_delta*/ 15,
ThreadGoalAccountingMode::ActiveOnly,
)
@@ -359,10 +371,12 @@ mod tests {
};
assert_eq!(crate::ThreadGoalStatus::BudgetLimited, goal.status);
assert_eq!(20, goal.tokens_used);
assert_eq!(10, goal.time_used_seconds);
let outcome = runtime
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 5,
/*token_delta*/ 5,
ThreadGoalAccountingMode::ActiveOnly,
)
@@ -373,6 +387,7 @@ mod tests {
};
assert_eq!(crate::ThreadGoalStatus::BudgetLimited, goal.status);
assert_eq!(20, goal.tokens_used);
assert_eq!(10, goal.time_used_seconds);
}
#[tokio::test]
@@ -392,6 +407,7 @@ mod tests {
let active_only = runtime
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 30,
/*token_delta*/ 200,
ThreadGoalAccountingMode::ActiveOnly,
)
@@ -402,10 +418,12 @@ mod tests {
};
assert_eq!(crate::ThreadGoalStatus::Complete, goal.status);
assert_eq!(0, goal.tokens_used);
assert_eq!(0, goal.time_used_seconds);
let completing_turn = runtime
.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 30,
/*token_delta*/ 200,
ThreadGoalAccountingMode::ActiveOrComplete,
)
@@ -416,6 +434,7 @@ mod tests {
};
assert_eq!(crate::ThreadGoalStatus::Complete, goal.status);
assert_eq!(200, goal.tokens_used);
assert_eq!(30, goal.time_used_seconds);
}
#[tokio::test]
@@ -434,11 +453,13 @@ mod tests {
let first = runtime.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 4,
/*token_delta*/ 40,
ThreadGoalAccountingMode::ActiveOnly,
);
let second = runtime.account_thread_goal_usage(
thread_id,
/*time_delta_seconds*/ 6,
/*token_delta*/ 60,
ThreadGoalAccountingMode::ActiveOnly,
);
@@ -452,5 +473,6 @@ mod tests {
.expect("goal read should succeed")
.expect("goal should exist");
assert_eq!(100, goal.tokens_used);
assert_eq!(10, goal.time_used_seconds);
}
}