Preserve user query attribution in shared-session transport

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<AgentAttachment>,
+
+    /// 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<String>,
 }
+
+#[cfg(test)]
+#[path = "agent_prompt_test.rs"]
+mod attribution_tests;

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,
+    );
+}

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),