use std::sync::Arc;use codex_extension_api::FunctionCallError;use codex_extension_api::JsonToolOutput;use codex_extension_api::ToolCall;use codex_extension_api::ToolExecutor;use codex_extension_api::ToolName;use codex_extension_api::ToolOutput;use codex_extension_api::ToolSpec;use codex_protocol::ThreadId;use codex_protocol::protocol::ThreadGoal;use codex_protocol::protocol::ThreadGoalStatus;use codex_protocol::protocol::validate_thread_goal_objective;use serde::Deserialize;use serde::Serialize;use crate::accounting::BudgetLimitedGoalDisposition;use crate::accounting::GoalAccountingState;use crate::analytics::GoalAnalytics;use crate::analytics::GoalEventAttribution;use crate::events::GoalEventEmitter;use crate::metrics::GoalMetrics;use crate::spec::CREATE_GOAL_TOOL_NAME;use crate::spec::GET_GOAL_TOOL_NAME;use crate::spec::UPDATE_GOAL_TOOL_NAME;use crate::spec::create_create_goal_tool;use crate::spec::create_get_goal_tool;use crate::spec::create_update_goal_tool;#[derive(Clone)]pub(crate) struct GoalToolExecutor {    kind: GoalToolKind,    thread_id: ThreadId,    state_db: Arc<codex_state::StateRuntime>,    accounting_state: Arc<GoalAccountingState>,    analytics: GoalAnalytics,    event_emitter: GoalEventEmitter,    metrics: GoalMetrics,}#[derive(Clone, Copy)]enum GoalToolKind {    Get,    Create,    Update,}#[derive(Debug, Deserialize)]#[serde(rename_all = "snake_case")]pub struct CreateGoalRequest {    pub objective: String,    pub 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 GoalToolExecutor {    pub(crate) fn get(        thread_id: ThreadId,        state_db: Arc<codex_state::StateRuntime>,        accounting_state: Arc<GoalAccountingState>,        analytics: GoalAnalytics,        event_emitter: GoalEventEmitter,        metrics: GoalMetrics,    ) -> Self {        Self {            kind: GoalToolKind::Get,            thread_id,            state_db,            accounting_state,            analytics,            event_emitter,            metrics,        }    }    pub(crate) fn create(        thread_id: ThreadId,        state_db: Arc<codex_state::StateRuntime>,        accounting_state: Arc<GoalAccountingState>,        analytics: GoalAnalytics,        event_emitter: GoalEventEmitter,        metrics: GoalMetrics,    ) -> Self {        Self {            kind: GoalToolKind::Create,            thread_id,            state_db,            accounting_state,            analytics,            event_emitter,            metrics,        }    }    pub(crate) fn update(        thread_id: ThreadId,        state_db: Arc<codex_state::StateRuntime>,        accounting_state: Arc<GoalAccountingState>,        analytics: GoalAnalytics,        event_emitter: GoalEventEmitter,        metrics: GoalMetrics,    ) -> Self {        Self {            kind: GoalToolKind::Update,            thread_id,            state_db,            accounting_state,            analytics,            event_emitter,            metrics,        }    }}impl ToolExecutor<ToolCall> for GoalToolExecutor {    fn tool_name(&self) -> ToolName {        ToolName::plain(match self.kind {            GoalToolKind::Get => GET_GOAL_TOOL_NAME,            GoalToolKind::Create => CREATE_GOAL_TOOL_NAME,            GoalToolKind::Update => UPDATE_GOAL_TOOL_NAME,        })    }    fn spec(&self) -> ToolSpec {        match self.kind {            GoalToolKind::Get => create_get_goal_tool(),            GoalToolKind::Create => create_create_goal_tool(),            GoalToolKind::Update => create_update_goal_tool(),        }    }    fn handle(&self, invocation: ToolCall) -> codex_extension_api::ToolExecutorFuture<'_> {        Box::pin(async move {            match self.kind {                GoalToolKind::Get => self.handle_get(invocation).await,                GoalToolKind::Create => self.handle_create(invocation).await,                GoalToolKind::Update => self.handle_update(invocation).await,            }        })    }}impl GoalToolExecutor {    async fn handle_get(        &self,        invocation: ToolCall,    ) -> Result<Box<dyn ToolOutput>, FunctionCallError> {        let _ = invocation.function_arguments()?;        let goal = self            .state_db            .thread_goals()            .get_thread_goal(self.thread_id)            .await            .map(|goal| goal.map(protocol_goal_from_state))            .map_err(|err| {                FunctionCallError::RespondToModel(format!("failed to read goal: {err}"))            })?;        goal_response(goal, CompletionBudgetReport::Omit)    }    async fn handle_create(        &self,        invocation: ToolCall,    ) -> Result<Box<dyn ToolOutput>, FunctionCallError> {        let mut request: CreateGoalRequest = parse_arguments(invocation.function_arguments()?)?;        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)?;        let goal = self            .state_db            .thread_goals()            .insert_thread_goal(                self.thread_id,                request.objective.as_str(),                codex_state::ThreadGoalStatus::Active,                request.token_budget,            )            .await            .map_err(|err| FunctionCallError::RespondToModel(format!("failed to create goal: {err}")))?            .ok_or_else(|| {                FunctionCallError::RespondToModel(                    "cannot create a new goal because this thread has an unfinished goal; complete the existing goal first"                        .to_string(),                )            })?;        fill_empty_thread_preview_if_possible(self.state_db.as_ref(), self.thread_id, &goal).await;        let turn_id = self            .accounting_state            .mark_current_turn_goal_active(goal.goal_id.clone());        self.metrics.record_created();        self.analytics.created(            &goal,            GoalEventAttribution::Turn(invocation.turn_id.as_str()),        );        let goal = protocol_goal_from_state(goal);        self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone());        goal_response(Some(goal), CompletionBudgetReport::Omit)    }    async fn handle_update(        &self,        invocation: ToolCall,    ) -> Result<Box<dyn ToolOutput>, FunctionCallError> {        let args: UpdateGoalArgs = parse_arguments(invocation.function_arguments()?)?;        if !matches!(            args.status,            ThreadGoalStatus::Complete | ThreadGoalStatus::Blocked        ) {            return Err(FunctionCallError::RespondToModel(                "update_goal can only mark the existing goal complete or blocked; pause, resume, budget-limited, and usage-limited status changes are controlled by the user or system"                    .to_string(),            ));        }        self.account_active_goal_progress(            match args.status {                ThreadGoalStatus::Complete => codex_state::GoalAccountingMode::ActiveOrComplete,                ThreadGoalStatus::Blocked => codex_state::GoalAccountingMode::ActiveOrStopped,                ThreadGoalStatus::Active                | ThreadGoalStatus::Paused                | ThreadGoalStatus::UsageLimited                | ThreadGoalStatus::BudgetLimited => unreachable!("status validated above"),            },            invocation.call_id.as_str(),            BudgetLimitedGoalDisposition::ClearActive,        )        .await?;        let previous_status = self            .current_goal_status_for_metrics(/*expected_goal_id*/ None)            .await?;        let goal = self            .state_db            .thread_goals()            .update_thread_goal(                self.thread_id,                codex_state::GoalUpdate {                    objective: None,                    status: Some(state_status_from_protocol(args.status)),                    token_budget: None,                    expected_goal_id: None,                },            )            .await            .map_err(|err| {                FunctionCallError::RespondToModel(format!("failed to update goal: {err}"))            })?            .ok_or_else(|| {                FunctionCallError::RespondToModel(                    "cannot update goal because this thread has no goal".to_string(),                )            })?;        self.metrics            .record_terminal_if_status_changed(previous_status, &goal);        self.analytics.status_changed(            &goal,            previous_status,            GoalEventAttribution::Turn(invocation.turn_id.as_str()),        );        let goal = protocol_goal_from_state(goal);        let turn_id = self.accounting_state.clear_current_turn_goal();        self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone());        goal_response(            Some(goal),            if args.status == ThreadGoalStatus::Complete {                CompletionBudgetReport::Include            } else {                CompletionBudgetReport::Omit            },        )    }    fn emit_goal_updated_from_tool_call(        &self,        invocation: &ToolCall,        turn_id: Option<String>,        goal: ThreadGoal,    ) {        self.event_emitter            .thread_goal_updated(invocation.call_id.clone(), turn_id, goal);    }    async fn account_active_goal_progress(        &self,        mode: codex_state::GoalAccountingMode,        event_id: &str,        budget_limited_goal_disposition: BudgetLimitedGoalDisposition,    ) -> Result<Option<ThreadGoal>, FunctionCallError> {        let Some(turn_id) = self.accounting_state.current_turn_id() else {            return Ok(None);        };        let _accounting_permit = self            .accounting_state            .progress_accounting_permit()            .await            .map_err(|err| {                FunctionCallError::Fatal(format!(                    "goal progress accounting semaphore closed: {err}"                ))            })?;        let Some(snapshot) = self.accounting_state.progress_snapshot(turn_id.as_str()) else {            return Ok(None);        };        let previous_status = self            .current_goal_status_for_metrics(Some(snapshot.expected_goal_id.as_str()))            .await?;        let outcome = self            .state_db            .thread_goals()            .account_thread_goal_usage(                self.thread_id,                snapshot.time_delta_seconds,                snapshot.token_delta,                mode,                Some(snapshot.expected_goal_id.as_str()),            )            .await            .map_err(|err| {                FunctionCallError::RespondToModel(format!("failed to account goal progress: {err}"))            })?;        Ok(match outcome {            codex_state::GoalAccountingOutcome::Updated(goal) => {                self.metrics                    .record_terminal_if_status_changed(previous_status, &goal);                self.analytics                    .usage_accounted(&goal, GoalEventAttribution::Turn(turn_id.as_str()));                self.analytics.status_changed(                    &goal,                    previous_status,                    GoalEventAttribution::Turn(turn_id.as_str()),                );                self.accounting_state.mark_progress_accounted_for_status(                    turn_id.as_str(),                    &snapshot,                    goal.status,                    budget_limited_goal_disposition,                );                let goal = protocol_goal_from_state(goal);                self.event_emitter.thread_goal_updated(                    event_id.to_string(),                    Some(turn_id),                    goal.clone(),                );                Some(goal)            }            codex_state::GoalAccountingOutcome::Unchanged(_) => None,        })    }    async fn current_goal_status_for_metrics(        &self,        expected_goal_id: Option<&str>,    ) -> Result<Option<codex_state::ThreadGoalStatus>, FunctionCallError> {        let goal = self            .state_db            .thread_goals()            .get_thread_goal(self.thread_id)            .await            .map_err(|err| {                FunctionCallError::RespondToModel(format!(                    "failed to read goal metrics status: {err}"                ))            })?;        Ok(goal.and_then(|goal| {            expected_goal_id                .is_none_or(|expected_goal_id| goal.goal_id == expected_goal_id)                .then_some(goal.status)        }))    }}fn parse_arguments<T>(arguments: &str) -> Result<T, FunctionCallError>where    T: for<'de> Deserialize<'de>,{    serde_json::from_str(arguments)        .map_err(|err| FunctionCallError::RespondToModel(err.to_string()))}pub(crate) fn validate_goal_budget(value: Option<i64>) -> Result<(), String> {    if let Some(value) = value        && value <= 0    {        return Err("goal budgets must be positive when provided".to_string());    }    Ok(())}fn goal_response(    goal: Option<ThreadGoal>,    completion_budget_report: CompletionBudgetReport,) -> Result<Box<dyn ToolOutput>, FunctionCallError> {    let value = serde_json::to_value(GoalToolResponse::new(goal, completion_budget_report))        .map_err(|err| FunctionCallError::Fatal(err.to_string()))?;    Ok(Box::new(JsonToolOutput::new(value)))}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,        }    }}pub(crate) async fn fill_empty_thread_preview_if_possible(    state_db: &codex_state::StateRuntime,    thread_id: ThreadId,    goal: &codex_state::ThreadGoal,) {    if let Err(err) = state_db        .set_thread_preview_if_empty(thread_id, goal.objective.as_str())        .await    {        tracing::warn!(            "failed to set empty thread preview from goal objective for {thread_id}: {err}"        );    }}pub(crate) fn protocol_goal_from_state(goal: codex_state::ThreadGoal) -> ThreadGoal {    ThreadGoal {        thread_id: goal.thread_id,        objective: goal.objective,        status: protocol_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(),    }}fn protocol_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::Blocked => ThreadGoalStatus::Blocked,        codex_state::ThreadGoalStatus::UsageLimited => ThreadGoalStatus::UsageLimited,        codex_state::ThreadGoalStatus::BudgetLimited => ThreadGoalStatus::BudgetLimited,        codex_state::ThreadGoalStatus::Complete => ThreadGoalStatus::Complete,    }}pub(crate) fn state_status_from_protocol(    status: ThreadGoalStatus,) -> codex_state::ThreadGoalStatus {    match status {        ThreadGoalStatus::Active => codex_state::ThreadGoalStatus::Active,        ThreadGoalStatus::Paused => codex_state::ThreadGoalStatus::Paused,        ThreadGoalStatus::Blocked => codex_state::ThreadGoalStatus::Blocked,        ThreadGoalStatus::UsageLimited => codex_state::ThreadGoalStatus::UsageLimited,        ThreadGoalStatus::BudgetLimited => codex_state::ThreadGoalStatus::BudgetLimited,        ThreadGoalStatus::Complete => codex_state::ThreadGoalStatus::Complete,    }}fn completion_budget_report(goal: &ThreadGoal) -> Option<String> {    if goal.token_budget.is_none() && goal.time_used_seconds <= 0 {        None    } else {        Some(            "Goal achieved. Report final usage from this tool result's structured goal fields. If `goal.tokenBudget` is present, include token usage from `goal.tokensUsed` and `goal.tokenBudget`. If `goal.timeUsedSeconds` is greater than 0, summarize elapsed time in a concise, human-friendly form appropriate to the response language."                .to_string(),        )    }}