fix: handle masked content in generate response and stream processing (#19)

Working alone on this project so I need to approve myself.

Author: lenoxys

src/handlers/generate.rs

@@ -148,7 +148,21 @@ async fn handle_non_streaming_generate(
         return build_violation_response(response_body);
     }
 
-    // Return safe response
+    // If the response was allowed but PANW provided masked content, use it
+    if assessment.is_masked {
+        debug!("Using masked content for generate response");
+
+        response_body.response = assessment.final_content;
+
+        let json_bytes = serde_json::to_vec(&response_body).map_err(|e| {
+            error!("Failed to serialize modified response: {}", e);
+            ApiError::InternalError("Failed to serialize response".to_string())
+        })?;
+
+        return build_json_response(json_bytes.into());
+    }
+
+    // Return original (safe) response
     build_json_response(body_bytes)
 }
 

src/stream.rs

@@ -662,6 +662,31 @@ where
             return None;
         }
 
+        // If the assessment indicates masking, replace the pending buffer with the masked content
+        if assessment.is_masked {
+            // Clear existing pending chunks
+            buffer.pending_buffer.clear();
+
+            // Build a single masked JSON object and return it as the only chunk.
+            // Add a trailing newline to make NDJSON consumers happier.
+            let masked_json = serde_json::json!({
+                "created_at": chrono::Utc::now().to_rfc3339(),
+                "done": true,
+                "message": { "content": assessment.final_content, "role": "assistant" }
+            });
+
+            let mut vec = serde_json::to_vec(&masked_json).unwrap_or_else(|_| assessment.final_content.clone().into_bytes());
+            vec.push(b'\n');
+            let bytes = Bytes::from(vec);
+
+            // Mark the buffer as blocked/finished so no further inner stream data is processed
+            buffer.waiting_for_assessment = false;
+            buffer.accumulating = false;
+            buffer.blocked = true;
+
+            return Some(Ok(bytes));
+        }
+
         // Mark the content as safe by updating the read position and clearing code buffer
         buffer.commit(true);