diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 8025227c08..eb44767356 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -717,6 +717,8 @@ Experimental: use `memory/reset` to clear local memory artifacts and sqlite-back Use `thread/goal/set` to create or update the current goal for a materialized thread. Clients can set `budgetLimited` when they stop because a token budget is exhausted or nearly exhausted, `blocked` when progress is waiting on outside intervention, and `usageLimited` when usage availability stops further work. The system also sets `budgetLimited` when accounting crosses a configured token budget and `usageLimited` when a turn ends on a hard usage-limit error. +When `goals.max_goal_token_budget` is configured, new goals default to that limit, larger budgets are rejected, and setting `tokenBudget` to `null` resets the budget to the configured limit instead of removing it. + ```json { "method": "thread/goal/set", "id": 27, "params": { "threadId": "thr_123", diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index df0daaee38..7cf1074b58 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -18,6 +18,7 @@ use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistry; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ExtensionWarning; +use codex_goal_extension::GoalExtensionConfig; use codex_goal_extension::GoalService; use codex_http_client::HttpClientFactory; use codex_login::AuthManager; @@ -85,7 +86,10 @@ where codex_otel::global(), thread_manager, goal_service, - |config: &Config| config.features.enabled(codex_features::Feature::Goals), + |config: &Config| GoalExtensionConfig { + enabled: config.features.enabled(codex_features::Feature::Goals), + max_goal_token_budget: config.max_goal_token_budget, + }, ); } codex_git_attribution::install( diff --git a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs index 835ae75563..ae46a4a5b4 100644 --- a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs @@ -122,6 +122,10 @@ impl ThreadGoalRequestProcessor { let state_db = self.state_db_for_materialized_thread(thread_id).await?; self.reconcile_thread_goal_rollout(thread_id, &state_db) .await?; + let max_goal_token_budget = match self.thread_manager.get_thread(thread_id).await { + Ok(thread) => thread.config().await.max_goal_token_budget, + Err(_) => self.config.max_goal_token_budget, + }; let listener_command_tx = { let thread_state = self.thread_state_manager.thread_state(thread_id).await; @@ -145,6 +149,7 @@ impl ThreadGoalRequestProcessor { Some(token_budget) => GoalTokenBudgetUpdate::Set(token_budget), None => GoalTokenBudgetUpdate::Keep, }, + max_goal_token_budget, }, ) .await diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index cfc9e53bef..fa18a5e3be 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -2051,6 +2051,103 @@ async fn thread_resume_keeps_paused_goal_paused() -> Result<()> { Ok(()) } +#[tokio::test] +async fn thread_goal_set_enforces_configured_maximum_token_budget() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + mock_responses_config(&server.uri()).write(codex_home.path())?; + let config_path = codex_home.path().join("config.toml"); + let config = std::fs::read_to_string(&config_path)?; + let config = config.replace("personality = true\n", "personality = true\ngoals = true\n"); + std::fs::write( + config_path, + format!("{config}\n[goals]\nmax_goal_token_budget = 200\n"), + )?; + + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_managed_config() + .build_initialized() + .await?; + let start_id = mcp + .send_thread_start_request_with_auto_env(ThreadStartParams { + model: Some("gpt-5.2-codex".to_string()), + config: Some( + [("goals.max_goal_token_budget".to_string(), json!(100))] + .into_iter() + .collect(), + ), + ..Default::default() + }) + .await?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; + + let oversized_creation_id = mcp + .send_raw_request( + "thread/goal/set", + Some(json!({ + "threadId": thread.id, + "objective": "oversized goal", + "tokenBudget": 101, + })), + ) + .await?; + let creation_error: JSONRPCError = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(oversized_creation_id)), + ) + .await??; + assert_eq!( + creation_error.error.message, + "goal token budget 101 exceeds the maximum allowed goal token budget of 100" + ); + + let creation_id = mcp + .send_raw_request( + "thread/goal/set", + Some(json!({ + "threadId": thread.id, + "objective": "bounded goal", + })), + ) + .await?; + let creation: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(creation_id)).await??; + assert_eq!(creation.goal.token_budget, Some(100)); + + let clear_budget_id = mcp + .send_raw_request( + "thread/goal/set", + Some(json!({ "threadId": thread.id, "tokenBudget": null })), + ) + .await?; + let clear_budget: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(clear_budget_id)).await??; + assert_eq!(clear_budget.goal.token_budget, Some(100)); + + let oversized_update_id = mcp + .send_raw_request( + "thread/goal/set", + Some(json!({ + "threadId": thread.id, + "tokenBudget": 101, + })), + ) + .await?; + let update_error: JSONRPCError = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(oversized_update_id)), + ) + .await??; + assert_eq!( + update_error.error.message, + "goal token budget 101 exceeds the maximum allowed goal token budget of 100" + ); + + Ok(()) +} + #[tokio::test] async fn thread_goal_set_preserves_budget_limited_same_objective() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; diff --git a/codex-rs/config/src/config_toml.rs b/codex-rs/config/src/config_toml.rs index 397abfe233..821026255d 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -2,6 +2,7 @@ use std::collections::BTreeMap; use std::collections::HashMap; +use std::num::NonZeroU64; use std::path::Path; use crate::HooksToml; @@ -432,6 +433,9 @@ pub struct ConfigToml { /// Agent-related settings (thread limits, etc.). pub agents: Option, + /// Goal-related settings. + pub goals: Option, + /// Memories subsystem settings. pub memories: Option, @@ -683,6 +687,13 @@ where }) } +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct GoalsToml { + /// Maximum token budget allowed for a goal and default budget for new goals. + pub max_goal_token_budget: Option, +} + #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] #[schemars(deny_unknown_fields)] pub struct AgentsToml { diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 5d8c170fe6..c94db1eb96 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1199,6 +1199,18 @@ }, "type": "object" }, + "GoalsToml": { + "additionalProperties": false, + "properties": { + "max_goal_token_budget": { + "description": "Maximum token budget allowed for a goal and default budget for new goals.", + "format": "uint64", + "minimum": 1.0, + "type": "integer" + } + }, + "type": "object" + }, "GranularApprovalConfig": { "properties": { "mcp_elicitations": { @@ -5521,6 +5533,14 @@ "default": null, "description": "Compatibility-only settings retained so legacy `ghost_snapshot` config still loads." }, + "goals": { + "allOf": [ + { + "$ref": "#/definitions/GoalsToml" + } + ], + "description": "Goal-related settings." + }, "hide_agent_reasoning": { "default": false, "description": "When set to `true`, `AgentReasoning` events will be hidden from the UI/output. Defaults to `false`.", @@ -5913,4 +5933,4 @@ }, "title": "ConfigToml", "type": "object" -} +} \ No newline at end of file diff --git a/codex-rs/core/src/config/config_loader_tests.rs b/codex-rs/core/src/config/config_loader_tests.rs index 59658c208a..65ad3d0ffc 100644 --- a/codex-rs/core/src/config/config_loader_tests.rs +++ b/codex-rs/core/src/config/config_loader_tests.rs @@ -800,6 +800,36 @@ extra = true assert_eq!(nested.get("extra"), Some(&TomlValue::Boolean(true))); } +#[tokio::test] +async fn managed_goal_token_budget_overrides_user_config() -> anyhow::Result<()> { + let tmp = tempdir()?; + let managed_path = tmp.path().join("managed_config.toml"); + std::fs::write( + tmp.path().join(CONFIG_TOML_FILE), + "[goals]\nmax_goal_token_budget = 20000\n", + )?; + std::fs::write(&managed_path, "[goals]\nmax_goal_token_budget = 5000\n")?; + + let state = load_config_layers_state( + LOCAL_FS.as_ref(), + tmp.path(), + Some(AbsolutePathBuf::try_from(tmp.path())?), + &[] as &[(String, TomlValue)], + LoaderOverrides::with_managed_config_path_for_tests(managed_path), + &codex_config::NoopThreadConfigLoader, + ) + .await?; + + assert_eq!( + state + .effective_config() + .get("goals") + .and_then(|goals| goals.get("max_goal_token_budget")), + Some(&TomlValue::Integer(5_000)) + ); + Ok(()) +} + #[tokio::test] async fn returns_empty_when_all_layers_missing() { let tmp = tempdir().expect("tempdir"); diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index bc76cd0fb3..b341bcc26c 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -326,6 +326,28 @@ consolidation_model = "gpt-5.2" ); } +#[tokio::test] +async fn goal_max_token_budget_requires_positive_integer() { + let config_toml = toml::from_str::("[goals]\nmax_goal_token_budget = 25000\n") + .expect("positive goal token budget should deserialize"); + let config = Config::load_from_base_config_with_overrides( + config_toml, + ConfigOverrides::default(), + tempdir().expect("tempdir").abs(), + ) + .await + .expect("positive goal token budget should load"); + assert_eq!(config.max_goal_token_budget, Some(25_000)); + + for invalid in ["0", "-1", "1.5", "\"100\""] { + let config = format!("[goals]\nmax_goal_token_budget = {invalid}\n"); + assert!( + toml::from_str::(&config).is_err(), + "invalid goal token budget should be rejected: {invalid}" + ); + } +} + #[test] fn parses_bundled_skills_config() { let cfg: ConfigToml = toml::from_str( diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 5760aeae28..5ca2239e30 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -898,6 +898,9 @@ pub struct Config { /// User-defined role declarations keyed by role name. pub agent_roles: BTreeMap, + /// Maximum token budget allowed for a goal and default budget for new goals. + pub max_goal_token_budget: Option, + /// Memories subsystem settings. pub memories: MemoriesConfig, @@ -4193,6 +4196,19 @@ impl Config { agent_default_subagent_reasoning_effort, agent_max_depth, agent_roles, + max_goal_token_budget: cfg + .goals + .as_ref() + .and_then(|goals| goals.max_goal_token_budget) + .map(|max_goal_token_budget| { + i64::try_from(max_goal_token_budget.get()).map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "goals.max_goal_token_budget exceeds the maximum supported token budget", + ) + }) + }) + .transpose()?, memories: memories_config, agent_interrupt_message_enabled, codex_home, diff --git a/codex-rs/ext/goal/src/api.rs b/codex-rs/ext/goal/src/api.rs index 83d58f7f01..149c95d0a1 100644 --- a/codex-rs/ext/goal/src/api.rs +++ b/codex-rs/ext/goal/src/api.rs @@ -54,6 +54,7 @@ pub struct GoalSetRequest<'a> { pub objective: GoalObjectiveUpdate<'a>, pub status: Option, pub token_budget: GoalTokenBudgetUpdate, + pub max_goal_token_budget: Option, } #[derive(Clone, Debug)] @@ -150,6 +151,7 @@ impl GoalService { objective, status, token_budget, + max_goal_token_budget, } = request; let status = status.map(state_status_from_protocol); let objective = match objective { @@ -158,14 +160,16 @@ impl GoalService { }; let token_budget = match token_budget { GoalTokenBudgetUpdate::Keep => None, - GoalTokenBudgetUpdate::Set(token_budget) => Some(token_budget), + GoalTokenBudgetUpdate::Set(token_budget) => { + Some(token_budget.or(max_goal_token_budget)) + } }; if let Some(objective) = objective { validate_thread_goal_objective(objective).map_err(GoalServiceError::InvalidRequest)?; } if objective.is_some() || token_budget.is_some() { - validate_goal_budget(token_budget.flatten()) + validate_goal_budget(token_budget.flatten(), max_goal_token_budget) .map_err(GoalServiceError::InvalidRequest)?; } @@ -225,7 +229,7 @@ impl GoalService { thread_id, objective, status.unwrap_or(codex_state::ThreadGoalStatus::Active), - token_budget.flatten(), + token_budget.flatten().or(max_goal_token_budget), ) .await .map_err(|err| { diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index dcf6bad73d..923f436581 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -48,12 +48,7 @@ use crate::tool::GoalToolExecutor; #[derive(Clone, Debug)] pub struct GoalExtensionConfig { pub enabled: bool, -} - -impl GoalExtensionConfig { - fn from_enabled(enabled: bool) -> Self { - Self { enabled } - } + pub max_goal_token_budget: Option, } #[derive(Clone)] @@ -64,7 +59,7 @@ pub struct GoalExtension { metrics: GoalMetrics, thread_manager: Weak, goal_service: Arc, - goals_enabled: Arc bool + Send + Sync>, + goal_config: Arc GoalExtensionConfig + Send + Sync>, } impl std::fmt::Debug for GoalExtension { @@ -81,7 +76,7 @@ impl GoalExtension { metrics_client: Option, thread_manager: Weak, goal_service: Arc, - goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, + goal_config: impl Fn(&C) -> GoalExtensionConfig + Send + Sync + 'static, ) -> Self { Self { state_dbs, @@ -90,7 +85,7 @@ impl GoalExtension { metrics: GoalMetrics::new(metrics_client), thread_manager, goal_service, - goals_enabled: Arc::new(goals_enabled), + goal_config: Arc::new(goal_config), } } } @@ -101,15 +96,14 @@ where { fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> { Box::pin(async move { - let enabled = (self.goals_enabled)(input.config); + let config = (self.goal_config)(input.config); + let enabled = config.enabled; let tools_available_for_thread = input.persistent_thread_state_available && !matches!( input.session_source, SessionSource::SubAgent(SubAgentSource::Review) ); - input - .thread_store - .insert(GoalExtensionConfig::from_enabled(enabled)); + input.thread_store.insert(config); let accounting_state = input .thread_store .get_or_init::(GoalAccountingState::default); @@ -186,8 +180,9 @@ where _previous_config: &C, new_config: &C, ) { - let enabled = (self.goals_enabled)(new_config); - thread_store.insert(GoalExtensionConfig::from_enabled(enabled)); + let config = (self.goal_config)(new_config); + let enabled = config.enabled; + thread_store.insert(config); if let Some(runtime) = goal_runtime_handle(thread_store) { runtime.set_enabled(enabled); } @@ -427,6 +422,9 @@ where if !runtime.tools_visible() { return Vec::new(); } + let max_goal_token_budget = thread_store + .get::() + .and_then(|config| config.max_goal_token_budget); vec![ Arc::new(GoalToolExecutor::get( @@ -444,6 +442,7 @@ where self.analytics.clone(), self.event_emitter.clone(), self.metrics.clone(), + max_goal_token_budget, )), Arc::new(GoalToolExecutor::update( runtime.thread_id(), @@ -464,7 +463,7 @@ pub fn install_with_backend( metrics_client: Option, thread_manager: Weak, goal_service: Arc, - goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, + goal_config: impl Fn(&C) -> GoalExtensionConfig + Send + Sync + 'static, ) where C: Send + Sync + 'static, { @@ -475,7 +474,7 @@ pub fn install_with_backend( metrics_client, thread_manager, Arc::clone(&goal_service), - goals_enabled, + goal_config, )); registry.thread_lifecycle_contributor(extension.clone()); registry.config_contributor(extension.clone()); diff --git a/codex-rs/ext/goal/src/tool.rs b/codex-rs/ext/goal/src/tool.rs index d280f6e04b..f15a5e90fb 100644 --- a/codex-rs/ext/goal/src/tool.rs +++ b/codex-rs/ext/goal/src/tool.rs @@ -36,6 +36,7 @@ pub(crate) struct GoalToolExecutor { analytics: GoalAnalytics, event_emitter: GoalEventEmitter, metrics: GoalMetrics, + max_goal_token_budget: Option, } #[derive(Clone, Copy)] @@ -89,6 +90,7 @@ impl GoalToolExecutor { analytics, event_emitter, metrics, + max_goal_token_budget: None, } } @@ -99,6 +101,7 @@ impl GoalToolExecutor { analytics: GoalAnalytics, event_emitter: GoalEventEmitter, metrics: GoalMetrics, + max_goal_token_budget: Option, ) -> Self { Self { kind: GoalToolKind::Create, @@ -108,6 +111,7 @@ impl GoalToolExecutor { analytics, event_emitter, metrics, + max_goal_token_budget, } } @@ -127,6 +131,7 @@ impl GoalToolExecutor { analytics, event_emitter, metrics, + max_goal_token_budget: None, } } } @@ -185,7 +190,9 @@ impl GoalToolExecutor { request.objective = request.objective.trim().to_string(); validate_thread_goal_objective(&request.objective) .map_err(FunctionCallError::RespondToModel)?; - validate_goal_budget(request.token_budget).map_err(FunctionCallError::RespondToModel)?; + request.token_budget = request.token_budget.or(self.max_goal_token_budget); + validate_goal_budget(request.token_budget, self.max_goal_token_budget) + .map_err(FunctionCallError::RespondToModel)?; let goal = self .state_db @@ -397,12 +404,23 @@ where .map_err(|err| FunctionCallError::RespondToModel(err.to_string())) } -pub(crate) fn validate_goal_budget(value: Option) -> Result<(), String> { +pub(crate) fn validate_goal_budget( + value: Option, + max_goal_token_budget: Option, +) -> Result<(), String> { if let Some(value) = value && value <= 0 { return Err("goal budgets must be positive when provided".to_string()); } + if let Some(value) = value + && let Some(max_goal_token_budget) = max_goal_token_budget + && value > max_goal_token_budget + { + return Err(format!( + "goal token budget {value} exceeds the maximum allowed goal token budget of {max_goal_token_budget}" + )); + } Ok(()) } diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 2ab1b954b7..7ac890183e 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -27,6 +27,7 @@ use codex_extension_api::ToolPayload; use codex_extension_api::TurnErrorInput; use codex_extension_api::TurnStartInput; use codex_extension_api::TurnStopInput; +use codex_goal_extension::GoalExtensionConfig; use codex_goal_extension::GoalObjectiveUpdate; use codex_goal_extension::GoalRuntimeHandle; use codex_goal_extension::GoalService; @@ -97,6 +98,55 @@ async fn installed_goal_tools_create_goal_and_fill_empty_preview() -> anyhow::Re Ok(()) } +#[tokio::test] +async fn installed_goal_tools_apply_maximum_token_budget() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.thread_store.insert(GoalExtensionConfig { + enabled: true, + max_goal_token_budget: Some(100), + }); + let tools = harness.tools(); + let create_tool = tool_by_name(&tools, "create_goal"); + + let result = create_tool + .handle(tool_call( + "create_goal", + "call-oversized-goal", + json!({ "objective": "oversized goal", "token_budget": 101 }), + )) + .await; + let error = match result { + Ok(_) => panic!("goal budget above the configured maximum should fail"), + Err(error) => error, + }; + assert_eq!( + error, + FunctionCallError::RespondToModel( + "goal token budget 101 exceeds the maximum allowed goal token budget of 100" + .to_string() + ) + ); + assert_eq!( + runtime.thread_goals().get_thread_goal(thread_id).await?, + None + ); + + let invocation = tool_call( + "create_goal", + "call-default-goal-budget", + json!({ "objective": "default goal budget" }), + ); + let output = create_tool.handle(invocation.clone()).await?; + assert_eq!( + output.code_mode_result(&invocation.payload)["goal"]["tokenBudget"], + json!(100) + ); + Ok(()) +} + #[tokio::test] async fn goal_tools_hidden_for_ephemeral_threads() -> anyhow::Result<()> { let runtime = test_runtime().await?; @@ -940,6 +990,7 @@ async fn goal_service_external_set_active_resets_baseline_without_live_thread() objective: GoalObjectiveUpdate::Set("new objective"), status: Some(ThreadGoalStatus::Active), token_budget: GoalTokenBudgetUpdate::Keep, + max_goal_token_budget: None, }, ) .await?; @@ -1069,6 +1120,7 @@ async fn goal_service_sets_gets_and_clears_thread_goal() -> anyhow::Result<()> { objective: GoalObjectiveUpdate::Set(" ship goal API ownership "), status: None, token_budget: GoalTokenBudgetUpdate::Set(Some(123)), + max_goal_token_budget: None, }, ) .await?; @@ -1096,6 +1148,84 @@ async fn goal_service_sets_gets_and_clears_thread_goal() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn goal_service_enforces_maximum_token_budget_on_creation_and_updates() -> anyhow::Result<()> +{ + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let service = GoalService::new(); + + let goal = service + .set_thread_goal( + runtime.as_ref(), + GoalSetRequest { + thread_id, + objective: GoalObjectiveUpdate::Set("bounded goal"), + status: None, + token_budget: GoalTokenBudgetUpdate::Keep, + max_goal_token_budget: Some(100), + }, + ) + .await?; + assert_eq!(goal.goal.token_budget, Some(100)); + + let error = service + .set_thread_goal( + runtime.as_ref(), + GoalSetRequest { + thread_id, + objective: GoalObjectiveUpdate::Keep, + status: None, + token_budget: GoalTokenBudgetUpdate::Set(Some(101)), + max_goal_token_budget: Some(100), + }, + ) + .await + .expect_err("goal budget above the configured maximum should fail"); + assert_eq!( + error.to_string(), + "goal token budget 101 exceeds the maximum allowed goal token budget of 100" + ); + assert_eq!( + service + .get_thread_goal(runtime.as_ref(), thread_id) + .await? + .expect("goal should remain unchanged") + .token_budget, + Some(100) + ); + + let goal = service + .set_thread_goal( + runtime.as_ref(), + GoalSetRequest { + thread_id, + objective: GoalObjectiveUpdate::Keep, + status: None, + token_budget: GoalTokenBudgetUpdate::Set(Some(99)), + max_goal_token_budget: Some(100), + }, + ) + .await?; + assert_eq!(goal.goal.token_budget, Some(99)); + + let goal = service + .set_thread_goal( + runtime.as_ref(), + GoalSetRequest { + thread_id, + objective: GoalObjectiveUpdate::Keep, + status: None, + token_budget: GoalTokenBudgetUpdate::Set(None), + max_goal_token_budget: Some(100), + }, + ) + .await?; + assert_eq!(goal.goal.token_budget, Some(100)); + Ok(()) +} + async fn installed_tools( runtime: Arc, thread_id: ThreadId, @@ -1124,7 +1254,10 @@ async fn installed_tools_with_start( /*metrics_client*/ None, Weak::new(), goal_service, - |_| true, + |_| GoalExtensionConfig { + enabled: true, + max_goal_token_budget: None, + }, ); let registry = builder.build(); let session_store = ExtensionData::new("session-1"); @@ -1178,7 +1311,10 @@ impl GoalExtensionHarness { /*metrics_client*/ None, Weak::new(), Arc::clone(&goal_service), - |_| true, + |_| GoalExtensionConfig { + enabled: true, + max_goal_token_budget: None, + }, ); let registry = builder.build(); let session_store = ExtensionData::new("session-1"); diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index a231102479..6dbc6d5214 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -303,6 +303,7 @@ fn new_config(model: Option, arg0_paths: Arg0DispatchPaths) -> anyhow::R background_terminal_max_timeout: 300_000, ghost_snapshot: GhostSnapshotConfig::default(), multi_agent_v2: MultiAgentV2Config::default(), + max_goal_token_budget: None, token_budget: None, rollout_budget: None, current_time_reminder: None,