use crate::product::agent::codex::GoalUsageSettlementMode;
use crate::product::agent::codex::Session;
use crate::product::agent::codex::TurnContext;
use crate::product::agent::function_tool::FunctionCallError;
use crate::product::agent::tools::context::ToolInvocation;
use crate::product::agent::tools::context::ToolOutput;
use crate::product::agent::tools::context::ToolPayload;
use crate::product::agent::tools::handlers::parse_arguments;
use crate::product::agent::tools::registry::ToolHandler;
use crate::product::agent::tools::registry::ToolKind;
use crate::product::protocol::config_types::IdentityKind;
use crate::product::protocol::protocol::EventMsg;
use crate::product::protocol::protocol::ThreadGoal;
use crate::product::protocol::protocol::ThreadGoalStatus;
use crate::product::protocol::protocol::ThreadGoalUpdatedEvent;
use crate::product::protocol::protocol::validate_thread_goal_objective;
use async_trait::async_trait;
use serde::Deserialize;
use serde_json::json;
pub struct GoalHandler;
#[derive(Deserialize)]
struct CreateGoalArgs {
objective: String,
}
#[derive(Deserialize)]
struct UpdateGoalArgs {
status: String,
}
#[async_trait]
impl ToolHandler for GoalHandler {
fn kind(&self) -> ToolKind {
ToolKind::Function
}
async fn handle(&self, invocation: ToolInvocation) -> Result<ToolOutput, FunctionCallError> {
let ToolInvocation {
session,
turn,
tool_name,
payload,
..
} = invocation;
let arguments = match payload {
ToolPayload::Function { arguments } => arguments,
_ => {
return Err(FunctionCallError::RespondToModel(
"goal handler received unsupported payload".to_string(),
));
}
};
let content = match tool_name.as_str() {
"get_goal" => get_goal(session.as_ref(), turn.as_ref()).await?,
"create_goal" => create_goal(session.as_ref(), turn.as_ref(), arguments).await?,
"update_goal" => update_goal(session.as_ref(), turn.as_ref(), arguments).await?,
other => {
return Err(FunctionCallError::RespondToModel(format!(
"unsupported goal tool `{other}`"
)));
}
};
Ok(ToolOutput::Function {
content,
content_items: None,
success: Some(true),
})
}
}
async fn get_goal(
session: &Session,
turn_context: &TurnContext,
) -> Result<String, FunctionCallError> {
ensure_goal_tool_allowed(session, turn_context).await?;
session
.settle_goal_usage_for_turn_context(
turn_context,
GoalUsageSettlementMode::RefreshForDisplay,
)
.await;
let state_db = state_db(session)?;
let goal = state_db
.get_thread_goal(session.conversation_id)
.await
.map_err(tool_error)?
.map(|goal| {
let goal_id = goal.goal_id.clone();
(goal_id, protocol_goal_from_state(goal))
});
if let Some((goal_id, _)) = &goal {
turn_context
.goal_context
.set_expected_goal_id(goal_id.clone())
.await;
}
let goal = goal.map(|(_, goal)| goal);
serde_json::to_string(&json!({ "goal": goal })).map_err(|err| {
FunctionCallError::RespondToModel(format!("failed to serialize goal: {err}"))
})
}
async fn create_goal(
session: &Session,
turn_context: &TurnContext,
arguments: String,
) -> Result<String, FunctionCallError> {
ensure_goal_tool_allowed(session, turn_context).await?;
let args: CreateGoalArgs = parse_arguments(&arguments)?;
validate_thread_goal_objective(&args.objective).map_err(FunctionCallError::RespondToModel)?;
let state_db = state_db(session)?;
let goal = state_db
.insert_thread_goal_or_replace_completed(
session.conversation_id,
&args.objective,
crate::product::state::ThreadGoalStatus::Active,
None,
)
.await
.map_err(tool_error)?
.ok_or_else(|| {
FunctionCallError::RespondToModel(
"an unfinished goal already exists; update or complete it instead of creating another"
.to_string(),
)
})?;
turn_context
.goal_context
.set_expected_goal_id(goal.goal_id.clone())
.await;
turn_context
.goal_context
.set_accounting_goal_id(goal.goal_id.clone())
.await;
session
.capture_goal_accounting_baseline_for_turn(turn_context)
.await;
emit_goal_updated(session, turn_context, &goal).await;
serde_json::to_string(&json!({ "goal": protocol_goal_from_state(goal) })).map_err(|err| {
FunctionCallError::RespondToModel(format!("failed to serialize goal: {err}"))
})
}
async fn update_goal(
session: &Session,
turn_context: &TurnContext,
arguments: String,
) -> Result<String, FunctionCallError> {
ensure_goal_tool_allowed(session, turn_context).await?;
let args: UpdateGoalArgs = parse_arguments(&arguments)?;
let status = match args.status.as_str() {
"complete" => crate::product::state::ThreadGoalStatus::Complete,
"blocked" => crate::product::state::ThreadGoalStatus::Blocked,
_ => {
return Err(FunctionCallError::RespondToModel(
"update_goal status must be `complete` or `blocked`".to_string(),
));
}
};
let state_db = state_db(session)?;
session
.settle_goal_usage_for_turn_context(
turn_context,
GoalUsageSettlementMode::RefreshForDisplay,
)
.await;
let expected_goal_id = turn_context.goal_context.expected_goal_id().await;
let Some(expected_goal_id) = expected_goal_id else {
let goal = state_db
.get_thread_goal(session.conversation_id)
.await
.map_err(tool_error)?;
let message = if goal.is_some() {
"call get_goal before updating the current goal"
} else {
"no goal is currently set"
};
return Err(FunctionCallError::RespondToModel(message.to_string()));
};
let goal = state_db
.update_thread_goal(
session.conversation_id,
crate::product::state::GoalUpdate {
objective: None,
status: Some(status),
token_budget: None,
expected_goal_id: Some(expected_goal_id),
},
)
.await
.map_err(tool_error)?
.ok_or_else(|| {
FunctionCallError::RespondToModel(
"goal changed before this update; call get_goal before updating the current goal"
.to_string(),
)
})?;
emit_goal_updated(session, turn_context, &goal).await;
serde_json::to_string(&json!({ "goal": protocol_goal_from_state(goal) })).map_err(|err| {
FunctionCallError::RespondToModel(format!("failed to serialize goal: {err}"))
})
}
async fn ensure_goal_tool_allowed(
session: &Session,
turn_context: &TurnContext,
) -> Result<(), FunctionCallError> {
if turn_context.identity.kind != IdentityKind::Programmer {
return Err(FunctionCallError::RespondToModel(
"goal tools are only available to the programmer identity".to_string(),
));
}
if session.state_db().is_none() {
return Err(FunctionCallError::RespondToModel(
"goal tools require a persisted session".to_string(),
));
}
Ok(())
}
fn state_db(
session: &Session,
) -> Result<crate::product::agent::state_db::StateDbHandle, FunctionCallError> {
session.state_db().ok_or_else(|| {
FunctionCallError::RespondToModel("goal tools require a persisted session".to_string())
})
}
fn tool_error(err: anyhow::Error) -> FunctionCallError {
FunctionCallError::RespondToModel(format!("goal tool failed: {err}"))
}
async fn emit_goal_updated(
session: &Session,
turn_context: &TurnContext,
goal: &crate::product::state::ThreadGoal,
) {
session
.send_event(
turn_context,
EventMsg::ThreadGoalUpdated(ThreadGoalUpdatedEvent {
thread_id: session.conversation_id,
turn_id: Some(turn_context.sub_id.clone()),
goal: protocol_goal_from_state(goal.clone()),
}),
)
.await;
}
fn protocol_goal_from_state(goal: crate::product::state::ThreadGoal) -> ThreadGoal {
ThreadGoal {
thread_id: goal.thread_id,
goal_id: goal.goal_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(),
}
}
fn protocol_goal_status_from_state(
status: crate::product::state::ThreadGoalStatus,
) -> ThreadGoalStatus {
match status {
crate::product::state::ThreadGoalStatus::Active => ThreadGoalStatus::Active,
crate::product::state::ThreadGoalStatus::Paused => ThreadGoalStatus::Paused,
crate::product::state::ThreadGoalStatus::Blocked => ThreadGoalStatus::Blocked,
crate::product::state::ThreadGoalStatus::UsageLimited => ThreadGoalStatus::UsageLimited,
crate::product::state::ThreadGoalStatus::BudgetLimited => ThreadGoalStatus::BudgetLimited,
crate::product::state::ThreadGoalStatus::Complete => ThreadGoalStatus::Complete,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::product::agent::codex::make_session_and_context;
use pretty_assertions::assert_eq;
async fn make_goal_tool_context() -> (
Session,
TurnContext,
crate::product::agent::state_db::StateDbHandle,
tempfile::TempDir,
) {
let (mut session, mut turn_context) = make_session_and_context().await;
let state_home = tempfile::tempdir().expect("create state temp dir");
let state_db = crate::product::state::StateRuntime::init(
state_home.path().to_path_buf(),
"test".to_string(),
None,
)
.await
.expect("state runtime should initialize");
session.services.state_db = Some(std::sync::Arc::clone(&state_db));
turn_context.identity.kind = IdentityKind::Programmer;
(session, turn_context, state_db, state_home)
}
fn response_goal(content: &str) -> ThreadGoal {
let value: serde_json::Value =
serde_json::from_str(content).expect("tool output should be JSON");
serde_json::from_value(value["goal"].clone()).expect("goal should deserialize")
}
#[tokio::test]
async fn create_goal_rejects_unfinished_goal() {
let (session, turn_context, state_db, _state_home) = make_goal_tool_context().await;
let first_goal = response_goal(
&create_goal(
&session,
&turn_context,
json!({ "objective": "finish current work" }).to_string(),
)
.await
.expect("first goal should be created"),
);
let err = create_goal(
&session,
&turn_context,
json!({ "objective": "start new work" }).to_string(),
)
.await
.expect_err("unfinished goal should block creation");
assert_eq!(
FunctionCallError::RespondToModel(
"an unfinished goal already exists; update or complete it instead of creating another"
.to_string()
),
err
);
let stored = state_db
.get_thread_goal(session.conversation_id)
.await
.expect("goal read should succeed")
.expect("goal should exist");
assert_eq!(first_goal, protocol_goal_from_state(stored));
}
#[tokio::test]
async fn create_goal_replaces_completed_goal() {
let (session, turn_context, state_db, _state_home) = make_goal_tool_context().await;
let first_goal = response_goal(
&create_goal(
&session,
&turn_context,
json!({ "objective": "finish current work" }).to_string(),
)
.await
.expect("first goal should be created"),
);
let completed_goal = response_goal(
&update_goal(
&session,
&turn_context,
json!({ "status": "complete" }).to_string(),
)
.await
.expect("goal should be completed"),
);
assert_eq!(ThreadGoalStatus::Complete, completed_goal.status);
let second_goal = response_goal(
&create_goal(
&session,
&turn_context,
json!({ "objective": "start new work" }).to_string(),
)
.await
.expect("completed goal should allow new creation"),
);
assert_ne!(first_goal.goal_id, second_goal.goal_id);
assert_eq!("start new work", second_goal.objective);
assert_eq!(ThreadGoalStatus::Active, second_goal.status);
assert_eq!(0, second_goal.tokens_used);
assert_eq!(0, second_goal.time_used_seconds);
let stored = state_db
.get_thread_goal(session.conversation_id)
.await
.expect("goal read should succeed")
.expect("goal should exist");
assert_eq!(second_goal, protocol_goal_from_state(stored));
assert_eq!(
Some(second_goal.goal_id),
turn_context.goal_context.expected_goal_id().await
);
}
}