diff --git a/apps/api/native/Cargo.toml b/apps/api/native/Cargo.toml index c884a31a095..3c8ffadc3f9 100644 --- a/apps/api/native/Cargo.toml +++ b/apps/api/native/Cargo.toml @@ -25,6 +25,7 @@ serde_json = "1.0.143" strsim = "0.11" texting_robots = "0.2.2" url = "2.5.7" +zip = { version = "8.6.0", default-features = false, features = ["deflate"] } tokio = "1.48.0" # Transitive dep (pdf-inspector -> unicode-normalization -> tinyvec). Pinned because # tinyvec 1.13.0 fails to compile (https://github.com/Lokathor/tinyvec/issues/225) diff --git a/apps/api/native/src/document.rs b/apps/api/native/src/document.rs index 9bde03f5fc4..ee75ba34f7a 100644 --- a/apps/api/native/src/document.rs +++ b/apps/api/native/src/document.rs @@ -1,5 +1,42 @@ use napi::bindgen_prelude::*; use napi_derive::napi; +use roxmltree::Document; +use std::collections::{HashMap, HashSet}; +use std::io::{Cursor, Read}; + +const MAX_OOXML_PART_BYTES: u64 = 128 * 1024 * 1024; +const MAX_OOXML_TOTAL_BYTES: u64 = 128 * 1024 * 1024; +const MAX_OOXML_PARTS: usize = 256; +const MAX_HEADER_FOOTER_REFERENCES: usize = 256; +const WORDPROCESSINGML_NAMESPACES: [&str; 2] = [ + "http://schemas.openxmlformats.org/wordprocessingml/2006/main", + "http://purl.oclc.org/ooxml/wordprocessingml/main", +]; + +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +enum RepeatingPartKind { + Header, + Footer, +} + +struct Relationship { + target: String, + kind: RepeatingPartKind, +} + +struct ZipReadBudget { + bytes_remaining: u64, + parts_remaining: usize, +} + +impl Default for ZipReadBudget { + fn default() -> Self { + Self { + bytes_remaining: MAX_OOXML_TOTAL_BYTES, + parts_remaining: MAX_OOXML_PARTS, + } + } +} /// Convert a document (doc, docx, odt/ods/odp, rtf, xls/xlsx, ppt/pptx, epub, /// csv, ...) to GitHub-Flavored Markdown using anydoc. The format is detected @@ -14,6 +51,241 @@ pub fn convert_document_to_markdown(data: &[u8], extension_hint: Option) .and_then(anydoc::Format::from_extension) }); - anydoc::to_markdown_bytes(data, format) - .map_err(|e| Error::new(Status::GenericFailure, e.to_string())) + let markdown = anydoc::to_markdown_bytes(data, format) + .map_err(|e| Error::new(Status::GenericFailure, e.to_string()))?; + + if format != Some(anydoc::Format::Docx) { + return Ok(markdown); + } + + let (headers, footers) = extract_docx_headers_and_footers(data); + if headers.is_empty() && footers.is_empty() { + return Ok(markdown); + } + let sections = [headers, markdown.trim().to_string(), footers] + .into_iter() + .filter(|section| !section.is_empty()) + .collect::>(); + + Ok(format!("{}\n", sections.join("\n\n"))) +} + +fn extract_docx_headers_and_footers(data: &[u8]) -> (String, String) { + let Ok(mut archive) = zip::ZipArchive::new(Cursor::new(data)) else { + return (String::new(), String::new()); + }; + let mut budget = ZipReadBudget::default(); + + let main_part = read_zip_text(&mut archive, "_rels/.rels", &mut budget) + .and_then(|xml| office_document_target(&xml)) + .unwrap_or_else(|| "word/document.xml".to_string()); + let Some(document_xml) = read_zip_text(&mut archive, &main_part, &mut budget) else { + return (String::new(), String::new()); + }; + let relationship_part = relationship_part_for(&main_part); + let Some(relationships_xml) = read_zip_text(&mut archive, &relationship_part, &mut budget) else { + return (String::new(), String::new()); + }; + let relationships = parse_repeating_relationships(&relationships_xml); + let Ok(document) = Document::parse(&document_xml) else { + return (String::new(), String::new()); + }; + + let mut seen = HashSet::new(); + let mut headers = Vec::new(); + let mut footers = Vec::new(); + let mut reference_count = 0; + + for reference in document + .descendants() + .filter(|node| is_wordprocessing_node(*node)) + { + let reference_kind = match reference.tag_name().name() { + "headerReference" => RepeatingPartKind::Header, + "footerReference" => RepeatingPartKind::Footer, + _ => continue, + }; + reference_count += 1; + if reference_count > MAX_HEADER_FOOTER_REFERENCES { + break; + } + let Some(id) = reference + .attributes() + .find(|attribute| attribute.name() == "id") + .map(|attribute| attribute.value()) + else { + continue; + }; + let Some(relationship) = relationships.get(id) else { + continue; + }; + if relationship.kind != reference_kind { + continue; + } + let Some(part) = resolve_part(&main_part, &relationship.target) else { + continue; + }; + if !seen.insert((reference_kind, part.clone())) { + continue; + } + let Some(xml) = read_zip_text(&mut archive, &part, &mut budget) else { + continue; + }; + let Some(text) = extract_wordprocessing_text(&xml) else { + continue; + }; + + match reference_kind { + RepeatingPartKind::Header => headers.push(text), + RepeatingPartKind::Footer => footers.push(text), + } + } + + (headers.join("\n\n"), footers.join("\n\n")) +} + +fn read_zip_text( + archive: &mut zip::ZipArchive, + name: &str, + budget: &mut ZipReadBudget, +) -> Option { + if budget.parts_remaining == 0 || budget.bytes_remaining == 0 { + return None; + } + let mut file = archive.by_name(name.trim_start_matches('/')).ok()?; + if file.size() > MAX_OOXML_PART_BYTES || file.size() > budget.bytes_remaining { + return None; + } + budget.parts_remaining -= 1; + let mut bytes = Vec::with_capacity(file.size() as usize); + let read_limit = MAX_OOXML_PART_BYTES.min(budget.bytes_remaining); + let read_result = file.by_ref().take(read_limit + 1).read_to_end(&mut bytes); + budget.bytes_remaining = budget.bytes_remaining.saturating_sub(bytes.len() as u64); + if read_result.is_err() { + return None; + } + if bytes.len() as u64 > read_limit { + budget.bytes_remaining = 0; + return None; + } + String::from_utf8(bytes).ok() +} + +fn office_document_target(xml: &str) -> Option { + let document = Document::parse(xml).ok()?; + let relationship = document.descendants().find(|node| { + node.is_element() + && node.tag_name().name() == "Relationship" + && node + .attribute("Type") + .is_some_and(|value| value.ends_with("/officeDocument")) + })?; + normalize_part(relationship.attribute("Target")?) +} + +fn parse_repeating_relationships(xml: &str) -> HashMap { + let Ok(document) = Document::parse(xml) else { + return HashMap::new(); + }; + + document + .descendants() + .filter_map(|node| { + if !node.is_element() + || node.tag_name().name() != "Relationship" + || node + .attribute("TargetMode") + .is_some_and(|value| value.eq_ignore_ascii_case("external")) + { + return None; + } + let kind = match node.attribute("Type")? { + value if value.ends_with("/header") => RepeatingPartKind::Header, + value if value.ends_with("/footer") => RepeatingPartKind::Footer, + _ => return None, + }; + Some(( + node.attribute("Id")?.to_string(), + Relationship { + target: node.attribute("Target")?.to_string(), + kind, + }, + )) + }) + .collect() +} + +fn relationship_part_for(part: &str) -> String { + match part.rsplit_once('/') { + Some((parent, file_name)) => format!("{parent}/_rels/{file_name}.rels"), + None => format!("_rels/{part}.rels"), + } +} + +fn resolve_part(base: &str, target: &str) -> Option { + if target.starts_with('/') { + return normalize_part(target); + } + let parent = base + .rsplit_once('/') + .map(|(parent, _)| parent) + .unwrap_or(""); + normalize_part(&format!("{parent}/{target}")) +} + +fn normalize_part(part: &str) -> Option { + let mut normalized = Vec::new(); + for component in part.trim_start_matches('/').split('/') { + match component { + "" | "." => {} + ".." => normalized.pop().map(|_| ())?, + value => normalized.push(value), + } + } + (!normalized.is_empty()).then(|| normalized.join("/")) +} + +fn extract_wordprocessing_text(xml: &str) -> Option { + let document = Document::parse(xml).ok()?; + let paragraphs = document + .descendants() + .filter(|node| { + is_wordprocessing_node(*node) + && node.tag_name().name() == "p" + && !node + .ancestors() + .skip(1) + .any(|ancestor| is_wordprocessing_node(ancestor) && ancestor.tag_name().name() == "p") + }) + .filter_map(|paragraph| { + let mut text = String::new(); + for node in paragraph + .descendants() + .filter(|node| is_wordprocessing_node(*node)) + { + match node.tag_name().name() { + "t" => { + if let Some(value) = node.text() { + text.push_str(value); + } + } + "tab" => text.push('\t'), + "br" | "cr" => text.push('\n'), + _ => {} + } + } + let text = text.trim(); + (!text.is_empty()).then(|| text.to_string()) + }) + .collect::>(); + + (!paragraphs.is_empty()).then(|| paragraphs.join("\n\n")) +} + +fn is_wordprocessing_node(node: roxmltree::Node<'_, '_>) -> bool { + node.is_element() + && node + .tag_name() + .namespace() + .is_some_and(|namespace| WORDPROCESSINGML_NAMESPACES.contains(&namespace)) } diff --git a/apps/api/src/__tests__/snips/v2/parse.test.ts b/apps/api/src/__tests__/snips/v2/parse.test.ts index e840d7c9aa5..83d84922b38 100644 --- a/apps/api/src/__tests__/snips/v2/parse.test.ts +++ b/apps/api/src/__tests__/snips/v2/parse.test.ts @@ -21,7 +21,7 @@ import { config } from "../../../config"; import { getRedisConnection } from "../../../services/queue-service"; const DOCX_FIXTURE_BASE64 = - "UEsDBBQAAAAIAKtlbVzXeYTq8QAAALgBAAATAAAAW0NvbnRlbnRfVHlwZXNdLnhtbH2QzU7DMBCE730Ky9cqccoBIZSkB36OwKE8wMreJFb9J69b2rdn00KREOVozXwz62nXB+/EHjPZGDq5qhspMOhobBg7+b55ru6koALBgIsBO3lEkut+0W6OCUkwHKiTUynpXinSE3qgOiYMrAwxeyj8zKNKoLcworppmlulYygYSlXmDNkvhGgfcYCdK+LpwMr5loyOpHg4e+e6TkJKzmoorKt9ML+Kqq+SmsmThyabaMkGqa6VzOL1jh/0lSfK1qB4g1xewLNRfcRslIl65xmu/0/649o4DFbjhZ/TUo4aiXh77+qL4sGG71+06jR8/wlQSwMEFAAAAAgAq2VtXCAbhuqyAAAALgEAAAsAAABfcmVscy8ucmVsc43Puw6CMBQG4J2naM4uBQdjDIXFmLAafICmPZRGeklbL7y9HRzEODie23fyN93TzOSOIWpnGdRlBQStcFJbxeAynDZ7IDFxK/nsLDJYMELXFs0ZZ57yTZy0jyQjNjKYUvIHSqOY0PBYOo82T0YXDE+5DIp6Lq5cId1W1Y6GTwPagpAVS3rJIPSyBjIsHv/h3ThqgUcnbgZt+vHlayPLPChMDB4uSCrf7TKzQHNKuorZvgBQSwMEFAAAAAgAq2VtXCCNfXOwAAAA7AAAABEAAAB3b3JkL2RvY3VtZW50LnhtbDWOMQvCMBCFd3/FkV1THURKGwfFVQcF19icWmjuQi5a/fcmBZeP93jw3TXbjx/gjVF6plYtF5UCpI5dT49WXc6H+UaBJEvODkzYqi+K2ppZM9aOu5dHSpANJPXYqmdKodZauid6KwsOSHm7c/Q25RofeuToQuQORfIBP+hVVa21tz0pMwPI1hu7b4lTCSYjFiRzslEQ9sfdFS5hYOvgjJIaXbbCODFMGv33lPT/0/wAUEsBAhQDFAAAAAgAq2VtXNd5hOrxAAAAuAEAABMAAAAAAAAAAAAAAIABAAAAAFtDb250ZW50X1R5cGVzXS54bWxQSwECFAMUAAAACACrZW1cIBuG6rIAAAAuAQAACwAAAAAAAAAAAAAAgAEiAQAAX3JlbHMvLnJlbHNQSwECFAMUAAAACACrZW1cII19c7AAAADsAAAAEQAAAAAAAAAAAAAAgAH9AQAAd29yZC9kb2N1bWVudC54bWxQSwUGAAAAAAMAAwC5AAAA3AIAAAAA"; + "UEsDBBQAAAAIAGJCNV0B3PB/6gAAAIsCAAATAAAAW0NvbnRlbnRfVHlwZXNdLnhtbLWSu07DMBSGXyXyWtUnZWBASToAKzDwAkf2cWLhm+zTUt4epykdUKELjPZ/+X5L7rYH75o95WJj6MVGtmI7dK8fiUpTlVB6MTGnO4CiJvJYZEwUqmJi9sj1mEdIqN5wJLhp21tQMTAFXvPcIYbugQzuHDePh3q9UDK5Ipr7xTizeoEpOauQqw77oL9R1ieCrMmjp0w2lVU1CLhImJWfAafcc312tpqaF8z8hL664D1mDTqqna9J+XvNhZ3RGKvonJ/bUo6KSrFh9E6eFY82rK7tmAg15c3fz1iKr/JNjPwv/KX4iw/H7zZ8AlBLAwQUAAAACABiQjVdXzOVUpUAAAAHAQAACwAAAF9yZWxzLy5yZWxzjc87DsIwDAbgq0Q+QJ0yMKCmXVi6Ii4QJW5T0TzkhNftycBAEQOjf//6LHfDw6/iRpyXGBS0jYSh70606lKD7JaURW2ErMCVkg6I2TjyOjcxUaibKbLXpY48Y9LmomfCnZR75E8DtqYYrQIebQvi/Ez0jx2naTF0jObqKZQfJ74aVdY8U1Fwj2zRvuOmsoB9h5sX+xdQSwMEFAAAAAgAYkI1Xf2I2NHjAAAAkQEAABEAAAB3b3JkL2RvY3VtZW50LnhtbI2QQU8DIRCF/wrhB5TVg4fN7jbRptGTjdHEK8LQJQGGzFDX/nthrem1l0fIfO/Ng2H7E4P4BmKPaZR3m05up2HpLZpThFREHSful1HOpeReKTYzRM0bzJDqzCFFXeqVjmpBspnQALNPxxjUfdc9qKh9kpcYuiUGnfMGdpcCfyEEQZfakGefWbaCX2jP7cxNqEmZDpoYxO716VN85IDainfgIh4rOqgGNKVVVxuDKYfVO4O2QG/ggCAZEBU+ZxilBadPoUhBvbejpBf7vJJSNZdDLLe59ivZXOq6Vv2/Ql3/e/oFUEsDBBQAAAAIAGJCNV2/NrkppgAAAIMBAAAcAAAAd29yZC9fcmVscy9kb2N1bWVudC54bWwucmVsc62QywrCMBBFfyXkAzqtCxfSx0ak3Yo/ENJJU2weJKPo3xvUioUuXLi8M8O5hymbm5nYFUMcna14keW8qcsjToLSIOrRR5YubKy4JvI7gCg1GhEz59GmjXLBCEoxDOCFPIsBYZPnWwjfDL5ksq6veOj6FkWPgbPT3eMvBU6pUeLeyYtBSys9oGeiCANSYj5zkSUQh3WJg3P0Xwk1E98Sr/yRgMV76wdQSwMEFAAAAAgAYkI1XRIFb3mKAAAAqQAAABAAAAB3b3JkL2hlYWRlcjEueG1sHc5NCsMgEAXgq4gHiLaLLoIxi3aRXbtooVuJ0yj4x4w0PX41mw8ebxiemn8xsC8g+Zwmfhokn7XaR2eRtSbRuE/c1VpGIWh1EA0NuUBq3SdjNLVF3MSe0RbMKxD5tMUgzlJeRDQ+8f6tdLBT9cMgAbvdr2/2KiEby55AlS1gLKAS/aSLh+WwjdF/UEsDBBQAAAAIAGJCNV0TuxHRigAAAKkAAAAQAAAAd29yZC9mb290ZXIxLnhtbB3OPQrDMAwF4KsYHyB2O3QITjK0dG2HFrqaRPmB2DKSqHv82lk+eDwhnht+YVdfIN4wdvrUWD30LrezkCpN5DZ3ehVJrTE8rhA8N5gglm5GCl5KpMVkpCkRjsC8xSXs5mztxQS/RV2/pQpVpH96YlC3x/Wj3mlHP6kXsKg7ogA5U0+qdJgOy5j+D1BLAQIUAxQAAAAIAGJCNV0B3PB/6gAAAIsCAAATAAAAAAAAAAAAAACAAQAAAABbQ29udGVudF9UeXBlc10ueG1sUEsBAhQDFAAAAAgAYkI1XV8zlVKVAAAABwEAAAsAAAAAAAAAAAAAAIABGwEAAF9yZWxzLy5yZWxzUEsBAhQDFAAAAAgAYkI1Xf2I2NHjAAAAkQEAABEAAAAAAAAAAAAAAIAB2QEAAHdvcmQvZG9jdW1lbnQueG1sUEsBAhQDFAAAAAgAYkI1Xb82uSmmAAAAgwEAABwAAAAAAAAAAAAAAIAB6wIAAHdvcmQvX3JlbHMvZG9jdW1lbnQueG1sLnJlbHNQSwECFAMUAAAACABiQjVdEgVveYoAAACpAAAAEAAAAAAAAAAAAAAAgAHLAwAAd29yZC9oZWFkZXIxLnhtbFBLAQIUAxQAAAAIAGJCNV0TuxHRigAAAKkAAAAQAAAAAAAAAAAAAACAAYMEAAB3b3JkL2Zvb3RlcjEueG1sUEsFBgAAAAAGAAYAfwEAADsFAAAAAA=="; const htmlFixture = ` @@ -572,7 +572,12 @@ describe("/v2/parse", () => { identity, ); - expect(result.markdown).toMatch(/Parse DOCX Upload Test/i); + const markdown = result.markdown!; + expect(markdown).toContain("Parse DOCX Upload Test Header"); + expect(markdown).toContain("Parse DOCX Upload Test Body"); + expect(markdown).toContain("Parse DOCX Upload Test Footer"); + expect(markdown.indexOf("Header")).toBeLessThan(markdown.indexOf("Body")); + expect(markdown.indexOf("Body")).toBeLessThan(markdown.indexOf("Footer")); expect(result.metadata.creditsUsed).toBe(1); }, scrapeTimeout, diff --git a/apps/api/src/controllers/v2/agent.ts b/apps/api/src/controllers/v2/agent.ts index a153bb0be0f..906c58d90b7 100644 --- a/apps/api/src/controllers/v2/agent.ts +++ b/apps/api/src/controllers/v2/agent.ts @@ -1,4 +1,5 @@ import { v7 as uuidv7 } from "uuid"; +import { AGENT_REQUEST_CREDITS_SHARDS } from "../../lib/request-credits-store"; import { Response } from "express"; import { AgentRequest, @@ -205,7 +206,7 @@ export async function agentController( target_hint: req.body.urls?.[0] ?? req.body.prompt ?? "", zeroDataRetention: false, // not supported for agent api_key_id: req.acuc?.api_key_id ?? null, - creditsShards: 8, + creditsShards: AGENT_REQUEST_CREDITS_SHARDS, }); const passthrough = await fetch( diff --git a/apps/api/src/controllers/v2/batch-scrape.ts b/apps/api/src/controllers/v2/batch-scrape.ts index a44e5ab4f8d..a0ec17177a0 100644 --- a/apps/api/src/controllers/v2/batch-scrape.ts +++ b/apps/api/src/controllers/v2/batch-scrape.ts @@ -50,7 +50,11 @@ import { calculateThreatScanCredits } from "../../lib/scrape-billing"; import { billTeam } from "../../services/billing/credit_billing"; import { emitRejectedScrapeActivityEvents } from "../../lib/siem-logging"; import { UnsupportedSiteError } from "../../lib/error"; -import { requestCreditsShards } from "../../lib/request-credits-store"; +import { + AGENT_REQUEST_CREDITS_SHARDS, + initializeRequestCredits, + requestCreditsShards, +} from "../../lib/request-credits-store"; export async function batchScrapeController( req: RequestWithAuth<{}, BatchScrapeResponse, BatchScrapeRequest>, @@ -351,6 +355,35 @@ export async function batchScrapeController( ), creditsShards: requestCreditsShards(urls.length), }); + } else if (!req.body.appendToId) { + // An agent-started batch is recorded under the agent's request rather + // than as a request of its own, so logRequest never created a credit + // row for it. Its status is still polled: v2 reads credits under the + // saved agent request id, v1 under the batch id, and its children record + // credits under the agent request id. Make sure both rows exist. The + // agent row is created once with the agent controller's shard count; + // for an agent that came through that controller this is a no-op, for + // one that reaches the API first through interop it is the only writer. + const agentRequestId = req.body.__agentInterop?.requestId; + if (agentRequestId) { + await initializeRequestCredits( + agentRequestId, + AGENT_REQUEST_CREDITS_SHARDS, + ).catch(error => { + logger.warn("Failed to initialize Bigtable request credits", { + error, + requestId: agentRequestId, + shards: AGENT_REQUEST_CREDITS_SHARDS, + }); + }); + } + const creditsShards = requestCreditsShards(urls.length); + await initializeRequestCredits(id, creditsShards).catch(error => { + logger.warn("Failed to initialize Bigtable request credits", { + error, + shards: creditsShards, + }); + }); } const sc: StoredCrawl = req.body.appendToId diff --git a/apps/api/src/controllers/v2/scrape-alexandria.ts b/apps/api/src/controllers/v2/scrape-alexandria.ts index 4d9cbddb1f6..ba19f4cc38d 100644 --- a/apps/api/src/controllers/v2/scrape-alexandria.ts +++ b/apps/api/src/controllers/v2/scrape-alexandria.ts @@ -30,8 +30,8 @@ const providerScrapeSchema = z.strictObject({ .number() .int() .positive() - .default(50000) - .transform(value => Math.min(value, 50000)), + .default(120_000) + .transform(value => Math.min(value, 120_000)), origin: z.string().default("api"), integration: z.string().nullable().optional(), __agentInterop: z diff --git a/apps/api/src/controllers/v2/search.ts b/apps/api/src/controllers/v2/search.ts index 982b2107e91..54bb66e5c5e 100644 --- a/apps/api/src/controllers/v2/search.ts +++ b/apps/api/src/controllers/v2/search.ts @@ -344,6 +344,9 @@ async function searchControllerInner( zeroDataRetention, api_key_id: req.acuc?.api_key_id ?? null, }); + // The rejection is surfaced where the promise is awaited below; this + // only stops it counting as unhandled until then. + logRequestPromise.catch(() => {}); } const toolsOnly = isToolsOnlySearch(req.body.sources, req.body.categories); @@ -359,6 +362,17 @@ async function searchControllerInner( zeroDataRetention, ) : 0; + // The request must be on record before anything is reserved, runs, or + // bills. A failed log fails the request here, ahead of the keyless + // reservation and the search, so there is nothing to refund or unbill. + const logStart = Date.now(); + await logRequestPromise; + const waited = Date.now() - logStart; + if (waited >= 5) + logger.warn("Had to wait for log request promise to complete", { + timeMs: waited, + }); + if (projectedKeylessCredits > 0) { const reservation = await reserveKeylessCredits( req.auth.team_id, @@ -446,15 +460,6 @@ async function searchControllerInner( const endTime = new Date().getTime(); const timeTakenInSeconds = (endTime - middlewareStartTime) / 1000; - // Wait for the parent log before inserting the child search log. - const logStart = Date.now(); - await logRequestPromise; - const waited = Date.now() - logStart; - if (waited >= 5) - logger.warn("Had to wait for log request promise to complete", { - timeMs: waited, - }); - logSearch( { id: jobId, diff --git a/apps/api/src/lib/request-credits-store.ts b/apps/api/src/lib/request-credits-store.ts index f1e4e87a633..88fcecc0688 100644 --- a/apps/api/src/lib/request-credits-store.ts +++ b/apps/api/src/lib/request-credits-store.ts @@ -81,6 +81,14 @@ export function requestCreditsShardForJob( ); } +/** + * Shard count for an agent request's credit row. An agent's job count is not + * known up front, and every writer that may create the row for an agent id + * (the agent controller, and agent-interop calls that reach the API first) + * must agree on it, because the row is created once and never resized. + */ +export const AGENT_REQUEST_CREDITS_SHARDS = 8; + export function requestCreditsShards(maxJobs: number): number { if (!Number.isFinite(maxJobs) || maxJobs <= 0) return 8; if (maxJobs <= 1_000) return 8; diff --git a/apps/api/src/services/alexandria/client.ts b/apps/api/src/services/alexandria/client.ts index 299fd9970ea..44a946d5a88 100644 --- a/apps/api/src/services/alexandria/client.ts +++ b/apps/api/src/services/alexandria/client.ts @@ -4,8 +4,8 @@ import type { ExchangeResponse } from "./contracts"; const dispatcher = new Agent({ connectTimeout: 10000, - headersTimeout: 50000, - bodyTimeout: 50000, + headersTimeout: 120_000, + bodyTimeout: 120_000, }); export async function exchangeRequest(input: { diff --git a/apps/api/src/services/logging/log_job.test.ts b/apps/api/src/services/logging/log_job.test.ts index 91e6516ecf3..315687dfa2e 100644 --- a/apps/api/src/services/logging/log_job.test.ts +++ b/apps/api/src/services/logging/log_job.test.ts @@ -233,12 +233,12 @@ describe("logSearch", () => { ); }); - it("reports serialization failures without losing the PostgreSQL attempt", async () => { + it("fails the log call on a serialization failure before touching PostgreSQL", async () => { const search = makeSearch({ options: { unsupported: 1n } }); - await expect(logSearch(search)).resolves.toBeUndefined(); + await expect(logSearch(search)).rejects.toThrow(); - expect(values).toHaveBeenCalledOnce(); + expect(values).not.toHaveBeenCalled(); expect(publishMessage).not.toHaveBeenCalled(); expect(metricInc).toHaveBeenCalledWith({ table: "searches", @@ -578,10 +578,11 @@ describe("logRequest", () => { const gaxOpts = publishes[0].options.gaxOpts; // A bare `timeout` would collapse the retry budget to one attempt. expect(gaxOpts.timeout).toBeUndefined(); + // The caller waits on the publish, so the whole budget is 30 s. expect(gaxOpts.retry.backoffSettings).toMatchObject({ - initialRpcTimeoutMillis: 15_000, - maxRpcTimeoutMillis: 15_000, - totalTimeoutMillis: 300_000, + initialRpcTimeoutMillis: 10_000, + maxRpcTimeoutMillis: 10_000, + totalTimeoutMillis: 30_000, }); expect(gaxOpts.retry.retryCodes).toBeUndefined(); @@ -633,12 +634,14 @@ describe("logRequest", () => { expect(values).not.toHaveBeenCalled(); }); - it("keeps the database write when Pub/Sub fails", async () => { + it("fails the log call and skips PostgreSQL when Pub/Sub fails", async () => { publishMessage.mockRejectedValueOnce(new Error("Pub/Sub unavailable")); - await expect(logRequest(makeRequest(null))).resolves.toBeUndefined(); + await expect(logRequest(makeRequest(null))).rejects.toThrow( + "Pub/Sub unavailable", + ); - expect(values).toHaveBeenCalled(); + expect(values).not.toHaveBeenCalled(); expect(logger.error).toHaveBeenCalledWith( "Failed to publish log to Pub/Sub", expect.objectContaining({ error: expect.any(Error) }), @@ -678,16 +681,22 @@ describe("logRequest", () => { ); }); - it("does not hold the caller on a slow publish", async () => { + it("holds the caller until the publish is acknowledged", async () => { const publication = deferred(); publishMessage.mockReturnValueOnce(publication.promise); + let finished = false; + const logging = logRequest(makeRequest(null)).then(() => { + finished = true; + }); - await expect(logRequest(makeRequest(null))).resolves.toBeUndefined(); - expect(values).toHaveBeenCalled(); - expect(metricInc).not.toHaveBeenCalled(); + await new Promise(resolve => setImmediate(resolve)); + expect(publishMessage).toHaveBeenCalledOnce(); + expect(values).not.toHaveBeenCalled(); + expect(finished).toBe(false); publication.resolve("message-id"); - await publication.promise; + await logging; + expect(values).toHaveBeenCalled(); expect(metricInc).toHaveBeenCalledWith({ table: "requests", outcome: "published", @@ -759,31 +768,42 @@ describe("logRequest", () => { expect(published.origin).toBe("api"); }); - it("drops publishes beyond the outstanding cap instead of queueing them", async () => { + it("refuses a publish beyond the outstanding cap and fails that log call", async () => { vi.resetModules(); const fresh = await import("./log_job.js"); const publication = deferred(); publishMessage.mockImplementation(async () => publication.promise); config.PUBSUB_MAX_OUTSTANDING_MESSAGES = 2; try { - await fresh.logRequest(makeRequest(null)); - await fresh.logRequest(makeRequest(null)); - await fresh.logRequest(makeRequest(null)); + const first = fresh.logRequest(makeRequest(null)); + const second = fresh.logRequest(makeRequest(null)); + await new Promise(resolve => setImmediate(resolve)); + await expect(fresh.logRequest(makeRequest(null))).rejects.toThrow( + "backlog is full", + ); + publication.resolve("message-id"); + await Promise.all([first, second]); } finally { config.PUBSUB_MAX_OUTSTANDING_MESSAGES = 10_000; publication.resolve("message-id"); await fresh.shutdownPubSubLogging(); } - // The database write is never held back by the publisher. - expect(values).toHaveBeenCalledTimes(3); + // The refused row never reaches PostgreSQL either. + expect(values).toHaveBeenCalledTimes(2); expect(publishMessage).toHaveBeenCalledTimes(2); expect(metricInc).toHaveBeenCalledWith({ table: "requests", outcome: "dropped", }); + // A refusal is the rate-limited warning, not a per-row error. + expect(logger.error).not.toHaveBeenCalled(); + expect(metricInc).not.toHaveBeenCalledWith({ + table: "requests", + outcome: "failed", + }); expect(logger.warn).toHaveBeenCalledWith( - "Dropping Pub/Sub log: publisher backlog is full", + "Refusing Pub/Sub log: publisher backlog is full", expect.objectContaining({ table: "requests", logId: makeRequest(null).id, @@ -802,8 +822,9 @@ describe("logRequest", () => { }); publishMessage.mockRejectedValueOnce(new Error("Pub/Sub unavailable")); - await logRequest(makeRequest(null)); - await new Promise(resolve => setImmediate(resolve)); + await expect(logRequest(makeRequest(null))).rejects.toThrow( + "Pub/Sub unavailable", + ); expect(metricInc).toHaveBeenCalledWith({ table: "requests", outcome: "failed", @@ -815,7 +836,9 @@ describe("logRequest", () => { config.PUBSUB_MAX_OUTSTANDING_MESSAGES = 1; try { publishMessage.mockRejectedValueOnce(new Error("Pub/Sub unavailable")); - await logRequest(makeRequest(null)); + await expect(logRequest(makeRequest(null))).rejects.toThrow( + "Pub/Sub unavailable", + ); await logRequest(makeRequest(null)); expect(publishMessage).toHaveBeenCalledTimes(2); expect(metricInc).not.toHaveBeenCalledWith({ @@ -854,7 +877,8 @@ describe("shutdownPubSubLogging deadline", () => { const fresh = await import("./log_job.js"); const publication = deferred(); publishMessage.mockReturnValueOnce(publication.promise); - await fresh.logSearch(makeSearch()); + const logging = fresh.logSearch(makeSearch()); + await new Promise(resolve => setImmediate(resolve)); const shutdown = fresh.shutdownPubSubLogging(); await new Promise(resolve => setImmediate(resolve)); @@ -863,6 +887,7 @@ describe("shutdownPubSubLogging deadline", () => { publication.resolve("message-id"); await shutdown; + await logging; expect(close).toHaveBeenCalledOnce(); }); @@ -871,7 +896,8 @@ describe("shutdownPubSubLogging deadline", () => { const fresh = await import("./log_job.js"); const publication = deferred(); publishMessage.mockReturnValueOnce(publication.promise); - await fresh.logSearch(makeSearch()); + const logging = fresh.logSearch(makeSearch()); + await new Promise(resolve => setImmediate(resolve)); vi.useFakeTimers(); const shutdown = fresh.shutdownPubSubLogging(); @@ -895,7 +921,7 @@ describe("shutdownPubSubLogging deadline", () => { ); publication.resolve("message-id"); - await publication.promise; + await logging; }); it("reports publication failures during shutdown and still closes", async () => { @@ -903,11 +929,13 @@ describe("shutdownPubSubLogging deadline", () => { const fresh = await import("./log_job.js"); const publication = deferred(); publishMessage.mockReturnValueOnce(publication.promise); - await fresh.logSearch(makeSearch()); + const logging = fresh.logSearch(makeSearch()); + await new Promise(resolve => setImmediate(resolve)); const shutdown = fresh.shutdownPubSubLogging(); publication.reject(new Error("Pub/Sub unavailable")); await shutdown; + await expect(logging).rejects.toThrow("Pub/Sub unavailable"); expect(close).toHaveBeenCalledOnce(); expect(metricInc).toHaveBeenCalledWith({ table: "searches", @@ -938,9 +966,11 @@ describe("shutdownPubSubLogging deadline", () => { flush.mockReturnValueOnce(flushing.promise); const shutdown = fresh.shutdownPubSubLogging(); - await fresh.logSearch(makeSearch()); + await expect(fresh.logSearch(makeSearch())).rejects.toThrow( + "shutting down", + ); expect(publishMessage).toHaveBeenCalledOnce(); - expect(values).toHaveBeenCalledTimes(2); + expect(values).toHaveBeenCalledTimes(1); expect(metricInc).toHaveBeenCalledWith({ table: "searches", outcome: "failed", diff --git a/apps/api/src/services/logging/log_job.ts b/apps/api/src/services/logging/log_job.ts index d4c57370bd8..4666044c676 100644 --- a/apps/api/src/services/logging/log_job.ts +++ b/apps/api/src/services/logging/log_job.ts @@ -151,17 +151,20 @@ let pubSubShutdown: Promise | undefined; // A retry can deliver a batch twice when the first attempt was persisted but // its response was lost; the ClickHouse tables dedupe on row id, not message // id, so those copies collapse. +// +// The log call waits for the publish, so the budget is what a caller can be +// held for while Pub/Sub is unreachable before the log call fails. const PUBSUB_PUBLISH_OPTIONS: PublishOptions = { gaxOpts: { retry: { backoffSettings: { initialRetryDelayMillis: 250, retryDelayMultiplier: 2, - maxRetryDelayMillis: 15_000, - initialRpcTimeoutMillis: 15_000, + maxRetryDelayMillis: 5_000, + initialRpcTimeoutMillis: 10_000, rpcTimeoutMultiplier: 1, - maxRpcTimeoutMillis: 15_000, - totalTimeoutMillis: 300_000, + maxRpcTimeoutMillis: 10_000, + totalTimeoutMillis: 30_000, }, }, }, @@ -212,6 +215,20 @@ function getTopic(client: PubSub, table: string): Topic { return topic; } +/** The publisher already holds as many unacknowledged rows as it may. */ +class PubSubBacklogFullError extends Error { + constructor() { + super("Pub/Sub log publisher backlog is full"); + this.name = "PubSubBacklogFullError"; + } +} + +/** + * Publishes one job-log row to its Pub/Sub topic and waits for the + * acknowledgement. Pub/Sub is the log's store of record, so a publish that + * fails, is refused by a full backlog, or arrives during shutdown throws to + * the caller; nothing else is written for that row. + */ async function publishLog(table: string, data: any, logger: Logger) { const startedAt = Date.now(); try { @@ -252,7 +269,7 @@ async function publishLog(table: string, data: any, logger: Logger) { const now = Date.now(); if (now - lastDropWarningAt >= 60_000) { lastDropWarningAt = now; - logger.warn("Dropping Pub/Sub log: publisher backlog is full", { + logger.warn("Refusing Pub/Sub log: publisher backlog is full", { table, logId: data.id, payloadBytes: payload.length, @@ -261,7 +278,7 @@ async function publishLog(table: string, data: any, logger: Logger) { droppedTotal, }); } - return; + throw new PubSubBacklogFullError(); } const publication = getTopic(client, table).publishMessage({ @@ -284,14 +301,18 @@ async function publishLog(table: string, data: any, logger: Logger) { } }); } catch (error) { - pubsubLogPublishTotal.inc({ table, outcome: "failed" }); - - logger.error("Failed to publish log to Pub/Sub", { - error, - table, - logId: data.id, - durationMs: Date.now() - startedAt, - }); + // A backlog refusal counted itself as "dropped" and warned, rate-limited, + // above; only a genuine publish failure is an error here. + if (!(error instanceof PubSubBacklogFullError)) { + pubsubLogPublishTotal.inc({ table, outcome: "failed" }); + logger.error("Failed to publish log to Pub/Sub", { + error, + table, + logId: data.id, + durationMs: Date.now() - startedAt, + }); + } + throw error; } } @@ -388,37 +409,36 @@ async function robustInsert( canonicalLog: "log_job/robustInsert", }); + if (config.USE_DB_AUTHENTICATION !== true) { + logger.info( + "Skipping database insertion due to USE_DB_AUTHENTICATION being off", + ); + return; + } + + // The single point where a row leaves for both stores: clean it once so + // PostgreSQL and ClickHouse receive identical, accepted values. + data = sanitizeLogData({ + ...data, + created_at: data.created_at ?? new Date(), + }); + + // Pub/Sub is the store of record: publish first and wait for it. A failed + // publish fails the log call, and the PostgreSQL copy is not attempted. + await publishLog(table, data, logger); + + // The PostgreSQL copy is on its way out; its failure is logged, not thrown. const attempts: { error: any; timeMs: number; backoffMs: number }[] = []; try { - const inserted = await withSpan("log_job.postgres.insert", async span => { + await withSpan("log_job.postgres.insert", async span => { setSpanAttributes(span, { "db.system": "postgresql", "log_job.table": table, "log_job.id": data.id, "log_job.force": force, + "log_job.postgres.enabled": true, }); - - if (config.USE_DB_AUTHENTICATION !== true) { - logger.info( - "Skipping database insertion due to USE_DB_AUTHENTICATION being off", - ); - setSpanAttributes(span, { - "log_job.postgres.enabled": false, - "log_job.postgres.outcome": "skipped", - }); - return false; - } - - setSpanAttributes(span, { "log_job.postgres.enabled": true }); const target = tableMap[table]; - // The single point where a row leaves for both stores: clean it once so - // PostgreSQL and ClickHouse receive identical, accepted values. - data = sanitizeLogData({ - ...data, - created_at: data.created_at ?? new Date(), - }); - // Publish in the background. Customer responses must not wait for Pub/Sub. - void publishLog(table, data, logger); const maxAttempts = force ? 10 : 1; for (let i = 0; i < maxAttempts; i++) { @@ -456,10 +476,8 @@ async function robustInsert( lastAttempt?.error ?? new Error("Database insert was not attempted") ); } - return true; }); - if (!inserted) return; if (attempts.length === 1) { logger.debug("Inserted into database successfully", { attempts }); } else {