mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
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:
@@ -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",
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<AgentsToml>,
|
||||
|
||||
/// Goal-related settings.
|
||||
pub goals: Option<GoalsToml>,
|
||||
|
||||
/// Memories subsystem settings.
|
||||
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)]
|
||||
#[schemars(deny_unknown_fields)]
|
||||
pub struct AgentsToml {
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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");
|
||||
|
||||
@@ -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]
|
||||
fn parses_bundled_skills_config() {
|
||||
let cfg: ConfigToml = toml::from_str(
|
||||
|
||||
@@ -898,6 +898,9 @@ pub struct Config {
|
||||
/// User-defined role declarations keyed by role name.
|
||||
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.
|
||||
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,
|
||||
|
||||
@@ -54,6 +54,7 @@ pub struct GoalSetRequest<'a> {
|
||||
pub objective: GoalObjectiveUpdate<'a>,
|
||||
pub status: Option<ThreadGoalStatus>,
|
||||
pub token_budget: GoalTokenBudgetUpdate,
|
||||
pub max_goal_token_budget: Option<i64>,
|
||||
}
|
||||
|
||||
#[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| {
|
||||
|
||||
@@ -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<i64>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -64,7 +59,7 @@ pub struct GoalExtension<C> {
|
||||
metrics: GoalMetrics,
|
||||
thread_manager: Weak<ThreadManager>,
|
||||
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> {
|
||||
@@ -81,7 +76,7 @@ impl<C> GoalExtension<C> {
|
||||
metrics_client: Option<MetricsClient>,
|
||||
thread_manager: Weak<ThreadManager>,
|
||||
goal_service: Arc<GoalService>,
|
||||
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<C> GoalExtension<C> {
|
||||
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>(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::<GoalExtensionConfig>()
|
||||
.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<C>(
|
||||
metrics_client: Option<MetricsClient>,
|
||||
thread_manager: Weak<ThreadManager>,
|
||||
goal_service: Arc<GoalService>,
|
||||
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<C>(
|
||||
metrics_client,
|
||||
thread_manager,
|
||||
Arc::clone(&goal_service),
|
||||
goals_enabled,
|
||||
goal_config,
|
||||
));
|
||||
registry.thread_lifecycle_contributor(extension.clone());
|
||||
registry.config_contributor(extension.clone());
|
||||
|
||||
@@ -36,6 +36,7 @@ pub(crate) struct GoalToolExecutor {
|
||||
analytics: GoalAnalytics,
|
||||
event_emitter: GoalEventEmitter,
|
||||
metrics: GoalMetrics,
|
||||
max_goal_token_budget: Option<i64>,
|
||||
}
|
||||
|
||||
#[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<i64>,
|
||||
) -> 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<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
|
||||
&& 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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -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<codex_state::StateRuntime>,
|
||||
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");
|
||||
|
||||
@@ -303,6 +303,7 @@ fn new_config(model: Option<String>, 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,
|
||||
|
||||
Reference in New Issue
Block a user