Add configurable goal token budget limits (#37878)

## What changed

- Add `goals.max_goal_token_budget` as a positive-integer configuration setting.
- Use the configured maximum as the default budget for new goals and when `tokenBudget` is reset to `null`.
- Reject goal creation and updates whose token budget exceeds the configured maximum, including requests through goal tools and `thread/goal/set`.
- Respect managed configuration precedence and per-thread configuration overrides.

## Testing

- Cover configuration parsing and managed overrides.
- Cover defaulting, resetting, and rejecting oversized budgets through the goal service, goal tools, and app-server API.

GitOrigin-RevId: f8d7e6418cdc237d454c8cf47bb32ba0d44a60cf
This commit is contained in:
Eric Traut authored and copyberry committed 2026-08-10 21:00:00 +00:00
1 parent a603d7ca5c
commit a9dee37f9c
14 files changed
+391 -26

No files matched your search

+2
View File
@@ -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. 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 ```json
{ "method": "thread/goal/set", "id": 27, "params": { { "method": "thread/goal/set", "id": 27, "params": {
"threadId": "thr_123", "threadId": "thr_123",
+5 -1
View File
@@ -18,6 +18,7 @@ use codex_extension_api::ExtensionEventSink;
use codex_extension_api::ExtensionRegistry; use codex_extension_api::ExtensionRegistry;
use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::ExtensionWarning; use codex_extension_api::ExtensionWarning;
use codex_goal_extension::GoalExtensionConfig;
use codex_goal_extension::GoalService; use codex_goal_extension::GoalService;
use codex_http_client::HttpClientFactory; use codex_http_client::HttpClientFactory;
use codex_login::AuthManager; use codex_login::AuthManager;
@@ -85,7 +86,10 @@ where
codex_otel::global(), codex_otel::global(),
thread_manager, thread_manager,
goal_service, 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( codex_git_attribution::install(
@@ -122,6 +122,10 @@ impl ThreadGoalRequestProcessor {
let state_db = self.state_db_for_materialized_thread(thread_id).await?; let state_db = self.state_db_for_materialized_thread(thread_id).await?;
self.reconcile_thread_goal_rollout(thread_id, &state_db) self.reconcile_thread_goal_rollout(thread_id, &state_db)
.await?; .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 listener_command_tx = {
let thread_state = self.thread_state_manager.thread_state(thread_id).await; 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), Some(token_budget) => GoalTokenBudgetUpdate::Set(token_budget),
None => GoalTokenBudgetUpdate::Keep, None => GoalTokenBudgetUpdate::Keep,
}, },
max_goal_token_budget,
}, },
) )
.await .await
@@ -2051,6 +2051,103 @@ async fn thread_resume_keeps_paused_goal_paused() -> Result<()> {
Ok(()) 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] #[tokio::test]
async fn thread_goal_set_preserves_budget_limited_same_objective() -> Result<()> { async fn thread_goal_set_preserves_budget_limited_same_objective() -> Result<()> {
let server = create_mock_responses_server_repeating_assistant("Done").await; let server = create_mock_responses_server_repeating_assistant("Done").await;
+11
View File
@@ -2,6 +2,7 @@
use std::collections::BTreeMap; use std::collections::BTreeMap;
use std::collections::HashMap; use std::collections::HashMap;
use std::num::NonZeroU64;
use std::path::Path; use std::path::Path;
use crate::HooksToml; use crate::HooksToml;
@@ -432,6 +433,9 @@ pub struct ConfigToml {
/// Agent-related settings (thread limits, etc.). /// Agent-related settings (thread limits, etc.).
pub agents: Option<AgentsToml>, pub agents: Option<AgentsToml>,
/// Goal-related settings.
pub goals: Option<GoalsToml>,
/// Memories subsystem settings. /// Memories subsystem settings.
pub memories: Option<MemoriesToml>, pub memories: Option<MemoriesToml>,
@@ -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<NonZeroU64>,
}
#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)]
#[schemars(deny_unknown_fields)] #[schemars(deny_unknown_fields)]
pub struct AgentsToml { pub struct AgentsToml {
+21 -1
View File
@@ -1199,6 +1199,18 @@
}, },
"type": "object" "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": { "GranularApprovalConfig": {
"properties": { "properties": {
"mcp_elicitations": { "mcp_elicitations": {
@@ -5521,6 +5533,14 @@
"default": null, "default": null,
"description": "Compatibility-only settings retained so legacy `ghost_snapshot` config still loads." "description": "Compatibility-only settings retained so legacy `ghost_snapshot` config still loads."
}, },
"goals": {
"allOf": [
{
"$ref": "#/definitions/GoalsToml"
}
],
"description": "Goal-related settings."
},
"hide_agent_reasoning": { "hide_agent_reasoning": {
"default": false, "default": false,
"description": "When set to `true`, `AgentReasoning` events will be hidden from the UI/output. Defaults to `false`.", "description": "When set to `true`, `AgentReasoning` events will be hidden from the UI/output. Defaults to `false`.",
@@ -5913,4 +5933,4 @@
}, },
"title": "ConfigToml", "title": "ConfigToml",
"type": "object" "type": "object"
} }
@@ -800,6 +800,36 @@ extra = true
assert_eq!(nested.get("extra"), Some(&TomlValue::Boolean(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] #[tokio::test]
async fn returns_empty_when_all_layers_missing() { async fn returns_empty_when_all_layers_missing() {
let tmp = tempdir().expect("tempdir"); let tmp = tempdir().expect("tempdir");
+22
View File
@@ -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::<ConfigToml>("[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::<ConfigToml>(&config).is_err(),
"invalid goal token budget should be rejected: {invalid}"
);
}
}
#[test] #[test]
fn parses_bundled_skills_config() { fn parses_bundled_skills_config() {
let cfg: ConfigToml = toml::from_str( let cfg: ConfigToml = toml::from_str(
+16
View File
@@ -898,6 +898,9 @@ pub struct Config {
/// User-defined role declarations keyed by role name. /// User-defined role declarations keyed by role name.
pub agent_roles: BTreeMap<String, AgentRoleConfig>, pub agent_roles: BTreeMap<String, AgentRoleConfig>,
/// Maximum token budget allowed for a goal and default budget for new goals.
pub max_goal_token_budget: Option<i64>,
/// Memories subsystem settings. /// Memories subsystem settings.
pub memories: MemoriesConfig, pub memories: MemoriesConfig,
@@ -4193,6 +4196,19 @@ impl Config {
agent_default_subagent_reasoning_effort, agent_default_subagent_reasoning_effort,
agent_max_depth, agent_max_depth,
agent_roles, 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, memories: memories_config,
agent_interrupt_message_enabled, agent_interrupt_message_enabled,
codex_home, codex_home,
+7 -3
View File
@@ -54,6 +54,7 @@ pub struct GoalSetRequest<'a> {
pub objective: GoalObjectiveUpdate<'a>, pub objective: GoalObjectiveUpdate<'a>,
pub status: Option<ThreadGoalStatus>, pub status: Option<ThreadGoalStatus>,
pub token_budget: GoalTokenBudgetUpdate, pub token_budget: GoalTokenBudgetUpdate,
pub max_goal_token_budget: Option<i64>,
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
@@ -150,6 +151,7 @@ impl GoalService {
objective, objective,
status, status,
token_budget, token_budget,
max_goal_token_budget,
} = request; } = request;
let status = status.map(state_status_from_protocol); let status = status.map(state_status_from_protocol);
let objective = match objective { let objective = match objective {
@@ -158,14 +160,16 @@ impl GoalService {
}; };
let token_budget = match token_budget { let token_budget = match token_budget {
GoalTokenBudgetUpdate::Keep => None, 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 { if let Some(objective) = objective {
validate_thread_goal_objective(objective).map_err(GoalServiceError::InvalidRequest)?; validate_thread_goal_objective(objective).map_err(GoalServiceError::InvalidRequest)?;
} }
if objective.is_some() || token_budget.is_some() { 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)?; .map_err(GoalServiceError::InvalidRequest)?;
} }
@@ -225,7 +229,7 @@ impl GoalService {
thread_id, thread_id,
objective, objective,
status.unwrap_or(codex_state::ThreadGoalStatus::Active), status.unwrap_or(codex_state::ThreadGoalStatus::Active),
token_budget.flatten(), token_budget.flatten().or(max_goal_token_budget),
) )
.await .await
.map_err(|err| { .map_err(|err| {
+16 -17
View File
@@ -48,12 +48,7 @@ use crate::tool::GoalToolExecutor;
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct GoalExtensionConfig { pub struct GoalExtensionConfig {
pub enabled: bool, pub enabled: bool,
} pub max_goal_token_budget: Option<i64>,
impl GoalExtensionConfig {
fn from_enabled(enabled: bool) -> Self {
Self { enabled }
}
} }
#[derive(Clone)] #[derive(Clone)]
@@ -64,7 +59,7 @@ pub struct GoalExtension<C> {
metrics: GoalMetrics, metrics: GoalMetrics,
thread_manager: Weak<ThreadManager>, thread_manager: Weak<ThreadManager>,
goal_service: Arc<GoalService>, goal_service: Arc<GoalService>,
goals_enabled: Arc<dyn Fn(&C) -> bool + Send + Sync>, goal_config: Arc<dyn Fn(&C) -> GoalExtensionConfig + Send + Sync>,
} }
impl<C> std::fmt::Debug for GoalExtension<C> { impl<C> std::fmt::Debug for GoalExtension<C> {
@@ -81,7 +76,7 @@ impl<C> GoalExtension<C> {
metrics_client: Option<MetricsClient>, metrics_client: Option<MetricsClient>,
thread_manager: Weak<ThreadManager>, thread_manager: Weak<ThreadManager>,
goal_service: Arc<GoalService>, goal_service: Arc<GoalService>,
goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, goal_config: impl Fn(&C) -> GoalExtensionConfig + Send + Sync + 'static,
) -> Self { ) -> Self {
Self { Self {
state_dbs, state_dbs,
@@ -90,7 +85,7 @@ impl<C> GoalExtension<C> {
metrics: GoalMetrics::new(metrics_client), metrics: GoalMetrics::new(metrics_client),
thread_manager, thread_manager,
goal_service, 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, ()> { fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> {
Box::pin(async move { 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 let tools_available_for_thread = input.persistent_thread_state_available
&& !matches!( && !matches!(
input.session_source, input.session_source,
SessionSource::SubAgent(SubAgentSource::Review) SessionSource::SubAgent(SubAgentSource::Review)
); );
input input.thread_store.insert(config);
.thread_store
.insert(GoalExtensionConfig::from_enabled(enabled));
let accounting_state = input let accounting_state = input
.thread_store .thread_store
.get_or_init::<GoalAccountingState>(GoalAccountingState::default); .get_or_init::<GoalAccountingState>(GoalAccountingState::default);
@@ -186,8 +180,9 @@ where
_previous_config: &C, _previous_config: &C,
new_config: &C, new_config: &C,
) { ) {
let enabled = (self.goals_enabled)(new_config); let config = (self.goal_config)(new_config);
thread_store.insert(GoalExtensionConfig::from_enabled(enabled)); let enabled = config.enabled;
thread_store.insert(config);
if let Some(runtime) = goal_runtime_handle(thread_store) { if let Some(runtime) = goal_runtime_handle(thread_store) {
runtime.set_enabled(enabled); runtime.set_enabled(enabled);
} }
@@ -427,6 +422,9 @@ where
if !runtime.tools_visible() { if !runtime.tools_visible() {
return Vec::new(); return Vec::new();
} }
let max_goal_token_budget = thread_store
.get::<GoalExtensionConfig>()
.and_then(|config| config.max_goal_token_budget);
vec![ vec![
Arc::new(GoalToolExecutor::get( Arc::new(GoalToolExecutor::get(
@@ -444,6 +442,7 @@ where
self.analytics.clone(), self.analytics.clone(),
self.event_emitter.clone(), self.event_emitter.clone(),
self.metrics.clone(), self.metrics.clone(),
max_goal_token_budget,
)), )),
Arc::new(GoalToolExecutor::update( Arc::new(GoalToolExecutor::update(
runtime.thread_id(), runtime.thread_id(),
@@ -464,7 +463,7 @@ pub fn install_with_backend<C>(
metrics_client: Option<MetricsClient>, metrics_client: Option<MetricsClient>,
thread_manager: Weak<ThreadManager>, thread_manager: Weak<ThreadManager>,
goal_service: Arc<GoalService>, goal_service: Arc<GoalService>,
goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, goal_config: impl Fn(&C) -> GoalExtensionConfig + Send + Sync + 'static,
) where ) where
C: Send + Sync + 'static, C: Send + Sync + 'static,
{ {
@@ -475,7 +474,7 @@ pub fn install_with_backend<C>(
metrics_client, metrics_client,
thread_manager, thread_manager,
Arc::clone(&goal_service), Arc::clone(&goal_service),
goals_enabled, goal_config,
)); ));
registry.thread_lifecycle_contributor(extension.clone()); registry.thread_lifecycle_contributor(extension.clone());
registry.config_contributor(extension.clone()); registry.config_contributor(extension.clone());
+20 -2
View File
@@ -36,6 +36,7 @@ pub(crate) struct GoalToolExecutor {
analytics: GoalAnalytics, analytics: GoalAnalytics,
event_emitter: GoalEventEmitter, event_emitter: GoalEventEmitter,
metrics: GoalMetrics, metrics: GoalMetrics,
max_goal_token_budget: Option<i64>,
} }
#[derive(Clone, Copy)] #[derive(Clone, Copy)]
@@ -89,6 +90,7 @@ impl GoalToolExecutor {
analytics, analytics,
event_emitter, event_emitter,
metrics, metrics,
max_goal_token_budget: None,
} }
} }
@@ -99,6 +101,7 @@ impl GoalToolExecutor {
analytics: GoalAnalytics, analytics: GoalAnalytics,
event_emitter: GoalEventEmitter, event_emitter: GoalEventEmitter,
metrics: GoalMetrics, metrics: GoalMetrics,
max_goal_token_budget: Option<i64>,
) -> Self { ) -> Self {
Self { Self {
kind: GoalToolKind::Create, kind: GoalToolKind::Create,
@@ -108,6 +111,7 @@ impl GoalToolExecutor {
analytics, analytics,
event_emitter, event_emitter,
metrics, metrics,
max_goal_token_budget,
} }
} }
@@ -127,6 +131,7 @@ impl GoalToolExecutor {
analytics, analytics,
event_emitter, event_emitter,
metrics, metrics,
max_goal_token_budget: None,
} }
} }
} }
@@ -185,7 +190,9 @@ impl GoalToolExecutor {
request.objective = request.objective.trim().to_string(); request.objective = request.objective.trim().to_string();
validate_thread_goal_objective(&request.objective) validate_thread_goal_objective(&request.objective)
.map_err(FunctionCallError::RespondToModel)?; .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 let goal = self
.state_db .state_db
@@ -397,12 +404,23 @@ where
.map_err(|err| FunctionCallError::RespondToModel(err.to_string())) .map_err(|err| FunctionCallError::RespondToModel(err.to_string()))
} }
pub(crate) fn validate_goal_budget(value: Option<i64>) -> Result<(), String> { pub(crate) fn validate_goal_budget(
value: Option<i64>,
max_goal_token_budget: Option<i64>,
) -> Result<(), String> {
if let Some(value) = value if let Some(value) = value
&& value <= 0 && value <= 0
{ {
return Err("goal budgets must be positive when provided".to_string()); 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(()) Ok(())
} }
@@ -27,6 +27,7 @@ use codex_extension_api::ToolPayload;
use codex_extension_api::TurnErrorInput; use codex_extension_api::TurnErrorInput;
use codex_extension_api::TurnStartInput; use codex_extension_api::TurnStartInput;
use codex_extension_api::TurnStopInput; use codex_extension_api::TurnStopInput;
use codex_goal_extension::GoalExtensionConfig;
use codex_goal_extension::GoalObjectiveUpdate; use codex_goal_extension::GoalObjectiveUpdate;
use codex_goal_extension::GoalRuntimeHandle; use codex_goal_extension::GoalRuntimeHandle;
use codex_goal_extension::GoalService; use codex_goal_extension::GoalService;
@@ -97,6 +98,55 @@ async fn installed_goal_tools_create_goal_and_fill_empty_preview() -> anyhow::Re
Ok(()) 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] #[tokio::test]
async fn goal_tools_hidden_for_ephemeral_threads() -> anyhow::Result<()> { async fn goal_tools_hidden_for_ephemeral_threads() -> anyhow::Result<()> {
let runtime = test_runtime().await?; 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"), objective: GoalObjectiveUpdate::Set("new objective"),
status: Some(ThreadGoalStatus::Active), status: Some(ThreadGoalStatus::Active),
token_budget: GoalTokenBudgetUpdate::Keep, token_budget: GoalTokenBudgetUpdate::Keep,
max_goal_token_budget: None,
}, },
) )
.await?; .await?;
@@ -1069,6 +1120,7 @@ async fn goal_service_sets_gets_and_clears_thread_goal() -> anyhow::Result<()> {
objective: GoalObjectiveUpdate::Set(" ship goal API ownership "), objective: GoalObjectiveUpdate::Set(" ship goal API ownership "),
status: None, status: None,
token_budget: GoalTokenBudgetUpdate::Set(Some(123)), token_budget: GoalTokenBudgetUpdate::Set(Some(123)),
max_goal_token_budget: None,
}, },
) )
.await?; .await?;
@@ -1096,6 +1148,84 @@ async fn goal_service_sets_gets_and_clears_thread_goal() -> anyhow::Result<()> {
Ok(()) 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( async fn installed_tools(
runtime: Arc<codex_state::StateRuntime>, runtime: Arc<codex_state::StateRuntime>,
thread_id: ThreadId, thread_id: ThreadId,
@@ -1124,7 +1254,10 @@ async fn installed_tools_with_start(
/*metrics_client*/ None, /*metrics_client*/ None,
Weak::new(), Weak::new(),
goal_service, goal_service,
|_| true, |_| GoalExtensionConfig {
enabled: true,
max_goal_token_budget: None,
},
); );
let registry = builder.build(); let registry = builder.build();
let session_store = ExtensionData::new("session-1"); let session_store = ExtensionData::new("session-1");
@@ -1178,7 +1311,10 @@ impl GoalExtensionHarness {
/*metrics_client*/ None, /*metrics_client*/ None,
Weak::new(), Weak::new(),
Arc::clone(&goal_service), Arc::clone(&goal_service),
|_| true, |_| GoalExtensionConfig {
enabled: true,
max_goal_token_budget: None,
},
); );
let registry = builder.build(); let registry = builder.build();
let session_store = ExtensionData::new("session-1"); let session_store = ExtensionData::new("session-1");
@@ -303,6 +303,7 @@ fn new_config(model: Option<String>, arg0_paths: Arg0DispatchPaths) -> anyhow::R
background_terminal_max_timeout: 300_000, background_terminal_max_timeout: 300_000,
ghost_snapshot: GhostSnapshotConfig::default(), ghost_snapshot: GhostSnapshotConfig::default(),
multi_agent_v2: MultiAgentV2Config::default(), multi_agent_v2: MultiAgentV2Config::default(),
max_goal_token_budget: None,
token_budget: None, token_budget: None,
rollout_budget: None, rollout_budget: None,
current_time_reminder: None, current_time_reminder: None,