Add notify to code-mode (#14842)

Allows model to send an out-of-band notification.

The notification is injected as another tool call output for the same
call_id.
This commit is contained in:
pakrym-oai
2026-03-18 09:37:13 -07:00
committed by GitHub
Unverified
parent 7ae99576a6
commit 606d85055f
27 changed files with 323 additions and 77 deletions
@@ -30,6 +30,7 @@ Object.defineProperty(globalThis, '__codexContentItems', {
defineGlobal('exit', __codexRuntime.exit);
defineGlobal('image', __codexRuntime.image);
defineGlobal('load', __codexRuntime.load);
defineGlobal('notify', __codexRuntime.notify);
defineGlobal('store', __codexRuntime.store);
defineGlobal('text', __codexRuntime.text);
defineGlobal('tools', __codexRuntime.tools);
@@ -14,5 +14,6 @@
- `image(imageUrl: string)`: Appends an image item and returns it. `image_url` can be an HTTPS URL or a base64-encoded `data:` URL.
- `store(key: string, value: any)`: stores a serializable value under a string key for later `exec` calls in the same session.
- `load(key: string)`: returns the stored value for a string key, or `undefined` if it is missing.
- `notify(value: string | number | boolean | undefined | null)`: immediately injects an extra `custom_tool_call_output` for the current `exec` call. Values are stringified like `text(...)`.
- `ALL_TOOLS`: metadata for the enabled nested tools as `{ name, description }` entries.
- `yield_control()`: yields the accumulated output to the model immediately while the script keeps running.
@@ -43,6 +43,7 @@ impl CodeModeExecuteHandler {
&self,
session: std::sync::Arc<Session>,
turn: std::sync::Arc<TurnContext>,
call_id: String,
code: String,
) -> Result<FunctionToolOutput, FunctionCallError> {
let args = parse_freeform_args(&code)?;
@@ -62,6 +63,7 @@ impl CodeModeExecuteHandler {
let message = HostToNodeMessage::Start {
request_id: request_id.clone(),
cell_id: cell_id.clone(),
tool_call_id: call_id,
default_yield_time_ms: super::DEFAULT_EXEC_YIELD_TIME_MS,
enabled_tools,
stored_values,
@@ -198,6 +200,7 @@ impl ToolHandler for CodeModeExecuteHandler {
let ToolInvocation {
session,
turn,
call_id,
tool_name,
payload,
..
@@ -205,7 +208,7 @@ impl ToolHandler for CodeModeExecuteHandler {
match payload {
ToolPayload::Custom { input } if tool_name == PUBLIC_TOOL_NAME => {
self.execute(session, turn, input).await
self.execute(session, turn, call_id, input).await
}
_ => Err(FunctionCallError::RespondToModel(format!(
"{PUBLIC_TOOL_NAME} expects raw JavaScript source text"
+3
View File
@@ -110,6 +110,9 @@ async fn handle_node_message(
) -> Result<CodeModeSessionProgress, String> {
match message {
protocol::NodeToHostMessage::ToolCall { .. } => Err(protocol::unexpected_tool_call_error()),
protocol::NodeToHostMessage::Notify { .. } => Err(format!(
"unexpected {PUBLIC_TOOL_NAME} notify message in response path"
)),
protocol::NodeToHostMessage::Yielded { content_items, .. } => {
let mut delta_items = output_content_items_from_json_values(content_items)?;
delta_items = truncate_code_mode_result(delta_items, poll_max_output_tokens.flatten());
+9 -8
View File
@@ -13,7 +13,6 @@ use tracing::warn;
use super::CODE_MODE_RUNNER_SOURCE;
use super::PUBLIC_TOOL_NAME;
use super::protocol::CodeModeToolCall;
use super::protocol::HostToNodeMessage;
use super::protocol::NodeToHostMessage;
use super::protocol::message_request_id;
@@ -23,7 +22,7 @@ pub(super) struct CodeModeProcess {
pub(super) stdin: Arc<Mutex<tokio::process::ChildStdin>>,
pub(super) stdout_task: JoinHandle<()>,
pub(super) response_waiters: Arc<Mutex<HashMap<String, oneshot::Sender<NodeToHostMessage>>>>,
pub(super) tool_call_rx: Arc<Mutex<mpsc::UnboundedReceiver<CodeModeToolCall>>>,
pub(super) message_rx: Arc<Mutex<mpsc::UnboundedReceiver<NodeToHostMessage>>>,
}
impl CodeModeProcess {
@@ -92,7 +91,7 @@ pub(super) async fn spawn_code_mode_process(
String,
oneshot::Sender<NodeToHostMessage>,
>::new()));
let (tool_call_tx, tool_call_rx) = mpsc::unbounded_channel();
let (message_tx, message_rx) = mpsc::unbounded_channel();
tokio::spawn(async move {
let mut reader = BufReader::new(stderr);
@@ -135,12 +134,14 @@ pub(super) async fn spawn_code_mode_process(
}
};
match message {
NodeToHostMessage::ToolCall { tool_call } => {
let _ = tool_call_tx.send(tool_call);
message @ (NodeToHostMessage::ToolCall { .. }
| NodeToHostMessage::Notify { .. }) => {
let _ = message_tx.send(message);
}
message => {
let request_id = message_request_id(&message).to_string();
if let Some(waiter) = response_waiters.lock().await.remove(&request_id) {
if let Some(request_id) = message_request_id(&message)
&& let Some(waiter) = response_waiters.lock().await.remove(request_id)
{
let _ = waiter.send(message);
}
}
@@ -155,7 +156,7 @@ pub(super) async fn spawn_code_mode_process(
stdin,
stdout_task,
response_waiters,
tool_call_rx: Arc::new(Mutex::new(tool_call_rx)),
message_rx: Arc::new(Mutex::new(message_rx)),
})
}
+52 -4
View File
@@ -36,12 +36,20 @@ pub(super) struct CodeModeToolCall {
pub(super) input: Option<JsonValue>,
}
#[derive(Clone, Debug, Deserialize)]
pub(super) struct CodeModeNotify {
pub(super) cell_id: String,
pub(super) call_id: String,
pub(super) text: String,
}
#[derive(Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub(super) enum HostToNodeMessage {
Start {
request_id: String,
cell_id: String,
tool_call_id: String,
default_yield_time_ms: u64,
enabled_tools: Vec<EnabledTool>,
stored_values: HashMap<String, JsonValue>,
@@ -65,7 +73,7 @@ pub(super) enum HostToNodeMessage {
},
}
#[derive(Deserialize)]
#[derive(Debug, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub(super) enum NodeToHostMessage {
ToolCall {
@@ -80,6 +88,10 @@ pub(super) enum NodeToHostMessage {
request_id: String,
content_items: Vec<JsonValue>,
},
Notify {
#[serde(flatten)]
notify: CodeModeNotify,
},
Result {
request_id: String,
content_items: Vec<JsonValue>,
@@ -105,15 +117,51 @@ pub(super) fn build_source(
.replace("__CODE_MODE_USER_CODE_PLACEHOLDER__", user_code))
}
pub(super) fn message_request_id(message: &NodeToHostMessage) -> &str {
pub(super) fn message_request_id(message: &NodeToHostMessage) -> Option<&str> {
match message {
NodeToHostMessage::ToolCall { tool_call } => &tool_call.request_id,
NodeToHostMessage::ToolCall { .. } => None,
NodeToHostMessage::Yielded { request_id, .. }
| NodeToHostMessage::Terminated { request_id, .. }
| NodeToHostMessage::Result { request_id, .. } => request_id,
| NodeToHostMessage::Result { request_id, .. } => Some(request_id),
NodeToHostMessage::Notify { .. } => None,
}
}
pub(super) fn unexpected_tool_call_error() -> String {
format!("{PUBLIC_TOOL_NAME} received an unexpected tool call response")
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use super::CodeModeNotify;
use super::NodeToHostMessage;
use super::message_request_id;
#[test]
fn message_request_id_absent_for_notify() {
let message = NodeToHostMessage::Notify {
notify: CodeModeNotify {
cell_id: "1".to_string(),
call_id: "call-1".to_string(),
text: "hello".to_string(),
},
};
assert_eq!(None, message_request_id(&message));
}
#[test]
fn message_request_id_present_for_result() {
let message = NodeToHostMessage::Result {
request_id: "req-1".to_string(),
content_items: Vec::new(),
stored_values: HashMap::new(),
error_text: None,
max_output_tokens_per_exec_call: None,
};
assert_eq!(Some("req-1"), message_request_id(&message));
}
}
+41 -2
View File
@@ -233,7 +233,7 @@ function codeModeWorkerMain() {
throw new TypeError('image expects an http(s) or data URL');
}
function createCodeModeHelpers(context, state) {
function createCodeModeHelpers(context, state, toolCallId) {
const load = (key) => {
if (typeof key !== 'string') {
throw new TypeError('load key must be a string');
@@ -268,6 +268,21 @@ function codeModeWorkerMain() {
const yieldControl = () => {
parentPort.postMessage({ type: 'yield' });
};
const notify = (value) => {
const text = serializeOutputText(value);
if (text.trim().length === 0) {
throw new TypeError('notify expects non-empty text');
}
if (typeof toolCallId !== 'string' || toolCallId.length === 0) {
throw new TypeError('notify requires a valid tool call id');
}
parentPort.postMessage({
type: 'notify',
call_id: toolCallId,
text,
});
return text;
};
const exit = () => {
throw new CodeModeExitSignal();
};
@@ -276,6 +291,7 @@ function codeModeWorkerMain() {
exit,
image,
load,
notify,
output_image: image,
output_text: text,
store,
@@ -290,6 +306,7 @@ function codeModeWorkerMain() {
'exit',
'image',
'load',
'notify',
'output_text',
'output_image',
'store',
@@ -300,6 +317,7 @@ function codeModeWorkerMain() {
this.setExport('exit', helpers.exit);
this.setExport('image', helpers.image);
this.setExport('load', helpers.load);
this.setExport('notify', helpers.notify);
this.setExport('output_text', helpers.output_text);
this.setExport('output_image', helpers.output_image);
this.setExport('store', helpers.store);
@@ -316,6 +334,7 @@ function codeModeWorkerMain() {
exit: helpers.exit,
image: helpers.image,
load: helpers.load,
notify: helpers.notify,
store: helpers.store,
text: helpers.text,
tools: createGlobalToolsNamespace(callTool, enabledTools),
@@ -448,6 +467,7 @@ function codeModeWorkerMain() {
async function main() {
const start = workerData ?? {};
const toolCallId = start.tool_call_id;
const state = {
storedValues: cloneJsonValue(start.stored_values ?? {}),
};
@@ -457,7 +477,7 @@ function codeModeWorkerMain() {
const context = vm.createContext({
__codexContentItems: contentItems,
});
const helpers = createCodeModeHelpers(context, state);
const helpers = createCodeModeHelpers(context, state, toolCallId);
Object.defineProperty(context, '__codexRuntime', {
value: createBridgeRuntime(callTool, enabledTools, helpers),
configurable: true,
@@ -631,6 +651,9 @@ function sessionWorkerSource() {
}
function startSession(protocol, sessions, start) {
if (typeof start.tool_call_id !== 'string' || start.tool_call_id.length === 0) {
throw new TypeError('start requires a valid tool_call_id');
}
const maxOutputTokensPerExecCall =
start.max_output_tokens == null
? DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL
@@ -704,6 +727,22 @@ async function handleWorkerMessage(protocol, sessions, session, message) {
return;
}
if (message.type === 'notify') {
if (typeof message.text !== 'string' || message.text.trim().length === 0) {
throw new TypeError('notify requires non-empty text');
}
if (typeof message.call_id !== 'string' || message.call_id.length === 0) {
throw new TypeError('notify requires a valid call id');
}
await protocol.send({
type: 'notify',
cell_id: session.id,
call_id: message.call_id,
text: message.text,
});
return;
}
if (message.type === 'tool_call') {
void forwardToolCall(protocol, session, message);
return;
+63 -26
View File
@@ -1,13 +1,18 @@
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use tracing::error;
use tracing::warn;
use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ResponseInputItem;
use super::ExecContext;
use super::PUBLIC_TOOL_NAME;
use super::call_nested_tool;
use super::process::CodeModeProcess;
use super::process::write_message;
use super::protocol::HostToNodeMessage;
use super::protocol::NodeToHostMessage;
use crate::tools::parallel::ToolCallRuntime;
pub(crate) struct CodeModeWorker {
shutdown_tx: Option<oneshot::Sender<()>>,
@@ -29,39 +34,71 @@ impl CodeModeProcess {
) -> CodeModeWorker {
let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
let stdin = self.stdin.clone();
let tool_call_rx = self.tool_call_rx.clone();
let message_rx = self.message_rx.clone();
tokio::spawn(async move {
loop {
let tool_call = tokio::select! {
let next_message = tokio::select! {
_ = &mut shutdown_rx => break,
tool_call = async {
let mut tool_call_rx = tool_call_rx.lock().await;
tool_call_rx.recv().await
} => tool_call,
message = async {
let mut message_rx = message_rx.lock().await;
message_rx.recv().await
} => message,
};
let Some(tool_call) = tool_call else {
let Some(next_message) = next_message else {
break;
};
let exec = exec.clone();
let tool_runtime = tool_runtime.clone();
let stdin = stdin.clone();
tokio::spawn(async move {
let response = HostToNodeMessage::Response {
request_id: tool_call.request_id,
id: tool_call.id,
code_mode_result: call_nested_tool(
exec,
tool_runtime,
tool_call.name,
tool_call.input,
CancellationToken::new(),
)
.await,
};
if let Err(err) = write_message(&stdin, &response).await {
warn!("failed to write {PUBLIC_TOOL_NAME} tool response: {err}");
match next_message {
NodeToHostMessage::ToolCall { tool_call } => {
let exec = exec.clone();
let tool_runtime = tool_runtime.clone();
let stdin = stdin.clone();
tokio::spawn(async move {
let response = HostToNodeMessage::Response {
request_id: tool_call.request_id,
id: tool_call.id,
code_mode_result: call_nested_tool(
exec,
tool_runtime,
tool_call.name,
tool_call.input,
CancellationToken::new(),
)
.await,
};
if let Err(err) = write_message(&stdin, &response).await {
warn!("failed to write {PUBLIC_TOOL_NAME} tool response: {err}");
}
});
}
});
NodeToHostMessage::Notify { notify } => {
if notify.text.trim().is_empty() {
continue;
}
if exec
.session
.inject_response_items(vec![ResponseInputItem::CustomToolCallOutput {
call_id: notify.call_id.clone(),
name: Some(PUBLIC_TOOL_NAME.to_string()),
output: FunctionCallOutputPayload::from_text(notify.text),
}])
.await
.is_err()
{
warn!(
"failed to inject {PUBLIC_TOOL_NAME} notify message for cell {}: no active turn",
notify.cell_id
);
}
}
unexpected_message @ (NodeToHostMessage::Yielded { .. }
| NodeToHostMessage::Terminated { .. }
| NodeToHostMessage::Result { .. }) => {
error!(
"received unexpected {PUBLIC_TOOL_NAME} message in worker loop: {unexpected_message:?}"
);
break;
}
}
}
});