Add goal model tools (3 / 5) (#18075)

Adds the model-facing goal tools on top of the app-server API from PR 2.

## Why

Once goals are persisted and exposed to clients, the model needs a
small, constrained tool surface for goal workflows. The tool contract
should let the model inspect goals, create them only when explicitly
requested, and mark them complete without giving it broad control over
user/runtime-owned state.

## What changed

- Added `get_goal`, `create_goal`, and `update_goal` tool specs behind
the `goals` feature flag.
- Added core goal tool handlers that validate objectives and token
budgets before mutating persisted state.
- Constrained `create_goal` to create only when no goal exists, with
optional `token_budget` only when a budget is explicitly provided.
- Tightened the `create_goal` instructions so the model does not infer
goals from ordinary task requests.
- Constrained `update_goal` to expose only goal completion; pause,
resume, clear, and budget-limited transitions remain user- or
runtime-controlled.
- Registered the goal tools in the tool registry and kept them out of
review contexts where they should not appear.

## Verification

- Added tool-registry coverage for feature gating and tool availability.
- Added core session tests for create/get/update behavior, duplicate
goal rejection, budget validation, and completion-only updates.

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<i64>,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "snake_case")]
+struct UpdateGoalArgs {
+    status: ThreadGoalStatus,
+}
+
+#[derive(Debug, PartialEq, Serialize)]
+#[serde(rename_all = "camelCase")]
+struct GoalToolResponse {
+    goal: Option<ThreadGoal>,
+    remaining_tokens: Option<i64>,
+    completion_budget_report: Option<String>,
+}
+
+#[derive(Clone, Copy)]
+enum CompletionBudgetReport {
+    Include,
+    Omit,
+}
+
+impl GoalToolResponse {
+    fn new(goal: Option<ThreadGoal>, 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<Self::Output, FunctionCallError> {
+        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<FunctionToolOutput, FunctionCallError> {
+    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<FunctionToolOutput, FunctionCallError> {
+    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<FunctionToolOutput, FunctionCallError> {
+    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<ThreadGoal>,
+    completion_budget_report: CompletionBudgetReport,
+) -> Result<FunctionToolOutput, FunctionCallError> {
+    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<String> {
+    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,
+            }
+        );
+    }
+}

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;

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));
             }