mirror of
https://github.com/openai/codex.git
synced 2026-06-02 03:11:59 +00:00
Merge branch 'goal-mode-1-state' into goal-mode-2-app-server
This commit is contained in:
@@ -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
|
||||
);
|
||||
|
||||
@@ -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)?,
|
||||
})
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user