diff --git a/codex-rs/tui/src/security_review.rs b/codex-rs/tui/src/security_review.rs index 431343dce4..940f74c1c6 100644 --- a/codex-rs/tui/src/security_review.rs +++ b/codex-rs/tui/src/security_review.rs @@ -1125,6 +1125,33 @@ pub(crate) async fn run_security_review( bugs_markdown = format!("{table}\n\n{bugs_markdown}"); } bugs_markdown = fix_mermaid_blocks(&bugs_markdown); + if !bugs_markdown.trim().is_empty() { + record("Polishing bug markdown formatting.".to_string()); + let fix_prompt = build_fix_markdown_prompt(&bugs_markdown, None); + let polished_response = match call_model( + &client, + &request.provider, + &request.auth, + MARKDOWN_FIX_MODEL, + MARKDOWN_FIX_SYSTEM_PROMPT, + &fix_prompt, + metrics.clone(), + 0.0, + ) + .await + { + Ok(text) => text, + Err(err) => { + let message = format!("Failed to polish bug markdown: {err}"); + record(message.clone()); + return Err(SecurityReviewFailure { + message, + logs: logs.clone(), + }); + } + }; + bugs_markdown = fix_mermaid_blocks(&polished_response); + } let mut report_sections_prefix: Vec = Vec::new(); if matches!(request.mode, SecurityReviewMode::Full) { @@ -1148,7 +1175,7 @@ pub(crate) async fn run_security_review( } else { Some(format!("# Security Findings\n\n{}", bugs_markdown.trim())) }; - let report_markdown = match request.mode { + let mut report_markdown = match request.mode { SecurityReviewMode::Full => { let mut sections = report_sections_prefix.clone(); if let Some(section) = findings_section.clone() { @@ -1176,6 +1203,34 @@ pub(crate) async fn run_security_review( } }; + if let Some(current) = report_markdown.clone() { + record("Polishing final report markdown formatting.".to_string()); + let fix_prompt = build_fix_markdown_prompt(¤t, None); + let polished_response = match call_model( + &client, + &request.provider, + &request.auth, + MARKDOWN_FIX_MODEL, + MARKDOWN_FIX_SYSTEM_PROMPT, + &fix_prompt, + metrics.clone(), + 0.0, + ) + .await + { + Ok(text) => text, + Err(err) => { + let message = format!("Failed to polish final security report: {err}"); + record(message.clone()); + return Err(SecurityReviewFailure { + message, + logs: logs.clone(), + }); + } + }; + report_markdown = Some(fix_mermaid_blocks(&polished_response)); + } + let snapshot = SecurityReviewSnapshot { generated_at: OffsetDateTime::now_utc(), findings_summary: findings_summary.clone(), @@ -2403,7 +2458,7 @@ async fn combine_spec_markdown( SPEC_GENERATION_MODEL, SPEC_COMBINE_SYSTEM_PROMPT, &prompt, - metrics, + metrics.clone(), 0.1, ) .await @@ -2418,7 +2473,38 @@ async fn combine_spec_markdown( }; let sanitized = fix_mermaid_blocks(&response); - if let Err(e) = tokio_fs::write(combined_path, sanitized.as_bytes()).await { + let polish_message = "Polishing combined specification markdown formatting.".to_string(); + if let Some(tx) = progress_sender.as_ref() { + tx.send(AppEvent::SecurityReviewLog(polish_message.clone())); + } + logs.push(polish_message); + + let fix_prompt = build_fix_markdown_prompt(&sanitized, Some(SPEC_COMBINED_MARKDOWN_TEMPLATE)); + let polished_response = match call_model( + client, + provider, + auth, + MARKDOWN_FIX_MODEL, + MARKDOWN_FIX_SYSTEM_PROMPT, + &fix_prompt, + metrics.clone(), + 0.0, + ) + .await + { + Ok(text) => text, + Err(err) => { + let message = format!("Failed to polish combined specification markdown: {err}"); + if let Some(tx) = progress_sender.as_ref() { + tx.send(AppEvent::SecurityReviewLog(message.clone())); + } + logs.push(message.clone()); + return Err(SecurityReviewFailure { message, logs }); + } + }; + let polished = fix_mermaid_blocks(&polished_response); + + if let Err(e) = tokio_fs::write(combined_path, polished.as_bytes()).await { return Err(SecurityReviewFailure { message: format!( "Failed to write combined specification to {}: {e}", @@ -2437,7 +2523,7 @@ async fn combine_spec_markdown( } logs.push(done_message); - Ok((sanitized, logs)) + Ok((polished, logs)) } fn build_spec_prompt_text( @@ -2506,6 +2592,31 @@ fn slugify_label(input: &str) -> String { } } +fn build_fix_markdown_prompt(original_content: &str, template_hint: Option<&str>) -> String { + let mut prompt = String::from( + "Read the report below and fix the formatting issues. Write the corrected version as the output.\n\ +Make sure it looks professional and polished, but still concise and to the point.\n\n\ +Some common issues to fix:\n\ +- Unicode bullet points: •\n\ +- Extra backticks around code blocks (``` markers)\n\ +- Mermaid diagrams: nodes with unescaped characters like () or []\n\ +- Incorrect number continuation (e.g. 1. 1. 1.)\n", + ); + if let Some(template) = template_hint { + prompt + .push_str("\nWhen fixing, ensure the output conforms to this template:\n\n"); + } + prompt.push_str("\nOriginal Report:\n\n"); + prompt.push_str(original_content); + prompt.push_str( + "\n\n\n# Output\n- A valid markdown report\n\n# Important:\n- Do not add emojis, or any filler text in the output.\n- Do not add AI summary or thinking process in the output (usually at the beginning or end of the response)\n- Do not remove, rewrite, or replace any image/GIF/video embeds. If the input contains media embeds (e.g., ![alt](path) or or