diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 37a64fbe33..bf4c68299e 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -3521,6 +3521,52 @@ ], "type": "object" }, + "ThreadGoalBudgetParams": { + "oneOf": [ + { + "properties": { + "token_budget": { + "format": "int64", + "type": "integer" + }, + "type": { + "enum": [ + "tokens" + ], + "title": "TokensThreadGoalBudgetParamsType", + "type": "string" + } + }, + "required": [ + "token_budget", + "type" + ], + "title": "TokensThreadGoalBudgetParams", + "type": "object" + }, + { + "properties": { + "percent": { + "format": "double", + "type": "number" + }, + "type": { + "enum": [ + "fiveHourLimitPercent" + ], + "title": "FiveHourLimitPercentThreadGoalBudgetParamsType", + "type": "string" + } + }, + "required": [ + "percent", + "type" + ], + "title": "FiveHourLimitPercentThreadGoalBudgetParams", + "type": "object" + } + ] + }, "ThreadGoalStatus": { "enum": [ "active", diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index 82914f3a6f..b906256318 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -3045,6 +3045,16 @@ }, "ThreadGoal": { "properties": { + "budget": { + "anyOf": [ + { + "$ref": "#/definitions/ThreadGoalBudget" + }, + { + "type": "null" + } + ] + }, "createdAt": { "format": "int64", "type": "integer" @@ -3089,6 +3099,80 @@ ], "type": "object" }, + "ThreadGoalBudget": { + "oneOf": [ + { + "properties": { + "token_budget": { + "format": "int64", + "type": "integer" + }, + "type": { + "enum": [ + "tokens" + ], + "title": "TokensThreadGoalBudgetType", + "type": "string" + } + }, + "required": [ + "token_budget", + "type" + ], + "title": "TokensThreadGoalBudget", + "type": "object" + }, + { + "properties": { + "baseline_resets_at": { + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "baseline_used_percent": { + "format": "double", + "type": "number" + }, + "latest_resets_at": { + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "latest_used_percent": { + "format": "double", + "type": "number" + }, + "limit_id": { + "type": "string" + }, + "percent": { + "format": "double", + "type": "number" + }, + "type": { + "enum": [ + "fiveHourLimitPercent" + ], + "title": "FiveHourLimitPercentThreadGoalBudgetType", + "type": "string" + } + }, + "required": [ + "baseline_used_percent", + "latest_used_percent", + "limit_id", + "percent", + "type" + ], + "title": "FiveHourLimitPercentThreadGoalBudget", + "type": "object" + } + ] + }, "ThreadGoalClearedNotification": { "properties": { "threadId": { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index f856b43d66..fecbcbb69f 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -15265,6 +15265,16 @@ }, "ThreadGoal": { "properties": { + "budget": { + "anyOf": [ + { + "$ref": "#/definitions/v2/ThreadGoalBudget" + }, + { + "type": "null" + } + ] + }, "createdAt": { "format": "int64", "type": "integer" @@ -15309,6 +15319,126 @@ ], "type": "object" }, + "ThreadGoalBudget": { + "oneOf": [ + { + "properties": { + "token_budget": { + "format": "int64", + "type": "integer" + }, + "type": { + "enum": [ + "tokens" + ], + "title": "TokensThreadGoalBudgetType", + "type": "string" + } + }, + "required": [ + "token_budget", + "type" + ], + "title": "TokensThreadGoalBudget", + "type": "object" + }, + { + "properties": { + "baseline_resets_at": { + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "baseline_used_percent": { + "format": "double", + "type": "number" + }, + "latest_resets_at": { + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "latest_used_percent": { + "format": "double", + "type": "number" + }, + "limit_id": { + "type": "string" + }, + "percent": { + "format": "double", + "type": "number" + }, + "type": { + "enum": [ + "fiveHourLimitPercent" + ], + "title": "FiveHourLimitPercentThreadGoalBudgetType", + "type": "string" + } + }, + "required": [ + "baseline_used_percent", + "latest_used_percent", + "limit_id", + "percent", + "type" + ], + "title": "FiveHourLimitPercentThreadGoalBudget", + "type": "object" + } + ] + }, + "ThreadGoalBudgetParams": { + "oneOf": [ + { + "properties": { + "token_budget": { + "format": "int64", + "type": "integer" + }, + "type": { + "enum": [ + "tokens" + ], + "title": "TokensThreadGoalBudgetParamsType", + "type": "string" + } + }, + "required": [ + "token_budget", + "type" + ], + "title": "TokensThreadGoalBudgetParams", + "type": "object" + }, + { + "properties": { + "percent": { + "format": "double", + "type": "number" + }, + "type": { + "enum": [ + "fiveHourLimitPercent" + ], + "title": "FiveHourLimitPercentThreadGoalBudgetParamsType", + "type": "string" + } + }, + "required": [ + "percent", + "type" + ], + "title": "FiveHourLimitPercentThreadGoalBudgetParams", + "type": "object" + } + ] + }, "ThreadGoalClearedNotification": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index c17efe7a45..93c1f9d532 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -13151,6 +13151,16 @@ }, "ThreadGoal": { "properties": { + "budget": { + "anyOf": [ + { + "$ref": "#/definitions/ThreadGoalBudget" + }, + { + "type": "null" + } + ] + }, "createdAt": { "format": "int64", "type": "integer" @@ -13195,6 +13205,126 @@ ], "type": "object" }, + "ThreadGoalBudget": { + "oneOf": [ + { + "properties": { + "token_budget": { + "format": "int64", + "type": "integer" + }, + "type": { + "enum": [ + "tokens" + ], + "title": "TokensThreadGoalBudgetType", + "type": "string" + } + }, + "required": [ + "token_budget", + "type" + ], + "title": "TokensThreadGoalBudget", + "type": "object" + }, + { + "properties": { + "baseline_resets_at": { + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "baseline_used_percent": { + "format": "double", + "type": "number" + }, + "latest_resets_at": { + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "latest_used_percent": { + "format": "double", + "type": "number" + }, + "limit_id": { + "type": "string" + }, + "percent": { + "format": "double", + "type": "number" + }, + "type": { + "enum": [ + "fiveHourLimitPercent" + ], + "title": "FiveHourLimitPercentThreadGoalBudgetType", + "type": "string" + } + }, + "required": [ + "baseline_used_percent", + "latest_used_percent", + "limit_id", + "percent", + "type" + ], + "title": "FiveHourLimitPercentThreadGoalBudget", + "type": "object" + } + ] + }, + "ThreadGoalBudgetParams": { + "oneOf": [ + { + "properties": { + "token_budget": { + "format": "int64", + "type": "integer" + }, + "type": { + "enum": [ + "tokens" + ], + "title": "TokensThreadGoalBudgetParamsType", + "type": "string" + } + }, + "required": [ + "token_budget", + "type" + ], + "title": "TokensThreadGoalBudgetParams", + "type": "object" + }, + { + "properties": { + "percent": { + "format": "double", + "type": "number" + }, + "type": { + "enum": [ + "fiveHourLimitPercent" + ], + "title": "FiveHourLimitPercentThreadGoalBudgetParamsType", + "type": "string" + } + }, + "required": [ + "percent", + "type" + ], + "title": "FiveHourLimitPercentThreadGoalBudgetParams", + "type": "object" + } + ] + }, "ThreadGoalClearedNotification": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadGoalUpdatedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadGoalUpdatedNotification.json index 52a2e905a2..71414438fb 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadGoalUpdatedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadGoalUpdatedNotification.json @@ -3,6 +3,16 @@ "definitions": { "ThreadGoal": { "properties": { + "budget": { + "anyOf": [ + { + "$ref": "#/definitions/ThreadGoalBudget" + }, + { + "type": "null" + } + ] + }, "createdAt": { "format": "int64", "type": "integer" @@ -47,6 +57,80 @@ ], "type": "object" }, + "ThreadGoalBudget": { + "oneOf": [ + { + "properties": { + "token_budget": { + "format": "int64", + "type": "integer" + }, + "type": { + "enum": [ + "tokens" + ], + "title": "TokensThreadGoalBudgetType", + "type": "string" + } + }, + "required": [ + "token_budget", + "type" + ], + "title": "TokensThreadGoalBudget", + "type": "object" + }, + { + "properties": { + "baseline_resets_at": { + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "baseline_used_percent": { + "format": "double", + "type": "number" + }, + "latest_resets_at": { + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "latest_used_percent": { + "format": "double", + "type": "number" + }, + "limit_id": { + "type": "string" + }, + "percent": { + "format": "double", + "type": "number" + }, + "type": { + "enum": [ + "fiveHourLimitPercent" + ], + "title": "FiveHourLimitPercentThreadGoalBudgetType", + "type": "string" + } + }, + "required": [ + "baseline_used_percent", + "latest_used_percent", + "limit_id", + "percent", + "type" + ], + "title": "FiveHourLimitPercentThreadGoalBudget", + "type": "object" + } + ] + }, "ThreadGoalStatus": { "enum": [ "active", diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadGoal.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadGoal.ts index c68732324f..206b8e37ba 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadGoal.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadGoal.ts @@ -3,4 +3,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { ThreadGoalStatus } from "./ThreadGoalStatus"; -export type ThreadGoal = { threadId: string, objective: string, status: ThreadGoalStatus, tokenBudget: number | null, tokensUsed: number, timeUsedSeconds: number, createdAt: number, updatedAt: number, }; +export type ThreadGoal = { threadId: string, objective: string, status: ThreadGoalStatus, budget: ThreadGoalBudget | null, tokenBudget: number | null, tokensUsed: number, timeUsedSeconds: number, createdAt: number, updatedAt: number, }; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index c5a7d61f01..010f78494f 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -1675,6 +1675,7 @@ mod tests { thread_id: "goal-thread".to_string(), objective: Some("ship it".to_string()), status: None, + budget: None, token_budget: None, }, }; @@ -2795,6 +2796,7 @@ mod tests { thread_id: "thr_123".to_string(), objective: Some("ship goal mode".to_string()), status: Some(v2::ThreadGoalStatus::Active), + budget: None, token_budget: Some(Some(10_000)), }, }; @@ -2831,6 +2833,7 @@ mod tests { thread_id: "thr_123".to_string(), objective: "ship goal mode".to_string(), status: v2::ThreadGoalStatus::Active, + budget: None, token_budget: Some(10_000), tokens_used: 123, time_used_seconds: 45, diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index 963ac69000..b184520aba 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -3991,6 +3991,66 @@ v2_enum_from_core! { } } +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(tag = "type", rename_all = "camelCase")] +#[ts(tag = "type")] +#[ts(export_to = "v2/")] +pub enum ThreadGoalBudget { + Tokens { + #[ts(type = "number")] + token_budget: i64, + }, + FiveHourLimitPercent { + limit_id: String, + percent: f64, + baseline_used_percent: f64, + #[ts(type = "number | null")] + baseline_resets_at: Option, + latest_used_percent: f64, + #[ts(type = "number | null")] + latest_resets_at: Option, + }, +} + +impl From for ThreadGoalBudget { + fn from(value: codex_protocol::protocol::ThreadGoalBudget) -> Self { + match value { + codex_protocol::protocol::ThreadGoalBudget::Tokens { token_budget } => { + Self::Tokens { token_budget } + } + codex_protocol::protocol::ThreadGoalBudget::FiveHourLimitPercent { + limit_id, + percent, + baseline_used_percent, + baseline_resets_at, + latest_used_percent, + latest_resets_at, + } => Self::FiveHourLimitPercent { + limit_id, + percent, + baseline_used_percent, + baseline_resets_at, + latest_used_percent, + latest_resets_at, + }, + } + } +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(tag = "type", rename_all = "camelCase")] +#[ts(tag = "type")] +#[ts(export_to = "v2/")] +pub enum ThreadGoalBudgetParams { + Tokens { + #[ts(type = "number")] + token_budget: i64, + }, + FiveHourLimitPercent { + percent: f64, + }, +} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] @@ -3998,6 +4058,8 @@ pub struct ThreadGoal { pub thread_id: String, pub objective: String, pub status: ThreadGoalStatus, + #[ts(type = "ThreadGoalBudget | null")] + pub budget: Option, #[ts(type = "number | null")] pub token_budget: Option, #[ts(type = "number")] @@ -4016,6 +4078,7 @@ impl From for ThreadGoal { thread_id: value.thread_id.to_string(), objective: value.objective, status: value.status.into(), + budget: value.budget.map(ThreadGoalBudget::from), token_budget: value.token_budget, tokens_used: value.tokens_used, time_used_seconds: value.time_used_seconds, @@ -4040,6 +4103,14 @@ pub struct ThreadGoalSetParams { serialize_with = "super::serde_helpers::serialize_double_option", skip_serializing_if = "Option::is_none" )] + #[ts(optional = nullable, type = "ThreadGoalBudgetParams | null")] + pub budget: Option>, + #[serde( + default, + deserialize_with = "super::serde_helpers::deserialize_double_option", + serialize_with = "super::serde_helpers::serialize_double_option", + skip_serializing_if = "Option::is_none" + )] #[ts(optional = nullable, type = "number | null")] pub token_budget: Option>, } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index dab47ec3a2..8be493b22a 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -153,7 +153,7 @@ Example with notification opt-out: - `thread/metadata/update` — patch stored thread metadata in sqlite; currently supports updating persisted `gitInfo` fields and returns the refreshed `thread`. - `thread/memoryMode/set` — experimental; set a thread’s persisted memory eligibility to `"enabled"` or `"disabled"` for either a loaded thread or a stored rollout; returns `{}` on success. - `memory/reset` — experimental; clear the current `CODEX_HOME/memories` directory and reset persisted memory stage data in sqlite while preserving existing thread memory modes; returns `{}` on success. -- `thread/goal/set` — create, replace, or update the single persisted goal for a materialized thread; returns the current goal and emits `thread/goal/updated`. Supplying a new `objective` replaces the goal and resets usage accounting. Supplying the current non-terminal objective or omitting `objective` updates the existing goal’s status and/or token budget while preserving usage. +- `thread/goal/set` — create, replace, or update the single persisted goal for a materialized thread; returns the current goal and emits `thread/goal/updated`. Supplying a new `objective` replaces the goal and resets usage accounting. Supplying the current non-terminal objective or omitting `objective` updates the existing goal’s status and/or budget while preserving usage. - `thread/goal/get` — fetch the current persisted goal for a materialized thread; returns `goal: null` when no goal exists. - `thread/goal/clear` — clear the current persisted goal for a materialized thread; returns whether a goal was removed and emits `thread/goal/cleared` when state changes. - `thread/goal/updated` — notification emitted whenever a thread goal changes; includes the full current goal. @@ -483,18 +483,19 @@ Experimental: use `memory/reset` to clear local memory artifacts and sqlite-back ### Example: Set and update a thread goal -Use `thread/goal/set` with an `objective` to create or replace the current goal for a materialized thread. Supplying a new objective resets `tokensUsed`, `timeUsedSeconds`, and `createdAt`. Supplying the current non-terminal objective, or omitting `objective`, updates the existing goal’s status or token budget while preserving usage history. Clients can set `budgetLimited` when they stop because a token budget is exhausted or nearly exhausted; the system also sets it when accounting crosses a configured token budget. +Use `thread/goal/set` with an `objective` to create or replace the current goal for a materialized thread. Supplying a new objective resets `tokensUsed`, `timeUsedSeconds`, and `createdAt`. Supplying the current non-terminal objective, or omitting `objective`, updates the existing goal’s status or budget while preserving usage history. Clients can use the legacy `tokenBudget` field or the structured `budget` field. Structured budgets support fixed tokens and 5h limit percentage points, where the server snapshots the current 5h usage as the baseline and stops the goal after usage increases by the requested percent. Clients can set `budgetLimited` when they stop because a budget is exhausted or nearly exhausted; the system also sets it when accounting crosses a configured budget. ```json { "method": "thread/goal/set", "id": 27, "params": { "threadId": "thr_123", "objective": "Keep improving the benchmark until p95 latency is under 120ms", - "tokenBudget": 200000 + "budget": { "type": "tokens", "tokenBudget": 200000 } } } { "id": 27, "result": { "goal": { "threadId": "thr_123", "objective": "Keep improving the benchmark until p95 latency is under 120ms", "status": "active", + "budget": { "type": "tokens", "tokenBudget": 200000 }, "tokenBudget": 200000, "tokensUsed": 0, "timeUsedSeconds": 0, @@ -505,6 +506,7 @@ Use `thread/goal/set` with an `objective` to create or replace the current goal "threadId": "thr_123", "objective": "Keep improving the benchmark until p95 latency is under 120ms", "status": "active", + "budget": { "type": "tokens", "tokenBudget": 200000 }, "tokenBudget": 200000, "tokensUsed": 0, "timeUsedSeconds": 0, @@ -522,6 +524,7 @@ Use `thread/goal/set` with an `objective` to create or replace the current goal "threadId": "thr_123", "objective": "Keep improving the benchmark until p95 latency is under 120ms", "status": "paused", + "budget": { "type": "tokens", "tokenBudget": 200000 }, "tokenBudget": 200000, "tokensUsed": 10000, "timeUsedSeconds": 60, @@ -530,18 +533,45 @@ Use `thread/goal/set` with an `objective` to create or replace the current goal } } } ``` +```json +{ "method": "thread/goal/set", "id": 29, "params": { + "threadId": "thr_456", + "objective": "Complete the migration without using more than ten percent of the 5h limit", + "budget": { "type": "fiveHourLimitPercent", "percent": 10 } +} } +{ "id": 29, "result": { "goal": { + "threadId": "thr_456", + "objective": "Complete the migration without using more than ten percent of the 5h limit", + "status": "active", + "budget": { + "type": "fiveHourLimitPercent", + "limitId": "codex", + "percent": 10, + "baselineUsedPercent": 42, + "baselineResetsAt": 1776800000, + "latestUsedPercent": 42, + "latestResetsAt": 1776800000 + }, + "tokenBudget": null, + "tokensUsed": 0, + "timeUsedSeconds": 0, + "createdAt": 1776272400, + "updatedAt": 1776272400 +} } } +``` + Use `thread/goal/get` to read the current goal without changing it. ```json -{ "method": "thread/goal/get", "id": 29, "params": { "threadId": "thr_123" } } -{ "id": 29, "result": { "goal": null } } +{ "method": "thread/goal/get", "id": 30, "params": { "threadId": "thr_123" } } +{ "id": 30, "result": { "goal": null } } ``` Use `thread/goal/clear` to remove the current goal. ```json -{ "method": "thread/goal/clear", "id": 30, "params": { "threadId": "thr_123" } } -{ "id": 30, "result": { "cleared": true } } +{ "method": "thread/goal/clear", "id": 31, "params": { "threadId": "thr_123" } } +{ "id": 31, "result": { "cleared": true } } { "method": "thread/goal/cleared", "params": { "threadId": "thr_123" } } ``` diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index edc142c840..94aaf0fb9b 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -165,6 +165,8 @@ use codex_app_server_protocol::ThreadDecrementElicitationResponse; use codex_app_server_protocol::ThreadForkParams; use codex_app_server_protocol::ThreadForkResponse; use codex_app_server_protocol::ThreadGoal; +use codex_app_server_protocol::ThreadGoalBudget; +use codex_app_server_protocol::ThreadGoalBudgetParams; use codex_app_server_protocol::ThreadGoalClearParams; use codex_app_server_protocol::ThreadGoalClearResponse; use codex_app_server_protocol::ThreadGoalClearedNotification; diff --git a/codex-rs/app-server/src/codex_message_processor/thread_goal_handlers.rs b/codex-rs/app-server/src/codex_message_processor/thread_goal_handlers.rs index 049e0af21c..5b415b92e6 100644 --- a/codex-rs/app-server/src/codex_message_processor/thread_goal_handlers.rs +++ b/codex-rs/app-server/src/codex_message_processor/thread_goal_handlers.rs @@ -82,18 +82,25 @@ impl CodexMessageProcessor { }; let status = params.status.map(thread_goal_status_to_state); let objective = params.objective.as_deref().map(str::trim); + let budget = match self.state_goal_budget_from_params(¶ms).await { + Ok(budget) => budget, + Err(message) => { + self.send_invalid_request_error(request_id, message).await; + return; + } + }; if let Some(objective) = objective { if let Err(message) = validate_thread_goal_objective(objective) { self.send_invalid_request_error(request_id, message).await; return; } - if let Err(message) = validate_goal_budget(params.token_budget.flatten()) { + if let Err(message) = validate_goal_budget(budget.as_ref()) { self.send_invalid_request_error(request_id, message).await; return; } - } else if let Some(token_budget) = params.token_budget - && let Err(message) = validate_goal_budget(token_budget) + } else if (params.budget.is_some() || params.token_budget.is_some()) + && let Err(message) = validate_goal_budget(budget.as_ref()) { self.send_invalid_request_error(request_id, message).await; return; @@ -109,13 +116,14 @@ impl CodexMessageProcessor { if let Some(goal) = goal.as_ref().filter(|goal| { goal.objective == objective && goal.status != codex_state::ThreadGoalStatus::Complete + && supports_thread_goal_update_budget(budget.as_ref()) }) { state_db .update_thread_goal( thread_id, codex_state::ThreadGoalUpdate { status, - token_budget: params.token_budget, + token_budget: legacy_token_budget_update(budget.as_ref()), expected_goal_id: Some(goal.goal_id.clone()), }, ) @@ -129,11 +137,11 @@ impl CodexMessageProcessor { }) } else { state_db - .replace_thread_goal( + .replace_thread_goal_with_budget( thread_id, objective, status.unwrap_or(codex_state::ThreadGoalStatus::Active), - params.token_budget.flatten(), + budget.clone().flatten(), ) .await } @@ -146,7 +154,7 @@ impl CodexMessageProcessor { thread_id, codex_state::ThreadGoalUpdate { status, - token_budget: params.token_budget, + token_budget: legacy_token_budget_update(budget.as_ref()), expected_goal_id: None, }, ) @@ -219,6 +227,53 @@ impl CodexMessageProcessor { .await; } + async fn state_goal_budget_from_params( + &self, + params: &ThreadGoalSetParams, + ) -> Result>, String> { + if let Some(budget) = params.budget.clone() { + return match budget { + Some(ThreadGoalBudgetParams::Tokens { token_budget }) => { + Ok(Some(Some(codex_state::ThreadGoalBudget::Tokens { + token_budget, + }))) + } + Some(ThreadGoalBudgetParams::FiveHourLimitPercent { percent }) => { + let (snapshot, snapshots_by_limit_id) = self + .fetch_account_rate_limits() + .await + .map_err(|err| err.message)?; + let snapshot = snapshots_by_limit_id + .get("codex") + .cloned() + .unwrap_or(snapshot); + let Some(window) = snapshot.primary else { + return Err( + "cannot set a five-hour-limit goal because current Codex usage data is unavailable" + .to_string(), + ); + }; + let limit_id = snapshot.limit_id.unwrap_or_else(|| "codex".to_string()); + Ok(Some(Some( + codex_state::ThreadGoalBudget::FiveHourLimitPercent { + limit_id, + percent, + baseline_used_percent: window.used_percent, + baseline_resets_at: window.resets_at, + latest_used_percent: window.used_percent, + latest_resets_at: window.resets_at, + }, + ))) + } + None => Ok(Some(None)), + }; + } + + Ok(params.token_budget.map(|token_budget| { + token_budget.map(|token_budget| codex_state::ThreadGoalBudget::Tokens { token_budget }) + })) + } + pub(super) async fn thread_goal_clear( &self, request_id: ConnectionRequestId, @@ -438,13 +493,42 @@ impl CodexMessageProcessor { } } -fn validate_goal_budget(value: Option) -> Result<(), String> { - if let Some(value) = value - && value <= 0 - { - return Err("goal budgets must be positive when provided".to_string()); +fn supports_thread_goal_update_budget( + budget: Option<&Option>, +) -> bool { + matches!( + budget, + None | Some(None) | Some(Some(codex_state::ThreadGoalBudget::Tokens { .. })) + ) +} + +fn legacy_token_budget_update( + budget: Option<&Option>, +) -> Option> { + budget.map(|budget| { + budget + .as_ref() + .and_then(codex_state::ThreadGoalBudget::token_budget) + }) +} + +fn validate_goal_budget( + value: Option<&Option>, +) -> Result<(), String> { + let Some(Some(value)) = value else { + return Ok(()); + }; + match value { + codex_state::ThreadGoalBudget::Tokens { token_budget } if *token_budget <= 0 => { + Err("goal budgets must be positive when provided".to_string()) + } + codex_state::ThreadGoalBudget::FiveHourLimitPercent { percent, .. } + if !percent.is_finite() || *percent <= 0.0 || *percent > 100.0 => + { + Err("five-hour-limit goal budgets must be a percent from 0 to 100".to_string()) + } + _ => Ok(()), } - Ok(()) } fn thread_goal_status_to_state(status: ThreadGoalStatus) -> codex_state::ThreadGoalStatus { @@ -470,6 +554,7 @@ pub(super) fn api_thread_goal_from_state(goal: codex_state::ThreadGoal) -> Threa thread_id: goal.thread_id.to_string(), objective: goal.objective, status: thread_goal_status_from_state(goal.status), + budget: goal.budget.map(api_thread_goal_budget_from_state), token_budget: goal.token_budget, tokens_used: goal.tokens_used, time_used_seconds: goal.time_used_seconds, @@ -477,3 +562,26 @@ pub(super) fn api_thread_goal_from_state(goal: codex_state::ThreadGoal) -> Threa updated_at: goal.updated_at.timestamp(), } } + +fn api_thread_goal_budget_from_state(budget: codex_state::ThreadGoalBudget) -> ThreadGoalBudget { + match budget { + codex_state::ThreadGoalBudget::Tokens { token_budget } => { + ThreadGoalBudget::Tokens { token_budget } + } + codex_state::ThreadGoalBudget::FiveHourLimitPercent { + limit_id, + percent, + baseline_used_percent, + baseline_resets_at, + latest_used_percent, + latest_resets_at, + } => ThreadGoalBudget::FiveHourLimitPercent { + limit_id, + percent, + baseline_used_percent, + baseline_resets_at, + latest_used_percent, + latest_resets_at, + }, + } +} diff --git a/codex-rs/app-server/src/transport_tests.rs b/codex-rs/app-server/src/transport_tests.rs index 1600b8be87..bd443dac52 100644 --- a/codex-rs/app-server/src/transport_tests.rs +++ b/codex-rs/app-server/src/transport_tests.rs @@ -23,6 +23,7 @@ fn thread_goal_updated_notification() -> ServerNotification { thread_id: "thread-1".to_string(), objective: "ship goal mode".to_string(), status: ThreadGoalStatus::Active, + budget: None, token_budget: None, tokens_used: 0, time_used_seconds: 0, diff --git a/codex-rs/core/src/goals.rs b/codex-rs/core/src/goals.rs index f1805bb750..df01267c7c 100644 --- a/codex-rs/core/src/goals.rs +++ b/codex-rs/core/src/goals.rs @@ -18,6 +18,7 @@ use codex_protocol::models::ResponseInputItem; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ThreadGoal; +use codex_protocol::protocol::ThreadGoalBudget; use codex_protocol::protocol::ThreadGoalStatus; use codex_protocol::protocol::ThreadGoalUpdatedEvent; use codex_protocol::protocol::TokenUsage; @@ -44,6 +45,12 @@ pub(crate) struct SetGoalRequest { pub(crate) struct CreateGoalRequest { pub(crate) objective: String, pub(crate) token_budget: Option, + pub(crate) budget: Option, +} + +pub(crate) enum GoalBudgetRequest { + Tokens(i64), + FiveHourLimitPercent(f64), } static CONTINUATION_PROMPT_TEMPLATE: LazyLock