Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions codex-rs/ext/goal/src/extension.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,13 +22,16 @@ use codex_extension_api::TurnStartInput;
use codex_extension_api::TurnStopInput;
use codex_otel::MetricsClient;
use codex_protocol::ThreadId;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::ThreadGoalStatus;
use codex_protocol::protocol::TokenUsageInfo;

use crate::accounting::BudgetLimitedGoalDisposition;
use crate::accounting::GoalAccountingState;
use crate::events::GoalEventEmitter;
use crate::metrics::GoalMetrics;
use crate::runtime::GoalRuntimeConfig;
use crate::runtime::GoalRuntimeHandle;
use crate::spec::UPDATE_GOAL_TOOL_NAME;
use crate::steering::budget_limit_steering_item;
Expand Down Expand Up @@ -85,6 +88,11 @@ where
{
async fn on_thread_start(&self, input: ThreadStartInput<'_, C>) {
let enabled = (self.goals_enabled)(input.config);
let tools_available_for_thread = input.persistent_thread_state_available
&& !matches!(
input.session_source,
SessionSource::SubAgent(SubAgentSource::Review)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Hide goal tools for guardian review sessions

When auto-approval review uses the reusable guardian reviewer, the spawned session is tagged as SessionSource::SubAgent(SubAgentSource::Other("guardian")) in codex-rs/core/src/guardian/review_session.rs:607, and build_guardian_review_session_config does not disable Feature::Goals. In that non-ephemeral review session this exact SubAgentSource::Review check still leaves goal tools visible, so reviewer agents can still see and call create/update goal tools despite this gate trying to suppress goal controls in review contexts.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is in sync with the old behaviour

);
input
.thread_store
.insert(GoalExtensionConfig::from_enabled(enabled));
Expand All @@ -102,7 +110,10 @@ where
self.metrics.clone(),
self.thread_manager.clone(),
accounting_state,
enabled,
GoalRuntimeConfig {
enabled,
tools_available_for_thread,
},
)
});
runtime.set_enabled(enabled);
Expand Down Expand Up @@ -330,7 +341,7 @@ where
let Some(runtime) = goal_runtime_handle(thread_store) else {
return Vec::new();
};
if !runtime.is_enabled() {
if !runtime.tools_visible() {
return Vec::new();
}

Expand Down
15 changes: 13 additions & 2 deletions codex-rs/ext/goal/src/runtime.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,11 @@ pub struct GoalRuntimeHandle {
inner: Arc<GoalRuntimeInner>,
}

pub(crate) struct GoalRuntimeConfig {
pub(crate) enabled: bool,
pub(crate) tools_available_for_thread: bool,
}

struct GoalRuntimeInner {
thread_id: ThreadId,
state_dbs: Arc<codex_state::StateRuntime>,
Expand All @@ -28,6 +33,7 @@ struct GoalRuntimeInner {
thread_manager: Weak<ThreadManager>,
accounting_state: Arc<GoalAccountingState>,
enabled: AtomicBool,
tools_available_for_thread: bool,
}

pub(crate) struct AccountedGoalProgress {
Expand Down Expand Up @@ -66,7 +72,7 @@ impl GoalRuntimeHandle {
metrics: GoalMetrics,
thread_manager: Weak<ThreadManager>,
accounting_state: Arc<GoalAccountingState>,
enabled: bool,
config: GoalRuntimeConfig,
) -> Self {
Self {
inner: Arc::new(GoalRuntimeInner {
Expand All @@ -76,7 +82,8 @@ impl GoalRuntimeHandle {
metrics,
thread_manager,
accounting_state,
enabled: AtomicBool::new(enabled),
enabled: AtomicBool::new(config.enabled),
tools_available_for_thread: config.tools_available_for_thread,
}),
}
}
Expand All @@ -89,6 +96,10 @@ impl GoalRuntimeHandle {
self.inner.enabled.load(Ordering::Relaxed)
}

pub(crate) fn tools_visible(&self) -> bool {
self.is_enabled() && self.inner.tools_available_for_thread
}

pub(crate) fn thread_id(&self) -> ThreadId {
self.inner.thread_id
}
Expand Down
55 changes: 53 additions & 2 deletions codex-rs/ext/goal/tests/goal_extension_backend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ use codex_protocol::config_types::Settings;
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;
Expand Down Expand Up @@ -83,6 +84,38 @@ async fn installed_goal_tools_create_goal_and_fill_empty_preview() -> anyhow::Re
Ok(())
}

#[tokio::test]
async fn goal_tools_hidden_for_ephemeral_threads() -> anyhow::Result<()> {
let runtime = test_runtime().await?;
let thread_id = test_thread_id()?;
let tools = installed_tools_with_start(
runtime,
thread_id,
SessionSource::Cli,
/*persistent_thread_state_available*/ false,
)
.await;

assert_eq!(Vec::<String>::new(), tool_names(&tools));
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_reject_duplicate_goal_creation() -> anyhow::Result<()> {
let runtime = test_runtime().await?;
Expand Down Expand Up @@ -877,6 +910,21 @@ async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Resul
async fn installed_tools(
runtime: Arc<codex_state::StateRuntime>,
thread_id: ThreadId,
) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
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 ToolExecutor<ToolCall>>> {
let mut builder = ExtensionRegistryBuilder::<()>::new();
install_with_backend(
Expand All @@ -889,13 +937,12 @@ async fn installed_tools(
let registry = builder.build();
let session_store = ExtensionData::new("session-1");
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,
persistent_thread_state_available,
session_store: &session_store,
thread_store: &thread_store,
})
Expand All @@ -909,6 +956,10 @@ async fn installed_tools(
.collect()
}

fn tool_names(tools: &[Arc<dyn ToolExecutor<ToolCall>>]) -> Vec<String> {
tools.iter().map(|tool| tool.tool_name().name).collect()
}

struct GoalExtensionHarness {
registry: codex_extension_api::ExtensionRegistry<()>,
session_store: ExtensionData,
Expand Down
Loading