Bound diagnostic payload allocation and trace provider restarts

2026-09-23 16:05:32 -05:00
parent 3bfd7c0a54
commit ea7117cc66
4 changed files with 228 additions and 53 deletions
@@ -6395,6 +6395,20 @@ impl BlocklistAIController {
        request_params: api::RequestParams,
        ctx: &mut ModelContext<Self>,
    ) {
+        #[cfg(not(target_family = "wasm"))]
+        remote_logging::log_model_event(
+            ctx,
+            RemoteLogRecord {
+                level: RemoteLogLevel::Info,
+                message: "Provider transition milestone".to_string(),
+                context: serde_json::json!({
+                    "event": "provider_transition_milestone",
+                    "phase": "preparation_started",
+                    "conversation_id": conversation_id,
+                    "stream_id": stream_id.as_str(),
+                }),
+            },
+        );
        let orchestration_models = LLMPreferences::as_ref(ctx)
            .get_base_llm_choices_for_agent_mode(ctx)
            .filter(|model| model.disable_reason.is_none())
@@ -6444,6 +6458,20 @@ impl BlocklistAIController {
                return;
            }
        };
+        #[cfg(not(target_family = "wasm"))]
+        remote_logging::log_model_event(
+            ctx,
+            RemoteLogRecord {
+                level: RemoteLogLevel::Info,
+                message: "Provider transition milestone".to_string(),
+                context: serde_json::json!({
+                    "event": "provider_transition_milestone",
+                    "phase": "preparation_completed",
+                    "conversation_id": conversation_id,
+                    "stream_id": stream_id.as_str(),
+                }),
+            },
+        );
        let PreparedProviderRun {
            base_profile,
            cli_monitor_profile,
@@ -7355,7 +7383,32 @@ impl BlocklistAIController {
                        );
                    });
                }
-                let _ = acknowledgement.send(result);
+                let model_turn_requested = lifecycle.as_ref().is_some_and(|lifecycle| {
+                    matches!(lifecycle.phase, ProviderLlmLifecyclePhase::Requested)
+                });
+                let projection_succeeded = result.is_ok();
+                let acknowledged = acknowledgement.send(result).is_ok();
+                if model_turn_requested {
+                    log::debug!(
+                        "Provider model request projection: conversation={conversation_id:?} stream={stream_id:?} succeeded={projection_succeeded} acknowledged={acknowledged}"
+                    );
+                    #[cfg(not(target_family = "wasm"))]
+                    remote_logging::log_model_event(
+                        ctx,
+                        RemoteLogRecord {
+                            level: RemoteLogLevel::Info,
+                            message: "Provider transition milestone".to_string(),
+                            context: serde_json::json!({
+                                "event": "provider_transition_milestone",
+                                "phase": "model_request_projection_acknowledged",
+                                "conversation_id": conversation_id,
+                                "stream_id": stream_id.as_str(),
+                                "projection_succeeded": projection_succeeded,
+                                "acknowledged": acknowledged,
+                            }),
+                        },
+                    );
+                }
            }
            ProviderDriveMessage::Checkpoint {
                checkpoint,
@@ -8571,6 +8624,20 @@ impl BlocklistAIController {
            request_usage_model.refresh_request_usage_async(ctx);
        });
        self.maybe_refresh_ai_overages(ctx);
+        #[cfg(not(target_family = "wasm"))]
+        remote_logging::log_model_event(
+            ctx,
+            RemoteLogRecord {
+                level: RemoteLogLevel::Info,
+                message: "Provider transition milestone".to_string(),
+                context: serde_json::json!({
+                    "event": "provider_transition_milestone",
+                    "phase": "cleanup_completed",
+                    "conversation_id": conversation_id,
+                    "stream_id": stream_id.as_str(),
+                }),
+            },
+        );
        self.start_next_queued_provider_run(conversation_id, ctx);
        OrchestrationEventService::handle(ctx).update(ctx, |service, ctx| {
            service.reconcile_local_groups_for_child(conversation_id, ctx);
@@ -8604,6 +8671,20 @@ impl BlocklistAIController {
        let Some(next) = next else {
            return;
        };
+        #[cfg(not(target_family = "wasm"))]
+        remote_logging::log_model_event(
+            ctx,
+            RemoteLogRecord {
+                level: RemoteLogLevel::Info,
+                message: "Provider transition milestone".to_string(),
+                context: serde_json::json!({
+                    "event": "provider_transition_milestone",
+                    "phase": "queued_follow_up_dequeued",
+                    "conversation_id": conversation_id,
+                    "stream_id": next.slot.stream_id.as_str(),
+                }),
+            },
+        );
        let mut request_params = next.request_params;
        if let Some(conversation) =
            BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id)
@@ -313,9 +313,9 @@ impl ResponseStream {
                }),
            },
        );
-        if let Some(raw_payload) =
-            remote_logging::raw_model_payload_context(ctx, raw_model_request_payload(params))
-        {
+        if let Some(raw_payload) = remote_logging::raw_model_payload_context(ctx, |max_chars| {
+            raw_model_request_payload(params, max_chars)
+        }) {
            remote_logging::log_model_event(
                ctx,
                RemoteLogRecord {
@@ -448,9 +448,9 @@ impl ResponseStream {
        raw_payload: impl AsRef<str>,
        ctx: &mut ModelContext<Self>,
    ) {
-        let Some(raw_payload) =
-            remote_logging::raw_model_payload_context(ctx, raw_payload.as_ref())
-        else {
+        let Some(raw_payload) = remote_logging::raw_model_payload_context(ctx, |max_chars| {
+            remote_logging::tail_limited_payload_context(raw_payload.as_ref(), max_chars)
+        }) else {
            return;
        };
        let mut context = self.common_remote_log_context("raw_model_response", request_id);
@@ -1073,26 +1073,49 @@ impl ResponseStream {
 }

 #[cfg(not(target_family = "wasm"))]
-fn raw_model_request_payload(params: &api::RequestParams) -> String {
-    let payload = serde_json::json!({
-        "model_id": params.model.as_str(),
-        "coding_model_id": params.coding_model.as_str(),
-        "cli_agent_model_id": params.cli_agent_model.as_str(),
-        "computer_use_model_id": params.computer_use_model.as_str(),
-        "input": format!("{:#?}", params.input),
-        "tool_results": format!("{:#?}", params.tool_results),
-        "tasks": format!("{:#?}", params.tasks),
-        "message_history": serde_json::to_value(&params.message_history)
-            .unwrap_or_else(|_| serde_json::json!(format!("{:#?}", params.message_history))),
-        "progressive_summary": &params.progressive_summary,
-        "tool_result_archive": serde_json::to_value(&params.tool_result_archive)
-            .unwrap_or_else(|_| serde_json::json!(format!("{:#?}", params.tool_result_archive))),
-        "global_rules": &params.global_rules,
-        "mcp_context": format!("{:#?}", params.mcp_context),
-        "session": {
+fn raw_model_request_payload(params: &api::RequestParams, max_chars: usize) -> serde_json::Value {
+    use serde::ser::{SerializeMap, Serializer};
+
+    // collect_str writes debug output directly through serde_json's escaping
+    // writer, without format! allocating another copy of potentially large input.
+    struct DebugPayload<T>(T);
+    impl<T: std::fmt::Debug> serde::Serialize for DebugPayload<T> {
+        fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
+            serializer.collect_str(&format_args!("{:#?}", self.0))
+        }
+    }
+
+    remote_logging::bounded_payload_context(max_chars, |writer| {
+        let mut serializer = serde_json::Serializer::pretty(writer);
+        let mut map = serializer.serialize_map(None)?;
+        map.serialize_entry("model_id", &params.model.as_str())?;
+        map.serialize_entry("coding_model_id", &params.coding_model.as_str())?;
+        map.serialize_entry("cli_agent_model_id", &params.cli_agent_model.as_str())?;
+        map.serialize_entry("computer_use_model_id", &params.computer_use_model.as_str())?;
+        map.serialize_entry("input", &DebugPayload(&params.input))?;
+        map.serialize_entry("tool_results", &DebugPayload(&params.tool_results))?;
+        map.serialize_entry("tasks", &DebugPayload(&params.tasks))?;
+        map.serialize_entry("message_history", &params.message_history)?;
+        map.serialize_entry("progressive_summary", &params.progressive_summary)?;
+        map.serialize_entry("tool_result_archive", &params.tool_result_archive)?;
+        map.serialize_entry("global_rules", &params.global_rules)?;
+        map.serialize_entry("mcp_context", &DebugPayload(&params.mcp_context))?;
+        map.serialize_entry("autonomy_level", &DebugPayload(&params.autonomy_level))?;
+        map.serialize_entry("isolation_level", &DebugPayload(&params.isolation_level))?;
+        map.serialize_entry(
+            "supported_tools_override",
+            &DebugPayload(&params.supported_tools_override),
+        )?;
+        map.serialize_entry("context_window_limit", &params.context_window_limit)?;
+        map.serialize_entry(
+            "session",
+            &serde_json::json!({
                "is_remote": params.session_context.is_remote(),
-        },
-        "features": {
+            }),
+        )?;
+        map.serialize_entry(
+            "features",
+            &serde_json::json!({
                "memory_enabled": params.is_memory_enabled,
                "warp_drive_context_enabled": params.warp_drive_context_enabled,
                "planning_enabled": params.planning_enabled,
@@ -1101,18 +1124,14 @@ fn raw_model_request_payload(params: &api::RequestParams) -> String {
                "ask_user_question_enabled": params.ask_user_question_enabled,
                "research_agent_enabled": params.research_agent_enabled,
                "orchestration_enabled": params.orchestration_enabled,
-        },
-        "autonomy_level": format!("{:?}", params.autonomy_level),
-        "isolation_level": format!("{:?}", params.isolation_level),
-        "supported_tools_override": format!("{:#?}", params.supported_tools_override),
-        "context_window_limit": params.context_window_limit,
-        "omitted_sensitive_fields": [
-            "api_keys",
-            "custom_model_providers",
-            "custom_model_routers",
-        ],
-    });
-    serde_json::to_string_pretty(&payload).unwrap_or_else(|_| format!("{payload:#?}"))
+            }),
+        )?;
+        map.serialize_entry(
+            "omitted_sensitive_fields",
+            &["api_keys", "custom_model_providers", "custom_model_routers"],
+        )?;
+        map.end()
+    })
 }

 #[cfg(not(target_family = "wasm"))]
@@ -189,7 +189,7 @@ pub(crate) fn sanitize_error(error: impl std::fmt::Display) -> String {

 pub(crate) fn raw_model_payload_context<M>(
    ctx: &ModelContext<M>,
-    raw_payload: impl AsRef<str>,
+    build_payload: impl FnOnce(usize) -> Value,
 ) -> Option<Value>
 where
    M: Entity,
@@ -203,10 +203,7 @@ where
    }

    let max_chars = (*settings.remote_logging_model_payload_max_chars.value()).max(1);
-    Some(tail_limited_payload_context(
-        raw_payload.as_ref(),
-        max_chars,
-    ))
+    Some(build_payload(max_chars))
 }

 fn enrich_context(mut context: Value) -> Value {
@@ -290,7 +287,59 @@ fn truncate_chars(value: &str, max_chars: usize) -> String {
    truncated
 }

-fn tail_limited_payload_context(payload: &str, max_chars: usize) -> Value {
+/// Serialize directly into a bounded UTF-8 tail instead of materializing the full
+/// request and its archive as a JSON Value and then another large String.
+pub(crate) fn bounded_payload_context(
+    max_chars: usize,
+    serialize: impl FnOnce(&mut dyn std::io::Write) -> serde_json::Result<()>,
+) -> Value {
+    struct TailWriter {
+        bytes: std::collections::VecDeque<u8>,
+        max_bytes: usize,
+        total_chars: usize,
+    }
+
+    impl std::io::Write for TailWriter {
+        fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
+            self.total_chars += bytes.iter().filter(|byte| **byte & 0xc0 != 0x80).count();
+            let retained = &bytes[bytes.len().saturating_sub(self.max_bytes)..];
+            let discard = (self.bytes.len() + retained.len()).saturating_sub(self.max_bytes);
+            self.bytes.drain(..discard);
+            self.bytes.extend(retained);
+            Ok(bytes.len())
+        }
+
+        fn flush(&mut self) -> std::io::Result<()> {
+            Ok(())
+        }
+    }
+
+    let mut writer = TailWriter {
+        bytes: std::collections::VecDeque::new(),
+        max_bytes: max_chars.saturating_mul(4),
+        total_chars: 0,
+    };
+    if let Err(error) = serialize(&mut writer) {
+        return json!({ "serialization_error": sanitize_error(error) });
+    }
+    // A byte-limited tail can begin in the middle of a multibyte character.
+    while writer.bytes.front().is_some_and(|byte| byte & 0xc0 == 0x80) {
+        writer.bytes.pop_front();
+    }
+    let bytes = writer.bytes.into_iter().collect::<Vec<_>>();
+    let payload = String::from_utf8_lossy(&bytes);
+    let mut context = tail_limited_payload_context(&payload, max_chars);
+    context["payload_total_chars"] = json!(writer.total_chars);
+    context["truncated"] = json!(writer.total_chars > max_chars);
+    context["truncation_strategy"] = json!(if writer.total_chars > max_chars {
+        Some("tail")
+    } else {
+        None
+    });
+    context
+}
+
+pub(crate) fn tail_limited_payload_context(payload: &str, max_chars: usize) -> Value {
    let total_chars = payload.chars().count();
    let truncated = total_chars > max_chars;
    let payload = if truncated {
@@ -38,3 +38,29 @@ fn raw_payload_cap_keeps_tail() {
        })
    );
 }
+
+#[test]
+fn streamed_payload_matches_existing_tail_for_unicode_and_split_writes() {
+    let payload = "start \"escaped\" é日本🙂 ending".repeat(100);
+    for limit in [0, 1, 3, 17, payload.chars().count(), payload.len()] {
+        for chunk_size in [1, 2, 7, payload.len()] {
+            let actual = super::bounded_payload_context(limit, |writer| {
+                for chunk in payload.as_bytes().chunks(chunk_size) {
+                    writer.write_all(chunk).unwrap();
+                }
+                Ok(())
+            });
+            assert_eq!(actual, tail_limited_payload_context(&payload, limit));
+        }
+    }
+}
+
+#[test]
+fn streamed_json_matches_serialized_payload_without_materializing_the_archive() {
+    let payload = json!({ "archive": ["large tool output".repeat(1000)], "input": "steer 🙂" });
+    let expected = serde_json::to_string_pretty(&payload).unwrap();
+    let actual = super::bounded_payload_context(100, |writer| {
+        serde_json::to_writer_pretty(writer, &payload)
+    });
+    assert_eq!(actual, tail_limited_payload_context(&expected, 100));
+}