mirror of
https://github.com/openai/codex.git
synced 2026-11-01 19:33:35 +00:00
Add automation extension tool
This commit is contained in:
Andrei Eternal
Generated
+23
@@ -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"
|
||||
|
||||
@@ -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" }
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -61,6 +61,7 @@ where
|
||||
} = dependencies;
|
||||
let mut builder = ExtensionRegistryBuilder::<Config>::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,
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
load("//:defs.bzl", "codex_rust_crate")
|
||||
|
||||
codex_rust_crate(
|
||||
name = "automation",
|
||||
crate_name = "codex_automation_extension",
|
||||
)
|
||||
@@ -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 }
|
||||
@@ -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<PathBuf>,
|
||||
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<Self> {
|
||||
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::<Vec<_>>();
|
||||
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<codex_state::StateRuntime>,
|
||||
}
|
||||
|
||||
impl AutomationExtension {
|
||||
fn new(state_db: Arc<codex_state::StateRuntime>) -> 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<Config> 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<Config> 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::<AutomationThreadContext>() 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<Arc<dyn codex_extension_api::ToolExecutor<codex_extension_api::ToolCall>>> {
|
||||
let Some(context) = thread_store.get::<AutomationThreadContext>() 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<Config>,
|
||||
state_db: Arc<codex_state::StateRuntime>,
|
||||
) {
|
||||
let extension = Arc::new(AutomationExtension::new(state_db));
|
||||
registry.thread_lifecycle_contributor(extension.clone());
|
||||
registry.config_contributor(extension.clone());
|
||||
registry.tool_contributor(extension);
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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<codex_state::StateRuntime>,
|
||||
context: AutomationThreadContext,
|
||||
}
|
||||
|
||||
impl AutomationUpdateTool {
|
||||
pub(crate) fn new(
|
||||
state_db: Arc<codex_state::StateRuntime>,
|
||||
context: AutomationThreadContext,
|
||||
) -> Self {
|
||||
Self { state_db, context }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
struct AutomationToolRequest {
|
||||
mode: AutomationToolMode,
|
||||
automation_id: Option<String>,
|
||||
kind: Option<AutomationToolKind>,
|
||||
name: Option<String>,
|
||||
prompt: Option<String>,
|
||||
rrule: Option<String>,
|
||||
model: Option<String>,
|
||||
reasoning_effort: Option<ReasoningEffort>,
|
||||
status: Option<AutomationToolStatus>,
|
||||
cwds: Option<Vec<PathBuf>>,
|
||||
thread_id: Option<String>,
|
||||
}
|
||||
|
||||
#[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<AutomationResponse>,
|
||||
automations: Vec<AutomationResponse>,
|
||||
deleted: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AutomationResponse {
|
||||
id: String,
|
||||
name: String,
|
||||
prompt: String,
|
||||
status: String,
|
||||
rrule: String,
|
||||
next_run_at: Option<i64>,
|
||||
last_run_at: Option<i64>,
|
||||
created_at: i64,
|
||||
updated_at: i64,
|
||||
model: Option<String>,
|
||||
reasoning_effort: Option<ReasoningEffort>,
|
||||
target: AutomationTargetResponse,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "camelCase", tag = "type")]
|
||||
enum AutomationTargetResponse {
|
||||
Cron { cwds: Vec<PathBuf> },
|
||||
Heartbeat { thread_id: String },
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolExecutor<ToolCall> 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<Box<dyn ToolOutput>, 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<Box<dyn ToolOutput>, 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<Box<dyn ToolOutput>, 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<Box<dyn ToolOutput>, 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<Box<dyn ToolOutput>, 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<Box<dyn ToolOutput>, 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<StateAutomationTarget, FunctionCallError> {
|
||||
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<AutomationToolStatus> for StateAutomationStatus {
|
||||
fn from(status: AutomationToolStatus) -> Self {
|
||||
match status {
|
||||
AutomationToolStatus::Active => Self::Active,
|
||||
AutomationToolStatus::Paused => Self::Paused,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<StateAutomation> 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<StateAutomationTarget> 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<Option<AutomationDispatchSettings>, 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<String>, field: &str) -> Result<String, FunctionCallError> {
|
||||
value
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| FunctionCallError::RespondToModel(format!("{field} is required")))
|
||||
}
|
||||
|
||||
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()))
|
||||
}
|
||||
|
||||
fn automation_response(
|
||||
automation: Option<AutomationResponse>,
|
||||
automations: Vec<AutomationResponse>,
|
||||
deleted: Option<bool>,
|
||||
) -> Result<Box<dyn ToolOutput>, 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}"))
|
||||
}
|
||||
@@ -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::<String>::new(), tool_names(&disabled.tools()));
|
||||
|
||||
let ephemeral = AutomationExtensionHarness::new(
|
||||
/*automations_enabled*/ true,
|
||||
/*persistent_thread_state_available*/ false,
|
||||
SessionSource::Cli,
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(Vec::<String>::new(), tool_names(&ephemeral.tools()));
|
||||
|
||||
let automation_thread = AutomationExtensionHarness::new(
|
||||
/*automations_enabled*/ true,
|
||||
/*persistent_thread_state_available*/ true,
|
||||
SessionSource::automation(),
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(Vec::<String>::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<codex_state::StateRuntime>,
|
||||
registry: codex_extension_api::ExtensionRegistry<Config>,
|
||||
session_store: ExtensionData,
|
||||
thread_store: ExtensionData,
|
||||
}
|
||||
|
||||
impl AutomationExtensionHarness {
|
||||
async fn new(
|
||||
automations_enabled: bool,
|
||||
persistent_thread_state_available: bool,
|
||||
session_source: SessionSource,
|
||||
) -> anyhow::Result<Self> {
|
||||
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::<Config>::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<Arc<dyn ToolExecutor<ToolCall>>> {
|
||||
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<dyn ToolExecutor<ToolCall>>]) -> Vec<String> {
|
||||
tools.iter().map(|tool| tool.tool_name().name).collect()
|
||||
}
|
||||
|
||||
fn tool_by_name<'a>(
|
||||
tools: &'a [Arc<dyn ToolExecutor<ToolCall>>],
|
||||
name: &str,
|
||||
) -> &'a Arc<dyn ToolExecutor<ToolCall>> {
|
||||
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> {
|
||||
ThreadId::from_string("11111111-1111-4111-8111-111111111111").map_err(anyhow::Error::msg)
|
||||
}
|
||||
Reference in New Issue
Block a user