clone/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-04-29 00:55:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add memory citation to agent message (#14821)

Browse Source 
Client side to come
This commit is contained in:
jif-oai
2026-03-18 10:03:38 +00:00
committed by GitHub
parent 0f9484dc8a
commit a265d6043e
43 changed files with 1420 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
codex-rs/app-server-protocol/src/protocol/thread_history.rs
View File

@@ -118,9 +118,11 @@ impl ThreadHistoryBuilder {
pub fn handle_event(&mut self, event: &EventMsg) {
match event {
EventMsg::UserMessage(payload) => self.handle_user_message(payload),
EventMsg::AgentMessage(payload) => {
self.handle_agent_message(payload.message.clone(), payload.phase.clone())
}
EventMsg::AgentMessage(payload) => self.handle_agent_message(
payload.message.clone(),
payload.phase.clone(),
payload.memory_citation.clone().map(Into::into),
),
EventMsg::AgentReasoning(payload) => self.handle_agent_reasoning(payload),
EventMsg::AgentReasoningRawContent(payload) => {
self.handle_agent_reasoning_raw_content(payload)
@@ -208,15 +210,23 @@ impl ThreadHistoryBuilder {
self.current_turn = Some(turn);
}
fn handle_agent_message(&mut self, text: String, phase: Option<MessagePhase>) {
fn handle_agent_message(
&mut self,
text: String,
phase: Option<MessagePhase>,
memory_citation: Option<crate::protocol::v2::MemoryCitation>,
) {
if text.is_empty() {
return;
}
let id = self.next_item_id();
self.ensure_turn()
.items
.push(ThreadItem::AgentMessage { id, text, phase });
self.ensure_turn().items.push(ThreadItem::AgentMessage {
id,
text,
phase,
memory_citation,
});
}
fn handle_agent_reasoning(&mut self, payload: &AgentReasoningEvent) {
@@ -1178,6 +1188,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "Hi there".into(),
phase: None,
memory_citation: None,
}),
EventMsg::AgentReasoning(AgentReasoningEvent {
text: "thinking".into(),
@@ -1194,6 +1205,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "Reply two".into(),
phase: None,
memory_citation: None,
}),
];
@@ -1229,6 +1241,7 @@ mod tests {
id: "item-2".into(),
text: "Hi there".into(),
phase: None,
memory_citation: None,
}
);
assert_eq!(
@@ -1260,6 +1273,7 @@ mod tests {
id: "item-5".into(),
text: "Reply two".into(),
phase: None,
memory_citation: None,
}
);
}
@@ -1318,6 +1332,7 @@ mod tests {
let events = vec![EventMsg::AgentMessage(AgentMessageEvent {
message: "Final reply".into(),
phase: Some(CoreMessagePhase::FinalAnswer),
memory_citation: None,
})];
let items = events
@@ -1332,6 +1347,7 @@ mod tests {
id: "item-1".into(),
text: "Final reply".into(),
phase: Some(MessagePhase::FinalAnswer),
memory_citation: None,
}
);
}
@@ -1354,6 +1370,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "interlude".into(),
phase: None,
memory_citation: None,
}),
EventMsg::AgentReasoning(AgentReasoningEvent {
text: "second summary".into(),
@@ -1399,6 +1416,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "Working...".into(),
phase: None,
memory_citation: None,
}),
EventMsg::TurnAborted(TurnAbortedEvent {
turn_id: Some("turn-1".into()),
@@ -1413,6 +1431,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "Second attempt complete.".into(),
phase: None,
memory_citation: None,
}),
];
@@ -1442,6 +1461,7 @@ mod tests {
id: "item-2".into(),
text: "Working...".into(),
phase: None,
memory_citation: None,
}
);
@@ -1464,6 +1484,7 @@ mod tests {
id: "item-4".into(),
text: "Second attempt complete.".into(),
phase: None,
memory_citation: None,
}
);
}
@@ -1480,6 +1501,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "A1".into(),
phase: None,
memory_citation: None,
}),
EventMsg::UserMessage(UserMessageEvent {
message: "Second".into(),
@@ -1490,6 +1512,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "A2".into(),
phase: None,
memory_citation: None,
}),
EventMsg::ThreadRolledBack(ThreadRolledBackEvent { num_turns: 1 }),
EventMsg::UserMessage(UserMessageEvent {
@@ -1501,6 +1524,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "A3".into(),
phase: None,
memory_citation: None,
}),
];
@@ -1529,6 +1553,7 @@ mod tests {
id: "item-2".into(),
text: "A1".into(),
phase: None,
memory_citation: None,
},
]
);
@@ -1546,6 +1571,7 @@ mod tests {
id: "item-4".into(),
text: "A3".into(),
phase: None,
memory_citation: None,
},
]
);
@@ -1563,6 +1589,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "A1".into(),
phase: None,
memory_citation: None,
}),
EventMsg::UserMessage(UserMessageEvent {
message: "Two".into(),
@@ -1573,6 +1600,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "A2".into(),
phase: None,
memory_citation: None,
}),
EventMsg::ThreadRolledBack(ThreadRolledBackEvent { num_turns: 99 }),
];
@@ -2209,6 +2237,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "still in b".into(),
phase: None,
memory_citation: None,
}),
EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: "turn-b".into(),
@@ -2263,6 +2292,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "still in b".into(),
phase: None,
memory_citation: None,
}),
];
@@ -2497,6 +2527,7 @@ mod tests {
EventMsg::AgentMessage(AgentMessageEvent {
message: "done".into(),
phase: None,
memory_citation: None,
}),
EventMsg::Error(ErrorEvent {
message: "rollback failed".into(),

63
codex-rs/app-server-protocol/src/protocol/v2.rs
View File

@@ -30,6 +30,8 @@ use codex_protocol::items::TurnItem as CoreTurnItem;
use codex_protocol::mcp::Resource as McpResource;
use codex_protocol::mcp::ResourceTemplate as McpResourceTemplate;
use codex_protocol::mcp::Tool as McpTool;
use codex_protocol::memory_citation::MemoryCitation as CoreMemoryCitation;
use codex_protocol::memory_citation::MemoryCitationEntry as CoreMemoryCitationEntry;
use codex_protocol::models::FileSystemPermissions as CoreFileSystemPermissions;
use codex_protocol::models::MacOsAutomationPermission as CoreMacOsAutomationPermission;
use codex_protocol::models::MacOsContactsPermission as CoreMacOsContactsPermission;
@@ -3568,6 +3570,44 @@ pub struct Turn {
pub error: Option<TurnError>,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
pub struct MemoryCitation {
pub entries: Vec<MemoryCitationEntry>,
pub thread_ids: Vec<String>,
}
impl From<CoreMemoryCitation> for MemoryCitation {
fn from(value: CoreMemoryCitation) -> Self {
Self {
entries: value.entries.into_iter().map(Into::into).collect(),
thread_ids: value.rollout_ids,
}
}
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
pub struct MemoryCitationEntry {
pub path: String,
pub line_start: u32,
pub line_end: u32,
pub note: String,
}
impl From<CoreMemoryCitationEntry> for MemoryCitationEntry {
fn from(value: CoreMemoryCitationEntry) -> Self {
Self {
path: value.path,
line_start: value.line_start,
line_end: value.line_end,
note: value.note,
}
}
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS, Error)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
@@ -4068,6 +4108,8 @@ pub enum ThreadItem {
text: String,
#[serde(default)]
phase: Option<MessagePhase>,
#[serde(default)]
memory_citation: Option<MemoryCitation>,
},
#[serde(rename_all = "camelCase")]
#[ts(rename_all = "camelCase")]
@@ -4317,6 +4359,7 @@ impl From<CoreTurnItem> for ThreadItem {
id: agent.id,
text,
phase: agent.phase,
memory_citation: agent.memory_citation.map(Into::into),
}
}
CoreTurnItem::Plan(plan) => ThreadItem::Plan {
@@ -7393,6 +7436,7 @@ mod tests {
},
],
phase: None,
memory_citation: None,
});
assert_eq!(
@@ -7401,6 +7445,7 @@ mod tests {
id: "agent-1".to_string(),
text: "Hello world".to_string(),
phase: None,
memory_citation: None,
}
);
@@ -7410,6 +7455,15 @@ mod tests {
text: "final".to_string(),
}],
phase: Some(MessagePhase::FinalAnswer),
memory_citation: Some(CoreMemoryCitation {
entries: vec![CoreMemoryCitationEntry {
path: "MEMORY.md".to_string(),
line_start: 1,
line_end: 2,
note: "summary".to_string(),
}],
rollout_ids: vec!["rollout-1".to_string()],
}),
});
assert_eq!(
@@ -7418,6 +7472,15 @@ mod tests {
id: "agent-2".to_string(),
text: "final".to_string(),
phase: Some(MessagePhase::FinalAnswer),
memory_citation: Some(MemoryCitation {
entries: vec![MemoryCitationEntry {
path: "MEMORY.md".to_string(),
line_start: 1,
line_end: 2,
note: "summary".to_string(),
}],
thread_ids: vec!["rollout-1".to_string()],
}),
}
);