From e2b1d4deb95482ad37d746643814b21f642bee5f Mon Sep 17 00:00:00 2001 From: Ben Stobaugh Date: Fri, 11 Sep 2026 12:49:57 -0700 Subject: [PATCH] Preserve user query attribution in shared-session protocol Carry optional opaque attribution envelopes and account for their wire size. Co-Authored-By: Warp Agent --- src/common/agent_prompt.rs | 13 +++++- src/common/agent_prompt_test.rs | 81 +++++++++++++++++++++++++++++++++ src/viewer.rs | 9 +++- 3 files changed, 101 insertions(+), 2 deletions(-) create mode 100644 src/common/agent_prompt_test.rs diff --git a/src/common/agent_prompt.rs b/src/common/agent_prompt.rs index 40465f3..3aad90a 100644 --- a/src/common/agent_prompt.rs +++ b/src/common/agent_prompt.rs @@ -98,7 +98,7 @@ impl std::str::FromStr for ServerConversationToken { } /// The data for an agent prompt request from a viewer. -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, Default, Serialize, Deserialize)] pub struct AgentPromptRequest { /// Unique identifier for this request. pub id: AgentPromptRequestId, @@ -113,4 +113,15 @@ pub struct AgentPromptRequest { /// Optional attachments (blocks, files, etc.) referenced in the prompt. #[serde(default)] pub attachments: Vec, + + /// Standard Base64 of a `warp.multi_agent.v1.UserQueryAttribution` protobuf. + /// The relay preserves these unverified query metadata bytes opaquely; they + /// must not affect authentication, authorization, or the requester identity. + /// Defaulted so requests already persisted in Redis remain readable. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub user_query_attribution_b64: Option, } + +#[cfg(test)] +#[path = "agent_prompt_test.rs"] +mod attribution_tests; diff --git a/src/common/agent_prompt_test.rs b/src/common/agent_prompt_test.rs new file mode 100644 index 0000000..3f13947 --- /dev/null +++ b/src/common/agent_prompt_test.rs @@ -0,0 +1,81 @@ +use super::*; +use crate::{sharer, viewer}; + +fn legacy_request() -> AgentPromptRequest { + serde_json::from_str(r#"{"id":"req-1","server_conversation_token":null,"prompt":"hello"}"#) + .unwrap() +} + +#[test] +fn legacy_request_has_no_attribution_and_omits_the_field_when_serialized() { + let request = legacy_request(); + assert!(request.user_query_attribution_b64.is_none()); + let json = serde_json::to_value(request).unwrap(); + assert!(json.get("user_query_attribution_b64").is_none()); +} + +#[test] +fn envelope_survives_storage_and_both_websocket_directions_opaquely() { + // Unknown and malformed envelopes belong to the decoder at the destination; + // the relay must preserve them, including empty values, without dropping input. + for envelope in ["CgIKABIKCgYKBHVzZXI=", "not-base64!", ""] { + let mut request = legacy_request(); + request.user_query_attribution_b64 = Some(envelope.to_owned()); + let stored = serde_json::to_string(&request).unwrap(); + let request: AgentPromptRequest = serde_json::from_str(&stored).unwrap(); + + let upstream = viewer::UpstreamMessage::SendAgentPrompt(request.clone()); + let upstream = viewer::UpstreamMessage::from_json(&upstream.to_json().unwrap()).unwrap(); + let viewer::UpstreamMessage::SendAgentPrompt(request) = upstream else { + panic!("expected prompt request"); + }; + assert_eq!( + request.user_query_attribution_b64.as_deref(), + Some(envelope) + ); + + let downstream = sharer::DownstreamMessage::AgentPromptRequested { + id: request.id.clone(), + participant_id: super::super::ParticipantId::new(), + request, + }; + let downstream = + sharer::DownstreamMessage::from_json(&downstream.to_json().unwrap()).unwrap(); + let sharer::DownstreamMessage::AgentPromptRequested { request, .. } = downstream else { + panic!("expected prompt delivery"); + }; + assert_eq!( + request.user_query_attribution_b64.as_deref(), + Some(envelope) + ); + } +} + +#[test] +fn session_byte_accounting_includes_encoded_envelope_and_inline_content() { + let mut request = legacy_request(); + request.prompt = "hello 👋".to_owned(); + request.attachments = vec![ + AgentAttachment::PlainText { + content: "résumé".to_owned(), + }, + AgentAttachment::FileReference { + attachment_id: "already-uploaded-file".to_owned(), + file_name: "file.txt".to_owned(), + }, + ]; + let baseline_bytes = "hello 👋".len() + "résumé".len(); + assert_eq!( + viewer::UpstreamMessage::SendAgentPrompt(request.clone()) + .num_bytes() + .as_u64(), + baseline_bytes as u64, + ); + request.user_query_attribution_b64 = Some("CgIKAA==".to_owned()); + assert_eq!( + viewer::UpstreamMessage::SendAgentPrompt(request) + .num_bytes() + .as_u64(), + (baseline_bytes + "CgIKAA==".len()) as u64, + ); +} diff --git a/src/viewer.rs b/src/viewer.rs index 5295378..c11ab41 100644 --- a/src/viewer.rs +++ b/src/viewer.rs @@ -413,8 +413,14 @@ impl UpstreamMessage { UpstreamMessage::ExecuteCommand { command, .. } => command.len().into(), UpstreamMessage::WriteToPty { bytes, .. } => bytes.len().into(), UpstreamMessage::SendAgentPrompt(request) => { - // Count prompt length + attachments + // Count prompt, inline attachments, and the encoded attribution + // exactly as transported, without interpreting the envelope. let prompt_bytes: Byte = request.prompt.len().into(); + let attribution_bytes: Byte = request + .user_query_attribution_b64 + .as_ref() + .map_or(0, String::len) + .into(); let attachments_bytes: Byte = request .attachments .iter() @@ -430,6 +436,7 @@ impl UpstreamMessage { .into(); prompt_bytes .add(attachments_bytes) + .and_then(|bytes| bytes.add(attribution_bytes)) .unwrap_or(u64::MAX.into()) } _ => Byte::from_u64(0),