diff --git a/codex-rs/core/src/goals.rs b/codex-rs/core/src/goals.rs new file mode 100644 index 000000000..18fa8da0a --- /dev/null +++ b/codex-rs/core/src/goals.rs @@ -0,0 +1,259 @@ +//! Core support for persisted thread goals. +//! +//! This module bridges core sessions and the state-db goal table. It validates +//! goal mutations, converts between state and protocol shapes, emits goal-update +//! events, and owns helper hooks used by goal lifecycle behavior. + +use crate::StateDbHandle; +use crate::session::session::Session; +use crate::session::turn_context::TurnContext; +use anyhow::Context; +use codex_features::Feature; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::ThreadGoal; +use codex_protocol::protocol::ThreadGoalStatus; +use codex_protocol::protocol::ThreadGoalUpdatedEvent; +use codex_rollout::state_db::reconcile_rollout; +use codex_thread_store::LocalThreadStore; + +pub(crate) struct SetGoalRequest { + pub(crate) objective: Option, + pub(crate) status: Option, + pub(crate) token_budget: Option>, +} + +pub(crate) struct CreateGoalRequest { + pub(crate) objective: String, + pub(crate) token_budget: Option, +} + +impl Session { + pub(crate) async fn get_thread_goal(&self) -> anyhow::Result> { + if !self.enabled(Feature::Goals) { + anyhow::bail!("goals feature is disabled"); + } + + let state_db = self.state_db_for_thread_goals().await?; + state_db + .get_thread_goal(self.conversation_id) + .await + .map(|goal| goal.map(protocol_goal_from_state)) + } + + pub(crate) async fn set_thread_goal( + &self, + turn_context: &TurnContext, + request: SetGoalRequest, + ) -> anyhow::Result { + if !self.enabled(Feature::Goals) { + anyhow::bail!("goals feature is disabled"); + } + + validate_goal_budget(request.token_budget.flatten())?; + let state_db = self.state_db_for_thread_goals().await?; + let goal = if let Some(objective) = request.objective { + let objective = objective.trim(); + if objective.is_empty() { + anyhow::bail!("goal objective must not be empty"); + } + state_db + .replace_thread_goal( + self.conversation_id, + objective, + request + .status + .map(state_goal_status_from_protocol) + .unwrap_or(codex_state::ThreadGoalStatus::Active), + request.token_budget.flatten(), + ) + .await? + } else { + let status = request.status.map(state_goal_status_from_protocol); + state_db + .update_thread_goal( + self.conversation_id, + codex_state::ThreadGoalUpdate { + status, + token_budget: request.token_budget, + expected_goal_id: None, + }, + ) + .await? + .ok_or_else(|| { + anyhow::anyhow!( + "cannot update goal for thread {}: no goal exists", + self.conversation_id + ) + })? + }; + + let goal = protocol_goal_from_state(goal); + self.send_event( + turn_context, + EventMsg::ThreadGoalUpdated(ThreadGoalUpdatedEvent { + thread_id: self.conversation_id, + turn_id: Some(turn_context.sub_id.clone()), + goal: goal.clone(), + }), + ) + .await; + Ok(goal) + } + + pub(crate) async fn create_thread_goal( + &self, + turn_context: &TurnContext, + request: CreateGoalRequest, + ) -> anyhow::Result { + if !self.enabled(Feature::Goals) { + anyhow::bail!("goals feature is disabled"); + } + + let CreateGoalRequest { + objective, + token_budget, + } = request; + validate_goal_budget(token_budget)?; + let objective = objective.trim(); + if objective.is_empty() { + anyhow::bail!("goal objective must not be empty"); + } + + let state_db = self.state_db_for_thread_goals().await?; + let goal = state_db + .insert_thread_goal( + self.conversation_id, + objective, + codex_state::ThreadGoalStatus::Active, + token_budget, + ) + .await? + .ok_or_else(|| { + anyhow::anyhow!( + "cannot create a new goal because thread {} already has a goal", + self.conversation_id + ) + })?; + + let goal = protocol_goal_from_state(goal); + self.send_event( + turn_context, + EventMsg::ThreadGoalUpdated(ThreadGoalUpdatedEvent { + thread_id: self.conversation_id, + turn_id: Some(turn_context.sub_id.clone()), + goal: goal.clone(), + }), + ) + .await; + Ok(goal) + } +} + +impl Session { + async fn state_db_for_thread_goals(&self) -> anyhow::Result { + let config = self.get_config().await; + if config.ephemeral { + anyhow::bail!("thread goals require a persisted thread; this thread is ephemeral"); + } + + self.try_ensure_rollout_materialized() + .await + .context("failed to materialize rollout before opening state db for thread goals")?; + + let state_db = if let Some(state_db) = self.state_db() { + state_db + } else if let Some(local_store) = self + .services + .thread_store + .as_any() + .downcast_ref::() + { + local_store.state_db().await.ok_or_else(|| { + anyhow::anyhow!( + "thread goals require a local persisted thread with a state database" + ) + })? + } else { + anyhow::bail!("thread goals require a local persisted thread with a state database"); + }; + + let thread_metadata_present = state_db + .get_thread(self.conversation_id) + .await + .context("failed to read thread metadata before reconciling thread goals")? + .is_some(); + if !thread_metadata_present { + let rollout_path = self + .current_rollout_path() + .await + .context("failed to locate rollout before reconciling thread goals")? + .ok_or_else(|| { + anyhow::anyhow!("thread goals require materialized thread metadata") + })?; + reconcile_rollout( + Some(&state_db), + rollout_path.as_path(), + config.model_provider_id.as_str(), + /*builder*/ None, + &[], + /*archived_only*/ None, + /*new_thread_memory_mode*/ None, + ) + .await; + let thread_metadata_present = state_db + .get_thread(self.conversation_id) + .await + .context("failed to read thread metadata after reconciling thread goals")? + .is_some(); + if !thread_metadata_present { + anyhow::bail!("thread metadata is unavailable after reconciling thread goals"); + } + } + + Ok(state_db) + } +} + +pub(crate) fn protocol_goal_from_state(goal: codex_state::ThreadGoal) -> ThreadGoal { + ThreadGoal { + thread_id: goal.thread_id, + objective: goal.objective, + status: protocol_goal_status_from_state(goal.status), + token_budget: goal.token_budget, + tokens_used: goal.tokens_used, + time_used_seconds: goal.time_used_seconds, + created_at: goal.created_at.timestamp(), + updated_at: goal.updated_at.timestamp(), + } +} + +pub(crate) fn protocol_goal_status_from_state( + status: codex_state::ThreadGoalStatus, +) -> ThreadGoalStatus { + match status { + codex_state::ThreadGoalStatus::Active => ThreadGoalStatus::Active, + codex_state::ThreadGoalStatus::Paused => ThreadGoalStatus::Paused, + codex_state::ThreadGoalStatus::BudgetLimited => ThreadGoalStatus::BudgetLimited, + codex_state::ThreadGoalStatus::Complete => ThreadGoalStatus::Complete, + } +} + +pub(crate) fn state_goal_status_from_protocol( + status: ThreadGoalStatus, +) -> codex_state::ThreadGoalStatus { + match status { + ThreadGoalStatus::Active => codex_state::ThreadGoalStatus::Active, + ThreadGoalStatus::Paused => codex_state::ThreadGoalStatus::Paused, + ThreadGoalStatus::BudgetLimited => codex_state::ThreadGoalStatus::BudgetLimited, + ThreadGoalStatus::Complete => codex_state::ThreadGoalStatus::Complete, + } +} + +pub(crate) fn validate_goal_budget(value: Option) -> anyhow::Result<()> { + if let Some(value) = value + && value <= 0 + { + anyhow::bail!("goal budgets must be positive when provided"); + } + Ok(()) +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 54fadc6fd..3e2d2ee52 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -37,6 +37,7 @@ pub mod file_watcher; mod flags; #[cfg(test)] mod git_info_tests; +mod goals; mod guardian; mod hook_runtime; mod installation_id; diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index 9d502ab1d..799af791e 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -24,6 +24,7 @@ pub(super) async fn spawn_review_thread( let _ = review_features.disable(Feature::WebSearchRequest); let _ = review_features.disable(Feature::WebSearchCached); let review_web_search_mode = WebSearchMode::Disabled; + let goal_tools_supported = !config.ephemeral && parent_turn_context.tools_config.goal_tools; let tools_config = ToolsConfig::new(&ToolsConfigParams { model_info: &review_model_info, available_models: &sess @@ -51,6 +52,7 @@ pub(super) async fn spawn_review_thread( .with_spawn_agent_usage_hint(config.multi_agent_v2.usage_hint_enabled) .with_spawn_agent_usage_hint_text(config.multi_agent_v2.usage_hint_text.clone()) .with_hide_spawn_agent_metadata(config.multi_agent_v2.hide_spawn_agent_metadata) + .with_goal_tools_allowed(goal_tools_supported) .with_max_concurrent_threads_per_session(config.agent_max_threads) .with_agent_type_description(crate::agent::role::spawn_tool_spec::build( &config.agent_roles, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 677f6d7a4..7046243d5 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -57,6 +57,7 @@ use crate::tasks::execute_user_shell_command; use crate::tools::ToolRouter; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; +use crate::tools::handlers::GoalHandler; use crate::tools::handlers::ShellHandler; use crate::tools::handlers::UnifiedExecHandler; use crate::tools::registry::ToolHandler; @@ -101,6 +102,7 @@ use codex_protocol::protocol::ResumedHistory; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SkillScope; use codex_protocol::protocol::Submission; +use codex_protocol::protocol::ThreadGoalStatus; use codex_protocol::protocol::ThreadRolledBackEvent; use codex_protocol::protocol::TokenCountEvent; use codex_protocol::protocol::TokenUsage; @@ -3348,6 +3350,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { session_configuration.cwd.clone(), "turn_id".to_string(), skills_outcome, + /*goal_tools_supported*/ true, ); let (mailbox, mailbox_rx) = crate::agent::Mailbox::new(); @@ -4703,6 +4706,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( session_configuration.cwd.clone(), "turn_id".to_string(), skills_outcome, + /*goal_tools_supported*/ true, )); let (mailbox, mailbox_rx) = crate::agent::Mailbox::new(); @@ -6852,6 +6856,222 @@ async fn sample_rollout( ) } +#[tokio::test] +async fn create_goal_tool_rejects_existing_goal() { + let (mut session, turn_context) = make_session_and_context().await; + let _ = session.features.enable(Feature::Goals); + let session = Arc::new(session); + upsert_goal_tool_test_thread(session.as_ref()).await; + let turn_context = Arc::new(turn_context); + let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); + let handler = GoalHandler; + + handler + .handle(ToolInvocation { + session: Arc::clone(&session), + turn: Arc::clone(&turn_context), + cancellation_token: CancellationToken::new(), + tracker: Arc::clone(&tracker), + call_id: "create-goal-1".to_string(), + tool_name: codex_tools::ToolName::plain("create_goal"), + source: ToolCallSource::Direct, + payload: ToolPayload::Function { + arguments: serde_json::json!({ + "objective": "Keep the watcher alive", + "token_budget": 123, + }) + .to_string(), + }, + }) + .await + .expect("initial create_goal should succeed"); + + let response = handler + .handle(ToolInvocation { + session: Arc::clone(&session), + turn: Arc::clone(&turn_context), + cancellation_token: CancellationToken::new(), + tracker, + call_id: "create-goal-2".to_string(), + tool_name: codex_tools::ToolName::plain("create_goal"), + source: ToolCallSource::Direct, + payload: ToolPayload::Function { + arguments: serde_json::json!({ + "objective": "Replace the watcher", + "token_budget": 456, + }) + .to_string(), + }, + }) + .await; + + let Err(FunctionCallError::RespondToModel(output)) = response else { + panic!("expected create_goal to reject an existing goal"); + }; + assert_eq!( + output, + "cannot create a new goal because this thread already has a goal; use update_goal only when the existing goal is complete" + ); + + let goal = session + .get_thread_goal() + .await + .expect("read thread goal") + .expect("goal should still exist"); + assert_eq!(goal.objective, "Keep the watcher alive"); + assert_eq!(goal.token_budget, Some(123)); +} + +#[tokio::test] +async fn update_goal_tool_rejects_pausing_goal() { + let (mut session, turn_context) = make_session_and_context().await; + let _ = session.features.enable(Feature::Goals); + let session = Arc::new(session); + upsert_goal_tool_test_thread(session.as_ref()).await; + let turn_context = Arc::new(turn_context); + let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); + let handler = GoalHandler; + + handler + .handle(ToolInvocation { + session: Arc::clone(&session), + turn: Arc::clone(&turn_context), + cancellation_token: CancellationToken::new(), + tracker: Arc::clone(&tracker), + call_id: "create-goal".to_string(), + tool_name: codex_tools::ToolName::plain("create_goal"), + source: ToolCallSource::Direct, + payload: ToolPayload::Function { + arguments: serde_json::json!({ + "objective": "Keep the watcher alive", + "token_budget": 123, + }) + .to_string(), + }, + }) + .await + .expect("initial create_goal should succeed"); + + let response = handler + .handle(ToolInvocation { + session: Arc::clone(&session), + turn: Arc::clone(&turn_context), + cancellation_token: CancellationToken::new(), + tracker, + call_id: "pause-goal".to_string(), + tool_name: codex_tools::ToolName::plain("update_goal"), + source: ToolCallSource::Direct, + payload: ToolPayload::Function { + arguments: serde_json::json!({ + "status": "paused", + }) + .to_string(), + }, + }) + .await; + + let Err(FunctionCallError::RespondToModel(output)) = response else { + panic!("expected update_goal to reject pausing a goal"); + }; + assert_eq!( + output, + "update_goal can only mark the existing goal complete; pause, resume, and budget-limited status changes are controlled by the user or system" + ); + + let goal = session + .get_thread_goal() + .await + .expect("read thread goal") + .expect("goal should still exist"); + assert_eq!(goal.status, ThreadGoalStatus::Active); +} + +#[tokio::test] +async fn update_goal_tool_marks_goal_complete() { + let (mut session, turn_context) = make_session_and_context().await; + let _ = session.features.enable(Feature::Goals); + let session = Arc::new(session); + upsert_goal_tool_test_thread(session.as_ref()).await; + let turn_context = Arc::new(turn_context); + let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); + let handler = GoalHandler; + + handler + .handle(ToolInvocation { + session: Arc::clone(&session), + turn: Arc::clone(&turn_context), + cancellation_token: CancellationToken::new(), + tracker: Arc::clone(&tracker), + call_id: "create-goal".to_string(), + tool_name: codex_tools::ToolName::plain("create_goal"), + source: ToolCallSource::Direct, + payload: ToolPayload::Function { + arguments: serde_json::json!({ + "objective": "Keep the watcher alive", + "token_budget": 123, + }) + .to_string(), + }, + }) + .await + .expect("initial create_goal should succeed"); + + handler + .handle(ToolInvocation { + session: Arc::clone(&session), + turn: Arc::clone(&turn_context), + cancellation_token: CancellationToken::new(), + tracker, + call_id: "complete-goal".to_string(), + tool_name: codex_tools::ToolName::plain("update_goal"), + source: ToolCallSource::Direct, + payload: ToolPayload::Function { + arguments: serde_json::json!({ + "status": "complete", + }) + .to_string(), + }, + }) + .await + .expect("update_goal should mark the goal complete"); + + let goal = session + .get_thread_goal() + .await + .expect("read thread goal") + .expect("goal should still exist"); + assert_eq!(goal.status, ThreadGoalStatus::Complete); +} + +async fn upsert_goal_tool_test_thread(session: &Session) { + let config = session.get_config().await; + let state_db = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); + let mut builder = codex_state::ThreadMetadataBuilder::new( + session.conversation_id, + config + .codex_home + .join("goal-tool-test-rollout.jsonl") + .to_path_buf(), + chrono::Utc::now(), + SessionSource::Exec, + ); + builder.cwd = config.cwd.to_path_buf(); + builder.model_provider = Some(config.model_provider_id.clone()); + builder.cli_version = Some(env!("CARGO_PKG_VERSION").to_string()); + builder.sandbox_policy = config.permissions.sandbox_policy.get().clone(); + builder.approval_mode = config.permissions.approval_policy.value(); + let metadata = builder.build(config.model_provider_id.as_str()); + state_db + .upsert_thread(&metadata) + .await + .expect("thread metadata should be upserted"); +} + #[tokio::test] async fn rejects_escalated_permissions_when_policy_not_on_request() { use crate::exec::ExecParams; diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index d2e6b5a21..11292c81c 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -180,6 +180,7 @@ impl TurnContext { .with_spawn_agent_usage_hint(config.multi_agent_v2.usage_hint_enabled) .with_spawn_agent_usage_hint_text(config.multi_agent_v2.usage_hint_text.clone()) .with_hide_spawn_agent_metadata(config.multi_agent_v2.hide_spawn_agent_metadata) + .with_goal_tools_allowed(self.tools_config.goal_tools) .with_max_concurrent_threads_per_session(config.agent_max_threads) .with_agent_type_description(crate::agent::role::spawn_tool_spec::build( &config.agent_roles, @@ -405,6 +406,7 @@ impl Session { cwd: AbsolutePathBuf, sub_id: String, skills_outcome: Arc, + goal_tools_supported: bool, ) -> TurnContext { let reasoning_effort = session_configuration.collaboration_mode.reasoning_effort(); let reasoning_summary = session_configuration @@ -441,6 +443,7 @@ impl Session { .with_spawn_agent_usage_hint(per_turn_config.multi_agent_v2.usage_hint_enabled) .with_spawn_agent_usage_hint_text(per_turn_config.multi_agent_v2.usage_hint_text.clone()) .with_hide_spawn_agent_metadata(per_turn_config.multi_agent_v2.hide_spawn_agent_metadata) + .with_goal_tools_allowed(goal_tools_supported) .with_max_concurrent_threads_per_session(per_turn_config.agent_max_threads) .with_agent_type_description(crate::agent::role::spawn_tool_spec::build( &per_turn_config.agent_roles, @@ -653,6 +656,7 @@ impl Session { .skills_for_config(&skills_input, fs) .await, ); + let goal_tools_supported = !per_turn_config.ephemeral && self.state_db().is_some(); let mut turn_context: TurnContext = Self::make_turn_context( self.conversation_id, Some(Arc::clone(&self.services.auth_manager)), @@ -679,6 +683,7 @@ impl Session { cwd, sub_id, skills_outcome, + goal_tools_supported, ); turn_context.realtime_active = self.conversation.running_state().await.is_some(); diff --git a/codex-rs/core/src/tools/handlers/goal.rs b/codex-rs/core/src/tools/handlers/goal.rs new file mode 100644 index 000000000..90911fb77 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/goal.rs @@ -0,0 +1,276 @@ +//! Built-in model tool handlers for persisted thread goals. +//! +//! The public tool contract intentionally splits goal creation from completion: +//! `create_goal` starts an active objective, while `update_goal` can only mark +//! the existing goal complete. + +use crate::function_tool::FunctionCallError; +use crate::goals::CreateGoalRequest; +use crate::goals::SetGoalRequest; +use crate::session::session::Session; +use crate::session::turn_context::TurnContext; +use crate::tools::context::FunctionToolOutput; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; +use crate::tools::handlers::parse_arguments; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; +use codex_protocol::protocol::ThreadGoal; +use codex_protocol::protocol::ThreadGoalStatus; +use codex_tools::CREATE_GOAL_TOOL_NAME; +use codex_tools::GET_GOAL_TOOL_NAME; +use codex_tools::UPDATE_GOAL_TOOL_NAME; +use serde::Deserialize; +use serde::Serialize; +use std::fmt::Write as _; + +pub struct GoalHandler; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "snake_case")] +struct CreateGoalArgs { + objective: String, + token_budget: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "snake_case")] +struct UpdateGoalArgs { + status: ThreadGoalStatus, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +struct GoalToolResponse { + goal: Option, + remaining_tokens: Option, + completion_budget_report: Option, +} + +#[derive(Clone, Copy)] +enum CompletionBudgetReport { + Include, + Omit, +} + +impl GoalToolResponse { + fn new(goal: Option, report_mode: CompletionBudgetReport) -> Self { + let remaining_tokens = goal.as_ref().and_then(|goal| { + goal.token_budget + .map(|budget| (budget - goal.tokens_used).max(0)) + }); + let completion_budget_report = match report_mode { + CompletionBudgetReport::Include => goal + .as_ref() + .filter(|goal| goal.status == ThreadGoalStatus::Complete) + .and_then(completion_budget_report), + CompletionBudgetReport::Omit => None, + }; + Self { + goal, + remaining_tokens, + completion_budget_report, + } + } +} + +impl ToolHandler for GoalHandler { + type Output = FunctionToolOutput; + + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + async fn handle(&self, invocation: ToolInvocation) -> Result { + let ToolInvocation { + session, + turn, + payload, + tool_name, + .. + } = invocation; + + let arguments = match payload { + ToolPayload::Function { arguments } => arguments, + _ => { + return Err(FunctionCallError::RespondToModel( + "goal handler received unsupported payload".to_string(), + )); + } + }; + + match tool_name.name.as_str() { + GET_GOAL_TOOL_NAME => handle_get_goal(session.as_ref()).await, + CREATE_GOAL_TOOL_NAME => { + handle_create_goal(session.as_ref(), turn.as_ref(), &arguments).await + } + UPDATE_GOAL_TOOL_NAME => { + handle_update_goal(session.as_ref(), turn.as_ref(), &arguments).await + } + other => Err(FunctionCallError::Fatal(format!( + "goal handler received unsupported tool: {other}" + ))), + } + } +} + +async fn handle_get_goal(session: &Session) -> Result { + let goal = session + .get_thread_goal() + .await + .map_err(|err| FunctionCallError::RespondToModel(format_goal_error(err)))?; + goal_response(goal, CompletionBudgetReport::Omit) +} + +async fn handle_create_goal( + session: &Session, + turn_context: &TurnContext, + arguments: &str, +) -> Result { + let args: CreateGoalArgs = parse_arguments(arguments)?; + let goal = session + .create_thread_goal( + turn_context, + CreateGoalRequest { + objective: args.objective, + token_budget: args.token_budget, + }, + ) + .await + .map_err(|err| { + if err + .chain() + .any(|cause| cause.to_string().contains("already has a goal")) + { + FunctionCallError::RespondToModel( + "cannot create a new goal because this thread already has a goal; use update_goal only when the existing goal is complete" + .to_string(), + ) + } else { + FunctionCallError::RespondToModel(format_goal_error(err)) + } + })?; + goal_response(Some(goal), CompletionBudgetReport::Omit) +} + +async fn handle_update_goal( + session: &Session, + turn_context: &TurnContext, + arguments: &str, +) -> Result { + let args: UpdateGoalArgs = parse_arguments(arguments)?; + if args.status != ThreadGoalStatus::Complete { + return Err(FunctionCallError::RespondToModel( + "update_goal can only mark the existing goal complete; pause, resume, and budget-limited status changes are controlled by the user or system" + .to_string(), + )); + } + let goal = session + .set_thread_goal( + turn_context, + SetGoalRequest { + objective: None, + status: Some(ThreadGoalStatus::Complete), + token_budget: None, + }, + ) + .await + .map_err(|err| FunctionCallError::RespondToModel(format_goal_error(err)))?; + goal_response(Some(goal), CompletionBudgetReport::Include) +} + +fn format_goal_error(err: anyhow::Error) -> String { + let mut message = err.to_string(); + for cause in err.chain().skip(1) { + let _ = write!(message, ": {cause}"); + } + message +} + +fn goal_response( + goal: Option, + completion_budget_report: CompletionBudgetReport, +) -> Result { + let response = + serde_json::to_string_pretty(&GoalToolResponse::new(goal, completion_budget_report)) + .map_err(|err| FunctionCallError::Fatal(err.to_string()))?; + Ok(FunctionToolOutput::from_text(response, Some(true))) +} + +fn completion_budget_report(goal: &ThreadGoal) -> Option { + let mut parts = Vec::new(); + if let Some(budget) = goal.token_budget { + parts.push(format!("tokens used: {} of {budget}", goal.tokens_used)); + } + if goal.time_used_seconds > 0 { + parts.push(format!("time used: {} seconds", goal.time_used_seconds)); + } + if parts.is_empty() { + None + } else { + Some(format!( + "Goal achieved. Report final budget usage to the user: {}.", + parts.join("; ") + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use codex_protocol::ThreadId; + use pretty_assertions::assert_eq; + + #[test] + fn completed_budgeted_goal_response_reports_final_usage() { + let goal = ThreadGoal { + thread_id: ThreadId::new(), + objective: "Keep optimizing".to_string(), + status: ThreadGoalStatus::Complete, + token_budget: Some(10_000), + tokens_used: 3_250, + time_used_seconds: 75, + created_at: 1, + updated_at: 2, + }; + + let response = GoalToolResponse::new(Some(goal.clone()), CompletionBudgetReport::Include); + + assert_eq!( + response, + GoalToolResponse { + goal: Some(goal), + remaining_tokens: Some(6_750), + completion_budget_report: Some( + "Goal achieved. Report final budget usage to the user: tokens used: 3250 of 10000; time used: 75 seconds." + .to_string() + ), + } + ); + } + + #[test] + fn completed_unbudgeted_goal_response_omits_budget_report() { + let goal = ThreadGoal { + thread_id: ThreadId::new(), + objective: "Write a poem".to_string(), + status: ThreadGoalStatus::Complete, + token_budget: None, + tokens_used: 120, + time_used_seconds: 0, + created_at: 1, + updated_at: 2, + }; + + let response = GoalToolResponse::new(Some(goal.clone()), CompletionBudgetReport::Include); + + assert_eq!( + response, + GoalToolResponse { + goal: Some(goal), + remaining_tokens: None, + completion_budget_report: None, + } + ); + } +} diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 757b0d94b..f96b49ad4 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -1,6 +1,7 @@ pub(crate) mod agent_jobs; pub(crate) mod apply_patch; mod dynamic; +mod goal; mod list_dir; mod mcp; mod mcp_resource; @@ -36,6 +37,7 @@ pub use apply_patch::ApplyPatchHandler; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::protocol::AskForApproval; pub use dynamic::DynamicToolHandler; +pub use goal::GoalHandler; pub use list_dir::ListDirHandler; pub use mcp::McpHandler; pub use mcp_resource::McpResourceHandler; diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index e556cab30..ebbc38b8b 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -80,6 +80,7 @@ pub(crate) fn build_specs_with_discoverable_tools( use crate::tools::handlers::CodeModeExecuteHandler; use crate::tools::handlers::CodeModeWaitHandler; use crate::tools::handlers::DynamicToolHandler; + use crate::tools::handlers::GoalHandler; use crate::tools::handlers::ListDirHandler; use crate::tools::handlers::McpHandler; use crate::tools::handlers::McpResourceHandler; @@ -148,6 +149,7 @@ pub(crate) fn build_specs_with_discoverable_tools( let plan_handler = Arc::new(PlanHandler); let apply_patch_handler = Arc::new(ApplyPatchHandler); let dynamic_tool_handler = Arc::new(DynamicToolHandler); + let goal_handler = Arc::new(GoalHandler); let view_image_handler = Arc::new(ViewImageHandler); let mcp_handler = Arc::new(McpHandler); let mcp_resource_handler = Arc::new(McpResourceHandler); @@ -208,6 +210,9 @@ pub(crate) fn build_specs_with_discoverable_tools( ToolHandlerKind::FollowupTaskV2 => { builder.register_handler(handler.name, Arc::new(FollowupTaskHandlerV2)); } + ToolHandlerKind::Goal => { + builder.register_handler(handler.name, goal_handler.clone()); + } ToolHandlerKind::ListAgentsV2 => { builder.register_handler(handler.name, Arc::new(ListAgentsHandlerV2)); } diff --git a/codex-rs/tools/src/goal_tool.rs b/codex-rs/tools/src/goal_tool.rs new file mode 100644 index 000000000..489fd8db3 --- /dev/null +++ b/codex-rs/tools/src/goal_tool.rs @@ -0,0 +1,112 @@ +//! Responses API tool definitions for persisted thread goals. +//! +//! These specs expose goal read/update primitives to the model while keeping +//! usage accounting system-managed. + +use crate::JsonSchema; +use crate::ResponsesApiTool; +use crate::ToolSpec; +use serde_json::json; +use std::collections::BTreeMap; + +pub const GET_GOAL_TOOL_NAME: &str = "get_goal"; +pub const CREATE_GOAL_TOOL_NAME: &str = "create_goal"; +pub const UPDATE_GOAL_TOOL_NAME: &str = "update_goal"; + +pub fn create_get_goal_tool() -> ToolSpec { + ToolSpec::Function(ResponsesApiTool { + name: GET_GOAL_TOOL_NAME.to_string(), + description: "Get the current goal for this thread, including status, budgets, token and elapsed-time usage, and remaining token budget." + .to_string(), + strict: false, + defer_loading: None, + parameters: JsonSchema::object(BTreeMap::new(), Some(Vec::new()), Some(false.into())), + output_schema: None, + }) +} + +pub fn create_create_goal_tool() -> ToolSpec { + let properties = BTreeMap::from([ + ( + "objective".to_string(), + JsonSchema::string(Some( + "Required. The concrete objective to start pursuing. This starts a new active goal only when no goal is currently defined; if a goal already exists, this tool fails." + .to_string(), + )), + ), + ( + "token_budget".to_string(), + JsonSchema::integer(Some( + "Optional positive token budget for the new active goal.".to_string(), + )), + ), + ]); + + ToolSpec::Function(ResponsesApiTool { + name: CREATE_GOAL_TOOL_NAME.to_string(), + description: format!( + r#"Create a goal only when explicitly requested by the user or system/developer instructions; do not infer goals from ordinary tasks. +Set token_budget only when an explicit token budget is requested. Fails if a goal exists; use {UPDATE_GOAL_TOOL_NAME} only for status."# + ), + strict: false, + defer_loading: None, + parameters: JsonSchema::object( + properties, + /*required*/ Some(vec!["objective".to_string()]), + Some(false.into()), + ), + output_schema: None, + }) +} + +pub fn create_update_goal_tool() -> ToolSpec { + let properties = BTreeMap::from([( + "status".to_string(), + JsonSchema::string_enum( + vec![json!("complete")], + Some( + "Required. Set to complete only when the objective is achieved and no required work remains." + .to_string(), + ), + ), + )]); + + ToolSpec::Function(ResponsesApiTool { + name: UPDATE_GOAL_TOOL_NAME.to_string(), + description: r#"Update the existing goal. +Use this tool only to mark the goal achieved. +Set status to `complete` only when the objective has actually been achieved and no required work remains. +Do not mark a goal complete merely because its budget is nearly exhausted or because you are stopping work. +You cannot use this tool to pause, resume, or budget-limit a goal; those status changes are controlled by the user or system. +When marking a budgeted goal achieved with status `complete`, report the final token usage from the tool result to the user."# + .to_string(), + strict: false, + defer_loading: None, + parameters: JsonSchema::object( + properties, + /*required*/ Some(vec!["status".to_string()]), + Some(false.into()), + ), + output_schema: None, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn update_goal_tool_only_exposes_complete_status() { + let ToolSpec::Function(tool) = create_update_goal_tool() else { + panic!("update_goal should be a function tool"); + }; + let status = tool + .parameters + .properties + .as_ref() + .and_then(|properties| properties.get("status")) + .expect("status property should exist"); + + assert_eq!(status.enum_values, Some(vec![json!("complete")])); + } +} diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index fe140e479..516bda685 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -6,6 +6,7 @@ mod agent_tool; mod apply_patch_tool; mod code_mode; mod dynamic_tool; +mod goal_tool; mod image_detail; mod json_schema; mod local_tool; @@ -51,6 +52,12 @@ pub use code_mode::create_wait_tool; pub use code_mode::tool_spec_to_code_mode_tool_definition; pub use codex_protocol::ToolName; pub use dynamic_tool::parse_dynamic_tool; +pub use goal_tool::CREATE_GOAL_TOOL_NAME; +pub use goal_tool::GET_GOAL_TOOL_NAME; +pub use goal_tool::UPDATE_GOAL_TOOL_NAME; +pub use goal_tool::create_create_goal_tool; +pub use goal_tool::create_get_goal_tool; +pub use goal_tool::create_update_goal_tool; pub use image_detail::can_request_original_image_detail; pub use image_detail::normalize_output_image_detail; pub use image_detail::sanitize_original_image_detail; diff --git a/codex-rs/tools/src/tool_config.rs b/codex-rs/tools/src/tool_config.rs index 8f27578c5..7520beeae 100644 --- a/codex-rs/tools/src/tool_config.rs +++ b/codex-rs/tools/src/tool_config.rs @@ -101,6 +101,7 @@ pub struct ToolsConfig { pub code_mode_only_enabled: bool, pub can_request_original_image_detail: bool, pub collab_tools: bool, + pub goal_tools: bool, pub multi_agent_v2: bool, pub hide_spawn_agent_metadata: bool, pub spawn_agent_usage_hint: bool, @@ -140,6 +141,7 @@ impl ToolsConfig { let include_code_mode = features.enabled(Feature::CodeMode); let include_code_mode_only = include_code_mode && features.enabled(Feature::CodeModeOnly); let include_collab_tools = features.enabled(Feature::Collab); + let include_goal_tools = features.enabled(Feature::Goals); let include_multi_agent_v2 = features.enabled(Feature::MultiAgentV2); let include_agent_jobs = features.enabled(Feature::SpawnCsv); let include_default_mode_request_user_input = @@ -218,6 +220,7 @@ impl ToolsConfig { code_mode_only_enabled: include_code_mode_only, can_request_original_image_detail: include_original_image_detail, collab_tools: include_collab_tools, + goal_tools: include_goal_tools, multi_agent_v2: include_multi_agent_v2, hide_spawn_agent_metadata: false, spawn_agent_usage_hint: true, @@ -254,6 +257,11 @@ impl ToolsConfig { self } + pub fn with_goal_tools_allowed(mut self, allowed: bool) -> Self { + self.goal_tools = self.goal_tools && allowed; + self + } + pub fn with_max_concurrent_threads_per_session( mut self, max_concurrent_threads_per_session: Option, diff --git a/codex-rs/tools/src/tool_registry_plan.rs b/codex-rs/tools/src/tool_registry_plan.rs index 6b024658f..f7c10dc2c 100644 --- a/codex-rs/tools/src/tool_registry_plan.rs +++ b/codex-rs/tools/src/tool_registry_plan.rs @@ -26,8 +26,10 @@ use crate::create_apply_patch_json_tool; use crate::create_close_agent_tool_v1; use crate::create_close_agent_tool_v2; use crate::create_code_mode_tool; +use crate::create_create_goal_tool; use crate::create_exec_command_tool; use crate::create_followup_task_tool; +use crate::create_get_goal_tool; use crate::create_image_generation_tool; use crate::create_list_agents_tool; use crate::create_list_dir_tool; @@ -49,6 +51,7 @@ use crate::create_spawn_agents_on_csv_tool; use crate::create_test_sync_tool; use crate::create_tool_search_tool; use crate::create_tool_suggest_tool; +use crate::create_update_goal_tool; use crate::create_update_plan_tool; use crate::create_view_image_tool; use crate::create_wait_agent_tool_v1; @@ -215,6 +218,26 @@ pub fn build_tool_registry_plan( config.code_mode_enabled, ); plan.register_handler("update_plan", ToolHandlerKind::Plan); + if config.goal_tools { + plan.push_spec( + create_get_goal_tool(), + /*supports_parallel_tool_calls*/ false, + config.code_mode_enabled, + ); + plan.register_handler("get_goal", ToolHandlerKind::Goal); + plan.push_spec( + create_create_goal_tool(), + /*supports_parallel_tool_calls*/ false, + config.code_mode_enabled, + ); + plan.register_handler("create_goal", ToolHandlerKind::Goal); + plan.push_spec( + create_update_goal_tool(), + /*supports_parallel_tool_calls*/ false, + config.code_mode_enabled, + ); + plan.register_handler("update_goal", ToolHandlerKind::Goal); + } plan.push_spec( create_request_user_input_tool(request_user_input_tool_description( diff --git a/codex-rs/tools/src/tool_registry_plan_tests.rs b/codex-rs/tools/src/tool_registry_plan_tests.rs index c08666a2c..24a1ea9ce 100644 --- a/codex-rs/tools/src/tool_registry_plan_tests.rs +++ b/codex-rs/tools/src/tool_registry_plan_tests.rs @@ -104,6 +104,15 @@ fn test_full_toolset_specs_for_gpt5_codex_unified_exec_web_search() { ] { expected.insert(spec.name().to_string(), spec); } + if config.goal_tools { + for spec in [ + create_get_goal_tool(), + create_create_goal_tool(), + create_update_goal_tool(), + ] { + expected.insert(spec.name().to_string(), spec); + } + } let collab_specs = if config.multi_agent_v2 { vec![ create_spawn_agent_tool_v2(spawn_agent_tool_options(&config)), @@ -186,6 +195,51 @@ fn test_build_specs_collab_tools_enabled() { assert!(!properties.contains_key("fork_turns")); } +#[test] +fn goal_tools_require_goals_feature() { + let model_info = model_info(); + let available_models = Vec::new(); + let mut features = Features::with_defaults(); + let tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + available_models: &available_models, + features: &features, + image_generation_tool_auth_allowed: true, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + sandbox_policy: &SandboxPolicy::DangerFullAccess, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + }); + let (tools, _) = build_specs( + &tools_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ); + assert_lacks_tool_name(&tools, "get_goal"); + assert_lacks_tool_name(&tools, "create_goal"); + assert_lacks_tool_name(&tools, "update_goal"); + + features.enable(Feature::Goals); + let tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + available_models: &available_models, + features: &features, + image_generation_tool_auth_allowed: true, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + sandbox_policy: &SandboxPolicy::DangerFullAccess, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + }); + let (tools, _) = build_specs( + &tools_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[], + ); + assert_contains_tool_names(&tools, &["get_goal", "create_goal", "update_goal"]); +} + #[test] fn test_build_specs_multi_agent_v2_uses_task_names_and_hides_resume() { let model_info = model_info(); diff --git a/codex-rs/tools/src/tool_registry_plan_types.rs b/codex-rs/tools/src/tool_registry_plan_types.rs index 724f1bd15..d22335b61 100644 --- a/codex-rs/tools/src/tool_registry_plan_types.rs +++ b/codex-rs/tools/src/tool_registry_plan_types.rs @@ -18,6 +18,7 @@ pub enum ToolHandlerKind { CodeModeWait, DynamicTool, FollowupTaskV2, + Goal, ListAgentsV2, ListDir, Mcp,