Add the git attribution extension (#34769)

## What changed

- Add a `codex-git-attribution` context contributor that resolves the workspace's `commit_attribution_enabled` setting and supplies matching commit and pull request instructions through World State.
- Cache resolved policy by authentication generation, recover from unauthorized responses, and fail disabled on timeouts or settings errors with deferred retries.
- Handle enabled, disabled, and legacy attribution fragments so policy changes replace stale instructions without duplicating attribution.

## Testing

- Cover World State transitions, policy resolution after unauthorized responses, and authentication changes during an in-flight settings request.

GitOrigin-RevId: 6f733e708b5c960ba5e90751110910a03d1ba2df
This commit is contained in:
jay
2026-07-22 15:19:15 +00:00
committed by copyberry
parent e370d23691
commit ab816f3ca0
8 changed files with 643 additions and 0 deletions
+15
View File
@@ -3113,6 +3113,21 @@ dependencies = [
"tracing",
]
[[package]]
name = "codex-git-attribution"
version = "0.0.0"
dependencies = [
"codex-backend-client",
"codex-extension-api",
"codex-http-client",
"codex-login",
"codex-protocol",
"pretty_assertions",
"serde_json",
"tokio",
"wiremock",
]
[[package]]
name = "codex-git-utils"
version = "0.0.0"
+1
View File
@@ -52,6 +52,7 @@ members = [
"ext/connectors",
"ext/extension-api",
"ext/goal",
"ext/git-attribution",
"ext/guardian",
"ext/image-generation",
"ext/items",
+6
View File
@@ -0,0 +1,6 @@
load("//:defs.bzl", "codex_rust_crate")
codex_rust_crate(
name = "git-attribution",
crate_name = "codex_git_attribution",
)
+27
View File
@@ -0,0 +1,27 @@
[package]
edition.workspace = true
license.workspace = true
name = "codex-git-attribution"
version.workspace = true
[lib]
name = "codex_git_attribution"
path = "src/lib.rs"
doctest = false
[lints]
workspace = true
[dependencies]
codex-backend-client = { workspace = true }
codex-extension-api = { workspace = true }
codex-http-client = { workspace = true }
codex-login = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true, features = ["time"] }
[dev-dependencies]
codex-protocol = { workspace = true }
pretty_assertions = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt"] }
wiremock = { workspace = true }
@@ -0,0 +1,299 @@
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionRegistry;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::PreviousWorldStateSection;
use codex_extension_api::WorldStateContributionInput;
use codex_extension_api::WorldStateSectionContribution;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::ExternalAuth;
use codex_login::ExternalAuthFuture;
use codex_login::ExternalAuthRefreshContext;
use codex_protocol::ThreadId;
use pretty_assertions::assert_eq;
use serde_json::Value;
use tokio::sync::Notify;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
use super::install;
use super::policy::resolve_attribution_policy;
use super::world_state::DISABLED_INSTRUCTIONS;
use super::world_state::ENABLED_INSTRUCTIONS;
use super::world_state::LEGACY_COMMIT_ATTRIBUTION_INSTRUCTIONS;
async fn contribute(
registry: &ExtensionRegistry<String>,
thread_store: &ExtensionData,
) -> WorldStateSectionContribution {
let session_store = ExtensionData::new("session");
let turn_store = ExtensionData::new("turn");
let step_store = ExtensionData::new("step");
registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
thread_id: ThreadId::new(),
turn_id: "turn",
environments: &[],
ready_selected_capability_roots: &[],
executor_capability_discovery: None,
session_store: &session_store,
thread_store,
turn_store: &turn_store,
step_store: &step_store,
})
.await
.remove(0)
}
fn enterprise_auth_manager() -> Arc<AuthManager> {
AuthManager::from_auth_for_testing(enterprise_auth("workspace-123"))
}
fn http_client_factory() -> HttpClientFactory {
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
}
fn enterprise_auth(account_id: &str) -> CodexAuth {
CodexAuth::from_external_chatgpt_tokens("e30.e30.c2ln", account_id, Some("enterprise"))
.expect("fake ChatGPT auth should parse")
}
async fn mount_settings(server: &MockServer, response: ResponseTemplate) {
Mock::given(method("GET"))
.and(path("/backend-api/wham/settings/user"))
.respond_with(response)
.expect(1)
.mount(server)
.await;
}
struct StaticExternalAuth(CodexAuth);
impl ExternalAuth for StaticExternalAuth {
fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> {
Box::pin(async { Ok(self.0.clone()) })
}
fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> {
self.resolve()
}
}
async fn set_auth(auth_manager: &AuthManager, account_id: &str) {
auth_manager
.set_external_auth(Arc::new(StaticExternalAuth(enterprise_auth(account_id))))
.await
.expect("auth refresh should succeed");
}
#[tokio::test]
async fn installed_contributor_composes_policy_changes() {
let server = MockServer::start().await;
let base_url = format!("{}/backend-api", server.uri());
let api_key = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test"));
assert!(
!resolve_attribution_policy(&api_key, &base_url, &http_client_factory())
.await
.expect("API key policy should not time out")
.expect("API key policy should resolve")
.enabled
);
let auth_manager = enterprise_auth_manager();
let mut builder = ExtensionRegistryBuilder::<String>::new();
install(
&mut builder,
auth_manager.clone(),
base_url.clone(),
http_client_factory(),
);
let registry = builder.build();
let thread_store = ExtensionData::new("thread");
mount_settings(
&server,
ResponseTemplate::new(200)
.set_delay(Duration::from_secs(1))
.set_body_json(serde_json::json!({"commit_attribution_enabled": true})),
)
.await;
let unavailable = contribute(&registry, &thread_store).await;
assert_eq!(unavailable.snapshot(), &Value::Bool(false));
assert!(
unavailable
.render_diff(PreviousWorldStateSection::Absent)
.is_none()
);
assert!(!unavailable.has_retained_fragment_matcher());
assert!(
unavailable.matches_legacy_fragment("developer", LEGACY_COMMIT_ATTRIBUTION_INSTRUCTIONS)
);
assert_eq!(
unavailable
.render_diff(PreviousWorldStateSection::Unknown)
.map(|fragment| fragment.body().to_string()),
Some(DISABLED_INSTRUCTIONS.to_string())
);
assert_eq!(
contribute(&registry, &thread_store).await.snapshot(),
&Value::Bool(false)
);
server.verify().await;
server.reset().await;
mount_settings(&server, ResponseTemplate::new(503)).await;
set_auth(auth_manager.as_ref(), "workspace-456").await;
let failed = contribute(&registry, &thread_store).await;
assert_eq!(failed.snapshot(), &Value::Bool(false));
assert!(
failed
.render_diff(PreviousWorldStateSection::Absent)
.is_none()
);
server.verify().await;
server.reset().await;
mount_settings(
&server,
ResponseTemplate::new(200)
.set_body_json(serde_json::json!({"commit_attribution_enabled": true})),
)
.await;
set_auth(auth_manager.as_ref(), "workspace-789").await;
let enabled = contribute(&registry, &thread_store).await;
let enabled_fragment = enabled
.render_diff(PreviousWorldStateSection::Known(&Value::Bool(false)))
.expect("enabled policy should replace disabled policy");
assert_eq!(enabled_fragment.body(), ENABLED_INSTRUCTIONS);
let (start, end) = enabled_fragment.markers();
let rendered = format!("{start}{}{end}", enabled_fragment.body());
assert!(enabled.matches_legacy_fragment(enabled_fragment.role(), &rendered));
assert!(enabled.matches_retained_fragment(enabled_fragment.role(), &rendered));
assert!(!enabled.matches_legacy_fragment("developer", LEGACY_COMMIT_ATTRIBUTION_INSTRUCTIONS));
assert!(
contribute(&registry, &thread_store)
.await
.render_diff(PreviousWorldStateSection::Known(&Value::Bool(true)))
.is_none()
);
server.verify().await;
server.reset().await;
mount_settings(
&server,
ResponseTemplate::new(200)
.set_body_json(serde_json::json!({"commit_attribution_enabled": false})),
)
.await;
set_auth(auth_manager.as_ref(), "workspace-disabled").await;
let disabled = contribute(&registry, &thread_store).await;
assert_eq!(
disabled
.render_diff(PreviousWorldStateSection::Known(&Value::Bool(true)))
.map(|fragment| fragment.body().to_string()),
Some(DISABLED_INSTRUCTIONS.to_string())
);
assert!(
disabled
.render_diff(PreviousWorldStateSection::Absent)
.is_none()
);
server.verify().await;
}
#[tokio::test]
async fn policy_resolution_recovers_after_unauthorized() {
let server = MockServer::start().await;
let request_count = Arc::new(AtomicUsize::new(0));
Mock::given(method("GET"))
.and(path("/backend-api/wham/settings/user"))
.respond_with({
let request_count = request_count.clone();
move |_request: &wiremock::Request| {
if request_count.fetch_add(1, Ordering::SeqCst) == 0 {
ResponseTemplate::new(401)
} else {
ResponseTemplate::new(200)
.set_body_json(serde_json::json!({"commit_attribution_enabled": true}))
}
}
})
.expect(2)
.mount(&server)
.await;
let auth_manager = enterprise_auth_manager();
set_auth(auth_manager.as_ref(), "workspace-123").await;
let policy = resolve_attribution_policy(
&auth_manager,
&format!("{}/backend-api", server.uri()),
&http_client_factory(),
)
.await
.expect("policy resolution should not time out")
.expect("policy should resolve after auth recovery");
assert!(policy.enabled);
assert_eq!(request_count.load(Ordering::SeqCst), 2);
server.verify().await;
}
#[tokio::test]
async fn policy_resolution_retries_after_auth_refresh() {
let server = MockServer::start().await;
let request_started = Arc::new(Notify::new());
let request_count = Arc::new(AtomicUsize::new(0));
Mock::given(method("GET"))
.and(path("/backend-api/wham/settings/user"))
.respond_with({
let request_started = request_started.clone();
let request_count = request_count.clone();
move |_request: &wiremock::Request| match request_count.fetch_add(1, Ordering::SeqCst) {
0 => {
request_started.notify_one();
ResponseTemplate::new(200)
.set_delay(Duration::from_millis(100))
.set_body_json(serde_json::json!({
"commit_attribution_enabled": true,
}))
}
1 => ResponseTemplate::new(401),
_ => ResponseTemplate::new(200).set_body_json(serde_json::json!({
"commit_attribution_enabled": true,
})),
}
})
.expect(3)
.mount(&server)
.await;
let auth_manager = enterprise_auth_manager();
let resolve = tokio::spawn({
let auth_manager = auth_manager.clone();
let base_url = format!("{}/backend-api", server.uri());
async move {
resolve_attribution_policy(&auth_manager, &base_url, &http_client_factory())
.await
.ok()
.flatten()
}
});
tokio::time::timeout(Duration::from_secs(5), request_started.notified())
.await
.expect("first settings request should start");
set_auth(auth_manager.as_ref(), "workspace-456").await;
assert!(
resolve
.await
.expect("policy task should complete")
.expect("policy should resolve after refresh")
.enabled
);
server.verify().await;
}
+113
View File
@@ -0,0 +1,113 @@
mod policy;
mod world_state;
use std::sync::Arc;
use std::time::Instant;
use codex_extension_api::ContextContributor;
use codex_extension_api::ExtensionFuture;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::WorldStateContributionInput;
use codex_extension_api::WorldStateSectionContribution;
use codex_http_client::HttpClientFactory;
use codex_login::AuthManager;
use crate::policy::GitAttributionPolicy;
use crate::policy::GitAttributionRetry;
use crate::policy::POLICY_RETRY_DELAY;
use crate::policy::auth_generation;
use crate::policy::cached_attribution_policy;
use crate::policy::resolve_attribution_policy;
use crate::policy::retry_deferred;
use crate::world_state::git_attribution_world_state_section;
/// Contributes model instructions for agent-created git commits and pull requests.
#[derive(Clone)]
struct GitAttributionExtension {
auth_manager: Arc<AuthManager>,
base_url: String,
http_client_factory: HttpClientFactory,
}
impl ContextContributor for GitAttributionExtension {
fn contribute_world_state<'a>(
&'a self,
input: WorldStateContributionInput<'a>,
) -> ExtensionFuture<'a, Vec<WorldStateSectionContribution>> {
Box::pin(async move {
let enabled = loop {
let current_auth_generation = auth_generation(self.auth_manager.as_ref());
let policy = match cached_attribution_policy(
input.thread_store,
input.turn_store,
current_auth_generation,
) {
Some(policy) => policy,
None if retry_deferred(input.thread_store, current_auth_generation) => {
GitAttributionPolicy {
auth_generation: current_auth_generation,
enabled: false,
}
}
None => {
match resolve_attribution_policy(
&self.auth_manager,
&self.base_url,
&self.http_client_factory,
)
.await
{
Ok(Some(policy)) => {
input.thread_store.insert(policy.clone());
policy
}
Ok(None) => {
let policy = GitAttributionPolicy {
auth_generation: current_auth_generation,
enabled: false,
};
input.turn_store.insert(policy.clone());
policy
}
Err(_) => {
let auth_generation = auth_generation(self.auth_manager.as_ref());
if auth_generation == current_auth_generation {
input.thread_store.insert(GitAttributionRetry {
auth_generation,
retry_at: Instant::now() + POLICY_RETRY_DELAY,
});
}
GitAttributionPolicy {
auth_generation: current_auth_generation,
enabled: false,
}
}
}
}
};
if policy.auth_generation == auth_generation(self.auth_manager.as_ref()) {
break policy.enabled;
}
};
vec![git_attribution_world_state_section(enabled)]
})
}
}
/// Installs the git-attribution contributor into the extension registry.
pub fn install<C: Sync>(
registry: &mut ExtensionRegistryBuilder<C>,
auth_manager: Arc<AuthManager>,
base_url: String,
http_client_factory: HttpClientFactory,
) {
registry.prompt_contributor(Arc::new(GitAttributionExtension {
auth_manager,
base_url,
http_client_factory,
}));
}
#[cfg(test)]
#[path = "git_attribution_tests.rs"]
mod tests;
+106
View File
@@ -0,0 +1,106 @@
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;
use codex_backend_client::Client as BackendClient;
use codex_extension_api::ExtensionData;
use codex_http_client::HttpClientFactory;
use codex_login::AuthManager;
use tokio::time::timeout;
#[derive(Clone, Debug)]
pub(super) struct GitAttributionPolicy {
pub(super) auth_generation: u64,
pub(super) enabled: bool,
}
pub(super) struct GitAttributionRetry {
pub(super) auth_generation: u64,
pub(super) retry_at: Instant,
}
pub(super) fn retry_deferred(thread_store: &ExtensionData, auth_generation: u64) -> bool {
thread_store
.get::<GitAttributionRetry>()
.is_some_and(|retry| {
retry.auth_generation == auth_generation && retry.retry_at > Instant::now()
})
}
pub(super) fn cached_attribution_policy(
thread_store: &ExtensionData,
turn_store: &ExtensionData,
auth_generation: u64,
) -> Option<GitAttributionPolicy> {
thread_store
.get::<GitAttributionPolicy>()
.filter(|policy| policy.auth_generation == auth_generation)
.or_else(|| {
turn_store
.get::<GitAttributionPolicy>()
.filter(|policy| policy.auth_generation == auth_generation)
})
.map(|policy| policy.as_ref().clone())
}
#[cfg(not(test))]
const POLICY_RESOLUTION_TIMEOUT: Duration = Duration::from_secs(5);
#[cfg(test)]
const POLICY_RESOLUTION_TIMEOUT: Duration = Duration::from_millis(500);
pub(super) const POLICY_RETRY_DELAY: Duration = Duration::from_secs(30);
pub(super) async fn resolve_attribution_policy(
auth_manager: &Arc<AuthManager>,
base_url: &str,
http_client_factory: &HttpClientFactory,
) -> Result<Option<GitAttributionPolicy>, tokio::time::error::Elapsed> {
timeout(POLICY_RESOLUTION_TIMEOUT, async {
let mut recovery_generation = auth_generation(auth_manager);
let mut auth_recovery = auth_manager.unauthorized_recovery();
loop {
let auth_generation_at_start = auth_generation(auth_manager);
if auth_generation_at_start != recovery_generation {
auth_recovery = auth_manager.unauthorized_recovery();
recovery_generation = auth_generation_at_start;
}
let auth = auth_manager.auth().await;
if auth_generation(auth_manager) != auth_generation_at_start {
continue;
}
let enabled = match auth {
Some(auth) if auth.uses_codex_backend() => {
let client =
BackendClient::from_auth(base_url, &auth, http_client_factory.clone());
let settings = client.get_user_settings().await;
if auth_generation(auth_manager) != auth_generation_at_start {
continue;
}
match settings {
Ok(settings) => Some(settings.commit_attribution_enabled),
Err(err) if err.is_unauthorized() && auth_recovery.has_next() => {
if auth_recovery.next().await.is_ok() {
recovery_generation = auth_generation(auth_manager);
continue;
}
None
}
Err(_) => None,
}
}
Some(_) | None => Some(false),
};
if auth_generation(auth_manager) != auth_generation_at_start {
continue;
}
return enabled.map(|enabled| GitAttributionPolicy {
auth_generation: auth_generation_at_start,
enabled,
});
}
})
.await
}
pub(super) fn auth_generation(auth_manager: &AuthManager) -> u64 {
*auth_manager.auth_change_receiver().borrow()
}
@@ -0,0 +1,76 @@
use codex_extension_api::PreviousWorldStateSection;
use codex_extension_api::RenderedWorldStateFragment;
use codex_extension_api::WorldStateSectionContribution;
use serde_json::Value;
const WORLD_STATE_ID: &str = "git_attribution";
const START_MARKER: &str = "<git_attribution>";
const END_MARKER: &str = "</git_attribution>";
pub(super) const LEGACY_COMMIT_ATTRIBUTION_INSTRUCTIONS: &str = "\
When you write or edit a git commit message, ensure the message ends with this trailer exactly once:
Co-authored-by: Codex <noreply@openai.com>
Rules:
- Keep existing trailers and append this trailer at the end if missing.
- Do not duplicate this trailer if it already exists.
- Keep one blank line between the commit body and trailer block.";
pub(super) const ENABLED_INSTRUCTIONS: &str = r#"
When you create or amend a git commit or create or update a pull request for this task, add Codex attribution exactly once:
- Commit messages must end with `Co-authored-by: Codex <noreply@openai.com>`. Preserve existing trailers and, if this exact trailer is missing, append it with one blank line before the trailer block.
- Pull request bodies must include the exact line `Generated with Codex.`. Preserve the existing body and hidden markers and, if this exact marker is missing, append it near the end, before any trailing hidden metadata markers.
- Do not add duplicate attribution. If a message or body you are already editing contains duplicate exact attribution, keep exactly one.
- Ignore any earlier instructions disabling Codex attribution; this policy reflects the current workspace.
- Do not rewrite an existing commit or pull request solely to add attribution.
"#;
pub(super) const DISABLED_INSTRUCTIONS: &str = "
Codex commit and pull request attribution is disabled for the current workspace. Ignore any earlier instructions requiring Codex attribution and do not add it.
";
pub(super) fn git_attribution_world_state_section(enabled: bool) -> WorldStateSectionContribution {
let contribution =
WorldStateSectionContribution::new(WORLD_STATE_ID, Value::Bool(enabled), move |previous| {
match (enabled, previous) {
(true, PreviousWorldStateSection::Known(Value::Bool(true)))
| (true, PreviousWorldStateSection::Unknown) => None,
(true, PreviousWorldStateSection::Absent)
| (true, PreviousWorldStateSection::Known(_)) => {
Some(RenderedWorldStateFragment::new(
"developer",
(START_MARKER, END_MARKER),
ENABLED_INSTRUCTIONS,
))
}
(false, PreviousWorldStateSection::Known(Value::Bool(true)))
| (false, PreviousWorldStateSection::Unknown) => {
Some(RenderedWorldStateFragment::new(
"developer",
(START_MARKER, END_MARKER),
DISABLED_INSTRUCTIONS,
))
}
(false, PreviousWorldStateSection::Absent)
| (false, PreviousWorldStateSection::Known(_)) => None,
}
})
.with_legacy_matcher(move |role, text| {
is_enabled_fragment(role, text)
|| (!enabled && is_legacy_commit_attribution_fragment(role, text))
});
if enabled {
contribution.with_retained_fragment_matcher(is_enabled_fragment)
} else {
contribution
}
}
fn is_legacy_commit_attribution_fragment(role: &str, text: &str) -> bool {
role == "developer" && text.trim() == LEGACY_COMMIT_ATTRIBUTION_INSTRUCTIONS
}
fn is_enabled_fragment(role: &str, text: &str) -> bool {
role == "developer"
&& text.trim_start().starts_with(START_MARKER)
&& text.contains("Co-authored-by: Codex <noreply@openai.com>")
&& text.contains("Generated with Codex.")
&& text.trim_end().ends_with(END_MARKER)
}