audio-knife: request deferred parameters

CONTEXT.md

@@ -62,9 +62,14 @@ These three Microsoft offerings are distinct and must not all be called "Azure".
   either contains complete service parameters or declares that they will follow
   in a deferred params message.
 
-- **Deferred Params Message** — the message immediately following an initial
-  start message that declared deferred parameters. It carries the complete
-  service parameters for the same conversation.
+- **Deferred Params Message** — the first text message following an initial start
+  message that declared deferred parameters. It carries the complete service
+  parameters for the same conversation. Binary, Ping, and Pong messages received
+  while waiting are ignored; Close ends startup with an error.
+
+- **Deferred Params Request** — AudioKnife's server message telling a client to
+  send the deferred parameters. AudioKnife sends it after accepting a deferred
+  initial start and before it begins receiving the Deferred Params Message.
 
 - **Logical Start** — the complete conversation start presented to ContextSwitch.
   It may originate from one initial start message or be assembled from an initial

Cargo.toml

@@ -22,7 +22,7 @@ members = [
 ]
 
 [workspace.package]
-version = "3.7.0"
+version = "3.7.1"
 edition = "2024"
 license = "MIT"
 repository = "https://github.com/pragmatrix/context-switch"

audio-knife/src/main.rs

@@ -19,11 +19,11 @@ use axum::routing::get;
 use axum::serve::ListenerExt;
 use base64::Engine as _;
 use base64::engine::general_purpose;
-use futures_util::stream::SplitSink;
+use futures_util::stream::{SplitSink, SplitStream};
 use futures_util::{SinkExt, StreamExt};
 use reqwest::StatusCode;
 use serde::Deserialize;
-use serde_json::Value;
+use serde_json::{Value, json};
 use server_event_router::ServerEventRouter;
 use tokio::net::TcpListener;
 use tokio::sync::mpsc::{Receiver, Sender, UnboundedReceiver, channel, unbounded_channel};
@@ -39,6 +39,8 @@ use context_switch::{
     ServerEvent, audio,
 };
 
+use crate::mod_audio_fork::AudioForkEvent;
+
 const DEFAULT_PORT: u16 = 8123;
 
 #[tokio::main]
@@ -212,10 +214,11 @@ async fn ws(state: State, mut websocket: WebSocket) -> Result<()> {
     match websocket.recv().await {
         Some(msg) => {
             let msg = msg?;
+            let (mut ws_sender, mut ws_receiver) = websocket.split();
             let (session_state, conversation_span, cs_receiver) =
-                SessionState::start_session(state, msg, &mut websocket).await?;
+                SessionState::start_session(state, msg, &mut ws_sender, &mut ws_receiver).await?;
 
-            ws_session(session_state, cs_receiver, websocket)
+            ws_session(session_state, cs_receiver, ws_sender, ws_receiver)
                 .instrument(conversation_span)
                 .await
         }
@@ -229,9 +232,9 @@ async fn ws(state: State, mut websocket: WebSocket) -> Result<()> {
 async fn ws_session(
     mut session_state: SessionState,
     cs_receiver: UnboundedReceiver<ServerEvent>,
-    websocket: WebSocket,
+    ws_sender: SplitSink<WebSocket, Message>,
+    mut ws_receiver: SplitStream<WebSocket>,
 ) -> Result<()> {
-    let (ws_sender, mut ws_receiver) = websocket.split();
     let billing_collector = session_state.state.billing_collector.clone();
 
     // Channel from event_scheduler to websocket dispatcher. Currently unbounded, because it's not
@@ -336,7 +339,8 @@ impl SessionState {
     async fn start_session(
         state: State,
         msg: Message,
-        websocket: &mut WebSocket,
+        websocket_sender: &mut SplitSink<WebSocket, Message>,
+        websocket_receiver: &mut SplitStream<WebSocket>,
     ) -> Result<(Self, Span, UnboundedReceiver<ServerEvent>)> {
         let Message::Text(msg) = msg else {
             // What about Ping?
@@ -354,7 +358,14 @@ impl SessionState {
         let conversation_span = info_span!("conversation", cid = %short_conversation_id);
 
         if start_aux.defer_params {
-            let params = Self::receive_deferred_params(&start_aux.id, websocket)
+            let params_request = AudioForkEvent::json(json!({
+                "type": "sendParams",
+                "id": &start_aux.id,
+            }))?;
+            mod_audio_fork::dispatch_event(websocket_sender, params_request)
+                .instrument(conversation_span.clone())
+                .await?;
+            let params = Self::receive_deferred_params(&start_aux.id, websocket_receiver)
                 .instrument(conversation_span.clone())
                 .await?;
             json_value
@@ -439,18 +450,21 @@ impl SessionState {
 
     async fn receive_deferred_params(
         start_id: &ConversationId,
-        websocket: &mut WebSocket,
+        websocket: &mut SplitStream<WebSocket>,
     ) -> Result<Value> {
-        let msg = websocket
-            .recv()
-            .await
-            .context("WebSocket closed before deferred params message was received")??;
-        let Message::Text(msg) = msg else {
-            bail!("Expecting deferred params WebSocket message to be text");
+        let text = loop {
+            let msg = websocket
+                .next()
+                .await
+                .context("WebSocket closed before deferred params message was received")??;
+
+            if let Some(text) = Self::expect_text_ignoring_other_messages(msg)? {
+                break text;
+            }
         };
 
         let deferred: DeferredParamsMessage =
-            serde_json::from_value(Self::decode_json_value(msg.as_str())?)?;
+            serde_json::from_value(Self::decode_json_value(text.as_str())?)?;
 
         if deferred.r#type != "params" {
             bail!("Expecting deferred params WebSocket message to have type `params`");
@@ -465,6 +479,14 @@ impl SessionState {
         Ok(deferred.params)
     }
 
+    fn expect_text_ignoring_other_messages(msg: Message) -> Result<Option<String>> {
+        match msg {
+            Message::Text(text) => Ok(Some(text.to_string())),
+            Message::Close(_) => bail!("WebSocket closed while waiting for a text message"),
+            Message::Binary(_) | Message::Ping(_) | Message::Pong(_) => Ok(None),
+        }
+    }
+
     fn process_request(&mut self, pong_sender: &Sender<Pong>, msg: Message) -> Result<()> {
         match msg {
             Message::Text(msg) => {
@@ -632,7 +654,9 @@ async fn dispatch_server_event(
         ServerEvent::Audio { samples, .. } => {
             mod_audio_fork::dispatch_audio(socket, samples.into()).await
         }
-        ServerEvent::ClearAudio { .. } => mod_audio_fork::dispatch_kill_audio(socket).await,
+        ServerEvent::ClearAudio { .. } => {
+            mod_audio_fork::dispatch_event(socket, AudioForkEvent::kill_audio()).await
+        }
         ServerEvent::BillingRecords {
             service,
             scope,
@@ -694,6 +718,8 @@ async fn take_billing_records(
 
 #[cfg(test)]
 mod tests {
+    use axum::body::Bytes;
+    use axum::extract::ws::Message;
     use serde_json::json;
 
     use super::{SessionState, StartEventAuxiliary};
@@ -746,4 +772,31 @@ mod tests {
                 .contains("Deferred start must not contain inline params")
         );
     }
+
+    #[test]
+    fn text_expectation_ignores_non_text_messages() {
+        for message in [
+            Message::Binary(Bytes::new()),
+            Message::Ping(Bytes::new()),
+            Message::Pong(Bytes::new()),
+        ] {
+            assert!(
+                SessionState::expect_text_ignoring_other_messages(message)
+                    .unwrap()
+                    .is_none()
+            );
+        }
+    }
+
+    #[test]
+    fn text_expectation_rejects_close() {
+        let error =
+            SessionState::expect_text_ignoring_other_messages(Message::Close(None)).unwrap_err();
+
+        assert!(
+            error
+                .to_string()
+                .contains("WebSocket closed while waiting for a text message")
+        );
+    }
 }

audio-knife/src/mod_audio_fork.rs

@@ -2,7 +2,8 @@
 
 use anyhow::Result;
 use axum::extract::ws::{Message, WebSocket};
-use futures_util::{SinkExt, stream::SplitSink};
+use futures_util::SinkExt;
+use futures_util::stream::SplitSink;
 use serde::Serialize;
 use serde_json::Value;
 use tracing::debug;
@@ -53,11 +54,7 @@ pub async fn dispatch_json(
     dispatch_event(socket, AudioForkEvent::json(value)?).await
 }
 
-pub async fn dispatch_kill_audio(socket: &mut SplitSink<WebSocket, Message>) -> Result<()> {
-    dispatch_event(socket, AudioForkEvent::kill_audio()).await
-}
-
-async fn dispatch_event(
+pub(crate) async fn dispatch_event(
     socket: &mut SplitSink<WebSocket, Message>,
     event: AudioForkEvent,
 ) -> Result<()> {

docs/adr/0004-deferred-audio-knife-start-params.md

@@ -1,26 +1,45 @@
 # AudioKnife assembles deferred service parameters before starting a conversation
 
-mod_audio_fork limits its initial message to roughly 8 KiB, while Gemini and
-OpenAI dialog instructions embedded in service parameters can exceed that size.
-AudioKnife therefore accepts an opt-in two-message transport form and assembles
+mod_audio_fork limits its initial text message to 8191 bytes after UTF-8
+encoding, while Gemini and OpenAI dialog instructions embedded in service
+parameters can exceed that size.
+AudioKnife therefore accepts an opt-in deferred-parameter exchange and assembles
 it into one complete logical Start before passing it to ContextSwitch. This keeps
 the transport constraint out of the core protocol and service implementations.
 
 ## Wire contract
 
 - An ordinary initial Start remains unchanged and contains `params`.
 - A deferred initial Start contains `"deferParams": true` and omits `params`.
-- Its literal next WebSocket frame must be a text message containing
+- Before requesting the parameters, AudioKnife validates only the transport
+  fields needed for the deferred exchange: the message is a JSON object, its
+  type is `start`, it has a valid conversation ID, and it opts into deferral
+  without inline `params`. Validation of ContextSwitch Start fields remains with
+  ContextSwitch after the Logical Start has been assembled.
+- After accepting the deferred initial Start, AudioKnife sends this Deferred
+  Params Request and waits for the send to complete:
+  `{"type":"json","data":{"type":"sendParams","id":"<same conversation id>"}}`.
+- The Deferred Params Request acknowledges only that AudioKnife is ready to
+  receive parameters; it does not mean that the conversation has started.
+- The Deferred Params Request is an AudioKnife/mod_audio_fork transport message,
+  not a ContextSwitch `ServerEvent`. ContextSwitch observes only the assembled
+  Logical Start.
+- Only after sending the Deferred Params Request does AudioKnife begin receiving
+  the deferred params message. A params frame already buffered by the WebSocket
+  is accepted; AudioKnife does not attempt to detect whether the client sent it
+  before receiving the request.
+- Its first subsequent text WebSocket frame must contain
   `{"type":"params","id":"<same conversation id>","params":<complete JSON value>}`.
 - The deferred message supports the same plain JSON and `base64:`-prefixed JSON
   encodings as other AudioKnife client text messages.
-- AudioKnife rejects mixed inline and deferred parameters, non-text or malformed
-  second frames, the wrong event type, and mismatched conversation IDs through
-  its existing startup error path.
+- AudioKnife ignores Binary, Ping, and Pong frames while waiting for deferred
+  parameters. It rejects Close frames, mixed inline and deferred parameters,
+  malformed first text frames, the wrong event type, and mismatched conversation
+  IDs through its existing startup error path.
 - The deferred message inherits the WebSocket message-size limit. AudioKnife adds
-  no separate size limit, acknowledgement, timeout, retry, or chunking protocol.
-- Ping and Pong frames are not accepted between the initial Start and deferred
-  params message; support can be added if this occurs in practice.
+  no separate size limit, completion acknowledgement, timeout, retry, or
+  chunking protocol. The normal conversation `Started` or startup error follows
+  processing of the assembled Logical Start.
 
 ## Client implementation
 
@@ -29,11 +48,12 @@ initial-message limit must:
 
 1. Serialize the complete service parameters as one JSON value.
 2. Send an initial Start without `params` and with `"deferParams": true`.
-3. Immediately send one text WebSocket message with `type` set to `params`, the
-   Start conversation ID, and the complete serialized parameters.
-4. Only send audio or other client events after the deferred params message.
+3. Wait until AudioKnife tells the client to send the parameters.
+4. Send one text WebSocket message with `type` set to `params`, the Start
+  conversation ID, and the complete serialized parameters.
+5. Only send audio or other client events after the deferred params message.
 
-For example, a client sends these two messages in order:
+For example, the exchange starts with this client message:
 
 ```json
 {
@@ -46,6 +66,20 @@ For example, a client sends these two messages in order:
 }
 ```
 
+AudioKnife responds:
+
+```json
+{
+  "type": "json",
+  "data": {
+    "type": "sendParams",
+    "id": "conversation-id"
+  }
+}
+```
+
+After receiving this Deferred Params Request, the client responds:
+
 ```json
 {
   "type": "params",
@@ -58,8 +92,9 @@ For example, a client sends these two messages in order:
 
 Both messages may instead use AudioKnife's `base64:<encoded-json>` text
 encoding. The client must not defer only part of the parameters, combine inline
-and deferred parameters, split the deferred value over multiple messages, or
-send a Ping or Pong between the two messages.
+and deferred parameters, or split the deferred value over multiple messages.
+Binary, Ping, and Pong frames sent before the deferred params message are
+discarded. A Close frame ends startup with an error.
 
 Clients that do not need deferral continue to send a single Start with inline
 `params`. A client opting into deferral requires an AudioKnife version that