mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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:
committed by
GitHub
Unverified
parent
7ae99576a6
commit
606d85055f
@@ -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"
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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)),
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
|
||||
Reference in New Issue
Block a user