Skip to content

Commit 95d38b9

Browse files
Add exact steering and API-key attention waits
1 parent a3cf563 commit 95d38b9

16 files changed

Lines changed: 1490 additions & 50 deletions

‎.github/workflows/codeql.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ jobs:
1919
# A successful analysis job does not mean that no alerts were found.
2020
name: CodeQL (${{ matrix.language }})
2121
runs-on: ubuntu-latest
22-
timeout-minutes: 45
22+
timeout-minutes: ${{ matrix.language == 'rust' && 180 || 45 }}
2323
permissions:
2424
contents: read
2525
security-events: write

‎packages/codex/patches/0018-api-key-rate-limits.patch‎

Lines changed: 792 additions & 0 deletions
Large diffs are not rendered by default.

‎packages/codex/patches/series.json‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,11 @@
6767
{
6868
"path": "0017-native-policy-expiry-admission.patch",
6969
"sha256": "4406a07b8bcd9769595ba9295e10986311e8b3fd96bcfeda3d9dfe13a3095048"
70+
},
71+
{
72+
"path": "0018-api-key-rate-limits.patch",
73+
"sha256": "7214955c12f507f0534a7813205eeba656265aa6d919aa3704ec5a47f3ea7d64"
7074
}
7175
],
72-
"result_tree": "66f4dc7319790001e362f9dc425916708efe4729"
76+
"result_tree": "031267b70bbf9fda8d45193acc5d41776c136c5a"
7377
}

‎packages/codex/rust/native_interception.rs‎

Lines changed: 28 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6370,7 +6370,7 @@ pub(crate) mod tests {
63706370
"name":"update_plan", "namespace":"functions",
63716371
"arguments":"{\"plan\":[]}", "encrypted_function_args":[],
63726372
}));
6373-
let original = serde_json::to_value(&item).unwrap();
6373+
assert_eq!(item.id().unwrap().as_str(), "fc_server");
63746374
let upstream = Box::pin(stream::iter(vec![
63756375
Ok(ResponseEvent::Created {
63766376
response_id: Some("real-provider".into()),
@@ -6386,7 +6386,7 @@ pub(crate) mod tests {
63866386
let mut intercepted = Box::new(request).intercept(upstream);
63876387
assert!(matches!(
63886388
intercepted.next().await,
6389-
Some(Ok(ResponseEvent::Created { .. }))
6389+
Some(Ok(ResponseEvent::Created { response_id: Some(id) })) if id == "real-provider"
63906390
));
63916391
let edit = async {
63926392
let req = control.next().await.unwrap();
@@ -6404,16 +6404,21 @@ pub(crate) mod tests {
64046404
let Some(Ok(ResponseEvent::OutputItemDone(item))) = event else {
64056405
panic!("missing actual item")
64066406
};
6407-
let mut actual = serde_json::to_value(item).unwrap();
6407+
assert!(item.id().is_none());
6408+
assert!(matches!(&item, ResponseItem::FunctionCall { call_id, .. }
6409+
if call_id == "real-native-call"));
64086410
assert_eq!(
6409-
serde_json::from_str::<Value>(actual["arguments"].as_str().unwrap()).unwrap(),
6410-
json!({"plan":[],"explanation":"effective"})
6411+
serde_json::to_value(item).unwrap(),
6412+
json!({
6413+
"type": "function_call", "call_id": "real-native-call",
6414+
"name": "update_plan", "namespace": "functions",
6415+
"arguments": json!({"plan": [], "explanation": "effective"}).to_string(),
6416+
"encrypted_function_args": [],
6417+
})
64116418
);
6412-
actual["arguments"] = original["arguments"].clone();
6413-
assert_eq!(actual, original); // including encrypted flags, id, name, namespace and metadata
64146419
assert!(matches!(
64156420
intercepted.next().await,
6416-
Some(Ok(ResponseEvent::Created { .. }))
6421+
Some(Ok(ResponseEvent::Created { response_id: Some(id) })) if id == "next-provider-event"
64176422
));
64186423
assert!(intercepted.next().await.is_none());
64196424
}
@@ -6477,7 +6482,7 @@ pub(crate) mod tests {
64776482
"arguments":r#"{"task_name":"old","message":"opaque"}"#,
64786483
"encrypted_function_args":metadata,
64796484
}));
6480-
let original = serde_json::to_value(&item).unwrap();
6485+
assert_eq!(item.id().unwrap().as_str(), "fc_original");
64816486
let request = direct_request(control.clone(), "turn");
64826487
let edit = async {
64836488
let req = control.next().await.unwrap();
@@ -6488,13 +6493,20 @@ pub(crate) mod tests {
64886493
};
64896494
let (result, ()) = tokio::join!(request.before(item, None), edit);
64906495
if accepted {
6491-
let mut actual = serde_json::to_value(result.unwrap().0).unwrap();
6492-
assert_eq!(
6493-
serde_json::from_str::<Value>(actual["arguments"].as_str().unwrap()).unwrap(),
6494-
value
6495-
);
6496-
actual["arguments"] = original["arguments"].clone();
6497-
assert_eq!(actual, original);
6496+
let (item, changed, _) = result.unwrap();
6497+
assert!(changed);
6498+
assert!(item.id().is_none());
6499+
assert!(matches!(&item, ResponseItem::FunctionCall { call_id, .. }
6500+
if call_id == "original"));
6501+
let mut expected = json!({
6502+
"type": "function_call", "call_id": "original",
6503+
"name": "spawn_agent", "namespace": "collaboration",
6504+
"arguments": value.to_string(),
6505+
});
6506+
if !metadata.is_null() {
6507+
expected["encrypted_function_args"] = metadata;
6508+
}
6509+
assert_eq!(serde_json::to_value(item).unwrap(), expected);
64986510
} else {
64996511
assert!(
65006512
matches!(result, Err(ModelResponseError::InvalidRequest { message })

‎packages/codex/rust/ops.rs‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,7 @@ params!(TurnParams {
422422
#[serde(default)] mode: Mode,
423423
expected_turn_id: Option<String>,
424424
expected_previous_turn_id: Option<String>,
425+
#[serde(default)] exact_steering: bool,
425426
});
426427

427428
fn rejection(reason: NotSubmittedReason) -> Value {
@@ -561,6 +562,10 @@ async fn start_with_params<'a>(
561562
&mut Option<tokio::sync::MutexGuard<'a, ()>>,
562563
)>,
563564
) -> Result<Value> {
565+
ensure!(
566+
!p.exact_steering || (managed.is_some() && matches!(p.mode, Mode::Steer)),
567+
"exact steering requires a managed steer"
568+
);
564569
let input = if p.input.is_array() {
565570
TurnInput::UserInput {
566571
content: decode::<Vec<UserInput>>(p.input)?,
@@ -659,6 +664,7 @@ async fn start_with_params<'a>(
659664
let expected = p
660665
.expected_turn_id
661666
.context("steer requires expected_turn_id")?;
667+
let mut input_id = None;
662668
if managed.is_some()
663669
&& let TurnInput::UserInput { client_id, .. } = &mut request.input
664670
{
@@ -667,12 +673,19 @@ async fn start_with_params<'a>(
667673
.get::<crate::python_tools::PythonTools>()
668674
.context("Managed steering lost its tool owner")?
669675
.admission_ready(&expected)?;
670-
*client_id = Some(receipt.steering_client_id());
676+
let id = receipt.root_steering_client_id(p.exact_steering)?;
677+
if p.exact_steering {
678+
input_id = Some(id.clone());
679+
}
680+
*client_id = Some(id);
671681
}
672682
Ok(match thread.steer_turn(request, expected).await? {
673-
SteerSubmission::Steered { turn_id } => {
674-
json!({"status": "steered", "turn_id": turn_id})
675-
}
683+
SteerSubmission::Steered { turn_id } => match input_id {
684+
Some(input_id) => {
685+
json!({"status": "steered", "turn_id": turn_id, "input_id": input_id})
686+
}
687+
None => json!({"status": "steered", "turn_id": turn_id}),
688+
},
676689
SteerSubmission::NotSubmitted { reason } => rejection(reason),
677690
})
678691
}

‎packages/codex/rust/python_tools.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,7 @@ pub(crate) struct TurnAdmission {
216216
// Private namespace for SDK steers of this exact runtime, including same-ID recovery.
217217
steering_namespace: String,
218218
steering_tail: AtomicBool,
219+
exact_steering: OnceLock<bool>,
219220
provider_error: Mutex<Option<(Value, Value)>>,
220221
policy_failure: OnceLock<&'static str>,
221222
}
@@ -304,6 +305,18 @@ impl TurnAdmission {
304305
)
305306
}
306307

308+
pub(crate) fn root_steering_client_id(&self, exact: bool) -> Result<String> {
309+
ensure!(
310+
*self.exact_steering.get_or_init(|| exact) == exact,
311+
"Cannot mix exact and continuation steering on a native turn"
312+
);
313+
Ok(if exact {
314+
format!("agents-exact-steer:{}", codex_protocol::ThreadId::new())
315+
} else {
316+
self.steering_client_id()
317+
})
318+
}
319+
307320
fn record_steering(&self, client_id: Option<&str>) {
308321
if client_id.is_some_and(|id| id.starts_with(&self.steering_namespace)) {
309322
self.steering_tail.store(true, Ordering::Release);
@@ -2190,6 +2203,7 @@ impl TurnLifecycleContributor for PythonContributor {
21902203
codex_protocol::ThreadId::new()
21912204
),
21922205
steering_tail: AtomicBool::new(false),
2206+
exact_steering: OnceLock::new(),
21932207
provider_error: Mutex::default(),
21942208
policy_failure: OnceLock::new(),
21952209
});

‎src/agents/__init__.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,10 @@
5656
ReplaceArguments as ReplaceArguments,
5757
StopTurn as StopTurn,
5858
)
59+
from .rate_limits import (
60+
ApiRateLimits as ApiRateLimits,
61+
ApiRateLimitWindow as ApiRateLimitWindow,
62+
)
5963
from .schema import (
6064
NativeOutput as NativeOutput,
6165
OutputSpec as OutputSpec,
@@ -95,7 +99,12 @@
9599
ToolRef as ToolRef,
96100
WebSearch as WebSearch,
97101
)
98-
from .turn import PendingRequests as PendingRequests, Turn as Turn
102+
from .turn import (
103+
PendingRequests as PendingRequests,
104+
SteeringReceipt as SteeringReceipt,
105+
Turn as Turn,
106+
TurnNeedsAttention as TurnNeedsAttention,
107+
)
99108

100109
_MODULES = {
101110
"Allow": ".native_interception",
@@ -133,6 +142,8 @@
133142
"OutputValidationError": ".schema",
134143
"CheckpointStatus": ".state",
135144
"LocalStorage": ".state",
145+
"ApiRateLimits": ".rate_limits",
146+
"ApiRateLimitWindow": ".rate_limits",
136147
"InputImage": ".input",
137148
"InputText": ".input",
138149
"UserMessage": ".input",
@@ -166,6 +177,8 @@
166177
"SessionBusyError": ".session",
167178
"Turn": ".turn",
168179
"PendingRequests": ".turn",
180+
"SteeringReceipt": ".turn",
181+
"TurnNeedsAttention": ".turn",
169182
"Subagent": ".subagents",
170183
}
171184

‎src/agents/observation.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,12 +13,14 @@
1313
from dataclasses import dataclass, field, fields, replace
1414
from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypeVar, cast
1515

16+
from .rate_limits import ApiRateLimits
1617
from .usage import TokenUsageInfo
1718

1819
if TYPE_CHECKING:
1920
from .turn import Turn
2021

2122
__all__ = [
23+
"ApiRateLimitsUpdated",
2224
"ContextCompacted",
2325
"EventGap",
2426
"EventStream",
@@ -86,6 +88,14 @@ class UsageUpdated(_SubmissionEvent):
8688
type: Literal["usage_updated"] = field(default="usage_updated", init=False)
8789

8890

91+
@dataclass(frozen=True, slots=True, kw_only=True)
92+
class ApiRateLimitsUpdated(_SubmissionEvent):
93+
"""Rate-limit fields reported on a provider response for this turn."""
94+
95+
snapshot: ApiRateLimits
96+
type: Literal["api_rate_limits_updated"] = field(default="api_rate_limits_updated", init=False)
97+
98+
8999
@dataclass(frozen=True, slots=True, kw_only=True)
90100
class ContextCompacted(_SubmissionEvent):
91101
"""Core completed this native context-compaction item."""
@@ -164,6 +174,7 @@ class EventGap(_Event):
164174
| SubagentLifecycle
165175
| OutputTextDelta
166176
| UsageUpdated
177+
| ApiRateLimitsUpdated
167178
| ContextCompacted
168179
| ToolStarted
169180
| ToolFinished

‎src/agents/rate_limits.py‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
"""Rate limits reported by the API-key model provider."""
2+
3+
from __future__ import annotations
4+
5+
import re
6+
from dataclasses import dataclass
7+
from datetime import datetime, timezone
8+
from decimal import Context, Decimal, localcontext
9+
from email.utils import parsedate_to_datetime
10+
from typing import Any
11+
12+
__all__ = ["ApiRateLimitWindow", "ApiRateLimits"]
13+
14+
15+
@dataclass(frozen=True, slots=True, kw_only=True)
16+
class ApiRateLimitWindow:
17+
"""Reported counts and time until reset, measured at the provider response.
18+
19+
A missing field means the response did not supply a usable value.
20+
"""
21+
22+
limit: int | None = None
23+
remaining: int | None = None
24+
reset_after_seconds: Decimal | None = None
25+
26+
27+
@dataclass(frozen=True, slots=True, kw_only=True)
28+
class ApiRateLimits:
29+
"""Rate limits from the latest observed response with usable rate-limit data.
30+
31+
Retry advice can be a delay in seconds or an absolute UTC time. The
32+
provider reports it separately from the remaining request and token counts.
33+
"""
34+
35+
requests: ApiRateLimitWindow | None = None
36+
tokens: ApiRateLimitWindow | None = None
37+
project_tokens: ApiRateLimitWindow | None = None
38+
retry_after_seconds: int | None = None
39+
retry_at: datetime | None = None
40+
41+
42+
_DURATION_UNIT_SECONDS = {
43+
"h": Decimal(3600),
44+
"m": Decimal(60),
45+
"s": Decimal(1),
46+
"ms": Decimal("0.001"),
47+
"us": Decimal("0.000001"),
48+
"ns": Decimal("0.000000001"),
49+
}
50+
_DURATION_COMPONENT = re.compile(r"([0-9]+(?:\.[0-9]+)?)(ms|us|ns|h|m|s)")
51+
_DURATION = re.compile(r"(?:[0-9]+(?:\.[0-9]+)?(?:ms|us|ns|h|m|s))+")
52+
53+
54+
def _duration_seconds(value: Any) -> Decimal | None:
55+
if not isinstance(value, str) or len(value) > 64 or not _DURATION.fullmatch(value):
56+
return None
57+
# The input length bounds the needed precision. Keep the application's
58+
# decimal context from rounding a provider value.
59+
with localcontext(Context(prec=80)):
60+
return sum(
61+
(
62+
Decimal(amount) * _DURATION_UNIT_SECONDS[unit]
63+
for amount, unit in _DURATION_COMPONENT.findall(value)
64+
),
65+
start=Decimal(0),
66+
)
67+
68+
69+
def _count(value: Any) -> int | None:
70+
if type(value) is int and value >= 0:
71+
return value
72+
return None
73+
74+
75+
def _window(value: Any) -> ApiRateLimitWindow | None:
76+
if not isinstance(value, dict):
77+
return None
78+
limit = _count(value.get("limit"))
79+
remaining = _count(value.get("remaining"))
80+
reset = _duration_seconds(value.get("reset_after"))
81+
if limit is None and remaining is None and reset is None:
82+
return None
83+
return ApiRateLimitWindow(limit=limit, remaining=remaining, reset_after_seconds=reset)
84+
85+
86+
def _retry_after(value: Any) -> tuple[int | None, datetime | None]:
87+
if not isinstance(value, str) or len(value) > 64:
88+
return None, None
89+
if value.isascii() and value.isdecimal():
90+
return int(value), None
91+
try:
92+
at = parsedate_to_datetime(value)
93+
if at.tzinfo is None:
94+
at = at.replace(tzinfo=timezone.utc)
95+
return None, at.astimezone(timezone.utc)
96+
except (ValueError, OverflowError):
97+
return None, None
98+
99+
100+
def _api_rate_limits(value: Any) -> ApiRateLimits | None:
101+
"""Project the native wire snapshot without copying unrecognized fields."""
102+
if not isinstance(value, dict):
103+
return None
104+
requests = _window(value.get("requests"))
105+
tokens = _window(value.get("tokens"))
106+
project_tokens = _window(value.get("project_tokens"))
107+
delay, retry_at = _retry_after(value.get("retry_after"))
108+
if all(field is None for field in (requests, tokens, project_tokens, delay, retry_at)):
109+
return None
110+
return ApiRateLimits(
111+
requests=requests,
112+
tokens=tokens,
113+
project_tokens=project_tokens,
114+
retry_after_seconds=delay,
115+
retry_at=retry_at,
116+
)

0 commit comments

Comments
 (0)