mirror of
https://github.com/openai/codex.git
synced 2026-10-23 18:08:45 +00:00
Krish Chainani
aabf13307a
## What changed - Upload generated images to the attachment store and expose `fileId` in image-generation items when available, with inline `result` fallback on upload or decoding failure. - Preserve file references through legacy events, persisted rollouts, thread history, and task output summaries. Keep model-facing image preparation separate. - Add attachment request tags and mark generated-image uploads with `usage=image_gen`, preserving the thread's ephemeral setting. - Update public schemas and SDK types to accept either image representation, retain the Python `ImageGenerationThreadItem` alias, and prefer `fileId` when reading records containing both fields. ## Testing Add coverage for file-backed generation and rollout persistence, legacy history replay, inline and transitional serialization, task output visibility, and Python SDK alias compatibility. GitOrigin-RevId: f7a2bb5707e2f277348a547a988bc33c8e96d891
2104 lines
68 KiB
Rust
2104 lines
68 KiB
Rust
#![recursion_limit = "256"]
|
|
#![allow(clippy::expect_used)]
|
|
|
|
use codex_utils_absolute_path::test_support::PathExt;
|
|
use std::sync::Arc;
|
|
use std::sync::Mutex;
|
|
use std::sync::PoisonError;
|
|
use std::sync::Weak;
|
|
use std::time::Duration;
|
|
|
|
use codex_analytics::AnalyticsEventsClient;
|
|
use codex_extension_api::ExtensionData;
|
|
use codex_extension_api::ExtensionEventSink;
|
|
use codex_extension_api::ExtensionRegistryBuilder;
|
|
use codex_extension_api::ExtensionWarning;
|
|
use codex_extension_api::FunctionCallError;
|
|
use codex_extension_api::NoopTurnItemEmitter;
|
|
use codex_extension_api::ThreadResumeInput;
|
|
use codex_extension_api::ThreadStartInput;
|
|
use codex_extension_api::ThreadStopInput;
|
|
use codex_extension_api::ToolCall;
|
|
use codex_extension_api::ToolCallOutcome;
|
|
use codex_extension_api::ToolCallSource;
|
|
use codex_extension_api::ToolExecutor;
|
|
use codex_extension_api::ToolFinishInput;
|
|
use codex_extension_api::ToolPayload;
|
|
use codex_extension_api::TurnErrorInput;
|
|
use codex_extension_api::TurnStartInput;
|
|
use codex_extension_api::TurnStopInput;
|
|
use codex_goal_extension::GoalExtensionConfig;
|
|
use codex_goal_extension::GoalObjectiveUpdate;
|
|
use codex_goal_extension::GoalRuntimeHandle;
|
|
use codex_goal_extension::GoalService;
|
|
use codex_goal_extension::GoalServiceError;
|
|
use codex_goal_extension::GoalSetRequest;
|
|
use codex_goal_extension::GoalTokenBudgetUpdate;
|
|
use codex_goal_extension::install_with_backend;
|
|
use codex_protocol::ThreadId;
|
|
use codex_protocol::config_types::CollaborationMode;
|
|
use codex_protocol::config_types::ModeKind;
|
|
use codex_protocol::config_types::Settings;
|
|
use codex_protocol::error::CodexErr;
|
|
use codex_protocol::protocol::Event;
|
|
use codex_protocol::protocol::EventMsg;
|
|
use codex_protocol::protocol::SessionSource;
|
|
use codex_protocol::protocol::SubAgentSource;
|
|
use codex_protocol::protocol::ThreadGoalStatus;
|
|
use codex_protocol::protocol::TokenUsage;
|
|
use codex_protocol::protocol::TokenUsageInfo;
|
|
use codex_protocol::protocol::TruncationPolicy;
|
|
use pretty_assertions::assert_eq;
|
|
use serde_json::json;
|
|
use tempfile::TempDir;
|
|
|
|
#[tokio::test]
|
|
async fn installed_goal_tools_create_goal_and_fill_empty_preview() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let tools = installed_tools(runtime.clone(), thread_id).await;
|
|
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
let invocation = tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({
|
|
"objective": "ship goal extension backend",
|
|
"token_budget": 123,
|
|
}),
|
|
);
|
|
let output = create_tool.handle(invocation.clone()).await?;
|
|
let result = output.code_mode_result(&invocation.payload);
|
|
assert_eq!(
|
|
result,
|
|
json!({
|
|
"goal": {
|
|
"threadId": thread_id,
|
|
"objective": "ship goal extension backend",
|
|
"status": "active",
|
|
"tokenBudget": 123,
|
|
"tokensUsed": 0,
|
|
"timeUsedSeconds": 0,
|
|
"createdAt": result["goal"]["createdAt"],
|
|
"updatedAt": result["goal"]["updatedAt"],
|
|
},
|
|
"remainingTokens": 123,
|
|
"completionBudgetReport": serde_json::Value::Null,
|
|
})
|
|
);
|
|
|
|
let metadata = runtime
|
|
.get_thread(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("seeded thread metadata should exist"))?;
|
|
assert_eq!(
|
|
metadata.preview.as_deref(),
|
|
Some("ship goal extension backend")
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn installed_goal_tools_apply_maximum_token_budget() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.thread_store.insert(GoalExtensionConfig {
|
|
enabled: true,
|
|
max_goal_token_budget: Some(100),
|
|
});
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
|
|
let result = create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-oversized-goal",
|
|
json!({ "objective": "oversized goal", "token_budget": 101 }),
|
|
))
|
|
.await;
|
|
let error = match result {
|
|
Ok(_) => panic!("goal budget above the configured maximum should fail"),
|
|
Err(error) => error,
|
|
};
|
|
assert_eq!(
|
|
error,
|
|
FunctionCallError::RespondToModel(
|
|
"goal token budget 101 exceeds the maximum allowed goal token budget of 100"
|
|
.to_string()
|
|
)
|
|
);
|
|
assert_eq!(
|
|
runtime.thread_goals().get_thread_goal(thread_id).await?,
|
|
None
|
|
);
|
|
|
|
let invocation = tool_call(
|
|
"create_goal",
|
|
"call-default-goal-budget",
|
|
json!({ "objective": "default goal budget" }),
|
|
);
|
|
let output = create_tool.handle(invocation.clone()).await?;
|
|
assert_eq!(
|
|
output.code_mode_result(&invocation.payload)["goal"]["tokenBudget"],
|
|
json!(100)
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn ephemeral_goal_tools_preserve_specs_but_reject_execution() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
let tools = installed_tools_with_start(
|
|
runtime.clone(),
|
|
thread_id,
|
|
SessionSource::Cli,
|
|
/*persistent_thread_state_available*/ false,
|
|
)
|
|
.await;
|
|
|
|
let parent_tools = installed_tools(runtime, ThreadId::new()).await;
|
|
assert_eq!(
|
|
tools.iter().map(|tool| tool.spec()).collect::<Vec<_>>(),
|
|
parent_tools
|
|
.iter()
|
|
.map(|tool| tool.spec())
|
|
.collect::<Vec<_>>(),
|
|
);
|
|
for tool in tools {
|
|
let Err(error) = tool
|
|
.handle(tool_call(&tool.tool_name().name, "ephemeral", json!({})))
|
|
.await
|
|
else {
|
|
panic!("ephemeral goal execution should fail");
|
|
};
|
|
assert_eq!(
|
|
error,
|
|
FunctionCallError::RespondToModel(
|
|
"Goal tools require a persistent thread.".to_string()
|
|
)
|
|
);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn goal_tools_hidden_for_review_subagents() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
let tools = installed_tools_with_start(
|
|
runtime,
|
|
thread_id,
|
|
SessionSource::SubAgent(SubAgentSource::Review),
|
|
/*persistent_thread_state_available*/ true,
|
|
)
|
|
.await;
|
|
|
|
assert_eq!(Vec::<String>::new(), tool_names(&tools));
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn installed_goal_tools_only_replace_complete_goal() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime, thread_id).await?;
|
|
let tools = harness.tools();
|
|
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
let first = tool_call(
|
|
"create_goal",
|
|
"call-create-goal-1",
|
|
json!({ "objective": "first goal" }),
|
|
);
|
|
create_tool.handle(first).await?;
|
|
|
|
let second = tool_call(
|
|
"create_goal",
|
|
"call-create-goal-2",
|
|
json!({ "objective": "second goal" }),
|
|
);
|
|
let err = match create_tool.handle(second).await {
|
|
Ok(_) => panic!("duplicate create should fail"),
|
|
Err(err) => err,
|
|
};
|
|
|
|
assert_eq!(
|
|
err,
|
|
FunctionCallError::RespondToModel(
|
|
"cannot create a new goal because this thread has an unfinished goal; complete the existing goal first"
|
|
.to_string()
|
|
)
|
|
);
|
|
|
|
let update_tool = tool_by_name(&tools, "update_goal");
|
|
update_tool
|
|
.handle(tool_call(
|
|
"update_goal",
|
|
"call-complete-goal",
|
|
json!({ "status": "complete" }),
|
|
))
|
|
.await?;
|
|
|
|
let invocation = tool_call(
|
|
"create_goal",
|
|
"call-create-goal-3",
|
|
json!({ "objective": "replacement goal" }),
|
|
);
|
|
let output = create_tool.handle(invocation.clone()).await?;
|
|
let result = output.code_mode_result(&invocation.payload);
|
|
|
|
assert_eq!(json!("replacement goal"), result["goal"]["objective"]);
|
|
assert_eq!(json!("active"), result["goal"]["status"]);
|
|
assert_eq!(json!(0), result["goal"]["tokensUsed"]);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn create_goal_resets_baseline_before_turn_stop_accounting() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
// A missing baseline must not panic or seed accounting with invented token usage.
|
|
for contributor in harness.registry.turn_lifecycle_contributors() {
|
|
contributor
|
|
.on_turn_start(TurnStartInput {
|
|
turn_id: "missing-baseline",
|
|
collaboration_mode: &default_collaboration_mode(),
|
|
token_usage_at_turn_start: None,
|
|
session_store: &harness.session_store,
|
|
thread_store: &harness.thread_store,
|
|
turn_store: &ExtensionData::new("missing-baseline"),
|
|
})
|
|
.await;
|
|
}
|
|
harness
|
|
.start_turn(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 100, /*cached_input_tokens*/ 10,
|
|
/*output_tokens*/ 30, /*reasoning_output_tokens*/ 5,
|
|
/*total_tokens*/ 135,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 120, /*cached_input_tokens*/ 14,
|
|
/*output_tokens*/ 42, /*reasoning_output_tokens*/ 8,
|
|
/*total_tokens*/ 162,
|
|
),
|
|
)
|
|
.await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "ship goal extension backend" }),
|
|
))
|
|
.await?;
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 127, /*cached_input_tokens*/ 16,
|
|
/*output_tokens*/ 52, /*reasoning_output_tokens*/ 10,
|
|
/*total_tokens*/ 189,
|
|
),
|
|
)
|
|
.await;
|
|
harness.stop_turn("turn-1").await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(15, goal.tokens_used);
|
|
assert_eq!(ThreadGoalStatus::Active, protocol_status(goal.status));
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn tool_finish_accounts_active_goal_progress_and_emits_event() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "ship goal extension backend" }),
|
|
))
|
|
.await?;
|
|
harness.sink.clear();
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 20, /*cached_input_tokens*/ 5, /*output_tokens*/ 8,
|
|
/*reasoning_output_tokens*/ 2, /*total_tokens*/ 30,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.notify_tool_finish("turn-1", "call-shell", "shell")
|
|
.await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(23, goal.tokens_used);
|
|
|
|
assert_eq!(
|
|
vec![CapturedGoalEvent {
|
|
event_id: "call-shell".to_string(),
|
|
turn_id: Some("turn-1".to_string()),
|
|
status: ThreadGoalStatus::Active,
|
|
tokens_used: 23,
|
|
}],
|
|
harness.sink.goal_events()
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn parallel_tool_finish_accounts_active_goal_progress_once() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness
|
|
.start_turn(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 100, /*cached_input_tokens*/ 0,
|
|
/*output_tokens*/ 0, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 100,
|
|
),
|
|
)
|
|
.await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "ship goal extension backend" }),
|
|
))
|
|
.await?;
|
|
harness.sink.clear();
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 130, /*cached_input_tokens*/ 0,
|
|
/*output_tokens*/ 0, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 130,
|
|
),
|
|
)
|
|
.await;
|
|
|
|
tokio::join!(
|
|
harness.notify_tool_finish("turn-1", "call-shell-1", "shell"),
|
|
harness.notify_tool_finish("turn-1", "call-shell-2", "shell"),
|
|
);
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(30, goal.tokens_used);
|
|
|
|
assert_eq!(
|
|
vec![CapturedGoalEvent {
|
|
event_id: "call-shell-1".to_string(),
|
|
turn_id: Some("turn-1".to_string()),
|
|
status: ThreadGoalStatus::Active,
|
|
tokens_used: 30,
|
|
}],
|
|
harness.sink.goal_events()
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawned_descendant_usage_exhausts_root_goal_budget_once() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let child = harness.spawn_child(ThreadId::new()).await?;
|
|
child.start_turn("child-turn", &TokenUsage::default()).await;
|
|
let grandchild = child
|
|
.spawn_child_with_source(ThreadId::new(), SubAgentSource::Review)
|
|
.await?;
|
|
grandchild
|
|
.start_turn("grandchild-turn", &TokenUsage::default())
|
|
.await;
|
|
let tools = harness.tools();
|
|
tool_by_name(&tools, "create_goal")
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "account for the entire agent tree", "token_budget": 62 }),
|
|
))
|
|
.await?;
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 12, /*cached_input_tokens*/ 2, /*output_tokens*/ 4,
|
|
/*reasoning_output_tokens*/ 0, /*total_tokens*/ 16,
|
|
),
|
|
)
|
|
.await;
|
|
let first_child_usage = input_token_usage(/*input_tokens*/ 23);
|
|
child
|
|
.record_token_usage_with_last("child-turn", &first_child_usage, &first_child_usage)
|
|
.await;
|
|
child
|
|
.record_token_usage_with_last(
|
|
"child-turn",
|
|
&input_token_usage(/*input_tokens*/ 36),
|
|
&input_token_usage(/*input_tokens*/ 13),
|
|
)
|
|
.await;
|
|
let grandchild_usage = input_token_usage(/*input_tokens*/ 12);
|
|
grandchild
|
|
.record_token_usage_with_last("grandchild-turn", &grandchild_usage, &grandchild_usage)
|
|
.await;
|
|
harness
|
|
.notify_tool_finish("turn-1", "call-shell", "shell")
|
|
.await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(62, goal.tokens_used);
|
|
assert_eq!(codex_state::ThreadGoalStatus::BudgetLimited, goal.status);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn grandchild_usage_rolls_up_after_parent_runtime_unloads() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
tool_by_name(&tools, "create_goal")
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "account for evicted agent trees" }),
|
|
))
|
|
.await?;
|
|
|
|
let child = harness.spawn_child(ThreadId::new()).await?;
|
|
child.stop_thread().await;
|
|
let grandchild = child.spawn_child(ThreadId::new()).await?;
|
|
grandchild
|
|
.start_turn("grandchild-turn", &TokenUsage::default())
|
|
.await;
|
|
let usage = input_token_usage(/*input_tokens*/ 23);
|
|
grandchild
|
|
.record_token_usage_with_last("grandchild-turn", &usage, &usage)
|
|
.await;
|
|
harness
|
|
.notify_tool_finish("turn-1", "call-shell", "shell")
|
|
.await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(23, goal.tokens_used);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn subagent_usage_resets_when_root_goal_is_replaced() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
let child = harness.spawn_child(ThreadId::new()).await?;
|
|
child.start_turn("child-turn", &TokenUsage::default()).await;
|
|
|
|
let previous_usage = input_token_usage(/*input_tokens*/ 10);
|
|
child
|
|
.record_token_usage_with_last("child-turn", &previous_usage, &previous_usage)
|
|
.await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-first-goal",
|
|
json!({ "objective": "first goal" }),
|
|
))
|
|
.await?;
|
|
child
|
|
.record_token_usage_with_last(
|
|
"child-turn",
|
|
&input_token_usage(/*input_tokens*/ 35),
|
|
&input_token_usage(/*input_tokens*/ 25),
|
|
)
|
|
.await;
|
|
let completion = tool_call(
|
|
"update_goal",
|
|
"call-complete-first-goal",
|
|
json!({ "status": "complete" }),
|
|
);
|
|
let completed = tool_by_name(&tools, "update_goal")
|
|
.handle(completion.clone())
|
|
.await?;
|
|
assert_eq!(
|
|
json!(25),
|
|
completed.code_mode_result(&completion.payload)["goal"]["tokensUsed"]
|
|
);
|
|
|
|
child
|
|
.record_token_usage_with_last(
|
|
"child-turn",
|
|
&input_token_usage(/*input_tokens*/ 55),
|
|
&input_token_usage(/*input_tokens*/ 20),
|
|
)
|
|
.await;
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-second-goal",
|
|
json!({ "objective": "replacement goal" }),
|
|
))
|
|
.await?;
|
|
child
|
|
.record_token_usage_with_last(
|
|
"child-turn",
|
|
&input_token_usage(/*input_tokens*/ 64),
|
|
&input_token_usage(/*input_tokens*/ 9),
|
|
)
|
|
.await;
|
|
harness.stop_turn("turn-1").await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("replacement goal should exist"))?;
|
|
assert_eq!("replacement goal", goal.objective);
|
|
assert_eq!(9, goal.tokens_used);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn budget_limited_goal_keeps_accruing_until_turn_stop() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({
|
|
"objective": "ship goal extension backend",
|
|
"token_budget": 25,
|
|
}),
|
|
))
|
|
.await?;
|
|
harness.sink.clear();
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 20, /*cached_input_tokens*/ 5,
|
|
/*output_tokens*/ 10, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 30,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.notify_tool_finish("turn-1", "call-shell", "shell")
|
|
.await;
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 24, /*cached_input_tokens*/ 5,
|
|
/*output_tokens*/ 16, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 40,
|
|
),
|
|
)
|
|
.await;
|
|
harness.stop_turn("turn-1").await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(35, goal.tokens_used);
|
|
assert_eq!(codex_state::ThreadGoalStatus::BudgetLimited, goal.status);
|
|
|
|
assert_eq!(
|
|
vec![
|
|
CapturedGoalEvent {
|
|
event_id: "call-shell".to_string(),
|
|
turn_id: Some("turn-1".to_string()),
|
|
status: ThreadGoalStatus::BudgetLimited,
|
|
tokens_used: 25,
|
|
},
|
|
CapturedGoalEvent {
|
|
event_id: "turn-1:turn-stop".to_string(),
|
|
turn_id: Some("turn-1".to_string()),
|
|
status: ThreadGoalStatus::BudgetLimited,
|
|
tokens_used: 35,
|
|
},
|
|
],
|
|
harness.sink.goal_events()
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn budget_limited_goal_keeps_accounting_after_later_tool_finish() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({
|
|
"objective": "ship goal extension backend",
|
|
"token_budget": 25,
|
|
}),
|
|
))
|
|
.await?;
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 20, /*cached_input_tokens*/ 5,
|
|
/*output_tokens*/ 10, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 30,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.notify_tool_finish("turn-1", "call-shell-1", "shell")
|
|
.await;
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 24, /*cached_input_tokens*/ 5,
|
|
/*output_tokens*/ 16, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 40,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.notify_tool_finish("turn-1", "call-shell-2", "shell")
|
|
.await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(35, goal.tokens_used);
|
|
assert_eq!(codex_state::ThreadGoalStatus::BudgetLimited, goal.status);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn turn_error_usage_limit_accounts_progress_and_clears_accounting() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "ship goal extension backend" }),
|
|
))
|
|
.await?;
|
|
harness.sink.clear();
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 20, /*cached_input_tokens*/ 5, /*output_tokens*/ 8,
|
|
/*reasoning_output_tokens*/ 2, /*total_tokens*/ 30,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.notify_turn_error("turn-1", CodexErr::UsageNotIncluded)
|
|
.await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(23, goal.tokens_used);
|
|
assert_eq!(codex_state::ThreadGoalStatus::UsageLimited, goal.status);
|
|
assert_eq!(
|
|
vec![
|
|
CapturedGoalEvent {
|
|
event_id: "turn-1:usage-limit-progress".to_string(),
|
|
turn_id: Some("turn-1".to_string()),
|
|
status: ThreadGoalStatus::Active,
|
|
tokens_used: 23,
|
|
},
|
|
CapturedGoalEvent {
|
|
event_id: "turn-1:usage-limit".to_string(),
|
|
turn_id: Some("turn-1".to_string()),
|
|
status: ThreadGoalStatus::UsageLimited,
|
|
tokens_used: 23,
|
|
},
|
|
],
|
|
harness.sink.goal_events()
|
|
);
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 50, /*cached_input_tokens*/ 5,
|
|
/*output_tokens*/ 20, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 70,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.notify_tool_finish("turn-1", "call-shell-after-usage-limit", "shell")
|
|
.await;
|
|
harness.stop_turn("turn-1").await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(23, goal.tokens_used);
|
|
assert_eq!(codex_state::ThreadGoalStatus::UsageLimited, goal.status);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn turn_error_blocks_goal() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
tool_by_name(&tools, "create_goal")
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "ship goal extension backend" }),
|
|
))
|
|
.await?;
|
|
|
|
harness
|
|
.notify_turn_error("turn-1", CodexErr::Fatal("test error".to_string()))
|
|
.await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(codex_state::ThreadGoalStatus::Blocked, goal.status);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn failed_execution_turns_block_goal_unless_a_tool_succeeds() -> anyhow::Result<()> {
|
|
for (recovery_turn, status_only_turn, blocking_turn) in
|
|
[(None, None, 3), (Some(2), None, 5), (None, Some(2), 4)]
|
|
{
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
tool_by_name(&tools, "create_goal")
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "ship goal extension backend" }),
|
|
))
|
|
.await?;
|
|
|
|
for turn in 1..=blocking_turn {
|
|
let turn_id = format!("turn-{turn}");
|
|
if turn > 1 {
|
|
harness.start_turn(&turn_id, &TokenUsage::default()).await;
|
|
}
|
|
if status_only_turn != Some(turn) {
|
|
harness
|
|
.notify_tool_finish_with_outcome(
|
|
&turn_id,
|
|
&format!("call-exec-{turn}"),
|
|
"exec",
|
|
ToolCallOutcome::Failed {
|
|
handler_executed: true,
|
|
},
|
|
)
|
|
.await;
|
|
}
|
|
if recovery_turn == Some(turn) {
|
|
harness
|
|
.notify_tool_finish(&turn_id, "call-recovery", "shell")
|
|
.await;
|
|
}
|
|
harness.stop_turn(&turn_id).await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
let expected = if turn == blocking_turn {
|
|
codex_state::ThreadGoalStatus::Blocked
|
|
} else {
|
|
codex_state::ThreadGoalStatus::Active
|
|
};
|
|
assert_eq!(expected, goal.status);
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn usage_limit_budget_limited_goal_accounts_remaining_progress() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({
|
|
"objective": "ship goal extension backend",
|
|
"token_budget": 25,
|
|
}),
|
|
))
|
|
.await?;
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 20, /*cached_input_tokens*/ 5,
|
|
/*output_tokens*/ 10, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 30,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.notify_tool_finish("turn-1", "call-shell", "shell")
|
|
.await;
|
|
harness.sink.clear();
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 24, /*cached_input_tokens*/ 5,
|
|
/*output_tokens*/ 16, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 40,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.runtime_handle()
|
|
.usage_limit_active_goal_for_turn("turn-1")
|
|
.await
|
|
.map_err(anyhow::Error::msg)?;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(35, goal.tokens_used);
|
|
assert_eq!(codex_state::ThreadGoalStatus::UsageLimited, goal.status);
|
|
assert_eq!(
|
|
vec![
|
|
CapturedGoalEvent {
|
|
event_id: "turn-1:usage-limit-progress".to_string(),
|
|
turn_id: Some("turn-1".to_string()),
|
|
status: ThreadGoalStatus::BudgetLimited,
|
|
tokens_used: 35,
|
|
},
|
|
CapturedGoalEvent {
|
|
event_id: "turn-1:usage-limit".to_string(),
|
|
turn_id: Some("turn-1".to_string()),
|
|
status: ThreadGoalStatus::UsageLimited,
|
|
tokens_used: 35,
|
|
},
|
|
],
|
|
harness.sink.goal_events()
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn usage_limit_plan_turn_does_not_stop_goal() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "ship goal extension backend" }),
|
|
))
|
|
.await?;
|
|
|
|
harness
|
|
.start_turn_with_mode("turn-plan", ModeKind::Plan, &TokenUsage::default())
|
|
.await;
|
|
harness.sink.clear();
|
|
harness
|
|
.runtime_handle()
|
|
.usage_limit_active_goal_for_turn("turn-plan")
|
|
.await
|
|
.map_err(anyhow::Error::msg)?;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(codex_state::ThreadGoalStatus::Active, goal.status);
|
|
assert_eq!(Vec::<CapturedGoalEvent>::new(), harness.sink.goal_events());
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn usage_limit_stale_turn_does_not_stop_current_goal() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "ship goal extension backend" }),
|
|
))
|
|
.await?;
|
|
harness.stop_turn("turn-1").await;
|
|
harness.start_turn("turn-2", &TokenUsage::default()).await;
|
|
harness.sink.clear();
|
|
|
|
harness
|
|
.runtime_handle()
|
|
.usage_limit_active_goal_for_turn("turn-1")
|
|
.await
|
|
.map_err(anyhow::Error::msg)?;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(codex_state::ThreadGoalStatus::Active, goal.status);
|
|
assert_eq!(Vec::<CapturedGoalEvent>::new(), harness.sink.goal_events());
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn update_goal_can_stop_and_accounts_final_progress() -> anyhow::Result<()> {
|
|
for (status, token_budget, expected_status, state_status) in [
|
|
(
|
|
ThreadGoalStatus::Blocked,
|
|
100_i64,
|
|
ThreadGoalStatus::Blocked,
|
|
codex_state::ThreadGoalStatus::Blocked,
|
|
),
|
|
(
|
|
ThreadGoalStatus::Paused,
|
|
100,
|
|
ThreadGoalStatus::Paused,
|
|
codex_state::ThreadGoalStatus::Paused,
|
|
),
|
|
(
|
|
ThreadGoalStatus::Paused,
|
|
20,
|
|
ThreadGoalStatus::BudgetLimited,
|
|
codex_state::ThreadGoalStatus::BudgetLimited,
|
|
),
|
|
] {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "ship goal extension backend", "token_budget": token_budget }),
|
|
))
|
|
.await?;
|
|
harness.sink.clear();
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 20, /*cached_input_tokens*/ 5,
|
|
/*output_tokens*/ 8, /*reasoning_output_tokens*/ 2,
|
|
/*total_tokens*/ 30,
|
|
),
|
|
)
|
|
.await;
|
|
let update_tool = tool_by_name(&tools, "update_goal");
|
|
let invocation = tool_call(
|
|
"update_goal",
|
|
"call-update-goal",
|
|
json!({ "status": status }),
|
|
);
|
|
let output = update_tool.handle(invocation.clone()).await?;
|
|
let result = output.code_mode_result(&invocation.payload);
|
|
|
|
assert_eq!(
|
|
result,
|
|
json!({
|
|
"goal": {
|
|
"threadId": thread_id,
|
|
"objective": "ship goal extension backend",
|
|
"status": expected_status,
|
|
"tokenBudget": token_budget,
|
|
"tokensUsed": 23,
|
|
"timeUsedSeconds": 0,
|
|
"createdAt": result["goal"]["createdAt"],
|
|
"updatedAt": result["goal"]["updatedAt"],
|
|
},
|
|
"remainingTokens": (token_budget - 23).max(0),
|
|
"completionBudgetReport": serde_json::Value::Null,
|
|
})
|
|
);
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(23, goal.tokens_used);
|
|
assert_eq!(state_status, goal.status);
|
|
|
|
assert_eq!(
|
|
vec![
|
|
CapturedGoalEvent {
|
|
event_id: "call-update-goal".to_string(),
|
|
turn_id: Some("turn-1".to_string()),
|
|
status: if token_budget > 23 {
|
|
ThreadGoalStatus::Active
|
|
} else {
|
|
ThreadGoalStatus::BudgetLimited
|
|
},
|
|
tokens_used: 23,
|
|
},
|
|
CapturedGoalEvent {
|
|
event_id: "call-update-goal".to_string(),
|
|
turn_id: Some("turn-1".to_string()),
|
|
status: expected_status,
|
|
tokens_used: 23,
|
|
},
|
|
],
|
|
harness.sink.goal_events()
|
|
);
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 40, /*cached_input_tokens*/ 5,
|
|
/*output_tokens*/ 10, /*reasoning_output_tokens*/ 2,
|
|
/*total_tokens*/ 52,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.notify_tool_finish("turn-1", "call-shell", "shell")
|
|
.await;
|
|
harness.stop_turn("turn-1").await;
|
|
assert_eq!(
|
|
Some(goal),
|
|
runtime.thread_goals().get_thread_goal(thread_id).await?
|
|
);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn update_goal_rejects_resume_and_system_limit_statuses() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let tools = installed_tools(runtime, thread_id).await;
|
|
let update_tool = tool_by_name(&tools, "update_goal");
|
|
for status in ["active", "budgetLimited", "usageLimited"] {
|
|
let result = update_tool
|
|
.handle(tool_call(
|
|
"update_goal",
|
|
"call-update-goal",
|
|
json!({ "status": status }),
|
|
))
|
|
.await;
|
|
let error = match result {
|
|
Ok(_) => panic!("agent must not set {status}"),
|
|
Err(error) => error,
|
|
};
|
|
assert_eq!(
|
|
error,
|
|
FunctionCallError::RespondToModel(
|
|
"update_goal can only mark the existing goal complete, blocked, or paused at the user's explicit request; resume, budget-limited, and usage-limited status changes are controlled by the user or system".to_string()
|
|
)
|
|
);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn external_goal_mutation_start_accounts_active_goal_progress() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "ship goal extension backend" }),
|
|
))
|
|
.await?;
|
|
harness.sink.clear();
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 20, /*cached_input_tokens*/ 5, /*output_tokens*/ 8,
|
|
/*reasoning_output_tokens*/ 2, /*total_tokens*/ 30,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.runtime_handle()
|
|
.prepare_external_goal_mutation()
|
|
.await
|
|
.map_err(anyhow::Error::msg)?;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(23, goal.tokens_used);
|
|
assert_eq!(
|
|
vec![CapturedGoalEvent {
|
|
event_id: "turn-1:external-goal-mutation".to_string(),
|
|
turn_id: Some("turn-1".to_string()),
|
|
status: ThreadGoalStatus::Active,
|
|
tokens_used: 23,
|
|
}],
|
|
harness.sink.goal_events()
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn goal_service_external_set_active_preserves_concurrent_usage() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness
|
|
.start_turn(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 100, /*cached_input_tokens*/ 0,
|
|
/*output_tokens*/ 0, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 100,
|
|
),
|
|
)
|
|
.await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "old objective" }),
|
|
))
|
|
.await?;
|
|
let child = harness.spawn_child(ThreadId::new()).await?;
|
|
child.start_turn("child-turn", &TokenUsage::default()).await;
|
|
harness.sink.clear();
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 120, /*cached_input_tokens*/ 0,
|
|
/*output_tokens*/ 0, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 120,
|
|
),
|
|
)
|
|
.await;
|
|
let outcome = harness
|
|
.goal_service
|
|
.set_thread_goal(
|
|
runtime.as_ref(),
|
|
GoalSetRequest {
|
|
thread_id,
|
|
objective: GoalObjectiveUpdate::Set("new objective"),
|
|
status: Some(ThreadGoalStatus::Active),
|
|
token_budget: GoalTokenBudgetUpdate::Keep,
|
|
max_goal_token_budget: None,
|
|
},
|
|
std::future::ready(Ok(())),
|
|
)
|
|
.await?;
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 125, /*cached_input_tokens*/ 0,
|
|
/*output_tokens*/ 0, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 125,
|
|
),
|
|
)
|
|
.await;
|
|
let child_usage = input_token_usage(/*input_tokens*/ 23);
|
|
child
|
|
.record_token_usage_with_last("child-turn", &child_usage, &child_usage)
|
|
.await;
|
|
outcome.apply_runtime_effects(&harness.goal_service).await;
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 130, /*cached_input_tokens*/ 0,
|
|
/*output_tokens*/ 0, /*reasoning_output_tokens*/ 0,
|
|
/*total_tokens*/ 130,
|
|
),
|
|
)
|
|
.await;
|
|
harness
|
|
.notify_tool_finish("turn-1", "call-shell", "shell")
|
|
.await;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(53, goal.tokens_used);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn thread_stop_unregisters_goal_runtime_from_service() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
harness.start_turn("turn-1", &TokenUsage::default()).await;
|
|
|
|
let tools = harness.tools();
|
|
let create_tool = tool_by_name(&tools, "create_goal");
|
|
create_tool
|
|
.handle(tool_call(
|
|
"create_goal",
|
|
"call-create-goal",
|
|
json!({ "objective": "ship goal extension backend" }),
|
|
))
|
|
.await?;
|
|
harness.sink.clear();
|
|
|
|
harness
|
|
.record_token_usage(
|
|
"turn-1",
|
|
&token_usage(
|
|
/*input_tokens*/ 10, /*cached_input_tokens*/ 0, /*output_tokens*/ 0,
|
|
/*reasoning_output_tokens*/ 0, /*total_tokens*/ 10,
|
|
),
|
|
)
|
|
.await;
|
|
harness.stop_thread().await;
|
|
|
|
assert!(
|
|
harness
|
|
.goal_service
|
|
.clear_thread_goal(runtime.as_ref(), thread_id, std::future::ready(Ok(())))
|
|
.await?
|
|
);
|
|
assert_eq!(Vec::<CapturedGoalEvent>::new(), harness.sink.goal_events());
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
runtime
|
|
.thread_goals()
|
|
.replace_thread_goal(
|
|
thread_id,
|
|
"ship goal extension backend",
|
|
codex_state::ThreadGoalStatus::Active,
|
|
/*token_budget*/ None,
|
|
)
|
|
.await?;
|
|
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
|
|
|
|
harness.resume_thread().await;
|
|
let child = harness.spawn_child(ThreadId::new()).await?;
|
|
child.start_turn("child-turn", &TokenUsage::default()).await;
|
|
let usage = input_token_usage(/*input_tokens*/ 23);
|
|
child
|
|
.record_token_usage_with_last("child-turn", &usage, &usage)
|
|
.await;
|
|
tokio::time::sleep(Duration::from_millis(1_100)).await;
|
|
harness
|
|
.runtime_handle()
|
|
.prepare_external_goal_mutation()
|
|
.await
|
|
.map_err(anyhow::Error::msg)?;
|
|
|
|
let goal = runtime
|
|
.thread_goals()
|
|
.get_thread_goal(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
assert_eq!(ThreadGoalStatus::Active, protocol_status(goal.status));
|
|
assert_eq!(23, goal.tokens_used);
|
|
assert!(
|
|
goal.time_used_seconds >= 1,
|
|
"resumed idle accounting should add elapsed wall-clock time"
|
|
);
|
|
assert_eq!(
|
|
vec![CapturedGoalEvent {
|
|
event_id: format!("{thread_id}:external-goal-mutation"),
|
|
turn_id: None,
|
|
status: ThreadGoalStatus::Active,
|
|
tokens_used: 23,
|
|
}],
|
|
harness.sink.goal_events()
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn goal_service_sets_gets_and_clears_thread_goal() -> anyhow::Result<()> {
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let api = GoalService::new();
|
|
|
|
let request = GoalSetRequest {
|
|
thread_id,
|
|
objective: GoalObjectiveUpdate::Set(" ship goal API ownership "),
|
|
status: None,
|
|
token_budget: GoalTokenBudgetUpdate::Set(Some(123)),
|
|
max_goal_token_budget: None,
|
|
};
|
|
let failure = GoalServiceError::Internal("instruction recording failed".to_owned());
|
|
let result = api
|
|
.set_thread_goal(
|
|
runtime.as_ref(),
|
|
request,
|
|
std::future::ready(Err(failure.clone())),
|
|
)
|
|
.await;
|
|
assert_eq!(result.unwrap_err(), failure);
|
|
assert_eq!(
|
|
api.get_thread_goal(runtime.as_ref(), thread_id).await?,
|
|
None
|
|
);
|
|
let set = api
|
|
.set_thread_goal(runtime.as_ref(), request, std::future::ready(Ok(())))
|
|
.await?;
|
|
let get = api
|
|
.get_thread_goal(runtime.as_ref(), thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
|
|
let metadata = runtime
|
|
.get_thread(thread_id)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("seeded thread metadata should exist"))?;
|
|
|
|
assert_eq!(set.goal, get);
|
|
assert_eq!("ship goal API ownership", get.objective);
|
|
assert_eq!(ThreadGoalStatus::Active, get.status);
|
|
assert_eq!(Some(123), get.token_budget);
|
|
assert_eq!(Some("ship goal API ownership"), metadata.preview.as_deref());
|
|
|
|
assert_eq!(
|
|
api.clear_thread_goal(
|
|
runtime.as_ref(),
|
|
thread_id,
|
|
std::future::ready(Err(failure.clone()))
|
|
)
|
|
.await,
|
|
Err(failure),
|
|
);
|
|
assert_eq!(
|
|
api.get_thread_goal(runtime.as_ref(), thread_id).await?,
|
|
Some(get)
|
|
);
|
|
|
|
assert!(
|
|
api.clear_thread_goal(runtime.as_ref(), thread_id, std::future::ready(Ok(())))
|
|
.await?
|
|
);
|
|
assert_eq!(
|
|
None,
|
|
api.get_thread_goal(runtime.as_ref(), thread_id).await?
|
|
);
|
|
assert!(
|
|
!api.clear_thread_goal(runtime.as_ref(), thread_id, std::future::ready(Ok(())))
|
|
.await?
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn goal_service_enforces_maximum_token_budget_on_creation_and_updates() -> anyhow::Result<()>
|
|
{
|
|
let runtime = test_runtime().await?;
|
|
let thread_id = test_thread_id()?;
|
|
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
|
|
let service = GoalService::new();
|
|
|
|
let goal = service
|
|
.set_thread_goal(
|
|
runtime.as_ref(),
|
|
GoalSetRequest {
|
|
thread_id,
|
|
objective: GoalObjectiveUpdate::Set("bounded goal"),
|
|
status: None,
|
|
token_budget: GoalTokenBudgetUpdate::Keep,
|
|
max_goal_token_budget: Some(100),
|
|
},
|
|
std::future::ready(Ok(())),
|
|
)
|
|
.await?;
|
|
assert_eq!(goal.goal.token_budget, Some(100));
|
|
|
|
let error = service
|
|
.set_thread_goal(
|
|
runtime.as_ref(),
|
|
GoalSetRequest {
|
|
thread_id,
|
|
objective: GoalObjectiveUpdate::Keep,
|
|
status: None,
|
|
token_budget: GoalTokenBudgetUpdate::Set(Some(101)),
|
|
max_goal_token_budget: Some(100),
|
|
},
|
|
std::future::ready(Ok(())),
|
|
)
|
|
.await
|
|
.expect_err("goal budget above the configured maximum should fail");
|
|
assert_eq!(
|
|
error.to_string(),
|
|
"goal token budget 101 exceeds the maximum allowed goal token budget of 100"
|
|
);
|
|
assert_eq!(
|
|
service
|
|
.get_thread_goal(runtime.as_ref(), thread_id)
|
|
.await?
|
|
.expect("goal should remain unchanged")
|
|
.token_budget,
|
|
Some(100)
|
|
);
|
|
|
|
let goal = service
|
|
.set_thread_goal(
|
|
runtime.as_ref(),
|
|
GoalSetRequest {
|
|
thread_id,
|
|
objective: GoalObjectiveUpdate::Keep,
|
|
status: None,
|
|
token_budget: GoalTokenBudgetUpdate::Set(Some(99)),
|
|
max_goal_token_budget: Some(100),
|
|
},
|
|
std::future::ready(Ok(())),
|
|
)
|
|
.await?;
|
|
assert_eq!(goal.goal.token_budget, Some(99));
|
|
|
|
let goal = service
|
|
.set_thread_goal(
|
|
runtime.as_ref(),
|
|
GoalSetRequest {
|
|
thread_id,
|
|
objective: GoalObjectiveUpdate::Keep,
|
|
status: None,
|
|
token_budget: GoalTokenBudgetUpdate::Set(None),
|
|
max_goal_token_budget: Some(100),
|
|
},
|
|
std::future::ready(Ok(())),
|
|
)
|
|
.await?;
|
|
assert_eq!(goal.goal.token_budget, Some(100));
|
|
Ok(())
|
|
}
|
|
|
|
async fn installed_tools(
|
|
runtime: Arc<codex_state::StateRuntime>,
|
|
thread_id: ThreadId,
|
|
) -> Vec<Arc<dyn for<'call> ToolExecutor<ToolCall<'call>>>> {
|
|
installed_tools_with_start(
|
|
runtime,
|
|
thread_id,
|
|
SessionSource::Cli,
|
|
/*persistent_thread_state_available*/ true,
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn installed_tools_with_start(
|
|
runtime: Arc<codex_state::StateRuntime>,
|
|
thread_id: ThreadId,
|
|
session_source: SessionSource,
|
|
persistent_thread_state_available: bool,
|
|
) -> Vec<Arc<dyn for<'call> ToolExecutor<ToolCall<'call>>>> {
|
|
let mut builder = ExtensionRegistryBuilder::<()>::new();
|
|
let goal_service = Arc::new(GoalService::new());
|
|
install_with_backend(
|
|
&mut builder,
|
|
runtime,
|
|
AnalyticsEventsClient::disabled(),
|
|
/*metrics_client*/ None,
|
|
Weak::new(),
|
|
goal_service,
|
|
|_| GoalExtensionConfig {
|
|
enabled: true,
|
|
max_goal_token_budget: None,
|
|
},
|
|
);
|
|
let registry = builder.build();
|
|
let session_store = ExtensionData::new("session-1");
|
|
let thread_store = ExtensionData::new(thread_id.to_string());
|
|
for contributor in registry.thread_lifecycle_contributors() {
|
|
contributor
|
|
.on_thread_start(ThreadStartInput {
|
|
config: &(),
|
|
session_source: &session_source,
|
|
persistent_thread_state_available,
|
|
environments: &[],
|
|
mcp_resource_client: None,
|
|
extension_metrics: None,
|
|
session_store: &session_store,
|
|
thread_store: &thread_store,
|
|
})
|
|
.await;
|
|
}
|
|
|
|
registry
|
|
.tool_contributors()
|
|
.iter()
|
|
.flat_map(|contributor| contributor.tools(&session_store, &thread_store))
|
|
.collect()
|
|
}
|
|
|
|
fn tool_names(tools: &[Arc<dyn for<'call> ToolExecutor<ToolCall<'call>>>]) -> Vec<String> {
|
|
tools.iter().map(|tool| tool.tool_name().name).collect()
|
|
}
|
|
|
|
struct GoalExtensionHarness {
|
|
registry: Arc<codex_extension_api::ExtensionRegistry<()>>,
|
|
session_store: ExtensionData,
|
|
thread_store: ExtensionData,
|
|
goal_service: Arc<GoalService>,
|
|
sink: Arc<RecordingEventSink>,
|
|
}
|
|
|
|
impl GoalExtensionHarness {
|
|
async fn new(
|
|
runtime: Arc<codex_state::StateRuntime>,
|
|
thread_id: ThreadId,
|
|
) -> anyhow::Result<Self> {
|
|
let sink = Arc::new(RecordingEventSink::default());
|
|
let mut builder = ExtensionRegistryBuilder::<()>::with_event_sink(sink.clone());
|
|
let goal_service = Arc::new(GoalService::new());
|
|
install_with_backend(
|
|
&mut builder,
|
|
runtime,
|
|
AnalyticsEventsClient::disabled(),
|
|
/*metrics_client*/ None,
|
|
Weak::new(),
|
|
Arc::clone(&goal_service),
|
|
|_| GoalExtensionConfig {
|
|
enabled: true,
|
|
max_goal_token_budget: None,
|
|
},
|
|
);
|
|
let registry = Arc::new(builder.build());
|
|
let session_store = ExtensionData::new(thread_id.to_string());
|
|
let thread_store = ExtensionData::new(thread_id.to_string());
|
|
let session_source = SessionSource::Cli;
|
|
for contributor in registry.thread_lifecycle_contributors() {
|
|
contributor
|
|
.on_thread_start(ThreadStartInput {
|
|
config: &(),
|
|
session_source: &session_source,
|
|
persistent_thread_state_available: true,
|
|
environments: &[],
|
|
mcp_resource_client: None,
|
|
extension_metrics: None,
|
|
session_store: &session_store,
|
|
thread_store: &thread_store,
|
|
})
|
|
.await;
|
|
}
|
|
Ok(Self {
|
|
registry,
|
|
session_store,
|
|
thread_store,
|
|
goal_service,
|
|
sink,
|
|
})
|
|
}
|
|
|
|
async fn spawn_child(&self, thread_id: ThreadId) -> anyhow::Result<Self> {
|
|
let session_source = SubAgentSource::ThreadSpawn {
|
|
parent_thread_id: ThreadId::from_string(self.thread_store.level_id())?,
|
|
depth: 1,
|
|
agent_path: None,
|
|
agent_nickname: None,
|
|
agent_role: None,
|
|
};
|
|
self.spawn_child_with_source(thread_id, session_source)
|
|
.await
|
|
}
|
|
|
|
async fn spawn_child_with_source(
|
|
&self,
|
|
thread_id: ThreadId,
|
|
session_source: SubAgentSource,
|
|
) -> anyhow::Result<Self> {
|
|
let session_store = ExtensionData::new(self.session_store.level_id());
|
|
let thread_store = ExtensionData::new(thread_id.to_string());
|
|
let session_source = SessionSource::SubAgent(session_source);
|
|
for contributor in self.registry.thread_lifecycle_contributors() {
|
|
contributor
|
|
.on_thread_start(ThreadStartInput {
|
|
config: &(),
|
|
session_source: &session_source,
|
|
persistent_thread_state_available: true,
|
|
environments: &[],
|
|
mcp_resource_client: None,
|
|
extension_metrics: None,
|
|
session_store: &session_store,
|
|
thread_store: &thread_store,
|
|
})
|
|
.await;
|
|
}
|
|
Ok(Self {
|
|
registry: Arc::clone(&self.registry),
|
|
session_store,
|
|
thread_store,
|
|
goal_service: Arc::clone(&self.goal_service),
|
|
sink: Arc::clone(&self.sink),
|
|
})
|
|
}
|
|
|
|
fn tools(&self) -> Vec<Arc<dyn for<'call> ToolExecutor<ToolCall<'call>>>> {
|
|
self.registry
|
|
.tool_contributors()
|
|
.iter()
|
|
.flat_map(|contributor| contributor.tools(&self.session_store, &self.thread_store))
|
|
.collect()
|
|
}
|
|
|
|
async fn start_turn(&self, turn_id: &str, usage: &TokenUsage) {
|
|
self.start_turn_with_mode(turn_id, ModeKind::Default, usage)
|
|
.await;
|
|
}
|
|
|
|
async fn start_turn_with_mode(&self, turn_id: &str, mode: ModeKind, usage: &TokenUsage) {
|
|
let turn_store = ExtensionData::new(turn_id);
|
|
let mut collaboration_mode = default_collaboration_mode();
|
|
collaboration_mode.mode = mode;
|
|
for contributor in self.registry.turn_lifecycle_contributors() {
|
|
contributor
|
|
.on_turn_start(TurnStartInput {
|
|
turn_id,
|
|
collaboration_mode: &collaboration_mode,
|
|
token_usage_at_turn_start: Some(usage),
|
|
session_store: &self.session_store,
|
|
thread_store: &self.thread_store,
|
|
turn_store: &turn_store,
|
|
})
|
|
.await;
|
|
}
|
|
}
|
|
|
|
async fn stop_turn(&self, turn_id: &str) {
|
|
let turn_store = ExtensionData::new(turn_id);
|
|
for contributor in self.registry.turn_lifecycle_contributors() {
|
|
contributor
|
|
.on_turn_stop(TurnStopInput {
|
|
session_store: &self.session_store,
|
|
thread_store: &self.thread_store,
|
|
turn_store: &turn_store,
|
|
})
|
|
.await;
|
|
}
|
|
}
|
|
|
|
async fn record_token_usage(&self, turn_id: &str, usage: &TokenUsage) {
|
|
self.record_token_usage_with_last(turn_id, usage, &TokenUsage::default())
|
|
.await;
|
|
}
|
|
|
|
async fn record_token_usage_with_last(
|
|
&self,
|
|
turn_id: &str,
|
|
usage: &TokenUsage,
|
|
last_usage: &TokenUsage,
|
|
) {
|
|
let turn_store = ExtensionData::new(turn_id);
|
|
let token_usage = TokenUsageInfo {
|
|
total_token_usage: usage.clone(),
|
|
last_token_usage: last_usage.clone(),
|
|
model_context_window: None,
|
|
};
|
|
for contributor in self.registry.token_usage_contributors() {
|
|
contributor
|
|
.on_token_usage(
|
|
&self.session_store,
|
|
&self.thread_store,
|
|
&turn_store,
|
|
&token_usage,
|
|
)
|
|
.await;
|
|
}
|
|
}
|
|
|
|
async fn resume_thread(&self) {
|
|
for contributor in self.registry.thread_lifecycle_contributors() {
|
|
contributor
|
|
.on_thread_resume(ThreadResumeInput {
|
|
session_store: &self.session_store,
|
|
thread_store: &self.thread_store,
|
|
})
|
|
.await;
|
|
}
|
|
}
|
|
|
|
async fn stop_thread(&self) {
|
|
for contributor in self.registry.thread_lifecycle_contributors() {
|
|
contributor
|
|
.on_thread_stop(ThreadStopInput {
|
|
session_store: &self.session_store,
|
|
thread_store: &self.thread_store,
|
|
})
|
|
.await;
|
|
}
|
|
}
|
|
|
|
async fn notify_tool_finish(&self, turn_id: &str, call_id: &str, tool_name: &str) {
|
|
self.notify_tool_finish_with_outcome(
|
|
turn_id,
|
|
call_id,
|
|
tool_name,
|
|
ToolCallOutcome::Completed { success: true },
|
|
)
|
|
.await;
|
|
}
|
|
|
|
async fn notify_tool_finish_with_outcome(
|
|
&self,
|
|
turn_id: &str,
|
|
call_id: &str,
|
|
tool_name: &str,
|
|
outcome: ToolCallOutcome,
|
|
) {
|
|
let turn_store = ExtensionData::new(turn_id);
|
|
let tool_name = codex_extension_api::ToolName::plain(tool_name);
|
|
for contributor in self.registry.tool_lifecycle_contributors() {
|
|
contributor
|
|
.on_tool_finish(ToolFinishInput {
|
|
session_store: &self.session_store,
|
|
thread_store: &self.thread_store,
|
|
turn_store: &turn_store,
|
|
turn_id,
|
|
call_id,
|
|
tool_name: &tool_name,
|
|
source: ToolCallSource::Direct,
|
|
outcome,
|
|
})
|
|
.await;
|
|
}
|
|
}
|
|
|
|
async fn notify_turn_error(&self, turn_id: &str, error: CodexErr) {
|
|
let turn_store = ExtensionData::new(turn_id);
|
|
for contributor in self.registry.turn_lifecycle_contributors() {
|
|
contributor
|
|
.on_turn_error(TurnErrorInput {
|
|
turn_id,
|
|
error: error.to_codex_protocol_error(),
|
|
error_details: error.details(),
|
|
session_store: &self.session_store,
|
|
thread_store: &self.thread_store,
|
|
turn_store: &turn_store,
|
|
})
|
|
.await;
|
|
}
|
|
}
|
|
|
|
fn runtime_handle(&self) -> Arc<GoalRuntimeHandle> {
|
|
self.thread_store
|
|
.get::<GoalRuntimeHandle>()
|
|
.expect("goal runtime handle should exist")
|
|
}
|
|
}
|
|
|
|
fn tool_by_name<'a>(
|
|
tools: &'a [Arc<dyn for<'call> ToolExecutor<ToolCall<'call>>>],
|
|
name: &str,
|
|
) -> &'a Arc<dyn for<'call> ToolExecutor<ToolCall<'call>>> {
|
|
tools
|
|
.iter()
|
|
.find(|tool| tool.tool_name().namespace.is_none() && tool.tool_name().name == name)
|
|
.expect("requested goal tool should exist")
|
|
}
|
|
|
|
fn tool_call(tool_name: &str, call_id: &str, arguments: serde_json::Value) -> ToolCall<'static> {
|
|
ToolCall {
|
|
turn_id: "turn-1".to_string(),
|
|
call_id: call_id.to_string(),
|
|
tool_name: codex_extension_api::ToolName::plain(tool_name),
|
|
model: "gpt-test".to_string(),
|
|
codex_turn_metadata: None,
|
|
truncation_policy: TruncationPolicy::Bytes(1024),
|
|
source: ToolCallSource::Direct,
|
|
conversation_history: codex_extension_api::ConversationHistory::default(),
|
|
turn_item_emitter: Arc::new(NoopTurnItemEmitter),
|
|
attachment_store: codex_core::passthrough_image_store(),
|
|
ephemeral: false,
|
|
environments: Vec::new(),
|
|
payload: ToolPayload::Function {
|
|
arguments: arguments.to_string(),
|
|
},
|
|
}
|
|
}
|
|
|
|
async fn test_runtime() -> anyhow::Result<Arc<codex_state::StateRuntime>> {
|
|
let tempdir = TempDir::new()?;
|
|
codex_state::StateRuntime::init(
|
|
codex_state::SqliteConfig::new_for_testing(tempdir.keep().as_path().abs()),
|
|
"test-provider".to_string(),
|
|
)
|
|
.await
|
|
}
|
|
|
|
fn test_thread_id() -> anyhow::Result<ThreadId> {
|
|
ThreadId::from_string("11111111-1111-4111-8111-111111111111").map_err(anyhow::Error::msg)
|
|
}
|
|
|
|
async fn seed_thread_metadata(
|
|
runtime: &codex_state::StateRuntime,
|
|
thread_id: ThreadId,
|
|
) -> anyhow::Result<()> {
|
|
let builder = codex_state::ThreadMetadataBuilder::new(
|
|
thread_id,
|
|
runtime
|
|
.sqlite()
|
|
.home()
|
|
.join(format!("rollout-{thread_id}.jsonl")),
|
|
chrono::Utc::now(),
|
|
SessionSource::Cli,
|
|
);
|
|
runtime.upsert_thread(&builder.build("test-provider")).await
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
struct RecordingEventSink {
|
|
events: Mutex<Vec<Event>>,
|
|
}
|
|
|
|
impl RecordingEventSink {
|
|
fn goal_events(&self) -> Vec<CapturedGoalEvent> {
|
|
self.events()
|
|
.iter()
|
|
.filter_map(|event| match &event.msg {
|
|
EventMsg::ThreadGoalUpdated(updated) => Some(CapturedGoalEvent {
|
|
event_id: event.id.clone(),
|
|
turn_id: updated.turn_id.clone(),
|
|
status: updated.goal.status,
|
|
tokens_used: updated.goal.tokens_used,
|
|
}),
|
|
_ => None,
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn clear(&self) {
|
|
self.events().clear();
|
|
}
|
|
|
|
fn events(&self) -> std::sync::MutexGuard<'_, Vec<Event>> {
|
|
self.events.lock().unwrap_or_else(PoisonError::into_inner)
|
|
}
|
|
}
|
|
|
|
impl ExtensionEventSink for RecordingEventSink {
|
|
fn emit(&self, event: Event) {
|
|
self.events().push(event);
|
|
}
|
|
|
|
fn emit_warning(&self, _warning: ExtensionWarning) {
|
|
panic!("goal extension tests do not emit warnings");
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, PartialEq, Eq)]
|
|
struct CapturedGoalEvent {
|
|
event_id: String,
|
|
turn_id: Option<String>,
|
|
status: ThreadGoalStatus,
|
|
tokens_used: i64,
|
|
}
|
|
|
|
fn default_collaboration_mode() -> CollaborationMode {
|
|
CollaborationMode {
|
|
mode: ModeKind::Default,
|
|
settings: Settings {
|
|
model: "gpt-5".to_string(),
|
|
reasoning_effort: None,
|
|
developer_instructions: None,
|
|
},
|
|
}
|
|
}
|
|
|
|
fn token_usage(
|
|
input_tokens: i64,
|
|
cached_input_tokens: i64,
|
|
output_tokens: i64,
|
|
reasoning_output_tokens: i64,
|
|
total_tokens: i64,
|
|
) -> TokenUsage {
|
|
TokenUsage {
|
|
input_tokens,
|
|
cached_input_tokens,
|
|
cache_write_input_tokens: 0,
|
|
output_tokens,
|
|
reasoning_output_tokens,
|
|
total_tokens,
|
|
codex_rollout_budget_units: None,
|
|
}
|
|
}
|
|
|
|
fn input_token_usage(input_tokens: i64) -> TokenUsage {
|
|
TokenUsage {
|
|
input_tokens,
|
|
total_tokens: input_tokens,
|
|
..TokenUsage::default()
|
|
}
|
|
}
|
|
|
|
fn protocol_status(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,
|
|
}
|
|
}
|