From 2ec077f2662c675f5ea1d6722f6c057abc94ab16 Mon Sep 17 00:00:00 2001 From: Andrei Eternal Date: Tue, 9 Jun 2026 20:55:22 -0700 Subject: [PATCH] Add automation extension tool --- codex-rs/Cargo.lock | 23 + codex-rs/Cargo.toml | 2 + codex-rs/app-server/Cargo.toml | 1 + codex-rs/app-server/src/extensions.rs | 1 + codex-rs/ext/automation/BUILD.bazel | 6 + codex-rs/ext/automation/Cargo.toml | 34 ++ codex-rs/ext/automation/src/extension.rs | 155 +++++++ codex-rs/ext/automation/src/lib.rs | 8 + codex-rs/ext/automation/src/spec.rs | 105 +++++ codex-rs/ext/automation/src/tool.rs | 398 ++++++++++++++++++ .../automation/tests/automation_extension.rs | 251 +++++++++++ 11 files changed, 984 insertions(+) create mode 100644 codex-rs/ext/automation/BUILD.bazel create mode 100644 codex-rs/ext/automation/Cargo.toml create mode 100644 codex-rs/ext/automation/src/extension.rs create mode 100644 codex-rs/ext/automation/src/lib.rs create mode 100644 codex-rs/ext/automation/src/spec.rs create mode 100644 codex-rs/ext/automation/src/tool.rs create mode 100644 codex-rs/ext/automation/tests/automation_extension.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 59cea06edc..e24a2a8130 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1971,6 +1971,7 @@ dependencies = [ "codex-app-server-protocol", "codex-app-server-transport", "codex-arg0", + "codex-automation-extension", "codex-backend-client", "codex-chatgpt", "codex-cloud-config", @@ -2234,6 +2235,28 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "codex-automation-extension" +version = "0.0.0" +dependencies = [ + "anyhow", + "async-trait", + "chrono", + "codex-core", + "codex-extension-api", + "codex-features", + "codex-protocol", + "codex-state", + "codex-tools", + "codex-utils-absolute-path", + "pretty_assertions", + "serde", + "serde_json", + "tempfile", + "tokio", + "toml 0.9.11+spec-1.1.0", +] + [[package]] name = "codex-aws-auth" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index e5f5025e91..9d08080265 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -48,6 +48,7 @@ members = [ "execpolicy", "execpolicy-legacy", "ext/extension-api", + "ext/automation", "ext/goal", "ext/guardian", "ext/image-generation", @@ -174,6 +175,7 @@ codex-file-system = { path = "file-system" } codex-exec-server = { path = "exec-server" } codex-execpolicy = { path = "execpolicy" } codex-extension-api = { path = "ext/extension-api" } +codex-automation-extension = { path = "ext/automation" } codex-goal-extension = { path = "ext/goal" } codex-guardian = { path = "ext/guardian" } codex-image-generation-extension = { path = "ext/image-generation" } diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 3d08059a63..4519a56a64 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -41,6 +41,7 @@ codex-extension-api = { workspace = true } codex-external-agent-migration = { workspace = true } codex-external-agent-sessions = { workspace = true } codex-features = { workspace = true } +codex-automation-extension = { workspace = true } codex-goal-extension = { workspace = true } codex-guardian = { workspace = true } codex-git-utils = { workspace = true } diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index 8a6bc07f3d..e1db2b5592 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -61,6 +61,7 @@ where } = dependencies; let mut builder = ExtensionRegistryBuilder::::with_event_sink(event_sink); if let Some(state_db) = state_db { + codex_automation_extension::install(&mut builder, Arc::clone(&state_db)); codex_goal_extension::install_with_backend( &mut builder, state_db, diff --git a/codex-rs/ext/automation/BUILD.bazel b/codex-rs/ext/automation/BUILD.bazel new file mode 100644 index 0000000000..f5cb5bf68f --- /dev/null +++ b/codex-rs/ext/automation/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "automation", + crate_name = "codex_automation_extension", +) diff --git a/codex-rs/ext/automation/Cargo.toml b/codex-rs/ext/automation/Cargo.toml new file mode 100644 index 0000000000..038d0358ea --- /dev/null +++ b/codex-rs/ext/automation/Cargo.toml @@ -0,0 +1,34 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-automation-extension" +version.workspace = true + +[lib] +name = "codex_automation_extension" +path = "src/lib.rs" +test = false +doctest = false + +[lints] +workspace = true + +[dependencies] +async-trait = { workspace = true } +chrono = { workspace = true } +codex-core = { workspace = true } +codex-extension-api = { workspace = true } +codex-features = { workspace = true } +codex-protocol = { workspace = true } +codex-state = { workspace = true } +codex-tools = { workspace = true } +codex-utils-absolute-path = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } + +[dev-dependencies] +anyhow = { workspace = true } +pretty_assertions = { workspace = true } +tempfile = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +toml = { workspace = true } diff --git a/codex-rs/ext/automation/src/extension.rs b/codex-rs/ext/automation/src/extension.rs new file mode 100644 index 0000000000..b8df23bcea --- /dev/null +++ b/codex-rs/ext/automation/src/extension.rs @@ -0,0 +1,155 @@ +use std::path::PathBuf; +use std::sync::Arc; + +use async_trait::async_trait; +use codex_core::config::Config; +use codex_extension_api::ConfigContributor; +use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadStartInput; +use codex_extension_api::ToolContributor; +use codex_features::Feature; +use codex_protocol::ThreadId; +use codex_state::AutomationDispatchSettings; + +use crate::tool::AutomationUpdateTool; + +#[derive(Clone, Debug)] +pub(crate) struct AutomationThreadContext { + pub(crate) thread_id: ThreadId, + pub(crate) cwd: PathBuf, + pub(crate) workspace_roots: Vec, + pub(crate) dispatch_settings: AutomationDispatchSettings, + pub(crate) enabled: bool, + pub(crate) tools_available_for_thread: bool, +} + +impl AutomationThreadContext { + fn from_thread_start(input: &ThreadStartInput<'_, Config>) -> Option { + let thread_id = ThreadId::from_string(input.thread_store.level_id()).ok()?; + let enabled = input.config.features.enabled(Feature::Automations); + let tools_available_for_thread = input.persistent_thread_state_available + && !input.session_source.is_automation() + && !input.session_source.is_internal(); + Some(Self::from_config( + thread_id, + input.config, + enabled, + tools_available_for_thread, + )) + } + + fn from_config( + thread_id: ThreadId, + config: &Config, + enabled: bool, + tools_available_for_thread: bool, + ) -> Self { + let cwd = config.cwd.clone().into_path_buf(); + let workspace_roots = config + .effective_workspace_roots() + .into_iter() + .map(codex_utils_absolute_path::AbsolutePathBuf::into_path_buf) + .collect::>(); + let workspace_roots = if workspace_roots.is_empty() { + vec![cwd.clone()] + } else { + workspace_roots + }; + Self { + thread_id, + cwd, + workspace_roots: workspace_roots.clone(), + dispatch_settings: AutomationDispatchSettings { + workspace_roots, + approval_policy: *config.permissions.approval_policy.get(), + approvals_reviewer: config.approvals_reviewer, + permission_profile: config.permissions.effective_permission_profile(), + }, + enabled, + tools_available_for_thread, + } + } + + pub(crate) fn tools_visible(&self) -> bool { + self.enabled && self.tools_available_for_thread + } +} + +#[derive(Clone)] +struct AutomationExtension { + state_db: Arc, +} + +impl AutomationExtension { + fn new(state_db: Arc) -> Self { + Self { state_db } + } +} + +impl std::fmt::Debug for AutomationExtension { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AutomationExtension") + .finish_non_exhaustive() + } +} + +#[async_trait] +impl ThreadLifecycleContributor for AutomationExtension { + async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) { + let Some(context) = AutomationThreadContext::from_thread_start(&input) else { + return; + }; + input.thread_store.insert(context); + } +} + +impl ConfigContributor for AutomationExtension { + fn on_config_changed( + &self, + _session_store: &ExtensionData, + thread_store: &ExtensionData, + _previous_config: &Config, + new_config: &Config, + ) { + let Some(existing) = thread_store.get::() else { + return; + }; + thread_store.insert(AutomationThreadContext::from_config( + existing.thread_id, + new_config, + new_config.features.enabled(Feature::Automations), + existing.tools_available_for_thread, + )); + } +} + +impl ToolContributor for AutomationExtension { + fn tools( + &self, + _session_store: &ExtensionData, + thread_store: &ExtensionData, + ) -> Vec>> { + let Some(context) = thread_store.get::() else { + return Vec::new(); + }; + if !context.tools_visible() { + return Vec::new(); + } + vec![Arc::new(AutomationUpdateTool::new( + Arc::clone(&self.state_db), + context.as_ref().clone(), + ))] + } +} + +pub fn install( + registry: &mut ExtensionRegistryBuilder, + state_db: Arc, +) { + let extension = Arc::new(AutomationExtension::new(state_db)); + registry.thread_lifecycle_contributor(extension.clone()); + registry.config_contributor(extension.clone()); + registry.tool_contributor(extension); +} diff --git a/codex-rs/ext/automation/src/lib.rs b/codex-rs/ext/automation/src/lib.rs new file mode 100644 index 0000000000..cea022831c --- /dev/null +++ b/codex-rs/ext/automation/src/lib.rs @@ -0,0 +1,8 @@ +//! Extension crate for scheduled automations and heartbeats. + +mod extension; +mod spec; +mod tool; + +pub use extension::install; +pub use spec::AUTOMATION_UPDATE_TOOL_NAME; diff --git a/codex-rs/ext/automation/src/spec.rs b/codex-rs/ext/automation/src/spec.rs new file mode 100644 index 0000000000..303c653e3c --- /dev/null +++ b/codex-rs/ext/automation/src/spec.rs @@ -0,0 +1,105 @@ +//! Responses API tool definition for scheduled automations. + +use std::collections::BTreeMap; + +use codex_tools::JsonSchema; +use codex_tools::ResponsesApiTool; +use codex_tools::ToolSpec; +use serde_json::json; + +pub const AUTOMATION_UPDATE_TOOL_NAME: &str = "automation_update"; + +pub fn create_automation_update_tool() -> ToolSpec { + let properties = BTreeMap::from([ + ( + "mode".to_string(), + JsonSchema::string_enum( + vec![ + json!("list"), + json!("read"), + json!("create"), + json!("update"), + json!("delete"), + ], + Some("Required. Operation to perform on automations for this thread.".to_string()), + ), + ), + ( + "automation_id".to_string(), + JsonSchema::string(Some( + "Automation id. Required for read, update, and delete.".to_string(), + )), + ), + ( + "kind".to_string(), + JsonSchema::string_enum( + vec![json!("cron"), json!("heartbeat")], + Some("Automation target kind for create, or for retargeting an update.".to_string()), + ), + ), + ( + "name".to_string(), + JsonSchema::string(Some("Human-readable automation name.".to_string())), + ), + ( + "prompt".to_string(), + JsonSchema::string(Some("Prompt to submit when the automation fires.".to_string())), + ), + ( + "rrule".to_string(), + JsonSchema::string(Some( + "RFC 5545-style recurrence rule supported by Codex automations.".to_string(), + )), + ), + ( + "model".to_string(), + JsonSchema::string(Some("Optional model override for cron-created threads.".to_string())), + ), + ( + "reasoning_effort".to_string(), + JsonSchema::string(Some( + "Optional reasoning effort override for cron-created threads.".to_string(), + )), + ), + ( + "status".to_string(), + JsonSchema::string_enum( + vec![json!("ACTIVE"), json!("PAUSED")], + Some("Automation status.".to_string()), + ), + ), + ( + "cwds".to_string(), + JsonSchema::array( + JsonSchema::string(/*description*/ None), + Some( + "Cron working directories. Each path must be absolute and inside this thread's workspace roots." + .to_string(), + ), + ), + ), + ( + "thread_id".to_string(), + JsonSchema::string(Some( + "Heartbeat target thread id. Must be this current thread.".to_string(), + )), + ), + ]); + + ToolSpec::Function(ResponsesApiTool { + name: AUTOMATION_UPDATE_TOOL_NAME.to_string(), + description: r#"Create, list, read, update, or delete scheduled automations and heartbeats for the current thread. +Use cron automations for recurring work that should start a fresh thread in one or more working directories. +Use heartbeat automations for recurring follow-ups to this existing thread. +This tool manages definitions only; immediate runNow dispatch is handled by the app-server API in a later automation dispatch path."# + .to_string(), + strict: false, + defer_loading: None, + parameters: JsonSchema::object( + properties, + /*required*/ Some(vec!["mode".to_string()]), + Some(false.into()), + ), + output_schema: None, + }) +} diff --git a/codex-rs/ext/automation/src/tool.rs b/codex-rs/ext/automation/src/tool.rs new file mode 100644 index 0000000000..7cccbdb096 --- /dev/null +++ b/codex-rs/ext/automation/src/tool.rs @@ -0,0 +1,398 @@ +use std::path::PathBuf; +use std::sync::Arc; + +use async_trait::async_trait; +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::openai_models::ReasoningEffort; +use codex_state::Automation as StateAutomation; +use codex_state::AutomationCreateParams; +use codex_state::AutomationDispatchSettings; +use codex_state::AutomationStatus as StateAutomationStatus; +use codex_state::AutomationTarget as StateAutomationTarget; +use codex_state::AutomationUpdateParams; +use serde::Deserialize; +use serde::Serialize; + +use crate::extension::AutomationThreadContext; +use crate::spec::AUTOMATION_UPDATE_TOOL_NAME; +use crate::spec::create_automation_update_tool; + +#[derive(Clone)] +pub(crate) struct AutomationUpdateTool { + state_db: Arc, + context: AutomationThreadContext, +} + +impl AutomationUpdateTool { + pub(crate) fn new( + state_db: Arc, + context: AutomationThreadContext, + ) -> Self { + Self { state_db, context } + } +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "snake_case")] +struct AutomationToolRequest { + mode: AutomationToolMode, + automation_id: Option, + kind: Option, + name: Option, + prompt: Option, + rrule: Option, + model: Option, + reasoning_effort: Option, + status: Option, + cwds: Option>, + thread_id: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "lowercase")] +enum AutomationToolMode { + List, + Read, + Create, + Update, + Delete, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "lowercase")] +enum AutomationToolKind { + Cron, + Heartbeat, +} + +#[derive(Debug, Deserialize)] +enum AutomationToolStatus { + #[serde(rename = "ACTIVE")] + Active, + #[serde(rename = "PAUSED")] + Paused, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +struct AutomationToolResponse { + automation: Option, + automations: Vec, + deleted: Option, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +struct AutomationResponse { + id: String, + name: String, + prompt: String, + status: String, + rrule: String, + next_run_at: Option, + last_run_at: Option, + created_at: i64, + updated_at: i64, + model: Option, + reasoning_effort: Option, + target: AutomationTargetResponse, +} + +#[derive(Debug, PartialEq, Serialize)] +#[serde(rename_all = "camelCase", tag = "type")] +enum AutomationTargetResponse { + Cron { cwds: Vec }, + Heartbeat { thread_id: String }, +} + +#[async_trait] +impl ToolExecutor for AutomationUpdateTool { + fn tool_name(&self) -> ToolName { + ToolName::plain(AUTOMATION_UPDATE_TOOL_NAME) + } + + fn spec(&self) -> ToolSpec { + create_automation_update_tool() + } + + async fn handle(&self, invocation: ToolCall) -> Result, FunctionCallError> { + let request: AutomationToolRequest = parse_arguments(invocation.function_arguments()?)?; + match request.mode { + AutomationToolMode::List => self.handle_list().await, + AutomationToolMode::Read => self.handle_read(request).await, + AutomationToolMode::Create => self.handle_create(request).await, + AutomationToolMode::Update => self.handle_update(request).await, + AutomationToolMode::Delete => self.handle_delete(request).await, + } + } +} + +impl AutomationUpdateTool { + async fn handle_list(&self) -> Result, FunctionCallError> { + let automations = self + .state_db + .list_automations() + .await + .map_err(storage_error)? + .into_iter() + .filter(|automation| automation.owner_thread_id == self.context.thread_id) + .map(AutomationResponse::from) + .collect(); + automation_response(/*automation*/ None, automations, /*deleted*/ None) + } + + async fn handle_read( + &self, + request: AutomationToolRequest, + ) -> Result, FunctionCallError> { + let automation_id = required(request.automation_id, "automation_id")?; + let automation = self + .state_db + .get_automation(automation_id.as_str()) + .await + .map_err(storage_error)? + .filter(|automation| automation.owner_thread_id == self.context.thread_id) + .map(AutomationResponse::from); + automation_response(automation, Vec::new(), /*deleted*/ None) + } + + async fn handle_create( + &self, + request: AutomationToolRequest, + ) -> Result, FunctionCallError> { + let target = self.target_from_request(&request, /*existing*/ None)?; + let dispatch_settings = dispatch_settings_for_target(&self.context, &target)?; + let automation = self + .state_db + .create_automation(&AutomationCreateParams { + owner_thread_id: self.context.thread_id, + name: required(request.name, "name")?, + prompt: required(request.prompt, "prompt")?, + status: request + .status + .map(StateAutomationStatus::from) + .unwrap_or(StateAutomationStatus::Active), + rrule: request.rrule, + model: request.model, + reasoning_effort: request.reasoning_effort, + target, + dispatch_settings, + }) + .await + .map_err(mutation_error)?; + automation_response( + Some(AutomationResponse::from(automation)), + Vec::new(), + /*deleted*/ None, + ) + } + + async fn handle_update( + &self, + request: AutomationToolRequest, + ) -> Result, FunctionCallError> { + let automation_id = required(request.automation_id.clone(), "automation_id")?; + let Some(existing) = self + .state_db + .get_automation(automation_id.as_str()) + .await + .map_err(storage_error)? + .filter(|automation| automation.owner_thread_id == self.context.thread_id) + else { + return automation_response( + /*automation*/ None, + Vec::new(), + /*deleted*/ None, + ); + }; + + let target = self.target_from_request(&request, Some(&existing.target))?; + let dispatch_settings = dispatch_settings_for_target(&self.context, &target)?; + let updated = self + .state_db + .update_automation_if_owner( + &AutomationUpdateParams { + id: automation_id, + owner_thread_id: self.context.thread_id, + name: request.name.unwrap_or(existing.name), + prompt: request.prompt.unwrap_or(existing.prompt), + status: request.status.map(Into::into).unwrap_or(existing.status), + rrule: Some(request.rrule.unwrap_or(existing.rrule)), + model: request.model.or(existing.model), + reasoning_effort: request.reasoning_effort.or(existing.reasoning_effort), + target, + dispatch_settings, + }, + self.context.thread_id, + ) + .await + .map_err(mutation_error)? + .map(AutomationResponse::from); + automation_response(updated, Vec::new(), /*deleted*/ None) + } + + async fn handle_delete( + &self, + request: AutomationToolRequest, + ) -> Result, FunctionCallError> { + let automation_id = required(request.automation_id, "automation_id")?; + let deleted = self + .state_db + .delete_automation_if_owner(automation_id.as_str(), self.context.thread_id) + .await + .map_err(storage_error)?; + automation_response(/*automation*/ None, Vec::new(), Some(deleted)) + } + + fn target_from_request( + &self, + request: &AutomationToolRequest, + existing: Option<&StateAutomationTarget>, + ) -> Result { + match request.kind.as_ref() { + Some(AutomationToolKind::Cron) => { + let cwds = request + .cwds + .clone() + .unwrap_or_else(|| vec![self.context.cwd.clone()]); + Ok(StateAutomationTarget::Cron { cwds }) + } + Some(AutomationToolKind::Heartbeat) => { + let thread_id = match request.thread_id.as_deref() { + Some(thread_id) => ThreadId::from_string(thread_id) + .map_err(|err| FunctionCallError::RespondToModel(err.to_string()))?, + None => self.context.thread_id, + }; + Ok(StateAutomationTarget::Heartbeat { thread_id }) + } + None => match existing { + Some(target) => Ok(target.clone()), + None => Err(FunctionCallError::RespondToModel( + "kind is required when creating an automation".to_string(), + )), + }, + } + } +} + +impl From for StateAutomationStatus { + fn from(status: AutomationToolStatus) -> Self { + match status { + AutomationToolStatus::Active => Self::Active, + AutomationToolStatus::Paused => Self::Paused, + } + } +} + +impl From for AutomationResponse { + fn from(automation: StateAutomation) -> Self { + Self { + id: automation.id, + name: automation.name, + prompt: automation.prompt, + status: automation.status.as_str().to_string(), + rrule: automation.rrule, + next_run_at: automation + .next_run_at + .as_ref() + .map(chrono::DateTime::timestamp), + last_run_at: automation + .last_run_at + .as_ref() + .map(chrono::DateTime::timestamp), + created_at: automation.created_at.timestamp(), + updated_at: automation.updated_at.timestamp(), + model: automation.model, + reasoning_effort: automation.reasoning_effort, + target: automation.target.into(), + } + } +} + +impl From for AutomationTargetResponse { + fn from(target: StateAutomationTarget) -> Self { + match target { + StateAutomationTarget::Cron { cwds } => Self::Cron { cwds }, + StateAutomationTarget::Heartbeat { thread_id } => Self::Heartbeat { + thread_id: thread_id.to_string(), + }, + } + } +} + +fn dispatch_settings_for_target( + context: &AutomationThreadContext, + target: &StateAutomationTarget, +) -> Result, FunctionCallError> { + match target { + StateAutomationTarget::Cron { cwds } => { + if cwds_visible_under_roots(cwds, &context.workspace_roots) { + Ok(Some(context.dispatch_settings.clone())) + } else { + Err(FunctionCallError::RespondToModel( + "cron automation cwds must be inside this thread's workspace roots".to_string(), + )) + } + } + StateAutomationTarget::Heartbeat { thread_id } => { + if *thread_id == context.thread_id { + Ok(None) + } else { + Err(FunctionCallError::RespondToModel( + "heartbeat automations can only target the current thread".to_string(), + )) + } + } + } +} + +fn cwds_visible_under_roots(cwds: &[PathBuf], workspace_roots: &[PathBuf]) -> bool { + !cwds.is_empty() + && cwds.iter().all(|cwd| { + cwd.is_absolute() && workspace_roots.iter().any(|root| cwd.starts_with(root)) + }) +} + +fn required(value: Option, field: &str) -> Result { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + .ok_or_else(|| FunctionCallError::RespondToModel(format!("{field} is required"))) +} + +fn parse_arguments(arguments: &str) -> Result +where + T: for<'de> Deserialize<'de>, +{ + serde_json::from_str(arguments) + .map_err(|err| FunctionCallError::RespondToModel(err.to_string())) +} + +fn automation_response( + automation: Option, + automations: Vec, + deleted: Option, +) -> Result, FunctionCallError> { + let value = serde_json::to_value(AutomationToolResponse { + automation, + automations, + deleted, + }) + .map_err(|err| FunctionCallError::Fatal(err.to_string()))?; + Ok(Box::new(JsonToolOutput::new(value))) +} + +fn storage_error(err: impl std::fmt::Display) -> FunctionCallError { + FunctionCallError::RespondToModel(format!("failed to read automations: {err}")) +} + +fn mutation_error(err: impl std::fmt::Display) -> FunctionCallError { + FunctionCallError::RespondToModel(format!("failed to update automations: {err}")) +} diff --git a/codex-rs/ext/automation/tests/automation_extension.rs b/codex-rs/ext/automation/tests/automation_extension.rs new file mode 100644 index 0000000000..a658c95444 --- /dev/null +++ b/codex-rs/ext/automation/tests/automation_extension.rs @@ -0,0 +1,251 @@ +use std::sync::Arc; + +use codex_automation_extension::AUTOMATION_UPDATE_TOOL_NAME; +use codex_automation_extension::install; +use codex_core::config::Config; +use codex_core::config::ConfigBuilder; +use codex_core::config::ConfigOverrides; +use codex_extension_api::ConversationHistory; +use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::NoopTurnItemEmitter; +use codex_extension_api::ThreadStartInput; +use codex_extension_api::ToolCall; +use codex_extension_api::ToolExecutor; +use codex_extension_api::ToolName; +use codex_extension_api::ToolPayload; +use codex_protocol::ThreadId; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TruncationPolicy; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; + +#[tokio::test] +async fn automation_update_tool_visibility_is_gated() -> anyhow::Result<()> { + let visible = AutomationExtensionHarness::new( + /*automations_enabled*/ true, + /*persistent_thread_state_available*/ true, + SessionSource::Cli, + ) + .await?; + assert_eq!( + vec![AUTOMATION_UPDATE_TOOL_NAME.to_string()], + tool_names(&visible.tools()) + ); + + let disabled = AutomationExtensionHarness::new( + /*automations_enabled*/ false, + /*persistent_thread_state_available*/ true, + SessionSource::Cli, + ) + .await?; + assert_eq!(Vec::::new(), tool_names(&disabled.tools())); + + let ephemeral = AutomationExtensionHarness::new( + /*automations_enabled*/ true, + /*persistent_thread_state_available*/ false, + SessionSource::Cli, + ) + .await?; + assert_eq!(Vec::::new(), tool_names(&ephemeral.tools())); + + let automation_thread = AutomationExtensionHarness::new( + /*automations_enabled*/ true, + /*persistent_thread_state_available*/ true, + SessionSource::automation(), + ) + .await?; + assert_eq!(Vec::::new(), tool_names(&automation_thread.tools())); + Ok(()) +} + +#[tokio::test] +async fn automation_update_tool_creates_and_lists_thread_owned_automations() -> anyhow::Result<()> { + let harness = AutomationExtensionHarness::new( + /*automations_enabled*/ true, + /*persistent_thread_state_available*/ true, + SessionSource::Cli, + ) + .await?; + let tools = harness.tools(); + let tool = tool_by_name(&tools, AUTOMATION_UPDATE_TOOL_NAME); + let cwd = harness.cwd_string(); + + let create_invocation = tool_call( + AUTOMATION_UPDATE_TOOL_NAME, + "call-create-automation", + json!({ + "mode": "create", + "kind": "cron", + "name": "Daily plan", + "prompt": "Plan the day", + "cwds": [cwd], + "status": "ACTIVE", + }), + ); + let create_output = tool.handle(create_invocation.clone()).await?; + let create_result = create_output.code_mode_result(&create_invocation.payload); + let automation = create_result["automation"].clone(); + assert_eq!( + create_result, + json!({ + "automation": { + "id": automation["id"], + "name": "Daily plan", + "prompt": "Plan the day", + "status": "ACTIVE", + "rrule": "FREQ=HOURLY;INTERVAL=24;BYMINUTE=0", + "nextRunAt": automation["nextRunAt"], + "lastRunAt": null, + "createdAt": automation["createdAt"], + "updatedAt": automation["updatedAt"], + "model": null, + "reasoningEffort": null, + "target": { + "type": "cron", + "cwds": [harness.cwd_string()], + }, + }, + "automations": [], + "deleted": null, + }) + ); + let stored = harness + .state_db + .get_automation(automation["id"].as_str().expect("automation id")) + .await? + .expect("created automation should be stored"); + let workspace_roots = stored + .dispatch_settings + .expect("cron automation should store dispatch settings") + .workspace_roots; + assert!( + workspace_roots + .iter() + .any(|root| harness.cwd.path().starts_with(root)), + "cron cwd should be covered by stored dispatch roots: {workspace_roots:?}" + ); + + let list_invocation = tool_call( + AUTOMATION_UPDATE_TOOL_NAME, + "call-list-automations", + json!({ "mode": "list" }), + ); + let list_output = tool.handle(list_invocation.clone()).await?; + assert_eq!( + list_output.code_mode_result(&list_invocation.payload), + json!({ + "automation": null, + "automations": [automation], + "deleted": null, + }) + ); + Ok(()) +} + +struct AutomationExtensionHarness { + _codex_home: TempDir, + cwd: TempDir, + state_db: Arc, + registry: codex_extension_api::ExtensionRegistry, + session_store: ExtensionData, + thread_store: ExtensionData, +} + +impl AutomationExtensionHarness { + async fn new( + automations_enabled: bool, + persistent_thread_state_available: bool, + session_source: SessionSource, + ) -> anyhow::Result { + let codex_home = TempDir::new()?; + let cwd = TempDir::new()?; + let state_db = codex_state::StateRuntime::init( + codex_home.path().to_path_buf(), + "test-provider".into(), + ) + .await?; + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .harness_overrides(ConfigOverrides { + cwd: Some(cwd.path().to_path_buf()), + ..ConfigOverrides::default() + }) + .cli_overrides(vec![( + "features.automations".to_string(), + toml::Value::Boolean(automations_enabled), + )]) + .build() + .await?; + let mut builder = ExtensionRegistryBuilder::::new(); + install(&mut builder, Arc::clone(&state_db)); + let registry = builder.build(); + let session_store = ExtensionData::new("session-1"); + let thread_store = ExtensionData::new(test_thread_id()?.to_string()); + for contributor in registry.thread_lifecycle_contributors() { + contributor + .on_thread_start(ThreadStartInput { + config: &config, + session_source: &session_source, + persistent_thread_state_available, + session_store: &session_store, + thread_store: &thread_store, + }) + .await; + } + Ok(Self { + _codex_home: codex_home, + cwd, + state_db, + registry, + session_store, + thread_store, + }) + } + + fn tools(&self) -> Vec>> { + self.registry + .tool_contributors() + .iter() + .flat_map(|contributor| contributor.tools(&self.session_store, &self.thread_store)) + .collect() + } + + fn cwd_string(&self) -> String { + self.cwd.path().to_string_lossy().to_string() + } +} + +fn tool_names(tools: &[Arc>]) -> Vec { + tools.iter().map(|tool| tool.tool_name().name).collect() +} + +fn tool_by_name<'a>( + tools: &'a [Arc>], + name: &str, +) -> &'a Arc> { + tools + .iter() + .find(|tool| tool.tool_name().namespace.is_none() && tool.tool_name().name == name) + .unwrap_or_else(|| panic!("missing tool {name}")) +} + +fn tool_call(tool_name: &str, call_id: &str, arguments: serde_json::Value) -> ToolCall { + ToolCall { + turn_id: "turn-1".to_string(), + call_id: call_id.to_string(), + tool_name: ToolName::plain(tool_name), + model: "gpt-test".to_string(), + truncation_policy: TruncationPolicy::Bytes(1024), + conversation_history: ConversationHistory::default(), + turn_item_emitter: Arc::new(NoopTurnItemEmitter), + payload: ToolPayload::Function { + arguments: arguments.to_string(), + }, + } +} + +fn test_thread_id() -> anyhow::Result { + ThreadId::from_string("11111111-1111-4111-8111-111111111111").map_err(anyhow::Error::msg) +}