diff --git a/application/single_app/app.py b/application/single_app/app.py index 05fa0c8de..faf66d428 100644 --- a/application/single_app/app.py +++ b/application/single_app/app.py @@ -286,7 +286,7 @@ def start_background_tasks(): print("Background tasks disabled for this web process.") _background_tasks_started = True return - start_background_task_threads() + start_background_task_threads(app=app) _background_tasks_started = True diff --git a/application/single_app/background_tasks.py b/application/single_app/background_tasks.py index 464255dd1..aee7ef73e 100644 --- a/application/single_app/background_tasks.py +++ b/application/single_app/background_tasks.py @@ -684,14 +684,14 @@ def run_tabular_generated_output_scheduler_loop(): time.sleep(30) -def run_data_management_scheduler_loop(): - """Queue due Data Management backup jobs across scaled-out workers.""" +def run_data_management_scheduler_loop(app=None): + """Queue due backup and recoverable migration jobs across scaled-out workers.""" while True: lock_document = None try: lock_document = acquire_distributed_task_lock('data_management_scheduler_scan', lease_seconds=300) if lock_document: - check_due_data_management_jobs_once() + check_due_data_management_jobs_once(app=app) except Exception as exc: print(f"Error in Data Management scheduler check: {exc}") log_event(f"[DataManagement] Error in scheduler check: {exc}", level=logging.ERROR) @@ -736,7 +736,7 @@ def run_app_maintenance_loop(): time.sleep(max(int(sleep_seconds or 3600), 15)) -def start_background_task_threads(): +def start_background_task_threads(app=None): """Start all background task loops for the current process.""" task_specs = [ ('Logging timer background task started.', run_logging_timer_loop), @@ -747,7 +747,7 @@ def start_background_task_threads(): ('Workflow scheduler background task started.', run_workflow_scheduler_loop), ('File Sync scheduler background task started.', run_file_sync_scheduler_loop), ('Tabular generated-output scheduler background task started.', run_tabular_generated_output_scheduler_loop), - ('Data Management scheduler background task started.', run_data_management_scheduler_loop), + ('Data Management scheduler background task started.', lambda: run_data_management_scheduler_loop(app=app)), ('App maintenance background task started.', run_app_maintenance_loop), ] diff --git a/application/single_app/config.py b/application/single_app/config.py index 4c9b91a3e..33f12cf1e 100644 --- a/application/single_app/config.py +++ b/application/single_app/config.py @@ -97,7 +97,7 @@ EXECUTOR_TYPE = 'thread' EXECUTOR_MAX_WORKERS = 30 SESSION_TYPE = 'filesystem' -VERSION = "0.250.059" +VERSION = "0.250.078" IS_DEVELOPMENT = is_development_env_enabled() SESSION_COOKIE_SAMESITE = os.getenv('SESSION_COOKIE_SAMESITE', 'Lax') diff --git a/application/single_app/foundry_agent_runtime.py b/application/single_app/foundry_agent_runtime.py index 56189993e..4de9fc2be 100644 --- a/application/single_app/foundry_agent_runtime.py +++ b/application/single_app/foundry_agent_runtime.py @@ -42,7 +42,13 @@ FOUNDRY_INTERNAL_METADATA_KEYS = { "active_group_ids", "active_public_workspace_ids", + "document_context_requested", + "document_scope", + "group_id", + "selection_mode", + "selected_document_id", "selected_document_ids", + "user_id", } FOUNDRY_FILE_SEARCHABLE_CONTEXT_MAX_CHARS = 6000 FOUNDRY_FILE_SEARCHABLE_CONTEXT_HEADER = "Attached file searchable summary" @@ -391,6 +397,14 @@ async def execute_foundry_agent( ) -> FoundryAgentInvocationResult: """Invoke a Foundry agent using Semantic Kernel's AzureAIAgent abstraction.""" + message_history = _filter_foundry_document_context_messages( + message_history, + include_document_context=_coerce_bool( + foundry_settings.get("include_document_context"), + True, + ), + ) + agent_id = (foundry_settings.get("agent_id") or "").strip() if not agent_id: raise FoundryAgentInvocationError( @@ -491,6 +505,14 @@ async def execute_new_foundry_agent( ) -> FoundryAgentInvocationResult: """Invoke the new Foundry application runtime through its Responses protocol endpoint.""" + message_history = _filter_foundry_document_context_messages( + message_history, + include_document_context=_coerce_bool( + foundry_settings.get("include_document_context"), + True, + ), + ) + application_name = _resolve_new_foundry_application_name(foundry_settings) endpoint = _resolve_endpoint(foundry_settings, global_settings) responses_api_version = ( @@ -573,6 +595,14 @@ async def execute_new_foundry_agent_stream( ) -> AsyncIterator[FoundryAgentStreamMessage]: """Stream a new Foundry application response through the Responses API.""" + message_history = _filter_foundry_document_context_messages( + message_history, + include_document_context=_coerce_bool( + foundry_settings.get("include_document_context"), + True, + ), + ) + application_name = _resolve_new_foundry_application_name(foundry_settings) endpoint = _resolve_endpoint(foundry_settings, global_settings) responses_api_version = ( @@ -1222,10 +1252,45 @@ def _looks_like_document_context_message(text: str) -> bool: "chat-uploaded file", "selected document", "tabular analysis", + "mixed-source evidence handoff", + "evidence_envelopes", + "[workflow document search context]", ) return any(marker in normalized for marker in markers) +def _filter_foundry_document_context_messages( + message_history: List[ChatMessageContent], + include_document_context: bool = True, +) -> List[ChatMessageContent]: + """Remove document/evidence messages before Foundry transport when opted out.""" + if include_document_context: + return list(message_history or []) + + filtered_messages: List[ChatMessageContent] = [] + workflow_task_marker = "[workflow task]" + for message in list(message_history or []): + text = _extract_message_text(message).strip() + if not text: + continue + if not _looks_like_document_context_message(text): + filtered_messages.append(message) + continue + + marker_index = text.lower().rfind(workflow_task_marker) + if marker_index < 0: + continue + workflow_task = text[marker_index + len(workflow_task_marker):].strip() + if not workflow_task: + continue + filtered_messages.append(ChatMessageContent( + role=getattr(message, "role", "user"), + content=workflow_task, + metadata=getattr(message, "metadata", {}) or {}, + )) + return filtered_messages + + def _build_foundry_workflow_input_text( message_history: List[ChatMessageContent], max_context_chars: Optional[int] = None, @@ -1237,7 +1302,13 @@ def _build_foundry_workflow_input_text( if not text: continue if not include_document_context and _looks_like_document_context_message(text): - continue + workflow_task_marker = "[Workflow task]" + workflow_task_index = text.rfind(workflow_task_marker) + if workflow_task_index < 0: + continue + text = text[workflow_task_index + len(workflow_task_marker):].strip() + if not text: + continue role_value = getattr(message, "role", "user") role = str(role_value).strip().lower() or "user" if role.startswith("authorrole."): @@ -1800,6 +1871,8 @@ def _collect_foundry_response_file_inputs( foundry_settings: Dict[str, Any], metadata: Dict[str, Any], ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + if not _coerce_bool(foundry_settings.get("include_document_context"), True): + return [], [] if not _coerce_bool(foundry_settings.get("include_file_inputs"), True): return [], [] diff --git a/application/single_app/functions_assistant_table_exports.py b/application/single_app/functions_assistant_table_exports.py index 67bc76f54..b7cab0165 100644 --- a/application/single_app/functions_assistant_table_exports.py +++ b/application/single_app/functions_assistant_table_exports.py @@ -9,6 +9,58 @@ ASSISTANT_TABLE_EXPORT_PREVIEW_ROWS = 3 +CSV_FENCE_LANGUAGES = {'', 'csv', 'md', 'markdown', 'plaintext', 'text', 'text/csv', 'txt'} +CSV_OUTPUT_REQUEST_PATTERNS = ( + re.compile( + r'\b(?:build|create|download|export|generate|make|prepare|save)\b' + r'.{0,120}\b(?:a\s+)?(?:(?:single|combined|one)\s+)?csv(?:\s+(?:file|format|output))?\b' + ), + re.compile( + r'\b(?:respond|return|provide|output)\b' + r'(?:(?!\bfrom\b).){0,80}\b(?:as|in|to\s+)?(?:a\s+)?(?:(?:single|combined|one)\s+)?csv\b' + ), + re.compile( + r'\b(?:convert|format|put|turn)\b.{0,80}\b(?:as|in|into|to)\s+(?:a\s+)?(?:(?:single|combined|one)\s+)?csv\b' + ), + re.compile(r'\b(?:get|give)\s+(?:me\s+)?(?:a|the|one)\s+(?:(?:single|combined)\s+)?csv\b'), + re.compile( + r'\b(?:need|want)\s+(?:(?:a|the|one)\s+)?(?:direct\s+)?(?:(?:single|combined)\s+)?csv' + r'(?:\s+(?:file|format|output))?\b' + ), + re.compile( + r'\b(?:need|want)\b.{0,40}\b(?:results?|output|answer|response|fields?|rows?|data|this|that|it)\b' + r'.{0,40}\b(?:as|in)\s+(?:(?:single|combined|one)\s+)?csv\b' + ), + re.compile(r'\b(?:single|combined|one)\s+csv(?:\s+(?:file|format|output))?\b'), + re.compile(r'\bcsv\s+(?:output|version)\b'), + re.compile(r'^\s*(?:a\s+)?csv\s+file\s*(?:,?\s*please)?[.!?]?\s*$'), + re.compile( + r'\b(?:build|create|download|export|generate|make|prepare|save|turn)\b' + r'.{0,80}\bspreadsheet\b' + ), +) +CSV_PROSE_PREFIXES = ( + 'below ', + 'csv data', + 'for clarity', + 'for context', + 'here are ', + 'here is ', + 'i extracted ', + 'i found ', + 'in summary', + 'sure ', + 'the following ', + 'the requested ', + 'the export ', + 'this csv ', + 'this export ', + 'your csv ', +) +SIGNED_NUMBER_PATTERN = re.compile( + r'[+-]?(?:(?:\d{1,3}(?:,\d{3})+|\d+)(?:\.\d*)?|\.\d+)(?:e[+-]?\d+)?', + flags=re.IGNORECASE, +) TABLE_EXPORT_REQUEST_MARKERS = ( 'turn that into a csv', @@ -85,25 +137,83 @@ 'table for me', 'in table format', 'download table', - 'table file', - 'spreadsheet', - 'csv file', 'download csv', 'save csv', 'make a csv', 'make csv', 'create a csv', 'create csv', + 'csv version', ) +CSV_EXPLICIT_ROW_SCHEMA_PATTERNS = ( + re.compile(r'\b(?:one|a|each)\s+(?:row|record|line|entry|object)\s+(?:per|for|of)\b'), + re.compile(r'\b(?:one|a|each)\s+(?:file|document|source|record|item)\s+per\s+row\b'), + re.compile(r'\b(?:columns?|fields?)\s*(?::|=|are|should|must|include|with)\b'), + re.compile(r'\b(?:include|with|using)\s+(?:the\s+)?(?:columns?|fields?)\b'), +) + def assistant_table_export_requested(user_question: str) -> bool: """Return True when the user asked for table-shaped output or a CSV export.""" normalized_question = re.sub(r'\s+', ' ', str(user_question or '').strip().casefold()) if not normalized_question: return False - return any(marker in normalized_question for marker in TABLE_EXPORT_REQUEST_MARKERS) + csv_markers = tuple(marker for marker in TABLE_EXPORT_REQUEST_MARKERS if 'csv' in marker) + table_markers = tuple(marker for marker in TABLE_EXPORT_REQUEST_MARKERS if 'csv' not in marker) + question_clauses = [ + clause.strip() + for clause in re.split(r'[;.!?]+', normalized_question) + if clause.strip() + ] + + for question_clause in question_clauses: + if 'csv' not in question_clause: + continue + csv_request_negated = bool(re.search( + r"\b(?:do\s+not|don't|dont|never)\b.{0,80}\bcsv\b" + r'|\b(?:not|no|without)\s+(?:a\s+)?csv\b', + question_clause, + )) + if csv_request_negated: + continue + if any(marker in question_clause for marker in csv_markers): + return True + if any(pattern.search(question_clause) for pattern in CSV_OUTPUT_REQUEST_PATTERNS): + return True + + return ( + any(marker in normalized_question for marker in table_markers) + or any( + pattern.search(normalized_question) + for pattern in CSV_OUTPUT_REQUEST_PATTERNS + if 'spreadsheet' in pattern.pattern + ) + ) + + +def build_csv_output_clarification_guidance(user_question: str) -> str: + """Return model guidance for a CSV request that may need one schema clarification.""" + if not assistant_table_export_requested(user_question): + return '' + + normalized_question = re.sub(r'\s+', ' ', str(user_question or '').strip().casefold()) + if any(pattern.search(normalized_question) for pattern in CSV_EXPLICIT_ROW_SCHEMA_PATTERNS): + return ( + 'The user explicitly specified CSV row or column structure. Preserve that structure, ' + 'produce valid structured rows, and do not ask a clarification unless the requested ' + 'evidence itself is contradictory.' + ) + + return ( + 'The user requested a CSV artifact. If the authorized evidence and request do not establish ' + 'one stable row unit and column schema, ask exactly one concise clarification before generating ' + 'a file: whether each row should represent files, documents, or extracted records, and which ' + 'columns to include. Do not create an empty or guessed CSV. If the schema is clear from the ' + 'request or evidence, generate valid structured rows directly. If this conversation already ' + 'contains that clarification, use the user\'s latest answer instead of asking again.' + ) def build_assistant_table_csv_export(user_question: str, assistant_content: str) -> Optional[Dict[str, Any]]: @@ -129,15 +239,41 @@ def build_assistant_table_csv_export(user_question: str, assistant_content: str) } +def has_generated_tabular_csv_output(generated_outputs: List[Dict[str, Any]]) -> bool: + """Return whether a tabular CSV result already suppresses a duplicate table export.""" + for generated_output in generated_outputs or []: + if not isinstance(generated_output, dict): + continue + + capability = str(generated_output.get('capability') or '').strip().lower() + if generated_output.get('suppress_assistant_table_export') and ( + not capability or capability == 'tabular' + ): + return True + output_format = str(generated_output.get('output_format') or '').strip().lower() + file_name = str(generated_output.get('file_name') or '').strip().lower() + if (output_format == 'csv' or file_name.endswith('.csv')) and ( + not capability or capability == 'tabular' + ): + return True + + return False + + def extract_assistant_table_entries(assistant_content: str) -> List[Dict[str, str]]: - """Extract table rows from Markdown pipe tables or tab-separated assistant output.""" + """Extract table rows from Markdown, tab-separated, or CSV assistant output.""" normalized_content = str(assistant_content or '').replace('\r\n', '\n').replace('\r', '\n') if not normalized_content.strip(): return [] + fenced_csv_rows = _extract_fenced_comma_separated_table_entries(normalized_content) + if fenced_csv_rows: + return fenced_csv_rows + candidates = [ _extract_markdown_table_entries(normalized_content), _extract_tab_separated_table_entries(normalized_content), + _extract_comma_separated_table_entries(normalized_content), ] return max(candidates, key=len, default=[]) @@ -159,18 +295,44 @@ def build_assistant_table_csv(table_rows: List[Dict[str, Any]]) -> str: if not ordered_columns: ordered_columns = ['value'] + safe_columns = build_safe_csv_headers(ordered_columns) + output_buffer = io.StringIO() - writer = csv.DictWriter(output_buffer, fieldnames=ordered_columns, extrasaction='ignore') + writer = csv.DictWriter(output_buffer, fieldnames=safe_columns, extrasaction='ignore') writer.writeheader() for table_row in table_rows or []: serialized_row = {} if isinstance(table_row, dict): - for column_name in ordered_columns: - serialized_row[column_name] = _serialize_table_cell(table_row.get(column_name)) + for source_column, safe_column in zip(ordered_columns, safe_columns): + serialized_row[safe_column] = _serialize_table_cell(table_row.get(source_column)) writer.writerow(serialized_row) return output_buffer.getvalue() +def build_safe_csv_headers(header_cells: List[Any]) -> List[str]: + """Return non-empty, formula-safe, unique CSV headers in source order.""" + headers = [] + seen_headers = set() + for index, header_cell in enumerate(header_cells or []): + base_header = neutralize_csv_spreadsheet_formula(_clean_table_cell(header_cell)) or f'Column {index + 1}' + header = base_header + occurrence_count = 2 + while header.casefold() in seen_headers: + header = f'{base_header} {occurrence_count}' + occurrence_count += 1 + seen_headers.add(header.casefold()) + headers.append(header) + return headers + + +def neutralize_csv_spreadsheet_formula(value: Any) -> str: + """Prefix spreadsheet formula-like text while preserving signed numbers.""" + serialized_value = '' if value is None else str(value) + if _spreadsheet_formula_candidate(serialized_value): + return f"'{serialized_value}" + return serialized_value + + def _extract_markdown_table_entries(content: str) -> List[Dict[str, str]]: lines = content.split('\n') best_entries = [] @@ -215,6 +377,198 @@ def _extract_tab_separated_table_entries(content: str) -> List[Dict[str, str]]: return best_entries +def _extract_comma_separated_table_entries(content: str) -> List[Dict[str, str]]: + fenced_candidates = [] + unfenced_sections = [] + previous_end = 0 + fence_pattern = re.compile( + r'```[ \t]*(?P[^\n`]*)\n(?P.*?)```', + flags=re.IGNORECASE | re.DOTALL, + ) + + for fence_match in fence_pattern.finditer(content): + unfenced_sections.append(content[previous_end:fence_match.start()]) + language = str(fence_match.group('language') or '').strip().casefold() + if language in CSV_FENCE_LANGUAGES: + fenced_candidates.extend(_parse_csv_table_candidates(fence_match.group('body'), trusted=True)) + previous_end = fence_match.end() + unfenced_sections.append(content[previous_end:]) + + if fenced_candidates: + return max(fenced_candidates, key=len, default=[]) + + candidates = [] + for section in unfenced_sections: + candidates.extend(_parse_csv_table_candidates(section, trusted=False)) + + return max(candidates, key=len, default=[]) + + +def _extract_fenced_comma_separated_table_entries(content: str) -> List[Dict[str, str]]: + explicit_csv_candidates = [] + generic_fenced_candidates = [] + fence_pattern = re.compile( + r'```[ \t]*(?P[^\n`]*)\n(?P.*?)```', + flags=re.IGNORECASE | re.DOTALL, + ) + for fence_match in fence_pattern.finditer(content): + language = str(fence_match.group('language') or '').strip().casefold() + if language in CSV_FENCE_LANGUAGES: + parsed_candidates = _parse_csv_table_candidates(fence_match.group('body'), trusted=True) + if language in {'csv', 'text/csv'}: + explicit_csv_candidates.extend(parsed_candidates) + else: + generic_fenced_candidates.extend(parsed_candidates) + if explicit_csv_candidates: + return max(explicit_csv_candidates, key=len, default=[]) + return max(generic_fenced_candidates, key=len, default=[]) + + +def _parse_csv_table_candidates(content: str, trusted: bool = False) -> List[List[Dict[str, str]]]: + if ',' not in str(content or ''): + return [] + + try: + parsed_rows = list(csv.reader(io.StringIO(content), strict=True)) + except csv.Error: + return [] + + candidates = [] + current_rows = [] + current_width = None + for parsed_row in parsed_rows: + if not trusted and _is_csv_source_citation_row(parsed_row): + candidate_rows = _trim_trailing_csv_narration_rows(current_rows) + if len(candidate_rows) >= 2: + candidates.append(_build_csv_table_entries(candidate_rows, trusted=trusted)) + current_rows = [] + current_width = None + continue + + row_width = len(parsed_row) + if row_width < 2: + if len(current_rows) >= 2: + candidates.append(_build_csv_table_entries(current_rows, trusted=trusted)) + current_rows = [] + current_width = None + continue + + if current_width is not None and row_width != current_width: + if len(current_rows) >= 2: + candidates.append(_build_csv_table_entries(current_rows, trusted=trusted)) + current_rows = [] + + current_rows.append(parsed_row) + current_width = row_width + + if len(current_rows) >= 2: + candidates.append(_build_csv_table_entries(current_rows, trusted=trusted)) + + return [candidate for candidate in candidates if candidate] + + +def _build_csv_table_entries(parsed_rows: List[List[str]], trusted: bool = False) -> List[Dict[str, str]]: + if len(parsed_rows) < 2: + return [] + + if trusted: + header_index = 0 + else: + header_index = next( + ( + row_index + for row_index, parsed_row in enumerate(parsed_rows[:-1]) + if _is_likely_csv_header_row(parsed_row) + ), + None, + ) + if header_index is None: + return [] + + table_rows = parsed_rows[header_index:] + headers = _build_unique_headers(table_rows[0]) + if len(headers) < 2: + return [] + + entries = [] + for parsed_row in table_rows[1:]: + if not trusted and _is_csv_source_citation_row(parsed_row): + break + normalized_row = _coerce_row_length(parsed_row, len(headers)) + if not any(str(cell or '').strip() for cell in normalized_row): + continue + entries.append({ + header: _clean_csv_cell(normalized_row[index]) + for index, header in enumerate(headers) + }) + return entries + + +def _is_likely_csv_header_row(parsed_row: List[str]) -> bool: + cleaned_cells = [_clean_csv_cell(cell) for cell in parsed_row] + if len(cleaned_cells) < 2 or any(not cell for cell in cleaned_cells): + return False + + if _is_likely_csv_preamble_row(parsed_row) or _is_csv_source_citation_row(parsed_row): + return False + + return all('\n' not in cell for cell in cleaned_cells) + + +def _is_likely_csv_discourse_row(parsed_row: List[str]) -> bool: + if not parsed_row: + return False + + first_cell = _clean_csv_cell(parsed_row[0]).casefold() + return first_cell.startswith(CSV_PROSE_PREFIXES) + + +def _is_likely_csv_preamble_row(parsed_row: List[str]) -> bool: + if _is_likely_csv_discourse_row(parsed_row): + return True + + cleaned_cells = [_clean_csv_cell(cell) for cell in parsed_row] + combined_text = ', '.join(cell for cell in cleaned_cells if cell) + word_count = len(re.findall(r"\b[\w'-]+\b", combined_text)) + return ( + len(cleaned_cells) == 2 + and word_count >= 5 + and cleaned_cells[-1].rstrip().endswith(('.', '!')) + and not any(re.search(r'\d', cell) for cell in cleaned_cells) + ) + + +def _trim_trailing_csv_narration_rows(parsed_rows: List[List[str]]) -> List[List[str]]: + trimmed_rows = list(parsed_rows or []) + while len(trimmed_rows) > 1 and _is_likely_csv_trailing_narration_row(trimmed_rows[-1]): + trimmed_rows.pop() + return trimmed_rows + + +def _is_likely_csv_trailing_narration_row(parsed_row: List[str]) -> bool: + cleaned_cells = [_clean_csv_cell(cell) for cell in parsed_row] + if len(cleaned_cells) != 2 or not _is_likely_csv_discourse_row(parsed_row): + return False + + combined_text = ', '.join(cell for cell in cleaned_cells if cell) + word_count = len(re.findall(r"\b[\w'-]+\b", combined_text)) + return word_count >= 5 and cleaned_cells[-1].rstrip().endswith(('.', '!')) + + +def _is_csv_source_citation_row(parsed_row: List[str]) -> bool: + if not parsed_row: + return False + + first_cell = _clean_csv_cell(parsed_row[0]) + return bool(re.match( + r'^\s*(?:[-*+>\u2022\u25e6\u25aa\u2023]\s*)?' + r'(?:(?:\[\s*\d+\s*\]|\(\s*\d+\s*\)|\d+[.)])\s*)?' + r'[\[(]?\s*(?:sources?|citations?)\s*[\])]?\s*:', + first_cell, + flags=re.IGNORECASE, + )) + + def _parse_markdown_table_block(table_block: List[str]) -> List[Dict[str, str]]: split_rows = [ _split_markdown_table_line(line) @@ -321,17 +675,7 @@ def _is_markdown_separator_row(row: List[str]) -> bool: def _build_unique_headers(header_cells: List[str]) -> List[str]: - headers = [] - seen_headers = {} - for index, header_cell in enumerate(header_cells or []): - header = _clean_table_cell(header_cell) or f'Column {index + 1}' - normalized_header = header.casefold() - occurrence_count = seen_headers.get(normalized_header, 0) - seen_headers[normalized_header] = occurrence_count + 1 - if occurrence_count: - header = f'{header} {occurrence_count + 1}' - headers.append(header) - return headers + return build_safe_csv_headers(header_cells) def _coerce_row_length(row: List[str], target_length: int) -> List[str]: @@ -353,12 +697,25 @@ def _clean_table_cell(value: Any) -> str: return cleaned +def _clean_csv_cell(value: Any) -> str: + return str(value or '').replace('\r\n', '\n').replace('\r', '\n').strip() + + def _serialize_table_cell(value: Any) -> str: if value is None: return '' if isinstance(value, (dict, list)): - return str(value) - return str(value) + return neutralize_csv_spreadsheet_formula(str(value)) + return neutralize_csv_spreadsheet_formula(str(value)) + + +def _spreadsheet_formula_candidate(value: Any) -> bool: + normalized_value = str(value or '').lstrip() + if not normalized_value or normalized_value[0] not in ('=', '+', '-', '@'): + return False + if normalized_value[0] in ('+', '-') and SIGNED_NUMBER_PATTERN.fullmatch(normalized_value): + return False + return True def _build_assistant_table_export_file_name() -> str: diff --git a/application/single_app/functions_conversation_metadata.py b/application/single_app/functions_conversation_metadata.py index efd850e24..7bfe5cce4 100644 --- a/application/single_app/functions_conversation_metadata.py +++ b/application/single_app/functions_conversation_metadata.py @@ -139,8 +139,34 @@ def _extract_document_id_from_search_result(doc): return chunk_identifier -def _build_last_grounded_document_refs(document_map): +def _build_last_grounded_document_refs(document_map, source_continuity_refs=None): """Build the exact reusable grounded document set for the latest search-backed turn.""" + if isinstance(source_continuity_refs, list): + safe_continuity_refs = [] + for raw_ref in source_continuity_refs[:100]: + if not isinstance(raw_ref, dict): + continue + document_id = str(raw_ref.get('document_id') or '').strip() + scope = str(raw_ref.get('scope') or '').strip().lower() + scope_id = str(raw_ref.get('scope_id') or '').strip() + if not document_id or not scope or not scope_id: + continue + safe_ref = { + key: raw_ref.get(key) + for key in ( + 'document_id', 'scope', 'scope_id', 'group_id', + 'public_workspace_id', 'user_id', 'source_role', + 'requested_order', 'source_kind', 'engine', + 'source_version', 'status', 'coverage', + 'selection_origin', 'action_mode', + 'citation_count', 'artifact_count', + ) + if raw_ref.get(key) is not None + } + safe_continuity_refs.append(safe_ref) + if safe_continuity_refs: + return safe_continuity_refs + grounded_refs = [] for document_id, doc_info in document_map.items(): @@ -273,7 +299,8 @@ def collect_conversation_metadata(user_message, conversation_id, user_id, active image_gen_enabled=False, selected_documents=None, selected_agent=None, selected_agent_details=None, search_results=None, web_search_results=None, conversation_item=None, additional_participants=None, - active_group_ids=None, active_public_workspace_id=None, active_public_workspace_ids=None): + active_group_ids=None, active_public_workspace_id=None, active_public_workspace_ids=None, + source_continuity_refs=None): """ Collect comprehensive metadata for a conversation based on the user's interaction. @@ -736,8 +763,11 @@ def collect_conversation_metadata(user_message, conversation_id, user_id, active current_tags[semantic_key] = semantic_tag # Update the tags array conversation_item['tags'] = list(current_tags.values()) - if document_map: - conversation_item['last_grounded_document_refs'] = _build_last_grounded_document_refs(document_map) + if document_map or source_continuity_refs: + conversation_item['last_grounded_document_refs'] = _build_last_grounded_document_refs( + document_map, + source_continuity_refs=source_continuity_refs, + ) # --- Scope Lock Logic --- current_scope_locked = conversation_item.get('scope_locked') diff --git a/application/single_app/functions_data_management.py b/application/single_app/functions_data_management.py index 376b2f5dd..10049e4ef 100644 --- a/application/single_app/functions_data_management.py +++ b/application/single_app/functions_data_management.py @@ -2,6 +2,7 @@ """Data Management settings, schedules, and durable job records.""" import copy +import base64 import hashlib import json import logging @@ -11,17 +12,21 @@ import time import tempfile import uuid +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, as_completed, wait from datetime import datetime, timedelta, timezone +from threading import Event, Lock, Thread +from urllib.parse import urlparse from azure.core import MatchConditions from azure.core.credentials import AzureKeyCredential +from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError import azure.cosmos as azure_cosmos from azure.cosmos import PartitionKey from azure.cosmos.exceptions import CosmosResourceNotFoundError from azure.identity import DefaultAzureCredential from azure.search.documents import SearchClient from azure.search.documents.indexes import SearchIndexClient -from azure.search.documents.indexes.models import SearchIndex +from azure.search.documents.indexes.models import SearchField, SearchFieldDataType, SearchIndex from azure.storage.blob import BlobServiceClient, ContentSettings from cryptography.fernet import Fernet @@ -34,12 +39,59 @@ cosmos_settings_container, ) from functions_appinsights import log_event +from functions_cosmos_throughput import ( + CosmosThroughputError, + get_container_throughput, + get_database_throughput, + set_database_throughput, +) +from functions_data_management_migration_state import ( + MIGRATION_RESOURCE_STATUS_COMPLETED, + MIGRATION_RESOURCE_STATUS_FAILED, + build_transfer_metrics, + complete_migration_resource, + fail_migration_resource, + initialize_migration_state, + is_migration_resource_completed, + start_migration_resource, + update_migration_resource, +) +from functions_data_management_search_write_fence import ( + acquire_data_management_search_write_fence, + acquire_data_management_target_migration_coordinator, + release_data_management_search_write_fence, + release_data_management_target_migration_coordinator, + renew_data_management_search_write_fence, + renew_data_management_target_migration_coordinator, +) +from functions_migration_provenance import ( + BLOB_MIGRATION_STATUS_METADATA_KEY, + COSMOS_MIGRATION_PROVENANCE_FIELD, + SEARCH_MIGRATED_AT_FIELD, + SEARCH_MIGRATION_ID_FIELD, + SEARCH_MIGRATION_SOURCE_HASH_FIELD, + SEARCH_MIGRATION_SOURCE_VERSION_FIELD, + SEARCH_MIGRATION_STATUS_FIELD, + add_cosmos_migration_provenance, + add_search_migration_provenance, + create_migration_provenance_context, + get_blob_migration_provenance, + get_cosmos_migration_provenance, + get_search_migration_provenance, + is_successful_migration_record, + merge_blob_migration_metadata, + migration_record_matches_source, + should_skip_migration_record, +) DATA_MANAGEMENT_SETTINGS_ID = "backup_settings" DATA_MANAGEMENT_SETTINGS_TYPE = "data_management_settings" DATA_MANAGEMENT_JOB_TYPE = "data_management_job" DATA_MANAGEMENT_JOB_ITEM_TYPE = "data_management_job_item" +DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_TYPE = "data_management_migration_manifest_batch" +DATA_MANAGEMENT_MIRROR_DELETION_BATCH_TYPE = "data_management_mirror_deletion_batch" +DATA_MANAGEMENT_MIGRATION_LOCK_TYPE = "data_management_migration_lock" DATA_MANAGEMENT_OPERATION_BACKUP = "backup" DATA_MANAGEMENT_OPERATION_RESTORE = "restore" @@ -83,8 +135,40 @@ DATA_MANAGEMENT_DEFAULT_LEASE_SECONDS = 900 DATA_MANAGEMENT_DEFAULT_STALE_SECONDS = 1200 DATA_MANAGEMENT_DEFAULT_JOB_LIMIT = 25 +DATA_MANAGEMENT_DEFAULT_RECOVERY_JOB_LIMIT = 25 +DATA_MANAGEMENT_RECOVERY_QUEUE_DELAY_SECONDS = 60 +DATA_MANAGEMENT_RECOVERY_RESUBMIT_DELAY_SECONDS = 120 +DATA_MANAGEMENT_MIGRATION_LOCK_PREFIX = "data_management_migration_lock" DATA_MANAGEMENT_MIGRATION_CATALOG_LIMIT = 50 DATA_MANAGEMENT_MIGRATION_BATCH_SIZE = 500 +DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_SIZE = 100 +DATA_MANAGEMENT_SEARCH_KEYSET_PAGE_SIZE = 1000 +DATA_MANAGEMENT_SEARCH_SCOPE_FILTER_BATCH_SIZE = 100 +DATA_MANAGEMENT_BLOB_SERVICE_TIMEOUT_SECONDS = 120 +DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS = 30 +DATA_MANAGEMENT_MIGRATION_HEARTBEAT_INTERVAL_SECONDS = 2.0 +DATA_MANAGEMENT_MIGRATION_HEARTBEAT_POLL_SECONDS = 1.0 +DATA_MANAGEMENT_MIGRATION_LOCK_RECOVERY_GRACE_SECONDS = 120 +DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY = "new_only" +DATA_MANAGEMENT_MIGRATION_MODE_DELTA_UPSERT = "delta_upsert" +DATA_MANAGEMENT_MIGRATION_MODE_MIRROR = "mirror_with_deletions" +DATA_MANAGEMENT_MIGRATION_MODES = { + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + DATA_MANAGEMENT_MIGRATION_MODE_DELTA_UPSERT, + DATA_MANAGEMENT_MIGRATION_MODE_MIRROR, +} +DATA_MANAGEMENT_MIRROR_CONFIRMATION = "MIRROR WITH DELETIONS" +DATA_MANAGEMENT_SEARCH_WRITE_FREEZE_CONFIRMATION_ERROR = ( + "Confirm that external destination AI Search writers are frozen before migrating AI Search documents." +) +DATA_MANAGEMENT_MIGRATION_DEFAULT_PARALLEL_OPERATIONS = 8 +DATA_MANAGEMENT_MIGRATION_MAX_PARALLEL_OPERATIONS = 32 +DATA_MANAGEMENT_MIGRATION_DEFAULT_RETRY_COUNT = 5 +DATA_MANAGEMENT_MIGRATION_MAX_RETRY_COUNT = 10 +DATA_MANAGEMENT_MIGRATION_DEFAULT_SKIP_WITHIN_HOURS = 0 +DATA_MANAGEMENT_MIGRATION_MAX_SKIP_WITHIN_HOURS = 8760 +DATA_MANAGEMENT_MIGRATION_DEFAULT_DESTINATION_RU = 10000 +DATA_MANAGEMENT_MIGRATION_MAX_DESTINATION_RU = 10000 DATA_MANAGEMENT_KEY_VAULT_SCOPE_VALUE = "data-management" DATA_MANAGEMENT_ENCRYPTION_SECRET_NAME = "backup-encryption-key" DATA_MANAGEMENT_REDACTED_VALUE = "***REDACTED***" @@ -92,6 +176,8 @@ "We suggest not running backups, restores, or migrations during your operational business hours. " "These jobs run inside the App Service environment and can affect application performance." ) +DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME = {} +DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME_LOCK = Lock() DATA_MANAGEMENT_DEFAULT_SETTINGS = { "enabled": False, @@ -104,6 +190,8 @@ "target_cosmos_endpoint": "", "target_cosmos_key": "", "target_cosmos_database_name": DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME, + "target_cosmos_subscription_id": "", + "target_cosmos_resource_group": "", "target_ai_search_authentication_type": "managed_identity", "target_ai_search_endpoint": "", "target_ai_search_key": "", @@ -122,6 +210,11 @@ "include_source_blobs": True, "low_impact_mode": True, "max_parallel_operations": 1, + "migration_max_parallel_operations": DATA_MANAGEMENT_MIGRATION_DEFAULT_PARALLEL_OPERATIONS, + "migration_retry_count": DATA_MANAGEMENT_MIGRATION_DEFAULT_RETRY_COUNT, + "migration_skip_recent_within_hours": DATA_MANAGEMENT_MIGRATION_DEFAULT_SKIP_WITHIN_HOURS, + "migration_temporary_destination_ru_enabled": False, + "migration_temporary_destination_ru": DATA_MANAGEMENT_MIGRATION_DEFAULT_DESTINATION_RU, "next_full_backup_run_at": None, "next_partial_backup_run_at": None, "last_full_backup_completed_at": None, @@ -137,6 +230,14 @@ class DataManagementSettingsValidationError(ValueError): class DataManagementCosmosEditorError(ValueError): """Raised when Cosmos editor input is unsafe or incomplete.""" + +class DataManagementMigrationLeaseLostError(RuntimeError): + """Raised when a stale migration worker no longer owns its job lease.""" + + +class DataManagementMigrationCanceledError(RuntimeError): + """Raised when an administrator requests cooperative migration cancellation.""" + DATA_MANAGEMENT_FRONTEND_SECRET_FIELDS = { "backup_storage_connection_string", "encryption_key_reference", @@ -146,6 +247,7 @@ class DataManagementCosmosEditorError(ValueError): } DATA_MANAGEMENT_MIGRATION_TARGET_TYPES = {"users", "groups", "public_workspaces"} +DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER = ("users", "groups", "public_workspaces") DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS = { "users": [ @@ -498,6 +600,8 @@ def normalize_data_management_settings(payload=None, existing_settings=None, cur source["target_cosmos_authentication_type"] = "managed_identity" source["target_cosmos_endpoint"] = _safe_text(source.get("target_cosmos_endpoint")) source["target_cosmos_key"] = _safe_text(source.get("target_cosmos_key")) + source["target_cosmos_subscription_id"] = _safe_text(source.get("target_cosmos_subscription_id")) + source["target_cosmos_resource_group"] = _safe_text(source.get("target_cosmos_resource_group")) if source["target_cosmos_authentication_type"] == "managed_identity": source["target_cosmos_key"] = "" source["target_cosmos_database_name"] = DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME @@ -534,6 +638,34 @@ def normalize_data_management_settings(payload=None, existing_settings=None, cur source["include_source_blobs"] = False source["low_impact_mode"] = _safe_bool(source.get("low_impact_mode"), True) source["max_parallel_operations"] = _safe_int(source.get("max_parallel_operations"), default=1, minimum=1, maximum=5) + source["migration_max_parallel_operations"] = _safe_int( + source.get("migration_max_parallel_operations"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_PARALLEL_OPERATIONS, + minimum=1, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_PARALLEL_OPERATIONS, + ) + source["migration_retry_count"] = _safe_int( + source.get("migration_retry_count"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_RETRY_COUNT, + minimum=1, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_RETRY_COUNT, + ) + source["migration_skip_recent_within_hours"] = _safe_int( + source.get("migration_skip_recent_within_hours"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_SKIP_WITHIN_HOURS, + minimum=0, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_SKIP_WITHIN_HOURS, + ) + source["migration_temporary_destination_ru_enabled"] = _safe_bool( + source.get("migration_temporary_destination_ru_enabled"), + False, + ) + source["migration_temporary_destination_ru"] = _safe_int( + source.get("migration_temporary_destination_ru"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_DESTINATION_RU, + minimum=1000, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_DESTINATION_RU, + ) for key in ("next_full_backup_run_at", "next_partial_backup_run_at", "last_full_backup_completed_at", "last_partial_backup_completed_at"): parsed = _parse_iso_datetime(source.get(key)) @@ -698,16 +830,47 @@ def _normalize_data_management_settings_from_payload(settings=None): return normalize_data_management_settings(existing_settings=existing_settings, application_settings=application_settings) -def test_target_cosmos_connection(settings=None): +def test_target_cosmos_connection(settings=None, migration_plan=None): normalized_settings = _normalize_data_management_settings_from_payload(settings) target_database = _get_target_cosmos_database(normalized_settings) properties = target_database.read() - return { + normalized_plan = normalize_data_management_migration_plan({ + "migration_plan": migration_plan if isinstance(migration_plan, dict) else {}, + }) + selected_total = sum( + len((normalized_plan.get(target_type) or {}).get("ids") or []) + if (normalized_plan.get(target_type) or {}).get("mode") == "selected" + else 1 if (normalized_plan.get(target_type) or {}).get("mode") == "all" + else 0 + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER + ) + migration_access = None + if selected_total: + migration_access = _preflight_target_cosmos_migration_access( + normalized_settings, + normalized_plan, + ) + capacity = None + if normalized_settings.get("migration_temporary_destination_ru_enabled"): + inspected_capacity = _inspect_target_cosmos_migration_capacity( + normalized_settings, + normalized_plan, + ) + capacity = { + "target_ru": inspected_capacity.get("target_ru"), + "database_mode": inspected_capacity.get("database_mode"), + "database_current_ru": inspected_capacity.get("database_current_ru"), + "targets": inspected_capacity.get("targets"), + } + result = { "success": True, "target": "cosmos", "database_name": properties.get("id") or DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME, "authentication_type": normalized_settings.get("target_cosmos_authentication_type"), + "migration_access": migration_access, + "capacity": capacity, } + return result def test_target_search_connection(settings=None): @@ -715,7 +878,13 @@ def test_target_search_connection(settings=None): endpoint = _safe_text(normalized_settings.get("target_ai_search_endpoint")) if not endpoint: raise ValueError("Target Search endpoint is required.") - index_client = SearchIndexClient(endpoint=endpoint, credential=_get_target_ai_search_credential(normalized_settings)) + index_client = SearchIndexClient( + endpoint=endpoint, + credential=_get_target_ai_search_credential(normalized_settings), + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + ) existing_indexes = set(index_client.list_index_names()) expected_indexes = [artifact["index_name"] for artifact in DATA_MANAGEMENT_SEARCH_ARTIFACTS] return { @@ -784,7 +953,7 @@ def _normalize_migration_selection(selection=None): } -def _get_target_cosmos_database(settings): +def _get_target_cosmos_client(settings): endpoint = _safe_text((settings or {}).get("target_cosmos_endpoint")) if not endpoint: raise ValueError("Target Cosmos endpoint is required before running migration.") @@ -792,304 +961,5474 @@ def _get_target_cosmos_database(settings): key = _safe_text((settings or {}).get("target_cosmos_key")) if not key: raise ValueError("Target Cosmos account key is required when account key authentication is selected.") - client = azure_cosmos.CosmosClient(endpoint, credential=key, consistency_level="Session") - else: - client = azure_cosmos.CosmosClient(endpoint, credential=DefaultAzureCredential(), consistency_level="Session") - return client.create_database_if_not_exists(DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME) + return azure_cosmos.CosmosClient( + endpoint, + credential=key, + consistency_level="Session", + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + ) + return azure_cosmos.CosmosClient( + endpoint, + credential=DefaultAzureCredential(), + consistency_level="Session", + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + ) + + +def _get_target_cosmos_database(settings): + return _get_target_cosmos_client(settings).create_database_if_not_exists( + DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME + ) + + +def _get_existing_target_cosmos_database(settings): + """Open the destination database without creating resources during preview.""" + database = _get_target_cosmos_client(settings).get_database_client( + DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME + ) + database.read() + return database def _get_target_cosmos_container(target_database, container_name, partition_key_path): - return target_database.create_container_if_not_exists( + target_container = target_database.create_container_if_not_exists( id=container_name, partition_key=PartitionKey(path=partition_key_path), ) + _validate_target_cosmos_container_partition_key( + target_container, + container_name, + partition_key_path, + ) + return target_container + + +def _get_target_data_management_search_write_gate_container(target_database): + """Open the target SimpleChat job container used to coordinate Search writers.""" + return _get_target_cosmos_container( + target_database, + getattr( + app_config, + "cosmos_data_management_jobs_container_name", + "data_management_jobs", + ), + "/id", + ) -def _iter_selected_cosmos_records(container_definition, selection): - source_container = getattr(app_config, container_definition["container_attr"], None) - if not source_container: - return - if container_definition.get("documents") and not selection.get("include_documents"): - return - mode = selection.get("mode") - if mode == "all": - yield from _iter_cosmos_container_items(source_container) - return - if mode != "selected": - return - - ids = selection.get("ids") or [] - if container_definition.get("id_field") == "id": - for item_id in ids: - try: - yield _strip_cosmos_system_fields(source_container.read_item(item=item_id, partition_key=item_id)) - except CosmosResourceNotFoundError: - continue - return - - filter_field = container_definition.get("filter_field") - if not filter_field: - return - for selected_id in ids: - query = f"SELECT * FROM c WHERE c.{filter_field} = @selected_id" - for item in source_container.query_items( - query=query, - parameters=[{"name": "@selected_id", "value": selected_id}], - enable_cross_partition_query=True, - ): - yield _strip_cosmos_system_fields(item) +def _get_target_search_write_fence_lease_seconds(settings): + """Keep uncertain target Search writes quarantined beyond their bounded request lifetime.""" + return min( + _get_migration_lock_lease_seconds(settings), + DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS + + DATA_MANAGEMENT_MIGRATION_LOCK_RECOVERY_GRACE_SECONDS, + ) -def _copy_cosmos_records_to_target(target_database, target_type, selection): - copied = [] - for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type]: - target_container_name = getattr(app_config, container_definition["container_name_attr"], container_definition["name"]) - target_container = _get_target_cosmos_container( - target_database, - target_container_name, - container_definition["partition_key_path"], - ) - count = 0 - for item in _iter_selected_cosmos_records(container_definition, selection) or []: - target_container.upsert_item(item) - count += 1 - copied.append({ - "name": container_definition["name"], - "type": "cosmos_container", - "target_type": target_type, - "container_name": target_container_name, - "item_count": count, - "partition_key_path": container_definition["partition_key_path"], - }) - return copied +def _set_target_migration_coordinator_runtime(job_id, container, coordinator): + """Keep non-serializable target SDK state out of durable migration job records.""" + with DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME_LOCK: + DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME[_safe_text(job_id)] = { + "container": container, + "coordinator": coordinator, + } -def _escape_search_filter_value(value): - return _safe_text(value).replace("'", "''") +def _get_target_migration_coordinator_runtime(job_id): + with DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME_LOCK: + return DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME.get(_safe_text(job_id)) -def _build_search_filter(field_name, selection): - if selection.get("mode") == "all": - return None - if selection.get("mode") != "selected": - return "id eq '__no_migration_selection__'" - conditions = [f"{field_name} eq '{_escape_search_filter_value(item_id)}'" for item_id in selection.get("ids") or []] - if not conditions: - return "id eq '__no_migration_selection__'" - return " or ".join(conditions) +def _clear_target_migration_coordinator_runtime(job_id): + with DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME_LOCK: + return DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME.pop( + _safe_text(job_id), + None, + ) -def _get_target_ai_search_credential(settings): - if (settings or {}).get("target_ai_search_authentication_type") == "key": - key = _safe_text((settings or {}).get("target_ai_search_key")) - if not key: - raise ValueError("Target AI Search key is required when key authentication is selected.") - return AzureKeyCredential(key) - return DefaultAzureCredential() +def _acquire_target_migration_coordinator(job, target_database, settings): + """Acquire a destination-wide lease before this source performs any migration work.""" + target_container = _get_target_data_management_search_write_gate_container( + target_database + ) + coordinator = acquire_data_management_target_migration_coordinator( + target_container, + job.get("id"), + _get_migration_lock_lease_seconds(settings), + existing_lock=( + job.get("target_migration_coordinator") + if isinstance(job.get("target_migration_coordinator"), dict) + else None + ), + ) + job["target_migration_coordinator"] = copy.deepcopy(coordinator) + _set_target_migration_coordinator_runtime( + job.get("id"), + target_container, + coordinator, + ) + return coordinator -def _get_target_search_client(settings, index_name): - endpoint = _safe_text((settings or {}).get("target_ai_search_endpoint")) - if not endpoint: - raise ValueError("Target AI Search endpoint is required before running AI Search migration.") - return SearchClient(endpoint=endpoint, index_name=index_name, credential=_get_target_ai_search_credential(settings)) +def _renew_target_migration_coordinator(job, settings): + """Renew the destination-wide coordinator that protects independent source deployments.""" + runtime = _get_target_migration_coordinator_runtime(job.get("id")) + if not isinstance(runtime, dict): + return + coordinator = runtime.get("coordinator") + target_container = runtime.get("container") + if not isinstance(coordinator, dict) or target_container is None: + return + try: + renewed_coordinator = renew_data_management_target_migration_coordinator( + target_container, + coordinator, + _get_migration_lock_lease_seconds({ + "data_management_job_lease_seconds": ( + (settings or {}).get("data_management_job_lease_seconds") or + coordinator.get("lease_seconds") + ), + }), + ) + except Exception as exc: + raise DataManagementMigrationLeaseLostError( + "The target migration coordinator was lost or could not be renewed." + ) from exc + job["target_migration_coordinator"] = copy.deepcopy(renewed_coordinator) -def _ensure_target_search_index(settings, index_name, schema_file): - endpoint = _safe_text((settings or {}).get("target_ai_search_endpoint")) - if not endpoint: - raise ValueError("Target AI Search endpoint is required before running AI Search migration.") - index_client = SearchIndexClient(endpoint=endpoint, credential=_get_target_ai_search_credential(settings)) +def _release_target_migration_coordinator(job): + """Release a cleanly completed or drained-canceled destination coordinator lease.""" + runtime = _clear_target_migration_coordinator_runtime(job.get("id")) + if not isinstance(runtime, dict): + return False try: - index_client.get_index(index_name) - return "exists" + return release_data_management_target_migration_coordinator( + runtime.get("container"), + runtime.get("coordinator"), + ) except Exception: - schema = _get_search_schema(schema_file) - schema = {key: value for key, value in schema.items() if not key.startswith("@odata.")} - try: - index_client.create_or_update_index(SearchIndex.from_dict(schema)) - return "created" - except Exception as exc: - raise ValueError(f"Target AI Search index {index_name} is missing and could not be created: {exc}") from exc + return False -def _upload_search_documents_in_batches(search_client, documents): - uploaded = 0 - batch = [] - for document in documents: - batch.append(document) - if len(batch) >= DATA_MANAGEMENT_MIGRATION_BATCH_SIZE: - search_client.upload_documents(documents=batch) - uploaded += len(batch) - batch = [] - if batch: - search_client.upload_documents(documents=batch) - uploaded += len(batch) - return uploaded +def _validate_target_cosmos_container_partition_key(target_container, container_name, expected_path): + """Reject existing target containers whose partition key differs from SimpleChat's contract.""" + container_read = getattr(target_container, "read", None) + if not callable(container_read): + return + properties = container_read() + partition_key = properties.get("partitionKey") if isinstance(properties, dict) else None + actual_paths = partition_key.get("paths") if isinstance(partition_key, dict) else None + if not isinstance(actual_paths, list) or actual_paths != [expected_path]: + raise DataManagementSettingsValidationError( + f"Destination Cosmos container '{container_name}' has partition key paths " + f"{actual_paths!r}; expected [{expected_path!r}]." + ) -def _copy_ai_search_to_target(settings, migration_plan): - copied = [] - if not migration_plan.get("include_ai_search"): - return copied - search_mappings = [ - ("users", "user_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[0]), - ("groups", "group_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[1]), - ("public_workspaces", "public_workspace_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[2]), - ] - for target_type, field_name, artifact in search_mappings: - selection = migration_plan.get(target_type) or {} - if selection.get("mode") == "none" or not selection.get("include_documents"): - continue - source_client = CLIENTS.get(artifact["client_key"]) - if not source_client: - copied.append({"name": artifact["name"], "type": "ai_search_documents", "status": "skipped", "warning": "Source AI Search client is not initialized."}) +def _iter_migration_cosmos_container_definitions(migration_plan): + """Yield only destination containers that the selected migration can write.""" + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = (migration_plan or {}).get(target_type) or {} + if selection.get("mode") == "none": continue - index_status = _ensure_target_search_index(settings, artifact["index_name"], artifact["schema_file"]) - target_client = _get_target_search_client(settings, artifact["index_name"]) - search_filter = _build_search_filter(field_name, selection) - documents = _iter_search_documents_for_filter(source_client, search_filter) - uploaded = _upload_search_documents_in_batches(target_client, documents) - copied.append({ - "name": artifact["name"], - "type": "ai_search_documents", - "target_type": target_type, - "index_name": artifact["index_name"], - "item_count": uploaded, - "partial_filter": search_filter, - "index_status": index_status, - }) - return copied - - -def _iter_search_documents_for_filter(search_client, search_filter=None): - for result in search_client.search(search_text="*", filter=search_filter, include_total_count=True): - document = dict(result) - yield {key: value for key, value in document.items() if not key.startswith("@search.")} + for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type]: + if container_definition.get("documents") and not selection.get("include_documents"): + continue + yield target_type, container_definition -def _get_target_enhanced_citations_blob_client(settings): - auth_type = _safe_text((settings or {}).get("target_enhanced_citations_storage_authentication_type"), "managed_identity") - if auth_type == "connection_string": - connection_string = _safe_text((settings or {}).get("target_enhanced_citations_storage_connection_string")) - if not connection_string: - raise ValueError("Destination Enhanced Citations storage connection string is required when migrating source document blobs.") - return BlobServiceClient.from_connection_string(connection_string) - blob_endpoint = _safe_text((settings or {}).get("target_enhanced_citations_storage_blob_endpoint")) - if not blob_endpoint: - raise ValueError("Destination Enhanced Citations Blob endpoint is required when migrating source document blobs.") - return BlobServiceClient(account_url=blob_endpoint, credential=DefaultAzureCredential()) +def _target_cosmos_container_name(container_definition): + return getattr( + app_config, + container_definition["container_name_attr"], + container_definition["name"], + ) -def _document_blob_reference(document_item): - if not isinstance(document_item, dict): - return None - container_name = document_item.get("blob_container") or _blob_container_for_document(document_item) - blob_path = document_item.get("blob_path") or document_item.get("archived_blob_path") - if not blob_path and document_item.get("file_name"): - scope_id = document_item.get("public_workspace_id") or document_item.get("group_id") or document_item.get("user_id") - blob_path = f"{scope_id}/{document_item.get('file_name')}" if scope_id else "" - if not container_name or not blob_path: - return None - return container_name, blob_path +def _get_migration_parallel_operations(settings): + return _safe_int( + (settings or {}).get("migration_max_parallel_operations"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_PARALLEL_OPERATIONS, + minimum=1, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_PARALLEL_OPERATIONS, + ) -def _blob_container_for_document(document_item): - if document_item.get("public_workspace_id"): - return app_config.storage_account_public_documents_container_name - if document_item.get("group_id"): - return app_config.storage_account_group_documents_container_name - return app_config.storage_account_user_documents_container_name +def _get_migration_retry_count(settings): + return _safe_int( + (settings or {}).get("migration_retry_count"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_RETRY_COUNT, + minimum=1, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_RETRY_COUNT, + ) -def _iter_selected_document_records_for_blob_migration(migration_plan): - container_map = [ - ("users", app_config.cosmos_user_documents_container, "user_id"), - ("groups", app_config.cosmos_group_documents_container, "group_id"), - ("public_workspaces", app_config.cosmos_public_documents_container, "public_workspace_id"), +def _get_target_cosmos_account_name(endpoint): + normalized_endpoint = _safe_text(endpoint) + if not normalized_endpoint: + return "" + parsed_endpoint = urlparse( + normalized_endpoint if "://" in normalized_endpoint else f"https://{normalized_endpoint}" + ) + hostname = parsed_endpoint.hostname or "" + return hostname.split(".")[0].strip() + + +def _get_target_cosmos_management_settings(settings): + account_name = _get_target_cosmos_account_name((settings or {}).get("target_cosmos_endpoint")) + management_settings = { + "cosmos_throughput_subscription_id": _safe_text( + (settings or {}).get("target_cosmos_subscription_id") + ), + "cosmos_throughput_resource_group": _safe_text( + (settings or {}).get("target_cosmos_resource_group") + ), + "cosmos_throughput_account_name": account_name, + "cosmos_throughput_database_name": DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME, + } + missing_values = [ + field_name + for field_name, value in management_settings.items() + if not value ] - for target_type, container, filter_field in container_map: - selection = migration_plan.get(target_type) or {} - if selection.get("mode") == "none" or not selection.get("include_documents"): - continue - if selection.get("mode") == "all": - yield from _iter_cosmos_container_items(container) - continue - for selected_id in selection.get("ids") or []: - query = f"SELECT * FROM c WHERE c.{filter_field} = @selected_id" - for item in container.query_items( - query=query, - parameters=[{"name": "@selected_id", "value": selected_id}], - enable_cross_partition_query=True, - ): - yield _strip_cosmos_system_fields(item) - - -def _copy_source_blobs_to_target(settings, migration_plan): - if not migration_plan.get("include_source_blobs"): - return [] - source_client = _get_source_blob_service_client() - if not source_client: - return [{"name": "source_blobs", "type": "source_blobs", "status": "skipped", "warning": "Source Enhanced Citations storage is not configured."}] - target_client = _get_target_enhanced_citations_blob_client(settings) - copied = {} - warnings = [] - for document in _iter_selected_document_records_for_blob_migration(migration_plan): - reference = _document_blob_reference(document) - if not reference: - continue - container_name, blob_path = reference - artifact = copied.setdefault(container_name, {"name": container_name, "type": "source_blob_container", "container_name": container_name, "blob_count": 0, "bytes": 0}) - try: - target_container = target_client.create_container(container_name) - except Exception: - target_container = target_client.get_container_client(container_name) - try: - source_blob = source_client.get_blob_client(container=container_name, blob=blob_path) - blob_bytes = source_blob.download_blob().readall() - target_container.upload_blob(name=blob_path, data=blob_bytes, overwrite=True, content_settings=ContentSettings(content_type="application/octet-stream")) - artifact["blob_count"] += 1 - artifact["bytes"] += len(blob_bytes) - except Exception as exc: - warnings.append(f"{container_name}/{blob_path}: {exc}") - artifacts = list(copied.values()) - if warnings: - artifacts.append({"name": "source_blob_warnings", "type": "source_blobs", "status": "warning", "warning": "; ".join(warnings[:10])}) - return artifacts + if missing_values: + raise DataManagementSettingsValidationError( + "Destination Cosmos capacity management requires " + f"{', '.join(missing_values)}." + ) + return management_settings -def _migration_selection_summary(migration_plan): +def _build_migration_configuration_snapshot(settings, migration_plan): + """Build a durable state fingerprint without copying credential material.""" return { - target_type: { - "mode": (migration_plan.get(target_type) or {}).get("mode"), - "selected_count": len((migration_plan.get(target_type) or {}).get("ids") or []), - "include_documents": bool((migration_plan.get(target_type) or {}).get("include_documents")), - } - for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES + "migration_plan": copy.deepcopy(migration_plan), + "parallel_operations": _get_migration_parallel_operations(settings), + "retry_count": _get_migration_retry_count(settings), + "skip_recent_within_hours": _safe_int( + (settings or {}).get("migration_skip_recent_within_hours"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_SKIP_WITHIN_HOURS, + minimum=0, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_SKIP_WITHIN_HOURS, + ), + "temporary_destination_ru_enabled": _safe_bool( + (settings or {}).get("migration_temporary_destination_ru_enabled"), + False, + ), + "temporary_destination_ru": _safe_int( + (settings or {}).get("migration_temporary_destination_ru"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_DESTINATION_RU, + minimum=1000, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_DESTINATION_RU, + ), + "target_cosmos_endpoint": _safe_text((settings or {}).get("target_cosmos_endpoint")), + "target_cosmos_authentication_type": _safe_text( + (settings or {}).get("target_cosmos_authentication_type") + ), + "target_cosmos_subscription_id": _safe_text( + (settings or {}).get("target_cosmos_subscription_id") + ), + "target_cosmos_resource_group": _safe_text( + (settings or {}).get("target_cosmos_resource_group") + ), + "target_ai_search_endpoint": _safe_text((settings or {}).get("target_ai_search_endpoint")), + "target_ai_search_authentication_type": _safe_text( + (settings or {}).get("target_ai_search_authentication_type") + ), + "target_enhanced_citations_storage_endpoint": _safe_text( + (settings or {}).get("target_enhanced_citations_storage_blob_endpoint") + ), + "target_enhanced_citations_storage_authentication_type": _safe_text( + (settings or {}).get("target_enhanced_citations_storage_authentication_type") + ), } -def _job_lease_holder_id(): - return f"{socket.gethostname()}:{uuid.uuid4().hex}" - - -def _is_stale_job(job, stale_seconds=DATA_MANAGEMENT_DEFAULT_STALE_SECONDS): - last_heartbeat = _parse_iso_datetime((job or {}).get("last_heartbeat_at") or (job or {}).get("updated_at")) - if not last_heartbeat: - return True - return last_heartbeat <= _now_utc() - timedelta(seconds=stale_seconds) +def _migration_plan_scope_signature(migration_plan): + """Return the source scope fields that must match an incremental baseline.""" + signature = { + "include_ai_search": bool((migration_plan or {}).get("include_ai_search")), + "include_source_blobs": bool((migration_plan or {}).get("include_source_blobs")), + } + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = (migration_plan or {}).get(target_type) or {} + signature[target_type] = { + "mode": _safe_text(selection.get("mode"), "none"), + "ids": sorted(_dedupe_limited_strings(selection.get("ids"), limit=2000)), + "include_documents": bool(selection.get("include_documents")), + } + return signature -def _read_job(job_id): - return cosmos_data_management_jobs_container.read_item(item=job_id, partition_key=job_id) +def _migration_destination_signature(configuration, migration_plan): + """Return normalized destination identities without credential material.""" + return { + "cosmos": _safe_text((configuration or {}).get("target_cosmos_endpoint")).lower(), + "ai_search": ( + _safe_text((configuration or {}).get("target_ai_search_endpoint")).lower() + if (migration_plan or {}).get("include_ai_search") else "" + ), + "source_blobs": ( + _safe_text( + (configuration or {}).get("target_enhanced_citations_storage_endpoint") + ).lower() + if (migration_plan or {}).get("include_source_blobs") else "" + ), + } -def _replace_job(job): - return cosmos_data_management_jobs_container.replace_item( - item=job.get("id"), - body=job, - etag=job.get("_etag"), - match_condition=MatchConditions.IfNotModified, +def _validate_incremental_baseline_job(candidate_job, current_job, current_configuration, migration_plan): + """Return a compatible completed cutoff or reject an unsafe baseline.""" + candidate_state = candidate_job.get("migration_state") if isinstance(candidate_job, dict) else None + candidate_configuration = ( + candidate_state.get("configuration") + if isinstance(candidate_state, dict) and isinstance(candidate_state.get("configuration"), dict) + else {} + ) + candidate_plan = candidate_configuration.get("migration_plan") + candidate_plan = candidate_plan if isinstance(candidate_plan, dict) else {} + candidate_result = candidate_job.get("result") if isinstance(candidate_job, dict) else {} + reconciliation_resource = ( + (candidate_state.get("resources") or {}).get("reconciliation:cutover") + if isinstance(candidate_state, dict) else None + ) + reconciliation_result = ( + reconciliation_resource.get("result") + if isinstance(reconciliation_resource, dict) and + isinstance(reconciliation_resource.get("result"), dict) + else {} + ) + if ( + not isinstance(candidate_job, dict) or + candidate_job.get("id") == (current_job or {}).get("id") or + candidate_job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION or + candidate_job.get("status") not in { + DATA_MANAGEMENT_STATUS_COMPLETED, + DATA_MANAGEMENT_STATUS_COMPLETED_WITH_WARNINGS, + } or + not isinstance(candidate_state, dict) or + candidate_state.get("status") != MIGRATION_RESOURCE_STATUS_COMPLETED or + (isinstance(candidate_result, dict) and candidate_result.get("dry_run")) or + reconciliation_result.get("readiness") not in {"ready", "ready_with_warnings"} + ): + raise DataManagementSettingsValidationError( + "Incremental migration baseline must be a completed, reconciled, non-dry-run migration." + ) + if _migration_plan_scope_signature(candidate_plan) != _migration_plan_scope_signature(migration_plan): + raise DataManagementSettingsValidationError( + "Incremental migration baseline uses a different source scope." + ) + if ( + _migration_destination_signature(candidate_configuration, candidate_plan) != + _migration_destination_signature(current_configuration, migration_plan) + ): + raise DataManagementSettingsValidationError( + "Incremental migration baseline targets a different destination." + ) + source_cutoff_at = _safe_text(candidate_state.get("source_cutoff_at")) + if _parse_iso_datetime(source_cutoff_at) is None: + raise DataManagementSettingsValidationError( + "Incremental migration baseline does not contain a valid source watermark." + ) + return source_cutoff_at + + +def _get_completed_migration_baseline_candidates(limit=100): + query = ( + "SELECT * FROM c WHERE c.type = @type AND c.operation = @operation " + "AND (c.status = @completed OR c.status = @completed_with_warnings) " + "ORDER BY c.completed_at DESC" + ) + parameters = [ + {"name": "@type", "value": DATA_MANAGEMENT_JOB_TYPE}, + {"name": "@operation", "value": DATA_MANAGEMENT_OPERATION_MIGRATION}, + {"name": "@completed", "value": DATA_MANAGEMENT_STATUS_COMPLETED}, + { + "name": "@completed_with_warnings", + "value": DATA_MANAGEMENT_STATUS_COMPLETED_WITH_WARNINGS, + }, + ] + return list(cosmos_data_management_jobs_container.query_items( + query=query, + parameters=parameters, + enable_cross_partition_query=True, + max_item_count=limit, + ))[:limit] + + +def _resolve_data_management_migration_baseline(job, settings, migration_plan): + """Resolve one immutable prior cutoff for delta and mirror migration modes.""" + resolved_plan = copy.deepcopy(migration_plan) + migration_mode = resolved_plan.get("migration_mode", DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY) + if migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + resolved_plan["baseline_job_id"] = "" + resolved_plan["baseline_source_cutoff_at"] = "" + return resolved_plan + + existing_state = job.get("migration_state") if isinstance(job, dict) else None + existing_configuration = ( + existing_state.get("configuration") + if isinstance(existing_state, dict) and isinstance(existing_state.get("configuration"), dict) + else {} + ) + existing_plan = existing_configuration.get("migration_plan") + if isinstance(existing_plan, dict) and existing_plan.get("baseline_source_cutoff_at"): + resolved_plan["baseline_job_id"] = _safe_text(existing_plan.get("baseline_job_id")) + resolved_plan["baseline_source_cutoff_at"] = _safe_text( + existing_plan.get("baseline_source_cutoff_at") + ) + return resolved_plan + + current_configuration = _build_migration_configuration_snapshot(settings, resolved_plan) + requested_baseline_job_id = _safe_text(resolved_plan.get("baseline_job_id")) + if requested_baseline_job_id: + candidate_jobs = [_read_job(requested_baseline_job_id)] + else: + candidate_jobs = _get_completed_migration_baseline_candidates() + + incompatibility_messages = [] + for candidate_job in candidate_jobs: + try: + source_cutoff_at = _validate_incremental_baseline_job( + candidate_job, + job, + current_configuration, + resolved_plan, + ) + except DataManagementSettingsValidationError as exc: + incompatibility_messages.append(str(exc)) + if requested_baseline_job_id: + raise + continue + resolved_plan["baseline_job_id"] = _safe_text(candidate_job.get("id")) + resolved_plan["baseline_source_cutoff_at"] = source_cutoff_at + return resolved_plan + + detail = incompatibility_messages[0] if incompatibility_messages else "No completed migration was found." + raise DataManagementSettingsValidationError( + f"Incremental migration requires a compatible completed baseline. {detail}" + ) + + +def _initialize_data_management_migration_state(job, settings, migration_plan): + """Persist the job's stable migration identity before target writes begin.""" + configuration = _build_migration_configuration_snapshot(settings, migration_plan) + state = initialize_migration_state( + job.get("migration_state"), + str(job.get("id")), + configuration, + ) + job["migration_state"] = state + context = create_migration_provenance_context( + migration_id=state["migration_id"], + migrated_at_utc=state["source_cutoff_at"], + skip_within_hours=configuration["skip_recent_within_hours"], + ) + migration_mode = migration_plan.get( + "migration_mode", + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ) + baseline_source_cutoff_at = _safe_text(migration_plan.get("baseline_source_cutoff_at")) + state["incremental"] = { + "mode": migration_mode, + "baseline_job_id": _safe_text(migration_plan.get("baseline_job_id")), + "baseline_source_cutoff_at": baseline_source_cutoff_at, + "source_cutoff_at": state["source_cutoff_at"], + "deletion_policy": ( + "confirmed_migration_owned_only" + if migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR else + "disabled" + ), + } + state["watermarks"] = { + "baseline": baseline_source_cutoff_at, + "current_cutoff": state["source_cutoff_at"], + "cosmos": {"strategy": "_ts_content_hash_with_identity", "through": state["source_cutoff_at"]}, + "ai_search": {"strategy": "id_keyset_with_observed_source_hash", "through": None}, + "source_blobs": {"strategy": "stable_etag_with_path", "through": None}, + } + context.update({ + "migration_mode": migration_mode, + "baseline_job_id": state["incremental"]["baseline_job_id"], + "baseline_source_cutoff_at": baseline_source_cutoff_at, + }) + return state, context + + +def _set_document_path_value(document, partition_key_path, value): + """Set a simple or nested Cosmos partition-key value in a probe document.""" + path_parts = [part for part in _safe_text(partition_key_path).strip("/").split("/") if part] + if not path_parts: + raise DataManagementSettingsValidationError("Destination Cosmos container has an invalid partition key path.") + current = document + for path_part in path_parts[:-1]: + nested = current.get(path_part) + if not isinstance(nested, dict): + nested = {} + current[path_part] = nested + current = nested + current[path_parts[-1]] = value + return value + + +def _verify_target_cosmos_container_access(target_container, partition_key_path): + """Prove the destination has create, read, and delete data-plane access.""" + probe_id = f"simplechat-migration-preflight-{uuid.uuid4().hex}" + probe_document = { + "id": probe_id, + "type": "simplechat_migration_preflight", + "created_at": _now_iso(), + } + partition_key_value = _set_document_path_value(probe_document, partition_key_path, probe_id) + target_container.create_item(body=probe_document) + try: + target_container.read_item(item=probe_id, partition_key=partition_key_value) + finally: + target_container.delete_item(item=probe_id, partition_key=partition_key_value) + + +def _preflight_target_cosmos_migration_access(settings, migration_plan): + """Create planned containers and verify their destination data-plane permissions.""" + target_database = _get_target_cosmos_database(settings) + container_results = [] + for target_type, container_definition in _iter_migration_cosmos_container_definitions(migration_plan): + source_container = getattr(app_config, container_definition["container_attr"], None) + if source_container is None: + raise DataManagementSettingsValidationError( + f"Source Cosmos container '{container_definition['name']}' is not initialized." + ) + source_probe = source_container.query_items( + query="SELECT TOP 1 * FROM c", + enable_cross_partition_query=True, + ) + next(iter(source_probe), None) + target_container_name = _target_cosmos_container_name(container_definition) + target_container = _get_target_cosmos_container( + target_database, + target_container_name, + container_definition["partition_key_path"], + ) + _verify_target_cosmos_container_access( + target_container, + container_definition["partition_key_path"], + ) + container_results.append({ + "target_type": target_type, + "container_name": target_container_name, + "partition_key_path": container_definition["partition_key_path"], + "source_read_verified": True, + "destination_write_verified": True, + }) + return { + "database_name": DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME, + "container_count": len(container_results), + "containers": container_results, + } + + +def _inspect_target_cosmos_migration_capacity(settings, migration_plan): + """Read destination capacity targets before an optional temporary RU boost.""" + management_settings = _get_target_cosmos_management_settings(settings) + database_throughput = get_database_throughput(management_settings) + targets = [] + if database_throughput.get("is_scalable"): + targets.append({ + "scope": "database", + "container_name": "", + "mode": database_throughput.get("mode"), + "current_ru": database_throughput.get("current_ru"), + }) + else: + seen_container_names = set() + for _, container_definition in _iter_migration_cosmos_container_definitions(migration_plan): + container_name = _target_cosmos_container_name(container_definition) + if container_name in seen_container_names: + continue + seen_container_names.add(container_name) + container_throughput = get_container_throughput(management_settings, container_name) + if container_throughput.get("is_scalable"): + targets.append({ + "scope": "container", + "container_name": container_name, + "mode": container_throughput.get("mode"), + "current_ru": container_throughput.get("current_ru"), + }) + if not targets: + raise CosmosThroughputError( + "Destination Cosmos throughput is not scalable at the database or selected container level." + ) + return { + "target_ru": _safe_int( + (settings or {}).get("migration_temporary_destination_ru"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_DESTINATION_RU, + minimum=1000, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_DESTINATION_RU, + ), + "management_settings": management_settings, + "database_mode": database_throughput.get("mode"), + "database_current_ru": database_throughput.get("current_ru"), + "targets": targets, + } + + +def _get_data_management_job_lease_seconds(settings): + return _safe_int( + (settings or {}).get("data_management_job_lease_seconds"), + default=DATA_MANAGEMENT_DEFAULT_LEASE_SECONDS, + minimum=60, + maximum=7200, + ) + + +def _migration_resource_metrics(resource): + if not isinstance(resource, dict): + return {} + result = resource.get("result") + if isinstance(result, dict) and result: + return result + progress = resource.get("progress") + return progress if isinstance(progress, dict) else {} + + +def _update_migration_state_totals(state): + totals = { + "copied_count": 0, + "skipped_count": 0, + "failed_count": 0, + "collision_count": 0, + "processed_count": 0, + "bytes": 0, + "request_units": 0.0, + } + for resource in (state.get("resources") or {}).values(): + metrics = _migration_resource_metrics(resource) + for field_name in ( + "copied_count", + "skipped_count", + "failed_count", + "collision_count", + "processed_count", + "bytes", + ): + totals[field_name] += _safe_int(metrics.get(field_name), default=0, minimum=0) + try: + totals["request_units"] += max(0.0, float(metrics.get("request_units") or 0.0)) + except (TypeError, ValueError): + continue + totals["request_units"] = round(totals["request_units"], 3) + state["totals"] = totals + return totals + + +def _migration_progress_marker(progress): + marker_fields = ( + "copied_count", + "created_count", + "updated_count", + "unchanged_count", + "skipped_count", + "failed_count", + "collision_count", + "missing_count", + "not_applicable_count", + "processed_count", + "bytes", + "observed_bytes", + "request_units", + "source_read_count", + "destination_accepted_count", + "destination_failed_count", + "completed_count", + "services_completed", + ) + return tuple((field_name, (progress or {}).get(field_name)) for field_name in marker_fields) + + +def _persist_migration_checkpoint( + job, + state, + settings, + resource_name, + progress, + message, + allow_cancel_requested=False, +): + """Save a batch checkpoint and renew the worker lease without timeline noise.""" + _assert_migration_job_lease( + job, + allow_cancel_requested=allow_cancel_requested, + ) + previous_resource = (state.get("resources") or {}).get(resource_name) or {} + previous_progress = ( + previous_resource.get("progress") + if isinstance(previous_resource.get("progress"), dict) + else {} + ) + now = _now_utc() + if _migration_progress_marker(progress) != _migration_progress_marker(previous_progress): + progress["last_progress_at"] = now.isoformat() + state["last_progress_at"] = now.isoformat() + job["last_progress_at"] = now.isoformat() + else: + progress["last_progress_at"] = ( + previous_progress.get("last_progress_at") or + state.get("last_progress_at") or + job.get("last_progress_at") + ) + update_migration_resource(state, resource_name, progress) + _update_migration_state_totals(state) + job.update({ + "migration_state": state, + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "lease_expires_at": ( + now + timedelta(seconds=_get_data_management_job_lease_seconds(settings)) + ).isoformat(), + "last_message": message, + }) + _save_data_management_job(job) + return job.get("migration_state") if isinstance(job.get("migration_state"), dict) else state + + +def _persist_migration_heartbeat(job, settings, message): + """Renew worker and coordinator leases without advancing resource progress.""" + _assert_migration_job_lease(job) + now = _now_utc() + job.update({ + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "lease_expires_at": ( + now + timedelta(seconds=_get_data_management_job_lease_seconds(settings)) + ).isoformat(), + "last_message": message, + }) + return _save_data_management_job(job) + + +def _complete_migration_resource_checkpoint(job, state, settings, resource_name, result, message): + """Complete a resource after its final durable checkpoint is written.""" + _assert_migration_job_lease(job) + complete_migration_resource(state, resource_name, result=result) + _update_migration_state_totals(state) + now = _now_utc() + result["last_progress_at"] = now.isoformat() + state["last_progress_at"] = now.isoformat() + job.update({ + "migration_state": state, + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "last_progress_at": now.isoformat(), + "lease_expires_at": ( + now + timedelta(seconds=_get_data_management_job_lease_seconds(settings)) + ).isoformat(), + "last_message": message, + }) + _save_data_management_job(job) + return job.get("migration_state") if isinstance(job.get("migration_state"), dict) else state + + +def _fail_migration_resource_checkpoint(job, state, settings, resource_name, error_message): + """Persist a failed resource before allowing the job to be retried.""" + _assert_migration_job_lease(job) + fail_migration_resource(state, resource_name, error_message) + _update_migration_state_totals(state) + now = _now_utc() + job.update({ + "migration_state": state, + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "lease_expires_at": ( + now + timedelta(seconds=_get_data_management_job_lease_seconds(settings)) + ).isoformat(), + "last_message": f"Migration resource failed: {resource_name}", + }) + _save_data_management_job(job) + return job.get("migration_state") if isinstance(job.get("migration_state"), dict) else state + + +def _persist_migration_state( + job, + state, + settings, + message, + allow_cancel_requested=False, +): + """Persist non-resource migration state changes and renew the worker lease.""" + _assert_migration_job_lease( + job, + allow_cancel_requested=allow_cancel_requested, + ) + now = _now_utc() + job.update({ + "migration_state": state, + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "lease_expires_at": ( + now + timedelta(seconds=_get_data_management_job_lease_seconds(settings)) + ).isoformat(), + "last_message": message, + }) + _save_data_management_job(job) + return job.get("migration_state") if isinstance(job.get("migration_state"), dict) else state + + +def _preflight_target_ai_search_migration_access(settings, migration_plan): + """Ensure target indexes are writable and have provenance fields before copying.""" + _validate_target_ai_search_migration_write_safety(settings, migration_plan) + target_database = None + target_search_write_gate_verified = False + results = [] + search_mappings = [ + ("users", DATA_MANAGEMENT_SEARCH_ARTIFACTS[0]), + ("groups", DATA_MANAGEMENT_SEARCH_ARTIFACTS[1]), + ("public_workspaces", DATA_MANAGEMENT_SEARCH_ARTIFACTS[2]), + ] + for target_type, artifact in search_mappings: + selection = (migration_plan or {}).get(target_type) or {} + if selection.get("mode") == "none" or not selection.get("include_documents"): + continue + if target_database is None: + target_database = _get_target_cosmos_database(settings) + target_search_write_gate_container = _get_target_data_management_search_write_gate_container( + target_database + ) + _verify_target_cosmos_container_access( + target_search_write_gate_container, + "/id", + ) + target_search_write_gate_verified = True + source_client = CLIENTS.get(artifact["client_key"]) + if not source_client: + raise DataManagementSettingsValidationError( + f"Source AI Search client '{artifact['name']}' is not initialized." + ) + next(iter(source_client.search( + search_text="*", + top=1, + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )), None) + index_status = _ensure_target_search_index( + settings, + artifact["index_name"], + artifact["schema_file"], + ) + target_client = _get_target_search_client(settings, artifact["index_name"]) + probe_id = f"simplechat-migration-preflight-{uuid.uuid4().hex}" + probe_document = { + "id": probe_id, + SEARCH_MIGRATION_ID_FIELD: probe_id, + SEARCH_MIGRATED_AT_FIELD: _now_iso(), + SEARCH_MIGRATION_STATUS_FIELD: "preflight", + } + probe_written = False + probe_succeeded = False + try: + upload_results = list(target_client.upload_documents( + documents=[probe_document], + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )) + if not upload_results or not _get_search_result_succeeded(upload_results[0]): + raise DataManagementSettingsValidationError( + f"Destination AI Search index '{artifact['index_name']}' rejected a write probe." + ) + probe_written = True + read_results = list(target_client.search( + search_text="*", + filter=f"id eq '{_escape_search_filter_value(probe_id)}'", + select=["id"], + top=1, + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )) + if not any(_safe_text(dict(item).get("id")) == probe_id for item in read_results): + raise DataManagementSettingsValidationError( + f"Destination AI Search index '{artifact['index_name']}' did not return its write probe." + ) + probe_succeeded = True + finally: + if probe_written: + try: + delete_results = list(target_client.delete_documents( + documents=[{"id": probe_id}], + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )) + if delete_results and not _get_search_result_succeeded(delete_results[0]): + raise DataManagementSettingsValidationError( + f"Destination AI Search index '{artifact['index_name']}' rejected write-probe cleanup." + ) + except Exception: + if probe_succeeded: + raise + results.append({ + "target_type": target_type, + "index_name": artifact["index_name"], + "source_read_verified": True, + "destination_write_verified": True, + "index_status": index_status, + }) + return { + "index_count": len(results), + "indexes": results, + "target_search_write_gate_verified": target_search_write_gate_verified, + } + + +def _preflight_target_blob_migration_access(settings, migration_plan): + """Verify source and target Blob clients before the long-running copy phase.""" + if not (migration_plan or {}).get("include_source_blobs"): + return {"enabled": False, "container_count": 0, "containers": []} + source_client = _get_source_blob_service_client() + if not source_client: + raise DataManagementSettingsValidationError( + "Source Enhanced Citations storage is not configured for blob migration." + ) + target_client = _get_target_enhanced_citations_blob_client(settings) + container_names = [] + for target_type, selection in (migration_plan or {}).items(): + if target_type not in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES: + continue + if not isinstance(selection, dict) or selection.get("mode") == "none" or not selection.get("include_documents"): + continue + if target_type == "users": + container_names.append(app_config.storage_account_user_documents_container_name) + elif target_type == "groups": + container_names.append(app_config.storage_account_group_documents_container_name) + elif target_type == "public_workspaces": + container_names.append(app_config.storage_account_public_documents_container_name) + results = [] + for container_name in sorted(set(container_names)): + source_container = source_client.get_container_client(container_name) + next(iter(source_container.list_blobs()), None) + try: + target_client.create_container(container_name) + except ResourceExistsError: + target_client.get_container_client(container_name).get_container_properties() + probe_blob_name = f"simplechat-migration-preflight/{uuid.uuid4().hex}" + target_blob = target_client.get_blob_client( + container=container_name, + blob=probe_blob_name, + ) + probe_written = False + probe_succeeded = False + try: + target_blob.upload_blob( + data=b"", + overwrite=False, + metadata={"simplechatMigrationPreflight": "true"}, + ) + probe_written = True + target_blob.get_blob_properties() + probe_succeeded = True + finally: + if probe_written: + try: + target_blob.delete_blob() + except Exception: + if probe_succeeded: + raise + results.append({ + "container_name": container_name, + "source_read_verified": True, + "destination_write_verified": True, + }) + return {"enabled": True, "container_count": len(results), "containers": results} + + +def _run_data_management_migration_preflight(job, migration_state, settings, migration_plan): + """Persist preflight evidence before any selected user data is written.""" + migration_state["preflight"] = { + "status": "in_progress", + "started_at": _now_iso(), + } + migration_state = _persist_migration_state( + job, + migration_state, + settings, + "Validating source and destination migration access", + ) + _assert_migration_job_lease(job) + cosmos_result = _preflight_target_cosmos_migration_access(settings, migration_plan) + _assert_migration_job_lease(job) + search_result = _preflight_target_ai_search_migration_access(settings, migration_plan) + _assert_migration_job_lease(job) + blob_result = _preflight_target_blob_migration_access(settings, migration_plan) + capacity_result = None + if _safe_bool((settings or {}).get("migration_temporary_destination_ru_enabled"), False): + capacity_result = _inspect_target_cosmos_migration_capacity(settings, migration_plan) + migration_state["preflight"] = { + "status": "completed", + "completed_at": _now_iso(), + "cosmos": cosmos_result, + "ai_search": search_result, + "source_blobs": blob_result, + "destination_capacity": capacity_result, + } + return _persist_migration_state( + job, + migration_state, + settings, + "Migration source and destination preflight completed", + ) + + +def _get_destination_capacity_current(management_settings, capacity_target): + if capacity_target.get("scope") == "container": + return get_container_throughput( + management_settings, + capacity_target.get("container_name"), + ) + return get_database_throughput(management_settings) + + +def _apply_temporary_destination_capacity(job, migration_state, settings, migration_plan): + """Temporarily increase only migration-targeted destination throughput to <=10,000 RU/s.""" + _assert_migration_job_lease(job) + if not _safe_bool((settings or {}).get("migration_temporary_destination_ru_enabled"), False): + migration_state["capacity"] = {"status": "not_requested", "restore_pending": False} + return _persist_migration_state( + job, + migration_state, + settings, + "Temporary destination Cosmos capacity boost is disabled", + ) + + capacity = migration_state.get("capacity") if isinstance(migration_state.get("capacity"), dict) else {} + if capacity.get("restore_pending") and capacity.get("targets"): + if capacity.get("status") in {"applying", "restore_pending"}: + _restore_temporary_destination_capacity(job, migration_state, settings) + capacity = migration_state.get("capacity") if isinstance(migration_state.get("capacity"), dict) else {} + if capacity.get("restore_pending"): + raise DataManagementSettingsValidationError( + "Destination Cosmos capacity restoration is still pending. Resolve the recorded restore error before retrying migration." + ) + return _persist_migration_state( + job, + migration_state, + settings, + "Recovered pending destination Cosmos capacity without reapplying a boost", + ) + else: + return migration_state + + inspection = _inspect_target_cosmos_migration_capacity(settings, migration_plan) + target_ru = inspection["target_ru"] + capacity = { + "status": "applying", + "restore_pending": True, + "target_ru": target_ru, + "management_settings": inspection["management_settings"], + "targets": [], + "started_at": _now_iso(), + } + migration_state["capacity"] = capacity + migration_state = _persist_migration_state( + job, + migration_state, + settings, + f"Preparing temporary destination Cosmos capacity up to {target_ru} RU/s", + ) + capacity = migration_state.setdefault("capacity", capacity) + + for target in inspection["targets"]: + _assert_migration_job_lease(job) + original_ru = _safe_int(target.get("current_ru"), default=0, minimum=0) + target_snapshot = { + "scope": target.get("scope"), + "container_name": target.get("container_name") or "", + "mode": target.get("mode"), + "original_ru": original_ru, + "target_ru": target_ru, + "changed": False, + "boost_attempted": False, + "restore_status": "not_required", + } + capacity["targets"].append(target_snapshot) + if original_ru and original_ru < target_ru: + target_snapshot["boost_attempted"] = True + target_snapshot["restore_status"] = "pending" + migration_state = _persist_migration_state( + job, + migration_state, + settings, + f"Recorded destination Cosmos capacity recovery snapshot for {target_snapshot['scope']}", + ) + capacity = migration_state.setdefault("capacity", capacity) + target_snapshot = capacity["targets"][-1] + _assert_migration_job_lease(job) + try: + scale_result = set_database_throughput( + inspection["management_settings"], + target_ru, + initiated_by=f"data_management_migration:{job.get('id')}", + reason="temporary_migration_capacity_boost", + decision={ + "scope": target_snapshot["scope"], + "container_name": target_snapshot["container_name"], + "target_mode": target_snapshot["mode"], + }, + ) + except Exception as exc: + target_snapshot["boost_error"] = str(exc)[:300] + migration_state["capacity"] = capacity + _persist_migration_state( + job, + migration_state, + settings, + f"Destination Cosmos capacity boost failed for {target_snapshot['scope']}", + ) + raise + target_snapshot["changed"] = True + target_snapshot["boosted_to_ru"] = scale_result.get("to_ru", target_ru) + target_snapshot["restore_status"] = "pending" + else: + target_snapshot["restore_status"] = "not_required" + migration_state = _persist_migration_state( + job, + migration_state, + settings, + f"Updated temporary destination Cosmos capacity for {target_snapshot['scope']}", + ) + capacity = migration_state.setdefault("capacity", capacity) + + capacity["status"] = "boosted" + capacity["completed_at"] = _now_iso() + return _persist_migration_state( + job, + migration_state, + settings, + f"Temporary destination Cosmos capacity is ready up to {target_ru} RU/s", + ) + + +def _restore_temporary_destination_capacity( + job, + migration_state, + settings, + allow_cancel_requested=False, +): + """Restore only the capacity values this migration changed, without clobbering outside updates.""" + _assert_migration_job_lease( + job, + allow_cancel_requested=allow_cancel_requested, + ) + capacity = migration_state.get("capacity") if isinstance(migration_state.get("capacity"), dict) else {} + if not capacity.get("restore_pending"): + return [], migration_state + management_settings = capacity.get("management_settings") or {} + restore_warnings = [] + for target in reversed(capacity.get("targets") or []): + _assert_migration_job_lease( + job, + allow_cancel_requested=allow_cancel_requested, + ) + if not (target.get("changed") or target.get("boost_attempted")): + continue + try: + current = _get_destination_capacity_current(management_settings, target) + current_ru = _safe_int(current.get("current_ru"), default=0, minimum=0) + boosted_to_ru = _safe_int(target.get("boosted_to_ru"), default=target.get("target_ru"), minimum=0) + original_ru = _safe_int(target.get("original_ru"), default=0, minimum=0) + if current_ru == original_ru: + target["restore_status"] = "not_changed" + continue + if current_ru != boosted_to_ru: + target["restore_status"] = "skipped_external_change" + restore_warnings.append( + f"Did not restore destination {target.get('scope')} capacity because it changed after the migration boost." + ) + continue + scale_result = set_database_throughput( + management_settings, + original_ru, + initiated_by=f"data_management_migration:{job.get('id')}", + reason="restore_temporary_migration_capacity", + decision={ + "scope": target.get("scope"), + "container_name": target.get("container_name") or "", + "target_mode": target.get("mode"), + }, + ) + target["restore_status"] = "restored" + target["restored_to_ru"] = scale_result.get("to_ru", original_ru) + except Exception as exc: + target["restore_status"] = "restore_failed" + restore_warnings.append( + f"Failed to restore destination {target.get('scope')} capacity: {str(exc)[:300]}" + ) + + unresolved_restore_statuses = {"pending", "restore_failed"} + capacity["restore_pending"] = any( + target.get("restore_status") in unresolved_restore_statuses + for target in capacity.get("targets") or [] + ) + capacity["restored_at"] = _now_iso() + if capacity["restore_pending"]: + capacity["status"] = "restore_pending" + else: + capacity["status"] = "restored_with_warnings" if restore_warnings else "restored" + capacity["restore_warnings"] = restore_warnings + migration_state["capacity"] = capacity + migration_state = _persist_migration_state( + job, + migration_state, + settings, + "Restored temporary destination Cosmos capacity", + allow_cancel_requested=allow_cancel_requested, + ) + return restore_warnings, migration_state + + +def _iter_selected_cosmos_records( + container_definition, + selection, + source_cutoff_epoch=None, + source_start_epoch=None, + include_source_version=False, +): + source_container = getattr(app_config, container_definition["container_attr"], None) + if not source_container: + return + if container_definition.get("documents") and not selection.get("include_documents"): + return + mode = selection.get("mode") + + def prepare_item(item): + cleaned_item = _strip_cosmos_system_fields(item) + if include_source_version: + return cleaned_item, _safe_text(item.get("_ts")) + return cleaned_item + + if mode == "all": + query = "SELECT * FROM c" + parameters = [] + conditions = [] + if source_start_epoch is not None: + conditions.append("c._ts >= @source_start_epoch") + parameters.append({"name": "@source_start_epoch", "value": source_start_epoch}) + if source_cutoff_epoch is not None: + conditions.append("c._ts <= @source_cutoff_epoch") + parameters = [{"name": "@source_cutoff_epoch", "value": source_cutoff_epoch}] + if source_start_epoch is not None: + parameters.insert(0, {"name": "@source_start_epoch", "value": source_start_epoch}) + if conditions: + query = f"SELECT * FROM c WHERE {' AND '.join(conditions)}" + for item in source_container.query_items( + query=query, + parameters=parameters, + enable_cross_partition_query=True, + ): + yield prepare_item(item) + return + if mode != "selected": + return + + ids = selection.get("ids") or [] + if container_definition.get("id_field") == "id": + for item_id in ids: + try: + item = source_container.read_item(item=item_id, partition_key=item_id) + except CosmosResourceNotFoundError: + continue + if ( + source_cutoff_epoch is not None and + item.get("_ts") is not None and + _safe_int(item.get("_ts"), default=0, minimum=0) > source_cutoff_epoch + ): + continue + if ( + source_start_epoch is not None and + item.get("_ts") is not None and + _safe_int(item.get("_ts"), default=0, minimum=0) < source_start_epoch + ): + continue + yield prepare_item(item) + return + + filter_field = container_definition.get("filter_field") + if not filter_field: + return + for selected_id in ids: + query = f"SELECT * FROM c WHERE c.{filter_field} = @selected_id" + parameters = [{"name": "@selected_id", "value": selected_id}] + if source_start_epoch is not None: + query += " AND c._ts >= @source_start_epoch" + parameters.append({"name": "@source_start_epoch", "value": source_start_epoch}) + if source_cutoff_epoch is not None: + query += " AND c._ts <= @source_cutoff_epoch" + parameters.append({"name": "@source_cutoff_epoch", "value": source_cutoff_epoch}) + for item in source_container.query_items( + query=query, + parameters=parameters, + enable_cross_partition_query=True, + ): + yield prepare_item(item) + + +def _build_cosmos_document_identity(document_id, partition_key_value): + if document_id is None or partition_key_value is None: + return "" + return json.dumps( + {"id": str(document_id), "partition_key": partition_key_value}, + ensure_ascii=True, + sort_keys=True, + default=str, + separators=(",", ":"), + ) + + +def _get_cosmos_document_identity(document, partition_key_path): + if not isinstance(document, dict): + return "" + return _build_cosmos_document_identity( + document.get("id"), + _get_document_path_value(document, partition_key_path), + ) + + +def _build_cosmos_source_hash(document): + """Hash source content without service fields or migration provenance.""" + canonical_document = { + key: value + for key, value in (document or {}).items() + if key not in { + COSMOS_MIGRATION_PROVENANCE_FIELD, + "_etag", + "_rid", + "_self", + "_attachments", + "_ts", + } + } + encoded = json.dumps( + canonical_document, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + default=_json_default, + ).encode("utf-8") + return f"sha256:{hashlib.sha256(encoded).hexdigest()}" + + +def _get_cosmos_partition_key_select_expression(partition_key_path): + path_parts = [part for part in _safe_text(partition_key_path).strip("/").split("/") if part] + if not path_parts or any(not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", part) for part in path_parts): + raise DataManagementSettingsValidationError("Destination Cosmos container has an unsupported partition key path.") + return "c." + ".".join(path_parts) + + +def _get_target_cosmos_document(target_container, document, partition_key_path): + """Read one exact destination item without allocating an account-wide key set.""" + document_id = _safe_text((document or {}).get("id")) + partition_key_value = _get_document_path_value(document, partition_key_path) + if not document_id or partition_key_value is None: + raise DataManagementSettingsValidationError( + "Migration source record is missing a required ID or partition key value." + ) + try: + return target_container.read_item( + item=document_id, + partition_key=partition_key_value, + ) + except (CosmosResourceNotFoundError, ResourceNotFoundError): + return None + + +def _classify_target_cosmos_document( + target_document, + provenance_context, + container_name, + migration_mode=DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + source_hash="", + source_version="", +): + """Allow only empty or migration-owned target identities to be written.""" + if target_document is None: + return "create" + target_provenance = get_cosmos_migration_provenance(target_document) + if is_successful_migration_record(target_provenance): + if ( + _safe_text(target_provenance.get("migrationId")) == + _safe_text((provenance_context or {}).get("migration_id")) + ): + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + return "resume_verified" + return "update" + if migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + return "unchanged" + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + return "unchanged" + return "update" + raise DataManagementSettingsValidationError( + f"Destination Cosmos container '{container_name}' contains an unowned record that conflicts with this migration." + ) + + +def _get_cosmos_response_request_units(headers): + for header_name, header_value in (headers or {}).items(): + if str(header_name).lower() != "x-ms-request-charge": + continue + try: + return max(0.0, float(header_value)) + except (TypeError, ValueError): + return 0.0 + return 0.0 + + +def _drain_migration_futures_with_heartbeat( + pending_futures, + consume_result, + persist_heartbeat, + cancel_event, + persist_stopped_results=None, +): + pending_by_future = ( + dict(pending_futures) + if isinstance(pending_futures, (dict, list, tuple)) + else {future: None for future in pending_futures} + ) + last_heartbeat = time.monotonic() + stop_error = None + while pending_by_future: + completed_futures, _pending_futures = wait( + set(pending_by_future), + timeout=DATA_MANAGEMENT_MIGRATION_HEARTBEAT_POLL_SECONDS, + return_when=FIRST_COMPLETED, + ) + for future in completed_futures: + context = pending_by_future.pop(future) + if future.cancelled(): + continue + try: + consume_result(future.result(), context) + except Exception as exc: + if stop_error is None: + stop_error = exc + cancel_event.set() + for pending_future in pending_by_future: + pending_future.cancel() + now = time.monotonic() + if ( + pending_by_future and + stop_error is None and + now - last_heartbeat >= DATA_MANAGEMENT_MIGRATION_HEARTBEAT_INTERVAL_SECONDS + ): + try: + persist_heartbeat() + except Exception as exc: + stop_error = exc + cancel_event.set() + for future in pending_by_future: + future.cancel() + last_heartbeat = now + if stop_error is not None: + if callable(persist_stopped_results): + try: + persist_stopped_results() + except Exception: + pass + raise stop_error + + +def _write_cosmos_migration_record( + target_container, + document, + partition_key_path, + provenance_context, + retry_count, + disposition="create", + target_document=None, + source_hash="", + source_version="", + cancel_event=None, +): + """Write one provenance-tagged Cosmos record with bounded transient retries.""" + writable_document = copy.deepcopy(document) + add_cosmos_migration_provenance( + writable_document, + provenance_context, + source_hash=source_hash, + source_version=source_version, + ) + payload_bytes = len(json.dumps(writable_document, default=_json_default).encode("utf-8")) + request_units = 0.0 + request_unit_lock = Lock() + + def response_hook(headers, _response): + nonlocal request_units + with request_unit_lock: + request_units += _get_cosmos_response_request_units(headers) + + started_at = time.perf_counter() + last_error = None + for attempt in range(1, retry_count + 1): + if cancel_event is not None and cancel_event.is_set(): + raise DataManagementMigrationCanceledError( + "Cosmos migration stopped after cancellation or lease loss." + ) + try: + if disposition == "update": + target_container.replace_item( + item=target_document, + body=writable_document, + etag=(target_document or {}).get("_etag"), + match_condition=MatchConditions.IfNotModified, + response_hook=response_hook, + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + ) + else: + target_container.create_item( + writable_document, + response_hook=response_hook, + retry_write=1, + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + ) + return { + "copied": True, + "created": disposition == "create", + "updated": disposition == "update", + "bytes": payload_bytes, + "request_units": request_units, + "attempt": attempt, + "elapsed_seconds": time.perf_counter() - started_at, + } + except Exception as exc: + last_error = exc + status_code = getattr(exc, "status_code", None) + if status_code in {409, 412}: + target_document = _get_target_cosmos_document( + target_container, + document, + partition_key_path, + ) + try: + disposition = _classify_target_cosmos_document( + target_document, + provenance_context, + "migration target", + migration_mode=(provenance_context or {}).get( + "migration_mode", + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ), + source_hash=source_hash, + source_version=source_version, + ) + except DataManagementSettingsValidationError as collision_error: + return { + "copied": False, + "skipped": False, + "collision": True, + "bytes": 0, + "request_units": request_units, + "attempt": attempt, + "elapsed_seconds": time.perf_counter() - started_at, + "error": str(collision_error), + } + if disposition in {"unchanged", "resume_verified"}: + return { + "copied": False, + "skipped": True, + "resume_verified": disposition == "resume_verified", + "bytes": 0, + "request_units": request_units, + "attempt": attempt, + "elapsed_seconds": time.perf_counter() - started_at, + } + if disposition == "update": + continue + retryable = status_code in {408, 429, 449, 500, 503} or status_code is None + if not retryable or attempt >= retry_count: + break + retry_delay = min(30, 2 ** (attempt - 1)) + if cancel_event is not None: + if cancel_event.wait(retry_delay): + raise DataManagementMigrationCanceledError( + "Cosmos migration stopped during retry backoff." + ) + else: + time.sleep(retry_delay) + + return { + "copied": False, + "skipped": False, + "collision": False, + "bytes": 0, + "request_units": request_units, + "attempt": retry_count, + "elapsed_seconds": time.perf_counter() - started_at, + "error": str(last_error or "Cosmos document write failed.")[:500], + } + + +def _copy_cosmos_records_to_target( + target_database, + target_type, + selection, + job, + migration_state, + provenance_context, + settings, +): + copied = [] + source_cutoff = _parse_iso_datetime(migration_state.get("source_cutoff_at")) + source_cutoff_epoch = int(source_cutoff.timestamp()) if source_cutoff else None + migration_mode = _safe_text( + (provenance_context or {}).get("migration_mode"), + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ) + baseline_cutoff = _parse_iso_datetime( + (provenance_context or {}).get("baseline_source_cutoff_at") + ) + source_start_epoch = ( + int(baseline_cutoff.timestamp()) + if baseline_cutoff and migration_mode != DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY + else None + ) + parallel_operations = _get_migration_parallel_operations(settings) + retry_count = _get_migration_retry_count(settings) + for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type]: + if container_definition.get("documents") and not selection.get("include_documents"): + continue + target_container_name = _target_cosmos_container_name(container_definition) + resource_name = f"cosmos:{target_type}:{container_definition['name']}" + completed_resource = migration_state.get("resources", {}).get(resource_name) + if is_migration_resource_completed(migration_state, resource_name): + previous_result = completed_resource.get("result") if isinstance(completed_resource, dict) else {} + copied.append({ + "name": container_definition["name"], + "type": "cosmos_container", + "target_type": target_type, + "container_name": target_container_name, + "partition_key_path": container_definition["partition_key_path"], + "status": "checkpoint_completed", + **(previous_result if isinstance(previous_result, dict) else {}), + }) + continue + + target_container = _get_target_cosmos_container( + target_database, + target_container_name, + container_definition["partition_key_path"], + ) + resource = start_migration_resource(migration_state, resource_name) + resource_started_at = resource.get("attempt_started_at") or resource.get("started_at") + previous_progress = resource.get("progress") if isinstance(resource.get("progress"), dict) else {} + copied_count = _safe_int(previous_progress.get("copied_count"), default=0, minimum=0) + skipped_count = _safe_int(previous_progress.get("skipped_count"), default=0, minimum=0) + prior_failed_count = _safe_int(previous_progress.get("failed_count"), default=0, minimum=0) + failed_count = 0 + collision_count = _safe_int(previous_progress.get("collision_count"), default=0, minimum=0) + destination_provenance_skip_count = _safe_int( + previous_progress.get("destination_provenance_skip_count"), + default=0, + minimum=0, + ) + created_count = _safe_int(previous_progress.get("created_count"), default=0, minimum=0) + updated_count = _safe_int(previous_progress.get("updated_count"), default=0, minimum=0) + unchanged_count = _safe_int(previous_progress.get("unchanged_count"), default=0, minimum=0) + source_read_count = _safe_int(previous_progress.get("source_read_count"), default=0, minimum=0) + retry_attempt_count = _safe_int(previous_progress.get("retry_attempt_count"), default=0, minimum=0) + byte_count = _safe_int(previous_progress.get("bytes"), default=0, minimum=0) + request_units = float(previous_progress.get("request_units") or 0.0) + errors = [] + append_manifest, flush_manifest = _create_migration_manifest_writer( + job.get("id"), + resource_name, + ) + transfer_cancel_event = Event() + + def persist_progress(allow_cancel_requested=False): + flush_manifest() + progress = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + request_units=request_units, + ) + progress.update({ + "parallel_operations": parallel_operations, + "retry_count": retry_count, + "destination_provenance_skip_count": destination_provenance_skip_count, + "collision_count": collision_count, + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "source_read_count": source_read_count, + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "retry_attempt_count": retry_attempt_count, + "prior_failed_count": prior_failed_count, + }) + return _persist_migration_checkpoint( + job, + migration_state, + settings, + resource_name, + progress, + f"Migrating Cosmos container {target_container_name}", + allow_cancel_requested=allow_cancel_requested, + ) + + def persist_heartbeat(): + return _persist_migration_heartbeat( + job, + settings, + f"Waiting for Cosmos writes to {target_container_name}", + ) + + def persist_stopped_results(): + return persist_progress(allow_cancel_requested=True) + + def consume_write_result(write_result, manifest_context): + nonlocal copied_count, skipped_count, failed_count + nonlocal collision_count, destination_provenance_skip_count + nonlocal created_count, updated_count, unchanged_count + nonlocal byte_count, request_units, retry_attempt_count + request_units += write_result.get("request_units", 0.0) + retry_attempt_count += max(0, _safe_int(write_result.get("attempt"), default=1) - 1) + if write_result.get("copied"): + copied_count += 1 + if write_result.get("updated"): + updated_count += 1 + else: + created_count += 1 + byte_count += write_result.get("bytes", 0) + elif write_result.get("skipped"): + if not write_result.get("resume_verified"): + skipped_count += 1 + unchanged_count += 1 + destination_provenance_skip_count += 1 + else: + failed_count += 1 + if write_result.get("collision"): + collision_count += 1 + errors.append(write_result.get("error", "Cosmos write failed.")) + status = ( + "updated" + if write_result.get("updated") else + "created" + if write_result.get("copied") else + "resume_verified" + if write_result.get("resume_verified") else + "unchanged" + if write_result.get("skipped") else + "collision" + if write_result.get("collision") else + "failed" + ) + append_manifest({ + **manifest_context, + "status": status, + "attempt": write_result.get("attempt"), + "bytes": write_result.get("bytes", 0), + "error": write_result.get("error"), + }) + + try: + pending_writes = [] + with ThreadPoolExecutor(max_workers=parallel_operations) as executor: + for item, source_version in _iter_selected_cosmos_records( + container_definition, + selection, + source_cutoff_epoch=source_cutoff_epoch, + source_start_epoch=source_start_epoch, + include_source_version=True, + ) or []: + source_read_count += 1 + source_hash = _build_cosmos_source_hash(item) + manifest_context = { + "service": "cosmos", + "resource_name": resource_name, + "target_type": target_type, + "source_identity": _hash_migration_manifest_identity( + target_type, + container_definition["name"], + item.get("id"), + _get_document_path_value( + item, + container_definition["partition_key_path"], + ), + ), + "destination_identity": _hash_migration_manifest_identity( + target_container_name, + item.get("id"), + _get_document_path_value( + item, + container_definition["partition_key_path"], + ), + ), + "source_version": source_version, + "source_hash": source_hash, + "_locator": { + "service": "cosmos", + "resource_name": resource_name, + "target_type": target_type, + "document_id": item.get("id"), + "partition_key": _get_document_path_value( + item, + container_definition["partition_key_path"], + ), + "container_name": target_container_name, + }, + } + target_document = _get_target_cosmos_document( + target_container, + item, + container_definition["partition_key_path"], + ) + disposition = _classify_target_cosmos_document( + target_document, + provenance_context, + target_container_name, + migration_mode=migration_mode, + source_hash=source_hash, + source_version=source_version, + ) + if disposition in {"unchanged", "resume_verified"}: + if disposition == "unchanged": + skipped_count += 1 + unchanged_count += 1 + destination_provenance_skip_count += 1 + append_manifest({**manifest_context, "status": disposition}) + if (copied_count + skipped_count + failed_count) % max(100, parallel_operations * 4) == 0: + migration_state = persist_progress() + continue + + if not pending_writes: + _assert_migration_job_lease(job) + pending_writes.append(( + executor.submit( + _write_cosmos_migration_record, + target_container, + item, + container_definition["partition_key_path"], + provenance_context, + retry_count, + disposition=disposition, + target_document=target_document, + source_hash=source_hash, + source_version=source_version, + cancel_event=transfer_cancel_event, + ), + manifest_context, + )) + if len(pending_writes) < parallel_operations: + continue + + _drain_migration_futures_with_heartbeat( + pending_writes, + consume_write_result, + persist_heartbeat, + transfer_cancel_event, + persist_stopped_results=persist_stopped_results, + ) + pending_writes = [] + migration_state = persist_progress() + + _drain_migration_futures_with_heartbeat( + pending_writes, + consume_write_result, + persist_heartbeat, + transfer_cancel_event, + persist_stopped_results=persist_stopped_results, + ) + + flush_manifest() + + result = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + request_units=request_units, + ) + result.update({ + "name": container_definition["name"], + "type": "cosmos_container", + "target_type": target_type, + "container_name": target_container_name, + "partition_key_path": container_definition["partition_key_path"], + "parallel_operations": parallel_operations, + "retry_count": retry_count, + "destination_provenance_skip_count": destination_provenance_skip_count, + "collision_count": collision_count, + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "source_read_count": source_read_count, + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "retry_attempt_count": retry_attempt_count, + "prior_failed_count": prior_failed_count, + }) + if failed_count: + error_summary = "; ".join(errors[:3]) + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + error_summary, + ) + raise RuntimeError( + f"Cosmos migration failed for container '{target_container_name}' after {failed_count} item write failure(s)." + ) + migration_state = _complete_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + result, + f"Completed Cosmos container {target_container_name}", + ) + copied.append(result) + except (DataManagementMigrationCanceledError, DataManagementMigrationLeaseLostError): + flush_manifest() + raise + except Exception as exc: + flush_manifest() + if not failed_count: + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + str(exc), + ) + raise + return copied + + +def _escape_search_filter_value(value): + return _safe_text(value).replace("'", "''") + + +def _build_search_filter(field_name, selection): + if selection.get("mode") == "all": + return None + if selection.get("mode") != "selected": + return "id eq '__no_migration_selection__'" + conditions = [f"{field_name} eq '{_escape_search_filter_value(item_id)}'" for item_id in selection.get("ids") or []] + if not conditions: + return "id eq '__no_migration_selection__'" + return " or ".join(conditions) + + +def _combine_search_filters(*search_filters): + normalized_filters = [ + _safe_text(search_filter) + for search_filter in search_filters + if _safe_text(search_filter) + ] + if not normalized_filters: + return None + return " and ".join(f"({search_filter})" for search_filter in normalized_filters) + + +def _build_search_scope_filter_batches(field_name, selection): + """Bound selected-scope filters so large migrations stay within Search limits.""" + if selection.get("mode") == "all": + return [None] + selected_ids = sorted({ + _safe_text(item_id) + for item_id in (selection.get("ids") or []) + if _safe_text(item_id) + }) + if selection.get("mode") != "selected" or not selected_ids: + return ["id eq '__no_migration_selection__'"] + return [ + _build_search_filter(field_name, { + "mode": "selected", + "ids": selected_ids[offset:offset + DATA_MANAGEMENT_SEARCH_SCOPE_FILTER_BATCH_SIZE], + }) + for offset in range(0, len(selected_ids), DATA_MANAGEMENT_SEARCH_SCOPE_FILTER_BATCH_SIZE) + ] + + +def _iter_search_document_pages( + search_client, + search_filter=None, + page_size=None, + last_id="", + select=None, +): + """Yield stable keyset pages without allowing the SDK to fall back to deep skip.""" + safe_page_size = _safe_int( + page_size, + default=DATA_MANAGEMENT_SEARCH_KEYSET_PAGE_SIZE, + minimum=1, + maximum=DATA_MANAGEMENT_SEARCH_KEYSET_PAGE_SIZE, + ) + cursor_id = _safe_text(last_id) + while True: + cursor_filter = ( + f"id gt '{_escape_search_filter_value(cursor_id)}'" + if cursor_id + else None + ) + search_kwargs = { + "search_text": "*", + "filter": _combine_search_filters(search_filter, cursor_filter), + "order_by": ["id asc"], + "top": safe_page_size, + "include_total_count": False, + } + if select: + search_kwargs["select"] = select + results = search_client.search( + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + **search_kwargs, + ) + page = [] + previous_id = cursor_id + for result in results: + document = { + key: value + for key, value in dict(result).items() + if not key.startswith("@search.") + } + document_id = _safe_text(document.get("id")) + if not document_id: + raise DataManagementSettingsValidationError( + "Source AI Search returned a document without an ID during keyset pagination." + ) + if previous_id and document_id <= previous_id: + raise DataManagementSettingsValidationError( + "Source AI Search did not honor the required strictly increasing ID order." + ) + previous_id = document_id + page.append(document) + if not page: + break + cursor_id = _safe_text(page[-1].get("id")) + yield page, cursor_id + if len(page) < safe_page_size: + break + + +def _get_target_ai_search_credential(settings): + if (settings or {}).get("target_ai_search_authentication_type") == "key": + key = _safe_text((settings or {}).get("target_ai_search_key")) + if not key: + raise ValueError("Target AI Search key is required when key authentication is selected.") + return AzureKeyCredential(key) + return DefaultAzureCredential() + + +def _get_target_search_client(settings, index_name): + endpoint = _safe_text((settings or {}).get("target_ai_search_endpoint")) + if not endpoint: + raise ValueError("Target AI Search endpoint is required before running AI Search migration.") + return SearchClient( + endpoint=endpoint, + index_name=index_name, + credential=_get_target_ai_search_credential(settings), + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + ) + + +def _normalize_search_service_endpoint(endpoint): + """Compare Search service endpoints without credentials or request paths.""" + normalized_endpoint = _safe_text(endpoint).strip().rstrip("/").lower() + parsed_endpoint = urlparse( + normalized_endpoint if "://" in normalized_endpoint else f"https://{normalized_endpoint}" + ) + return _safe_text(parsed_endpoint.netloc or normalized_endpoint).rstrip("/").lower() + + +def _get_search_client_service_endpoint(search_client): + """Read the SDK endpoint when the source client exposes one.""" + return _normalize_search_service_endpoint( + getattr(search_client, "_endpoint", None) or + getattr(search_client, "endpoint", None) + ) + + +def _migration_plan_includes_ai_search_documents(migration_plan): + """Return whether this plan will write selected document chunks to AI Search.""" + if not (migration_plan or {}).get("include_ai_search"): + return False + return any( + (migration_plan.get(target_type) or {}).get("mode") != "none" and + bool((migration_plan.get(target_type) or {}).get("include_documents")) + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES + ) + + +def _validate_target_ai_search_migration_write_safety(settings, migration_plan): + """Require a frozen, distinct target before non-conditional Search document writes.""" + if not _migration_plan_includes_ai_search_documents(migration_plan): + return + if not bool((migration_plan or {}).get("target_ai_search_writes_frozen")): + raise DataManagementSettingsValidationError( + DATA_MANAGEMENT_SEARCH_WRITE_FREEZE_CONFIRMATION_ERROR + ) + target_endpoint = _normalize_search_service_endpoint( + (settings or {}).get("target_ai_search_endpoint") + ) + if not target_endpoint: + return + for artifact in DATA_MANAGEMENT_SEARCH_ARTIFACTS: + source_client = CLIENTS.get(artifact["client_key"]) + source_endpoint = _get_search_client_service_endpoint(source_client) + if source_endpoint and source_endpoint == target_endpoint: + raise DataManagementSettingsValidationError( + "Target AI Search endpoint must differ from the selected source Search service." + ) + + +def _get_migration_search_provenance_fields(): + return [ + SearchField( + name=SEARCH_MIGRATION_ID_FIELD, + type=SearchFieldDataType.String, + searchable=False, + filterable=True, + sortable=True, + facetable=False, + retrievable=True, + ), + SearchField( + name=SEARCH_MIGRATED_AT_FIELD, + type=SearchFieldDataType.DateTimeOffset, + filterable=True, + sortable=True, + facetable=False, + retrievable=True, + ), + SearchField( + name=SEARCH_MIGRATION_STATUS_FIELD, + type=SearchFieldDataType.String, + searchable=False, + filterable=True, + sortable=True, + facetable=False, + retrievable=True, + ), + SearchField( + name=SEARCH_MIGRATION_SOURCE_HASH_FIELD, + type=SearchFieldDataType.String, + searchable=False, + filterable=True, + sortable=False, + facetable=False, + retrievable=True, + ), + SearchField( + name=SEARCH_MIGRATION_SOURCE_VERSION_FIELD, + type=SearchFieldDataType.String, + searchable=False, + filterable=True, + sortable=False, + facetable=False, + retrievable=True, + ), + ] + + +def _ensure_target_search_migration_provenance_fields(index): + """Add compatible queryable provenance fields without rebuilding an index.""" + existing_fields = { + field.name: field + for field in (getattr(index, "fields", None) or []) + if getattr(field, "name", None) + } + fields = list(getattr(index, "fields", None) or []) + added_fields = 0 + for required_field in _get_migration_search_provenance_fields(): + existing_field = existing_fields.get(required_field.name) + if existing_field is None: + fields.append(required_field) + added_fields += 1 + continue + if ( + str(getattr(existing_field, "type", "")) != str(required_field.type) or + getattr(existing_field, "filterable", False) is not True or + getattr(existing_field, "retrievable", False) is not True + ): + raise DataManagementSettingsValidationError( + f"Target AI Search index '{getattr(index, 'name', '')}' has incompatible " + f"migration provenance field '{required_field.name}'." + ) + if added_fields: + index.fields = fields + return added_fields + + +def _ensure_target_search_index(settings, index_name, schema_file): + endpoint = _safe_text((settings or {}).get("target_ai_search_endpoint")) + if not endpoint: + raise ValueError("Target AI Search endpoint is required before running AI Search migration.") + index_client = SearchIndexClient( + endpoint=endpoint, + credential=_get_target_ai_search_credential(settings), + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + ) + try: + index = index_client.get_index(index_name) + added_fields = _ensure_target_search_migration_provenance_fields(index) + if added_fields: + index_client.create_or_update_index(index) + return "updated_with_migration_provenance" + return "exists_with_migration_provenance" + except ResourceNotFoundError: + schema = _get_search_schema(schema_file) + schema = {key: value for key, value in schema.items() if not key.startswith("@odata.")} + try: + index = SearchIndex.from_dict(schema) + _ensure_target_search_migration_provenance_fields(index) + index_client.create_or_update_index(index) + return "created_with_migration_provenance" + except Exception as exc: + raise ValueError(f"Target AI Search index {index_name} is missing and could not be created: {exc}") from exc + + +def _get_target_search_documents_by_ids(search_client, document_ids): + """Read only the current bounded upload batch from the target index.""" + normalized_ids = [ + _safe_text(document_id) + for document_id in document_ids + if _safe_text(document_id) + ] + if not normalized_ids: + return {} + conditions = [ + f"id eq '{_escape_search_filter_value(document_id)}'" + for document_id in normalized_ids + ] + results = search_client.search( + search_text="*", + filter=" or ".join(conditions), + select=[ + "id", + SEARCH_MIGRATION_ID_FIELD, + SEARCH_MIGRATED_AT_FIELD, + SEARCH_MIGRATION_STATUS_FIELD, + SEARCH_MIGRATION_SOURCE_HASH_FIELD, + SEARCH_MIGRATION_SOURCE_VERSION_FIELD, + ], + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + ) + documents_by_id = {} + for result in results: + document = dict(result) + document_id = _safe_text(document.get("id")) + if document_id: + documents_by_id[document_id] = document + return documents_by_id + + +def _build_search_source_hash(document): + """Hash canonical source fields while excluding destination migration metadata.""" + excluded_fields = { + SEARCH_MIGRATION_ID_FIELD, + SEARCH_MIGRATED_AT_FIELD, + SEARCH_MIGRATION_STATUS_FIELD, + SEARCH_MIGRATION_SOURCE_HASH_FIELD, + SEARCH_MIGRATION_SOURCE_VERSION_FIELD, + } + canonical_document = { + key: value + for key, value in (document or {}).items() + if key not in excluded_fields and not key.startswith("@search.") + } + encoded = json.dumps( + canonical_document, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + default=_json_default, + ).encode("utf-8") + return f"sha256:{hashlib.sha256(encoded).hexdigest()}" + + +def _classify_target_search_document( + target_document, + provenance_context, + index_name, + migration_mode=DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + source_hash="", +): + """Allow only empty or migration-owned target keys to be uploaded.""" + if target_document is None: + return "create" + target_provenance = get_search_migration_provenance(target_document) + if is_successful_migration_record(target_provenance): + if ( + _safe_text(target_provenance.get("migrationId")) == + _safe_text((provenance_context or {}).get("migration_id")) + ): + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + ): + return "resume_verified" + return "update" + if migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + return "unchanged" + if migration_record_matches_source(target_provenance, source_hash=source_hash): + return "unchanged" + return "update" + raise DataManagementSettingsValidationError( + f"Destination AI Search index '{index_name}' contains an unowned document that conflicts with this migration." + ) + + +def _get_search_result_succeeded(result): + if isinstance(result, dict): + return bool(result.get("succeeded", result.get("status", False))) + return bool(getattr(result, "succeeded", getattr(result, "status", False))) + + +def _reclassify_pending_search_migration_entries( + search_client, + pending_entries, + provenance_context, + index_name, + migration_mode, +): + """Allow only verified response-lost writes to leave an ambiguous Search retry.""" + document_ids = [ + _safe_text((entry.get("document") or {}).get("id")) + for entry in pending_entries + ] + target_documents = _get_target_search_documents_by_ids(search_client, document_ids) + verified_entries = [] + retry_entries = [] + for entry in pending_entries: + document = entry.get("document") if isinstance(entry, dict) else {} + document_id = _safe_text((document or {}).get("id")) + target_document = target_documents.get(document_id) + if target_document is None: + retry_entries.append(entry) + continue + source_hash = _safe_text((document or {}).get(SEARCH_MIGRATION_SOURCE_HASH_FIELD)) + disposition = _classify_target_search_document( + target_document, + provenance_context, + index_name, + migration_mode=migration_mode, + source_hash=source_hash, + ) + if disposition != "resume_verified": + raise DataManagementSettingsValidationError( + f"Destination AI Search index '{index_name}' changed after an ambiguous " + f"migration write for document '{document_id}'. The migration stopped before retrying it." + ) + verified_entries.append(entry) + return verified_entries, retry_entries + + +def _upload_search_migration_batch( + search_client, + documents, + retry_count, + dispositions=None, + manifest_contexts=None, + cancel_event=None, + provenance_context=None, + index_name="", + migration_mode=DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, +): + """Upload one idempotent batch and surface partial indexing failures.""" + pending_entries = [ + { + "document": document, + "disposition": ( + dispositions[index] + if isinstance(dispositions, list) and index < len(dispositions) + else "create" + ), + "bytes": len(json.dumps(document, default=_json_default).encode("utf-8")), + "manifest_context": ( + manifest_contexts[index] + if isinstance(manifest_contexts, list) and index < len(manifest_contexts) + else {} + ), + } + for index, document in enumerate(documents) + ] + started_at = time.perf_counter() + last_error = "" + copied_count = 0 + created_count = 0 + updated_count = 0 + accepted_bytes = 0 + manifest_entries = [] + + def record_accepted_entry(entry, attempt): + nonlocal copied_count, created_count, updated_count, accepted_bytes + copied_count += 1 + accepted_bytes += entry["bytes"] + if entry["disposition"] == "update": + updated_count += 1 + else: + created_count += 1 + manifest_entries.append({ + **entry["manifest_context"], + "status": entry["disposition"] + "d", + "attempt": attempt, + "bytes": entry["bytes"], + }) + + for attempt in range(1, retry_count + 1): + if cancel_event is not None and cancel_event.is_set(): + raise DataManagementMigrationCanceledError( + "AI Search migration stopped after cancellation or lease loss." + ) + try: + results = list(search_client.upload_documents( + documents=[entry["document"] for entry in pending_entries], + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )) + if len(results) != len(pending_entries): + last_error = "AI Search returned an incomplete document result set." + else: + failed_entries = [] + for entry, result in zip(pending_entries, results): + if not _get_search_result_succeeded(result): + failed_entries.append(entry) + continue + record_accepted_entry(entry, attempt) + pending_entries = failed_entries + if not pending_entries: + return { + "copied": copied_count, + "created": created_count, + "updated": updated_count, + "failed": 0, + "bytes": accepted_bytes, + "attempt": attempt, + "manifest_entries": manifest_entries, + "elapsed_seconds": time.perf_counter() - started_at, + } + if not last_error: + last_error = f"AI Search returned {len(pending_entries)} failed document result(s)." + except DataManagementSettingsValidationError: + raise + except Exception as exc: + last_error = str(exc)[:500] + + if attempt < retry_count: + if provenance_context and index_name: + try: + verified_entries, pending_entries = _reclassify_pending_search_migration_entries( + search_client, + pending_entries, + provenance_context, + index_name, + migration_mode, + ) + except DataManagementSettingsValidationError: + raise + except Exception as exc: + last_error = ( + "AI Search migration could not verify unresolved documents after an " + f"ambiguous write: {str(exc)[:400]}" + ) + break + for entry in verified_entries: + record_accepted_entry(entry, attempt) + if not pending_entries: + return { + "copied": copied_count, + "created": created_count, + "updated": updated_count, + "failed": 0, + "bytes": accepted_bytes, + "attempt": attempt, + "manifest_entries": manifest_entries, + "elapsed_seconds": time.perf_counter() - started_at, + } + retry_delay = min(30, 2 ** (attempt - 1)) + if cancel_event is not None: + if cancel_event.wait(retry_delay): + return { + "copied": copied_count, + "created": created_count, + "updated": updated_count, + "failed": 0, + "bytes": accepted_bytes, + "attempt": attempt, + "manifest_entries": manifest_entries, + "interrupted": True, + "elapsed_seconds": time.perf_counter() - started_at, + } + else: + time.sleep(retry_delay) + + manifest_entries.extend({ + **entry["manifest_context"], + "status": "failed", + "attempt": retry_count, + "bytes": 0, + "error": last_error or "AI Search document indexing failed.", + } for entry in pending_entries) + return { + "copied": copied_count, + "created": created_count, + "updated": updated_count, + "failed": len(pending_entries), + "bytes": accepted_bytes, + "attempt": retry_count, + "manifest_entries": manifest_entries, + "elapsed_seconds": time.perf_counter() - started_at, + "error": last_error or "AI Search document indexing failed.", + } + + +def _copy_ai_search_to_target( + settings, + migration_plan, + job, + migration_state, + provenance_context, + target_search_write_fence=None, +): + copied = [] + if not migration_plan.get("include_ai_search"): + return copied + search_mappings = [ + ("users", "user_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[0]), + ("groups", "group_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[1]), + ("public_workspaces", "public_workspace_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[2]), + ] + parallel_operations = _get_migration_parallel_operations(settings) + retry_count = _get_migration_retry_count(settings) + migration_mode = _safe_text( + (provenance_context or {}).get("migration_mode"), + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ) + batch_size = min(50, DATA_MANAGEMENT_MIGRATION_BATCH_SIZE) + search_page_size = min( + DATA_MANAGEMENT_SEARCH_KEYSET_PAGE_SIZE, + max(batch_size, batch_size * parallel_operations), + ) + gate_container = None + gate_fence = None + if isinstance(target_search_write_fence, tuple) and len(target_search_write_fence) == 2: + gate_container, gate_fence = target_search_write_fence + + def renew_target_search_write_fence(): + if gate_container is None or gate_fence is None: + return + renew_data_management_search_write_fence( + gate_container, + gate_fence, + _get_target_search_write_fence_lease_seconds(settings), + ) + + for target_type, field_name, artifact in search_mappings: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none" or not selection.get("include_documents"): + continue + source_client = CLIENTS.get(artifact["client_key"]) + if not source_client: + copied.append({"name": artifact["name"], "type": "ai_search_documents", "status": "skipped", "warning": "Source AI Search client is not initialized."}) + continue + resource_name = f"ai_search:{target_type}:{artifact['index_name']}" + completed_resource = migration_state.get("resources", {}).get(resource_name) + if is_migration_resource_completed(migration_state, resource_name): + previous_result = completed_resource.get("result") if isinstance(completed_resource, dict) else {} + copied.append({ + "name": artifact["name"], + "type": "ai_search_documents", + "target_type": target_type, + "index_name": artifact["index_name"], + "status": "checkpoint_completed", + **(previous_result if isinstance(previous_result, dict) else {}), + }) + continue + + index_status = _ensure_target_search_index(settings, artifact["index_name"], artifact["schema_file"]) + target_client = _get_target_search_client(settings, artifact["index_name"]) + search_filters = _build_search_scope_filter_batches(field_name, selection) + resource = start_migration_resource(migration_state, resource_name) + resource_started_at = resource.get("attempt_started_at") or resource.get("started_at") + previous_progress = resource.get("progress") if isinstance(resource.get("progress"), dict) else {} + previous_cursor = previous_progress.get("keyset_cursor") + previous_cursor = previous_cursor if isinstance(previous_cursor, dict) else {} + copied_count = _safe_int(previous_progress.get("copied_count"), default=0, minimum=0) + skipped_count = _safe_int(previous_progress.get("skipped_count"), default=0, minimum=0) + prior_failed_count = _safe_int(previous_progress.get("failed_count"), default=0, minimum=0) + failed_count = 0 + collision_count = _safe_int(previous_progress.get("collision_count"), default=0, minimum=0) + destination_provenance_skip_count = _safe_int( + previous_progress.get("destination_provenance_skip_count"), + default=0, + minimum=0, + ) + byte_count = _safe_int(previous_progress.get("bytes"), default=0, minimum=0) + source_read_count = _safe_int(previous_progress.get("source_read_count"), default=0, minimum=0) + created_count = _safe_int(previous_progress.get("created_count"), default=0, minimum=0) + updated_count = _safe_int(previous_progress.get("updated_count"), default=0, minimum=0) + unchanged_count = _safe_int(previous_progress.get("unchanged_count"), default=0, minimum=0) + retry_attempt_count = _safe_int(previous_progress.get("retry_attempt_count"), default=0, minimum=0) + checkpoint_count = _safe_int(previous_progress.get("checkpoint_count"), default=0, minimum=0) + cursor_scope_batch_index = _safe_int( + previous_cursor.get("scope_batch_index"), + default=0, + minimum=0, + maximum=len(search_filters), + ) + cursor_last_id = _safe_text(previous_cursor.get("last_id")) + errors = [] + append_manifest, flush_manifest = _create_migration_manifest_writer( + job.get("id"), + resource_name, + ) + transfer_cancel_event = Event() + + def persist_progress(allow_cancel_requested=False): + renew_target_search_write_fence() + flush_manifest() + progress = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + ) + progress.update({ + "parallel_operations": parallel_operations, + "batch_size": batch_size, + "retry_count": retry_count, + "destination_provenance_skip_count": destination_provenance_skip_count, + "collision_count": collision_count, + "source_read_count": source_read_count, + "destination_accepted_count": copied_count, + "destination_failed_count": failed_count, + "prior_failed_count": prior_failed_count, + "retry_attempt_count": retry_attempt_count, + "checkpoint_count": checkpoint_count, + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "keyset_cursor": { + "scope_batch_index": cursor_scope_batch_index, + "scope_batch_count": len(search_filters), + "last_id": cursor_last_id, + }, + }) + return _persist_migration_checkpoint( + job, + migration_state, + settings, + resource_name, + progress, + f"Migrating AI Search index {artifact['index_name']}", + allow_cancel_requested=allow_cancel_requested, + ) + + def persist_heartbeat(): + result = _persist_migration_heartbeat( + job, + settings, + f"Waiting for AI Search uploads to {artifact['index_name']}", + ) + renew_target_search_write_fence() + return result + + def persist_stopped_results(): + return persist_progress(allow_cancel_requested=True) + + def prepare_upload_batch(source_batch): + nonlocal collision_count + source_ids = [_safe_text(document.get("id")) for document in source_batch] + if any(not document_id for document_id in source_ids): + raise DataManagementSettingsValidationError( + f"Source AI Search index '{artifact['index_name']}' contains a document without an ID." + ) + _assert_migration_job_lease(job) + renew_target_search_write_fence() + target_documents = _get_target_search_documents_by_ids( + target_client, + source_ids, + ) + upload_documents = [] + upload_dispositions = [] + upload_manifest_contexts = [] + deferred_outcomes = { + "skipped_count": 0, + "unchanged_count": 0, + "destination_provenance_skip_count": 0, + "manifest_entries": [], + } + for source_document, source_document_id in zip(source_batch, source_ids): + source_hash = _build_search_source_hash(source_document) + manifest_context = { + "service": "ai_search", + "resource_name": resource_name, + "target_type": target_type, + "source_identity": _hash_migration_manifest_identity( + artifact["index_name"], + source_document_id, + ), + "destination_identity": _hash_migration_manifest_identity( + "destination", + artifact["index_name"], + source_document_id, + ), + "source_hash": source_hash, + "_locator": { + "service": "ai_search", + "resource_name": resource_name, + "target_type": target_type, + "document_id": source_document_id, + "index_name": artifact["index_name"], + }, + } + try: + disposition = _classify_target_search_document( + target_documents.get(source_document_id), + provenance_context, + artifact["index_name"], + migration_mode=migration_mode, + source_hash=source_hash, + ) + except DataManagementSettingsValidationError: + collision_count += 1 + raise + if disposition in {"unchanged", "resume_verified"}: + if disposition == "unchanged": + deferred_outcomes["skipped_count"] += 1 + deferred_outcomes["unchanged_count"] += 1 + deferred_outcomes["destination_provenance_skip_count"] += 1 + deferred_outcomes["manifest_entries"].append({ + **manifest_context, + "status": disposition, + }) + continue + upload_documents.append(add_search_migration_provenance( + copy.deepcopy(source_document), + provenance_context, + source_hash=source_hash, + )) + upload_dispositions.append(disposition) + upload_manifest_contexts.append(manifest_context) + return ( + upload_documents, + upload_dispositions, + upload_manifest_contexts, + deferred_outcomes, + ) + + def consume_batch_result(batch_result): + nonlocal copied_count, failed_count, byte_count, retry_attempt_count + nonlocal created_count, updated_count + copied_count += batch_result.get("copied", 0) + created_count += batch_result.get("created", 0) + updated_count += batch_result.get("updated", 0) + failed_count += batch_result.get("failed", 0) + byte_count += batch_result.get("bytes", 0) + retry_attempt_count += max(0, _safe_int(batch_result.get("attempt"), default=1) - 1) + if batch_result.get("error"): + errors.append(batch_result["error"]) + for manifest_entry in batch_result.get("manifest_entries") or []: + append_manifest(manifest_entry) + + def consume_batch_future(batch_result, _context): + consume_batch_result(batch_result) + + try: + with ThreadPoolExecutor(max_workers=parallel_operations) as executor: + for scope_batch_index, search_filter in enumerate(search_filters): + if scope_batch_index < cursor_scope_batch_index: + continue + last_id = cursor_last_id if scope_batch_index == cursor_scope_batch_index else "" + for source_page, page_last_id in _iter_search_document_pages( + source_client, + search_filter=search_filter, + page_size=search_page_size, + last_id=last_id, + ): + page_failed_count = failed_count + pending_batches = [] + page_outcomes = { + "skipped_count": 0, + "unchanged_count": 0, + "destination_provenance_skip_count": 0, + "manifest_entries": [], + } + for offset in range(0, len(source_page), batch_size): + source_batch = [] + source_batch.extend(source_page[offset:offset + batch_size]) + if not source_batch: + continue + ( + upload_documents, + upload_dispositions, + upload_manifest_contexts, + deferred_outcomes, + ) = prepare_upload_batch(source_batch) + for field_name in ( + "skipped_count", + "unchanged_count", + "destination_provenance_skip_count", + ): + page_outcomes[field_name] += deferred_outcomes[field_name] + page_outcomes["manifest_entries"].extend( + deferred_outcomes["manifest_entries"] + ) + if upload_documents: + pending_batches.append(executor.submit( + _upload_search_migration_batch, + target_client, + upload_documents, + retry_count, + upload_dispositions, + upload_manifest_contexts, + transfer_cancel_event, + provenance_context=provenance_context, + index_name=artifact["index_name"], + migration_mode=migration_mode, + )) + if len(pending_batches) < parallel_operations: + continue + _drain_migration_futures_with_heartbeat( + {future: None for future in pending_batches}, + consume_batch_future, + persist_heartbeat, + transfer_cancel_event, + persist_stopped_results=persist_stopped_results, + ) + pending_batches = [] + _drain_migration_futures_with_heartbeat( + {future: None for future in pending_batches}, + consume_batch_future, + persist_heartbeat, + transfer_cancel_event, + persist_stopped_results=persist_stopped_results, + ) + + if failed_count > page_failed_count: + migration_state = persist_progress() + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + "; ".join(errors[:3]) or "AI Search document indexing failed.", + ) + raise RuntimeError( + f"AI Search migration failed for index '{artifact['index_name']}' " + "before its keyset cursor could advance." + ) + + source_read_count += len(source_page) + skipped_count += page_outcomes["skipped_count"] + unchanged_count += page_outcomes["unchanged_count"] + destination_provenance_skip_count += page_outcomes[ + "destination_provenance_skip_count" + ] + for manifest_entry in page_outcomes["manifest_entries"]: + append_manifest(manifest_entry) + cursor_scope_batch_index = scope_batch_index + cursor_last_id = page_last_id + checkpoint_count += 1 + migration_state = persist_progress() + + cursor_scope_batch_index = scope_batch_index + 1 + cursor_last_id = "" + checkpoint_count += 1 + migration_state = persist_progress() + + result = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + ) + flush_manifest() + result.update({ + "name": artifact["name"], + "type": "ai_search_documents", + "target_type": target_type, + "index_name": artifact["index_name"], + "scope_filter_batch_count": len(search_filters), + "index_status": index_status, + "parallel_operations": parallel_operations, + "batch_size": batch_size, + "search_page_size": search_page_size, + "retry_count": retry_count, + "destination_provenance_skip_count": destination_provenance_skip_count, + "collision_count": collision_count, + "source_read_count": source_read_count, + "destination_accepted_count": copied_count, + "destination_failed_count": failed_count, + "prior_failed_count": prior_failed_count, + "retry_attempt_count": retry_attempt_count, + "checkpoint_count": checkpoint_count, + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "keyset_cursor": { + "scope_batch_index": len(search_filters), + "scope_batch_count": len(search_filters), + "last_id": "", + "completed": True, + }, + }) + if failed_count: + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + "; ".join(errors[:3]) or "AI Search document indexing failed.", + ) + raise RuntimeError( + f"AI Search migration failed for index '{artifact['index_name']}' after {failed_count} document failure(s)." + ) + migration_state = _complete_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + result, + f"Completed AI Search index {artifact['index_name']}", + ) + copied.append(result) + except (DataManagementMigrationCanceledError, DataManagementMigrationLeaseLostError): + flush_manifest() + raise + except Exception as exc: + flush_manifest() + if not failed_count: + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + str(exc), + ) + raise + return copied + + +def _get_target_enhanced_citations_blob_client(settings): + auth_type = _safe_text((settings or {}).get("target_enhanced_citations_storage_authentication_type"), "managed_identity") + if auth_type == "connection_string": + connection_string = _safe_text((settings or {}).get("target_enhanced_citations_storage_connection_string")) + if not connection_string: + raise ValueError("Destination Enhanced Citations storage connection string is required when migrating source document blobs.") + return BlobServiceClient.from_connection_string(connection_string) + blob_endpoint = _safe_text((settings or {}).get("target_enhanced_citations_storage_blob_endpoint")) + if not blob_endpoint: + raise ValueError("Destination Enhanced Citations Blob endpoint is required when migrating source document blobs.") + return BlobServiceClient(account_url=blob_endpoint, credential=DefaultAzureCredential()) + + +def _document_blob_reference(document_item): + if not isinstance(document_item, dict): + return None + container_name = document_item.get("blob_container") or _blob_container_for_document(document_item) + blob_path = document_item.get("blob_path") or document_item.get("archived_blob_path") + if not container_name or not blob_path: + return None + return container_name, blob_path + + +def _document_migration_scope_hash(document_item): + if not isinstance(document_item, dict): + return "" + for scope_type, field_name in ( + ("public_workspaces", "public_workspace_id"), + ("groups", "group_id"), + ("users", "user_id"), + ): + scope_id = _safe_text(document_item.get(field_name)) + if scope_id: + encoded = f"{scope_type}:{scope_id}".encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + return "" + + +def _blob_container_for_document(document_item): + if document_item.get("public_workspace_id"): + return app_config.storage_account_public_documents_container_name + if document_item.get("group_id"): + return app_config.storage_account_group_documents_container_name + return app_config.storage_account_user_documents_container_name + + +def _iter_selected_document_records_for_blob_migration(migration_plan, source_cutoff_epoch=None): + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none" or not selection.get("include_documents"): + continue + container_definition = next( + ( + definition + for definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type] + if definition.get("documents") + ), + None, + ) + if not container_definition: + continue + yield from _iter_selected_cosmos_records( + container_definition, + selection, + source_cutoff_epoch=source_cutoff_epoch, + ) or [] + + +def _get_blob_properties_or_none(blob_client): + try: + return blob_client.get_blob_properties() + except ResourceNotFoundError: + return None + + +def _get_blob_content_md5(blob_properties): + content_settings = getattr(blob_properties, "content_settings", None) + content_md5 = getattr(content_settings, "content_md5", None) + if isinstance(content_md5, (bytes, bytearray)): + return base64.b64encode(bytes(content_md5)).decode("ascii") + return _safe_text(content_md5) + + +def _build_blob_source_fingerprint(blob_properties): + """Build stable source content and version markers from service properties.""" + content_md5 = _get_blob_content_md5(blob_properties) + last_modified = getattr(blob_properties, "last_modified", None) + if isinstance(last_modified, datetime): + last_modified = last_modified.astimezone(timezone.utc).isoformat() + version_payload = { + "etag": _safe_text(getattr(blob_properties, "etag", None)), + "last_modified": _safe_text(last_modified), + "size": _safe_int(getattr(blob_properties, "size", 0), default=0, minimum=0), + "blob_type": _safe_text(getattr(blob_properties, "blob_type", None)), + } + encoded = json.dumps( + version_payload, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + source_hash = f"md5:{content_md5}" if content_md5 else "" + source_version = f"sha256:{hashlib.sha256(encoded).hexdigest()}" + return source_hash, source_version + + +def _get_source_blob_tags(source_blob_client, source_properties): + source_tags = getattr(source_properties, "tags", None) + if isinstance(source_tags, dict): + return source_tags + if _safe_int(getattr(source_properties, "tag_count", 0), default=0, minimum=0) <= 0: + return None + get_blob_tags = getattr(source_blob_client, "get_blob_tags", None) + if not callable(get_blob_tags): + return None + tags = get_blob_tags() + return tags if isinstance(tags, dict) else None + + +def _classify_target_blob( + target_properties, + provenance_context, + migration_mode, + source_hash, + source_version, +): + if target_properties is None: + return "create" + target_provenance = get_blob_migration_provenance( + getattr(target_properties, "metadata", None) + ) + if ( + _safe_text(target_provenance.get("migrationId")) == + _safe_text((provenance_context or {}).get("migration_id")) and + _safe_text(target_provenance.get("status")).lower() == "pending" + ): + return "update" + if is_successful_migration_record(target_provenance): + if ( + _safe_text(target_provenance.get("migrationId")) == + _safe_text((provenance_context or {}).get("migration_id")) + ): + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + return "resume_verified" + return "update" + if migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + return "unchanged" + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + return "unchanged" + return "update" + return "collision" + + +def _build_pending_blob_migration_metadata( + source_metadata, + provenance_context, + source_hash, + source_version, + scope_hash, +): + metadata = merge_blob_migration_metadata( + source_metadata, + provenance_context, + source_hash=source_hash, + source_version=source_version, + scope_hash=scope_hash, + ) + metadata[BLOB_MIGRATION_STATUS_METADATA_KEY] = "pending" + return metadata + + +def _iter_migration_blob_chunks( + source_download, + progress_callback=None, + cancel_event=None, +): + for chunk in source_download.chunks(): + if cancel_event is not None and cancel_event.is_set(): + raise DataManagementMigrationCanceledError( + "Blob transfer stopped after migration cancellation or lease loss." + ) + if progress_callback: + progress_callback(len(chunk)) + yield chunk + if cancel_event is not None and cancel_event.is_set(): + raise DataManagementMigrationCanceledError( + "Blob transfer stopped after migration cancellation or lease loss." + ) + + +def _copy_source_blob_migration_record( + source_blob_client, + target_blob_client, + provenance_context, + retry_count, + scope_hash="", + progress_callback=None, + cancel_event=None, +): + """Stream one blob to the target and stamp metadata only on success.""" + try: + source_properties = source_blob_client.get_blob_properties() + except ResourceNotFoundError: + return {"status": "missing", "bytes": 0} + + migration_mode = _safe_text( + (provenance_context or {}).get("migration_mode"), + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ) + source_hash, source_version = _build_blob_source_fingerprint(source_properties) + existing_target_properties = _get_blob_properties_or_none(target_blob_client) + disposition = _classify_target_blob( + existing_target_properties, + provenance_context, + migration_mode, + source_hash, + source_version, + ) + if disposition in {"unchanged", "resume_verified"}: + return { + "status": disposition, + "bytes": 0, + "source_hash": source_hash, + "source_version": source_version, + } + if disposition == "collision": + return { + "status": "collision", + "bytes": 0, + "error": "Destination blob exists without successful migration provenance.", + } + + source_metadata = getattr(source_properties, "metadata", None) or {} + pending_metadata = _build_pending_blob_migration_metadata( + source_metadata, + provenance_context, + source_hash, + source_version, + scope_hash, + ) + destination_metadata = merge_blob_migration_metadata( + source_metadata, + provenance_context, + source_hash=source_hash, + source_version=source_version, + scope_hash=scope_hash, + ) + source_size = _safe_int(getattr(source_properties, "size", 0), default=0, minimum=0) + content_settings = getattr(source_properties, "content_settings", None) + source_tags = _get_source_blob_tags(source_blob_client, source_properties) + source_blob_tier = getattr(source_properties, "blob_tier", None) + source_blob_type = getattr(source_properties, "blob_type", None) + source_etag = _safe_text(getattr(source_properties, "etag", None)) + started_at = time.perf_counter() + last_error = "" + + for attempt in range(1, retry_count + 1): + try: + if progress_callback: + progress_callback(0, reset=True) + if cancel_event is not None and cancel_event.is_set(): + raise DataManagementMigrationCanceledError( + "Blob transfer stopped after migration cancellation or lease loss." + ) + download_kwargs = { + "max_concurrency": 1, + "timeout": DATA_MANAGEMENT_BLOB_SERVICE_TIMEOUT_SECONDS, + } + if source_etag: + download_kwargs.update({ + "etag": source_etag, + "match_condition": MatchConditions.IfNotModified, + }) + source_download = source_blob_client.download_blob(**download_kwargs) + upload_kwargs = { + "data": _iter_migration_blob_chunks( + source_download, + progress_callback=progress_callback, + cancel_event=cancel_event, + ), + "length": source_size, + "overwrite": disposition == "update", + "metadata": pending_metadata, + "content_settings": content_settings, + "max_concurrency": 1, + "timeout": DATA_MANAGEMENT_BLOB_SERVICE_TIMEOUT_SECONDS, + } + if source_tags is not None: + upload_kwargs["tags"] = source_tags + if source_blob_tier is not None: + upload_kwargs["standard_blob_tier"] = source_blob_tier + if source_blob_type is not None: + upload_kwargs["blob_type"] = source_blob_type + if disposition == "update" and existing_target_properties is not None: + target_etag = getattr(existing_target_properties, "etag", None) + if target_etag: + upload_kwargs["etag"] = target_etag + upload_kwargs["match_condition"] = MatchConditions.IfNotModified + target_blob_client.upload_blob(**upload_kwargs) + destination_properties = _get_blob_properties_or_none(target_blob_client) + destination_size = _safe_int( + getattr(destination_properties, "size", source_size), + default=source_size, + minimum=0, + ) + destination_md5 = _get_blob_content_md5(destination_properties) + if destination_size != source_size or ( + source_hash and destination_md5 and source_hash != f"md5:{destination_md5}" + ): + raise RuntimeError( + "Destination blob verification did not match the source content properties." + ) + current_source_properties = source_blob_client.get_blob_properties() + current_source_etag = _safe_text(getattr(current_source_properties, "etag", None)) + if source_etag and current_source_etag != source_etag: + raise RuntimeError("Source blob changed while it was being transferred.") + destination_etag = _safe_text(getattr(destination_properties, "etag", None)) + set_metadata_kwargs = {"metadata": destination_metadata} + if destination_etag: + set_metadata_kwargs.update({ + "etag": destination_etag, + "match_condition": MatchConditions.IfNotModified, + }) + target_blob_client.set_blob_metadata(**set_metadata_kwargs) + return { + "status": "updated" if disposition == "update" else "copied", + "bytes": source_size, + "attempt": attempt, + "source_size": source_size, + "destination_size": destination_size, + "source_hash": source_hash, + "source_version": source_version, + "elapsed_seconds": time.perf_counter() - started_at, + } + except Exception as exc: + if isinstance(exc, (DataManagementMigrationCanceledError, DataManagementMigrationLeaseLostError)): + raise + last_error = str(exc)[:500] + status_code = getattr(exc, "status_code", None) + if isinstance(exc, ResourceExistsError) or status_code in {409, 412}: + concurrent_target_properties = _get_blob_properties_or_none(target_blob_client) + concurrent_disposition = _classify_target_blob( + concurrent_target_properties, + provenance_context, + migration_mode, + source_hash, + source_version, + ) + if concurrent_disposition in {"unchanged", "resume_verified"}: + return {"status": concurrent_disposition, "bytes": 0} + if concurrent_disposition == "update": + disposition = "update" + existing_target_properties = concurrent_target_properties + continue + if concurrent_disposition == "collision": + return { + "status": "collision", + "bytes": 0, + "error": "Destination blob changed without migration ownership during this migration.", + } + retryable = status_code in {408, 429, 500, 503} or status_code is None + if not retryable or attempt >= retry_count: + break + time.sleep(min(30, 2 ** (attempt - 1))) + + if disposition == "create": + failed_target_properties = _get_blob_properties_or_none(target_blob_client) + failed_provenance = get_blob_migration_provenance( + getattr(failed_target_properties, "metadata", None) + ) if failed_target_properties else {} + if ( + _safe_text(failed_provenance.get("migrationId")) == + _safe_text((provenance_context or {}).get("migration_id")) and + _safe_text(failed_provenance.get("status")).lower() == "pending" + ): + delete_kwargs = {} + failed_etag = _safe_text(getattr(failed_target_properties, "etag", None)) + if failed_etag: + delete_kwargs.update({ + "etag": failed_etag, + "match_condition": MatchConditions.IfNotModified, + }) + try: + target_blob_client.delete_blob(**delete_kwargs) + except Exception: + pass + + return { + "status": "failed", + "bytes": 0, + "attempt": retry_count, + "elapsed_seconds": time.perf_counter() - started_at, + "error": last_error or "Source blob copy failed.", + } + + +def _copy_source_blobs_to_target(settings, migration_plan, job, migration_state, provenance_context): + if not migration_plan.get("include_source_blobs"): + return [] + source_client = _get_source_blob_service_client() + if not source_client: + return [{"name": "source_blobs", "type": "source_blobs", "status": "skipped", "warning": "Source Enhanced Citations storage is not configured."}] + + resource_name = "source_blobs:selected_documents" + completed_resource = migration_state.get("resources", {}).get(resource_name) + if is_migration_resource_completed(migration_state, resource_name): + previous_result = completed_resource.get("result") if isinstance(completed_resource, dict) else {} + return [{ + "name": "source_blobs", + "type": "source_blobs", + "status": "checkpoint_completed", + **(previous_result if isinstance(previous_result, dict) else {}), + }] + + target_client = _get_target_enhanced_citations_blob_client(settings) + resource = start_migration_resource(migration_state, resource_name) + resource_started_at = resource.get("attempt_started_at") or resource.get("started_at") + source_cutoff = _parse_iso_datetime(migration_state.get("source_cutoff_at")) + source_cutoff_epoch = int(source_cutoff.timestamp()) if source_cutoff else None + parallel_operations = _get_migration_parallel_operations(settings) + retry_count = _get_migration_retry_count(settings) + migration_mode = _safe_text( + (provenance_context or {}).get("migration_mode"), + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ) + previous_progress = resource.get("progress") if isinstance(resource.get("progress"), dict) else {} + copied_count = _safe_int(previous_progress.get("copied_count"), default=0, minimum=0) + skipped_count = _safe_int(previous_progress.get("skipped_count"), default=0, minimum=0) + prior_failed_count = _safe_int(previous_progress.get("failed_count"), default=0, minimum=0) + prior_collision_count = _safe_int(previous_progress.get("collision_count"), default=0, minimum=0) + prior_missing_count = _safe_int(previous_progress.get("missing_count"), default=0, minimum=0) + failed_count = 0 + collision_count = 0 + missing_count = 0 + not_applicable_count = _safe_int(previous_progress.get("not_applicable_count"), default=0, minimum=0) + created_count = _safe_int(previous_progress.get("created_count"), default=0, minimum=0) + updated_count = _safe_int(previous_progress.get("updated_count"), default=0, minimum=0) + unchanged_count = _safe_int(previous_progress.get("unchanged_count"), default=0, minimum=0) + retry_attempt_count = _safe_int(previous_progress.get("retry_attempt_count"), default=0, minimum=0) + byte_count = _safe_int(previous_progress.get("bytes"), default=0, minimum=0) + container_stats = {} + errors = [] + append_manifest, flush_manifest = _create_migration_manifest_writer( + job.get("id"), + resource_name, + ) + transfer_progress = {} + transfer_progress_lock = Lock() + transfer_cancel_event = Event() + + def persist_progress(message="Migrating selected source document blobs"): + nonlocal migration_state + flush_manifest() + with transfer_progress_lock: + in_flight_bytes = sum(transfer_progress.values()) + progress = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + ) + progress.update({ + "missing_count": missing_count, + "not_applicable_count": not_applicable_count, + "parallel_operations": parallel_operations, + "retry_count": retry_count, + "collision_count": collision_count, + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "retry_attempt_count": retry_attempt_count, + "prior_failed_count": prior_failed_count, + "prior_collision_count": prior_collision_count, + "prior_missing_count": prior_missing_count, + "in_flight_bytes": in_flight_bytes, + "observed_bytes": byte_count + in_flight_bytes, + "observed_bytes_per_second": round( + (byte_count + in_flight_bytes) / + max(0.001, float(progress.get("elapsed_seconds") or 0.001)), + 3, + ), + "active_transfer_count": len(transfer_progress), + }) + migration_state = _persist_migration_checkpoint( + job, + migration_state, + settings, + resource_name, + progress, + message, + ) + return migration_state + + def make_transfer_progress_callback(reference_key): + def report_progress(byte_delta, reset=False): + with transfer_progress_lock: + if reset: + transfer_progress[reference_key] = 0 + else: + transfer_progress[reference_key] = ( + transfer_progress.get(reference_key, 0) + + _safe_int(byte_delta, default=0, minimum=0) + ) + return report_progress + + def consume_copy_future(result_container_name, reference_key, future): + nonlocal copied_count, skipped_count, failed_count, collision_count, missing_count, byte_count + nonlocal created_count, updated_count, unchanged_count, retry_attempt_count + with transfer_progress_lock: + transfer_progress.pop(reference_key, None) + copy_result = future.result() + artifact = container_stats[result_container_name] + artifact["blob_count"] += 1 + status = copy_result.get("status") + retry_attempt_count += max(0, _safe_int(copy_result.get("attempt"), default=1) - 1) + if status in {"copied", "updated"}: + copied_count += 1 + if status == "updated": + updated_count += 1 + artifact["updated_count"] += 1 + else: + created_count += 1 + artifact["created_count"] += 1 + byte_count += copy_result.get("bytes", 0) + artifact["copied_count"] += 1 + artifact["bytes"] += copy_result.get("bytes", 0) + elif status in {"unchanged", "resume_verified"}: + if status == "unchanged": + skipped_count += 1 + unchanged_count += 1 + artifact["skipped_count"] += 1 + if status == "unchanged": + artifact["unchanged_count"] += 1 + elif status == "missing": + missing_count += 1 + artifact["missing_count"] += 1 + else: + failed_count += 1 + if status == "collision": + collision_count += 1 + artifact["collision_count"] += 1 + artifact["failed_count"] += 1 + errors.append(copy_result.get("error", "Source blob copy failed.")) + append_manifest({ + "service": "source_blobs", + "resource_name": resource_name, + "source_identity": _hash_migration_manifest_identity( + "source", + result_container_name, + reference_key.split("\n", 1)[-1], + ), + "destination_identity": _hash_migration_manifest_identity( + "destination", + result_container_name, + reference_key.split("\n", 1)[-1], + ), + "status": status, + "attempt": copy_result.get("attempt"), + "bytes": copy_result.get("bytes", 0), + "source_version": copy_result.get("source_version"), + "source_hash": copy_result.get("source_hash"), + "error": copy_result.get("error"), + "_locator": { + "service": "source_blobs", + "resource_name": resource_name, + "container_name": result_container_name, + "blob_name": reference_key.split("\n", 1)[-1], + }, + }) + + def process_pending_copies(pending_copies): + pending_by_future = { + future: (result_container_name, reference_key) + for result_container_name, reference_key, future in pending_copies + } + last_heartbeat = 0.0 + while pending_by_future: + completed_futures, _pending_futures = wait( + set(pending_by_future), + timeout=1.0, + return_when=FIRST_COMPLETED, + ) + for future in completed_futures: + result_container_name, reference_key = pending_by_future.pop(future) + consume_copy_future(result_container_name, reference_key, future) + now = time.monotonic() + if pending_by_future and now - last_heartbeat >= 2.0: + try: + persist_progress("Migrating source blobs; worker heartbeat renewed") + except (DataManagementMigrationCanceledError, DataManagementMigrationLeaseLostError): + transfer_cancel_event.set() + for future in pending_by_future: + future.cancel() + raise + last_heartbeat = now + + try: + pending_copies = [] + pending_reference_keys = set() + with ThreadPoolExecutor(max_workers=parallel_operations) as executor: + for document in _iter_selected_document_records_for_blob_migration( + migration_plan, + source_cutoff_epoch=source_cutoff_epoch, + ): + reference = _document_blob_reference(document) + if not reference: + not_applicable_count += 1 + append_manifest({ + "service": "source_blobs", + "resource_name": resource_name, + "source_identity": _hash_migration_manifest_identity( + "no-source-blob", + document.get("id"), + ), + "status": "not_applicable", + "_locator": { + "service": "source_blobs", + "resource_name": resource_name, + "document_id": document.get("id"), + }, + }) + continue + container_name, blob_path = reference + reference_key = f"{container_name}\n{blob_path}" + if reference_key in pending_reference_keys: + continue + artifact = container_stats.setdefault(container_name, { + "name": container_name, + "type": "source_blob_container", + "container_name": container_name, + "blob_count": 0, + "copied_count": 0, + "created_count": 0, + "updated_count": 0, + "unchanged_count": 0, + "skipped_count": 0, + "missing_count": 0, + "collision_count": 0, + "failed_count": 0, + "bytes": 0, + }) + try: + target_client.create_container(container_name) + except ResourceExistsError: + pass + source_blob_client = source_client.get_blob_client( + container=container_name, + blob=blob_path, + ) + target_blob_client = target_client.get_blob_client( + container=container_name, + blob=blob_path, + ) + if not pending_copies: + _assert_migration_job_lease(job) + pending_reference_keys.add(reference_key) + progress_callback = make_transfer_progress_callback(reference_key) + pending_copies.append((container_name, reference_key, executor.submit( + _copy_source_blob_migration_record, + source_blob_client, + target_blob_client, + provenance_context, + retry_count, + _document_migration_scope_hash(document), + progress_callback, + transfer_cancel_event, + ))) + if len(pending_copies) < parallel_operations: + continue + process_pending_copies(pending_copies) + pending_copies = [] + pending_reference_keys.clear() + persist_progress() + + process_pending_copies(pending_copies) + + flush_manifest() + + result = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + ) + result.update({ + "name": "source_blobs", + "type": "source_blobs", + "missing_count": missing_count, + "not_applicable_count": not_applicable_count, + "parallel_operations": parallel_operations, + "retry_count": retry_count, + "collision_count": collision_count, + "container_count": len(container_stats), + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "retry_attempt_count": retry_attempt_count, + }) + if failed_count or missing_count: + failure_message = ( + f"{missing_count} persisted source blob reference(s) could not be resolved." + if missing_count else + "; ".join(errors[:3]) or "Source blob copy failed." + ) + persist_progress("Source blob migration requires remediation before retry") + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + failure_message, + ) + raise RuntimeError( + "Source blob migration requires remediation before retry: " + f"{failed_count} failed and {missing_count} unexpectedly missing blob(s)." + ) + migration_state = _complete_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + result, + "Completed selected source document blob migration", + ) + artifacts = list(container_stats.values()) + artifacts.append(result) + return artifacts + except Exception as exc: + flush_manifest() + if not failed_count: + _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + str(exc), + ) + raise + + +def _iter_target_cosmos_records(target_container, container_definition, selection): + """Enumerate the exact destination scope used by the migration plan.""" + mode = selection.get("mode") + if mode == "all": + yield from target_container.query_items( + query="SELECT * FROM c", + enable_cross_partition_query=True, + ) + return + if mode != "selected": + return + if container_definition.get("id_field") == "id": + for item_id in selection.get("ids") or []: + try: + yield target_container.read_item(item=item_id, partition_key=item_id) + except (CosmosResourceNotFoundError, ResourceNotFoundError): + continue + return + filter_field = container_definition.get("filter_field") + if not filter_field: + return + for selected_id in selection.get("ids") or []: + yield from target_container.query_items( + query=f"SELECT * FROM c WHERE c.{filter_field} = @selected_id", + parameters=[{"name": "@selected_id", "value": selected_id}], + enable_cross_partition_query=True, + ) + + +def _delete_target_cosmos_record(target_container, document, partition_key_path): + delete_kwargs = { + "item": document.get("id"), + "partition_key": _get_document_path_value(document, partition_key_path), + } + if document.get("_etag"): + delete_kwargs.update({ + "etag": document.get("_etag"), + "match_condition": MatchConditions.IfNotModified, + }) + target_container.delete_item(**delete_kwargs) + + +def _get_reconciliation_target_cosmos_container( + target_database, + container_name, + partition_key_path, + read_only, +): + if not read_only: + return _get_target_cosmos_container( + target_database, + container_name, + partition_key_path, + ) + get_container_client = getattr(target_database, "get_container_client", None) + if not callable(get_container_client): + return _get_target_cosmos_container( + target_database, + container_name, + partition_key_path, + ) + target_container = get_container_client(container_name) + try: + _validate_target_cosmos_container_partition_key( + target_container, + container_name, + partition_key_path, + ) + except (CosmosResourceNotFoundError, ResourceNotFoundError): + return None + return target_container + + +def _migration_blob_container_target_type(container_name): + container_mappings = { + getattr(app_config, "storage_account_user_documents_container_name", ""): "users", + getattr(app_config, "storage_account_group_documents_container_name", ""): "groups", + getattr(app_config, "storage_account_public_documents_container_name", ""): "public_workspaces", + } + return container_mappings.get(container_name) + + +def _selected_scope_hashes(target_type, selection): + if selection.get("mode") == "all": + return None + return { + hashlib.sha256(f"{target_type}:{scope_id}".encode("utf-8")).hexdigest() + for scope_id in _dedupe_limited_strings(selection.get("ids"), limit=2000) + } + + +def _new_reconciliation_report(service, include_blob_fields=False): + report = { + "service": service, + "matched_count": 0, + "missing_count": 0, + "destination_only_owned_count": 0, + "remaining_destination_only_owned_count": 0, + "destination_only_unowned_count": 0, + "conflict_count": 0, + "deleted_count": 0, + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "stale_count": 0, + "delete_candidate_count": 0, + "resources": [], + } + if include_blob_fields: + report.update({ + "unresolved_scope_count": 0, + "not_applicable_count": 0, + "source_missing_count": 0, + }) + return report + + +def _start_reconciliation_heartbeat_thread( + job, + settings, + state_holder, + stop_event, + failure_holder, + persistence_lock, +): + def run_heartbeat(): + while not stop_event.wait(2.0): + try: + with persistence_lock: + state = state_holder["state"] + state_holder["state"] = _persist_migration_state( + job, + state, + settings, + "Reconciliation worker heartbeat renewed", + ) + except Exception as exc: + failure_holder["error"] = exc + stop_event.set() + return + + heartbeat_thread = Thread(target=run_heartbeat, daemon=True) + heartbeat_thread.start() + return heartbeat_thread + + +def _raise_reconciliation_heartbeat_failure(stop_event, failure_holder): + if not stop_event.is_set(): + return + error = failure_holder.get("error") + if error: + raise error + raise DataManagementMigrationCanceledError("Reconciliation was canceled.") + + +def _append_reconciliation_resource(report, resource_report): + report["resources"].append(resource_report) + for field_name in report: + if field_name in {"service", "resources"} or str(field_name).startswith("_"): + continue + report[field_name] += _safe_int(resource_report.get(field_name), default=0, minimum=0) + + +def _emit_reconciliation_candidate(report, candidate, deletion_candidate_callback): + if callable(deletion_candidate_callback): + deletion_candidate_callback(_sanitize_mirror_deletion_candidate(candidate)) + + +def _reconcile_cosmos_migration( + target_database, + migration_plan, + migration_state, + provenance_context, + apply_deletions=False, + heartbeat_callback=None, + deletion_candidate_callback=None, + stop_event=None, + failure_holder=None, +): + """Reconcile Cosmos with bounded point reads and one destination scan.""" + report = _new_reconciliation_report("cosmos") + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none": + continue + for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type]: + if container_definition.get("documents") and not selection.get("include_documents"): + continue + container_name = _target_cosmos_container_name(container_definition) + target_container = _get_reconciliation_target_cosmos_container( + target_database, + container_name, + container_definition["partition_key_path"], + read_only=True, + ) + source_container = getattr(app_config, container_definition["container_attr"], None) + resource_report = { + "target_type": target_type, + "container_name": container_name, + "source_count": 0, + "target_count": 0, + "matched_count": 0, + "missing_count": 0, + "destination_only_owned_count": 0, + "remaining_destination_only_owned_count": 0, + "destination_only_unowned_count": 0, + "conflict_count": 0, + "deleted_count": 0, + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "stale_count": 0, + "delete_candidate_count": 0, + } + for item, source_version in _iter_selected_cosmos_records( + container_definition, + selection, + include_source_version=True, + ) or []: + if stop_event is not None: + _raise_reconciliation_heartbeat_failure(stop_event, failure_holder or {}) + resource_report["source_count"] += 1 + if heartbeat_callback and resource_report["source_count"] % 250 == 0: + heartbeat_callback( + f"Reconciling source Cosmos container {container_name}", + resource_report["source_count"], + ) + target_document = ( + _get_target_cosmos_document( + target_container, + item, + container_definition["partition_key_path"], + ) + if target_container is not None else None + ) + if target_document is None: + resource_report["missing_count"] += 1 + resource_report["create_count"] += 1 + continue + target_provenance = get_cosmos_migration_provenance(target_document) + if not is_successful_migration_record(target_provenance): + resource_report["conflict_count"] += 1 + continue + resource_report["matched_count"] += 1 + source_hash = _build_cosmos_source_hash(item) + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + resource_report["unchanged_count"] += 1 + else: + resource_report["stale_count"] += 1 + if migration_plan.get("migration_mode") != DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + resource_report["update_count"] += 1 + + target_records = ( + _iter_target_cosmos_records(target_container, container_definition, selection) + if target_container is not None else [] + ) + for target_document in target_records or []: + if stop_event is not None: + _raise_reconciliation_heartbeat_failure(stop_event, failure_holder or {}) + resource_report["target_count"] += 1 + if heartbeat_callback and resource_report["target_count"] % 250 == 0: + heartbeat_callback( + f"Reconciling destination Cosmos container {container_name}", + resource_report["target_count"], + ) + target_provenance = get_cosmos_migration_provenance(target_document) + if _source_cosmos_document_exists( + source_container, + target_document, + container_definition["partition_key_path"], + ): + continue + if is_successful_migration_record(target_provenance): + resource_report["destination_only_owned_count"] += 1 + resource_report["remaining_destination_only_owned_count"] += 1 + if migration_plan.get("migration_mode") == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR: + resource_report["delete_candidate_count"] += 1 + _emit_reconciliation_candidate(report, { + "service": "cosmos", + "target_type": target_type, + "container_name": container_name, + "document_id": target_document.get("id"), + "partition_key": _get_document_path_value( + target_document, + container_definition["partition_key_path"], + ), + "target_etag": target_document.get("_etag"), + }, deletion_candidate_callback) + else: + resource_report["destination_only_unowned_count"] += 1 + _append_reconciliation_resource(report, resource_report) + return report + + +def _source_cosmos_document_exists(source_container, target_document, partition_key_path): + if source_container is None: + return True + try: + source_container.read_item( + item=target_document.get("id"), + partition_key=_get_document_path_value(target_document, partition_key_path), + ) + return True + except (CosmosResourceNotFoundError, ResourceNotFoundError): + return False + + +def _iter_search_documents_keyset(search_client, search_filter=None, select=None): + for page, _last_id in _iter_search_document_pages( + search_client, + search_filter=search_filter, + select=select, + ): + yield from page + + +def _next_or_none(iterator): + try: + return next(iterator) + except StopIteration: + return None + + +def _reconcile_ai_search_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=None, + deletion_candidate_callback=None, + stop_event=None, + failure_holder=None, +): + """Reconcile Search by merging two ordered keyset streams.""" + report = _new_reconciliation_report("ai_search") + if not migration_plan.get("include_ai_search"): + return report + search_mappings = [ + ("users", "user_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[0]), + ("groups", "group_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[1]), + ("public_workspaces", "public_workspace_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[2]), + ] + provenance_select = [ + "id", + SEARCH_MIGRATION_ID_FIELD, + SEARCH_MIGRATED_AT_FIELD, + SEARCH_MIGRATION_STATUS_FIELD, + SEARCH_MIGRATION_SOURCE_HASH_FIELD, + SEARCH_MIGRATION_SOURCE_VERSION_FIELD, + ] + for target_type, field_name, artifact in search_mappings: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none" or not selection.get("include_documents"): + continue + source_client = CLIENTS.get(artifact["client_key"]) + if not source_client: + continue + target_client = _get_target_search_client(settings, artifact["index_name"]) + resource_report = { + "target_type": target_type, + "index_name": artifact["index_name"], + "source_count": 0, + "target_count": 0, + "matched_count": 0, + "missing_count": 0, + "destination_only_owned_count": 0, + "remaining_destination_only_owned_count": 0, + "destination_only_unowned_count": 0, + "conflict_count": 0, + "deleted_count": 0, + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "stale_count": 0, + "delete_candidate_count": 0, + } + for search_filter in _build_search_scope_filter_batches(field_name, selection): + source_iterator = iter(_iter_search_documents_keyset( + source_client, + search_filter=search_filter, + )) + try: + target_iterator = iter(_iter_search_documents_keyset( + target_client, + search_filter=search_filter, + select=provenance_select, + )) + source_document = _next_or_none(source_iterator) + target_document = _next_or_none(target_iterator) + except ResourceNotFoundError: + source_document = _next_or_none(source_iterator) + target_document = None + target_iterator = iter(()) + while source_document is not None or target_document is not None: + if stop_event is not None: + _raise_reconciliation_heartbeat_failure(stop_event, failure_holder or {}) + source_id = _safe_text((source_document or {}).get("id")) + target_id = _safe_text((target_document or {}).get("id")) + if target_document is None or (source_document is not None and source_id < target_id): + resource_report["source_count"] += 1 + resource_report["missing_count"] += 1 + resource_report["create_count"] += 1 + source_document = _next_or_none(source_iterator) + elif source_document is None or target_id < source_id: + resource_report["target_count"] += 1 + target_provenance = get_search_migration_provenance(target_document) + if is_successful_migration_record(target_provenance): + resource_report["destination_only_owned_count"] += 1 + resource_report["remaining_destination_only_owned_count"] += 1 + if migration_plan.get("migration_mode") == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR: + resource_report["delete_candidate_count"] += 1 + _emit_reconciliation_candidate(report, { + "service": "ai_search", + "target_type": target_type, + "index_name": artifact["index_name"], + "document_id": target_id, + }, deletion_candidate_callback) + else: + resource_report["destination_only_unowned_count"] += 1 + target_document = _next_or_none(target_iterator) + else: + resource_report["source_count"] += 1 + resource_report["target_count"] += 1 + target_provenance = get_search_migration_provenance(target_document) + if not is_successful_migration_record(target_provenance): + resource_report["conflict_count"] += 1 + else: + resource_report["matched_count"] += 1 + if migration_record_matches_source( + target_provenance, + source_hash=_build_search_source_hash(source_document), + ): + resource_report["unchanged_count"] += 1 + else: + resource_report["stale_count"] += 1 + if migration_plan.get("migration_mode") != DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + resource_report["update_count"] += 1 + source_document = _next_or_none(source_iterator) + target_document = _next_or_none(target_iterator) + processed_count = resource_report["source_count"] + resource_report["target_count"] + if heartbeat_callback and processed_count % 1000 == 0: + heartbeat_callback( + f"Reconciling AI Search index {artifact['index_name']}", + processed_count, + ) + _append_reconciliation_resource(report, resource_report) + return report + + +def _source_blob_reference_exists_in_plan( + migration_plan, + container_name, + blob_name, + scope_hash, +): + target_type = _migration_blob_container_target_type(container_name) + selection = migration_plan.get(target_type) or {} + allowed_scope_hashes = _selected_scope_hashes(target_type, selection) + if allowed_scope_hashes is not None and scope_hash not in allowed_scope_hashes: + return False + container_definition = next(( + definition + for definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS.get(target_type, []) + if definition.get("documents") + ), None) + source_container = ( + getattr(app_config, container_definition["container_attr"], None) + if container_definition else None + ) + if source_container is None: + return True + query = ( + "SELECT * FROM c WHERE c.blob_path = @blob_path " + "OR c.archived_blob_path = @blob_path" + ) + for document in source_container.query_items( + query=query, + parameters=[{"name": "@blob_path", "value": blob_name}], + enable_cross_partition_query=True, + ): + if allowed_scope_hashes is None or _document_migration_scope_hash(document) == scope_hash: + return True + return False + + +def _reconcile_blob_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=None, + deletion_candidate_callback=None, + stop_event=None, + failure_holder=None, +): + """Reconcile Blob paths with bounded point reads and a destination scan.""" + report = _new_reconciliation_report("source_blobs", include_blob_fields=True) + if not migration_plan.get("include_source_blobs"): + return report + target_service = _get_target_enhanced_citations_blob_client(settings) + source_service = _get_source_blob_service_client() + resource_reports = {} + for document_count, document in enumerate( + _iter_selected_document_records_for_blob_migration(migration_plan), + start=1, + ): + if stop_event is not None: + _raise_reconciliation_heartbeat_failure(stop_event, failure_holder or {}) + reference = _document_blob_reference(document) + if not reference: + report["not_applicable_count"] += 1 + continue + container_name, blob_name = reference + target_type = _migration_blob_container_target_type(container_name) + resource_report = resource_reports.setdefault(container_name, { + "target_type": target_type, + "container_name": container_name, + "source_count": 0, + "target_count": 0, + "matched_count": 0, + "missing_count": 0, + "destination_only_owned_count": 0, + "remaining_destination_only_owned_count": 0, + "destination_only_unowned_count": 0, + "conflict_count": 0, + "deleted_count": 0, + "unresolved_scope_count": 0, + "not_applicable_count": 0, + "source_missing_count": 0, + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "stale_count": 0, + "delete_candidate_count": 0, + }) + resource_report["source_count"] += 1 + source_blob = source_service.get_blob_client(container=container_name, blob=blob_name) + try: + source_properties = source_blob.get_blob_properties() + except ResourceNotFoundError: + resource_report["source_missing_count"] += 1 + continue + target_blob = target_service.get_blob_client(container=container_name, blob=blob_name) + target_properties = _get_blob_properties_or_none(target_blob) + if target_properties is None: + resource_report["missing_count"] += 1 + resource_report["create_count"] += 1 + continue + target_provenance = get_blob_migration_provenance( + getattr(target_properties, "metadata", None) + ) + if not is_successful_migration_record(target_provenance): + resource_report["conflict_count"] += 1 + continue + resource_report["matched_count"] += 1 + source_hash, source_version = _build_blob_source_fingerprint(source_properties) + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + resource_report["unchanged_count"] += 1 + else: + resource_report["stale_count"] += 1 + if migration_plan.get("migration_mode") != DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + resource_report["update_count"] += 1 + if heartbeat_callback and document_count % 250 == 0: + heartbeat_callback("Reconciling source Blob references", document_count) + + for container_name, resource_report in resource_reports.items(): + target_type = resource_report["target_type"] + selection = migration_plan.get(target_type) or {} + allowed_scope_hashes = _selected_scope_hashes(target_type, selection) + container_client = target_service.get_container_client(container_name) + try: + target_blob_items = container_client.list_blobs(include=["metadata"]) + except ResourceNotFoundError: + target_blob_items = [] + for target_count, blob_item in enumerate(target_blob_items, start=1): + if stop_event is not None: + _raise_reconciliation_heartbeat_failure(stop_event, failure_holder or {}) + metadata = ( + blob_item.get("metadata") + if isinstance(blob_item, dict) else getattr(blob_item, "metadata", None) + ) or {} + provenance = get_blob_migration_provenance(metadata) + blob_name = _safe_text( + blob_item.get("name") if isinstance(blob_item, dict) else getattr(blob_item, "name", None) + ) + if not blob_name: + continue + scope_hash = _safe_text(provenance.get("scopeHash")) + source_reference_exists = _source_blob_reference_exists_in_plan( + migration_plan, + container_name, + blob_name, + scope_hash, + ) + if is_successful_migration_record(provenance): + if allowed_scope_hashes is not None and not scope_hash: + resource_report["unresolved_scope_count"] += 1 + continue + if allowed_scope_hashes is not None and scope_hash not in allowed_scope_hashes: + continue + elif allowed_scope_hashes is not None and not source_reference_exists: + continue + resource_report["target_count"] += 1 + if source_reference_exists: + continue + if is_successful_migration_record(provenance): + resource_report["destination_only_owned_count"] += 1 + resource_report["remaining_destination_only_owned_count"] += 1 + if migration_plan.get("migration_mode") == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR: + resource_report["delete_candidate_count"] += 1 + _emit_reconciliation_candidate(report, { + "service": "source_blobs", + "target_type": target_type, + "container_name": container_name, + "blob_name": blob_name, + "target_etag": ( + blob_item.get("etag") + if isinstance(blob_item, dict) else getattr(blob_item, "etag", None) + ), + }, deletion_candidate_callback) + else: + resource_report["destination_only_unowned_count"] += 1 + if heartbeat_callback and target_count % 250 == 0: + heartbeat_callback(f"Reconciling Blob container {container_name}", target_count) + _append_reconciliation_resource(report, resource_report) + return report + + +def _source_cosmos_identity_exists(candidate): + container_definition = _get_migration_cosmos_container_definition( + candidate.get("target_type"), + candidate.get("container_name"), + ) + source_container = ( + getattr(app_config, container_definition["container_attr"], None) + if container_definition else None + ) + if source_container is None: + return True + try: + source_container.read_item( + item=candidate.get("document_id"), + partition_key=candidate.get("partition_key"), + ) + return True + except (CosmosResourceNotFoundError, ResourceNotFoundError): + return False + + +def _source_search_identity_exists(candidate): + artifact = _get_migration_search_artifact(candidate.get("target_type")) + source_client = CLIENTS.get(artifact["client_key"]) if artifact else None + document_id = _safe_text(candidate.get("document_id")) + if source_client is None or not document_id: + return True + results = source_client.search( + search_text="*", + filter=f"id eq '{_escape_search_filter_value(document_id)}'", + select=["id"], + top=1, + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + ) + return any(_safe_text(dict(result).get("id")) == document_id for result in results) + + +def _source_blob_identity_exists(candidate): + source_service = _get_source_blob_service_client() + if source_service is None: + return True + source_blob = source_service.get_blob_client( + container=candidate.get("container_name"), + blob=candidate.get("blob_name"), + ) + try: + source_blob.get_blob_properties() + return True + except ResourceNotFoundError: + return False + + +def _get_migration_cosmos_container_definition(target_type, container_name): + for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS.get( + _safe_text(target_type), + [], + ): + if _target_cosmos_container_name(container_definition) == _safe_text(container_name): + return container_definition + return None + + +def _get_migration_search_artifact(target_type): + artifacts_by_type = { + "users": DATA_MANAGEMENT_SEARCH_ARTIFACTS[0], + "groups": DATA_MANAGEMENT_SEARCH_ARTIFACTS[1], + "public_workspaces": DATA_MANAGEMENT_SEARCH_ARTIFACTS[2], + } + return artifacts_by_type.get(_safe_text(target_type)) + + +def _validate_mirror_deletion_candidate(candidate, settings=None, target_database=None): + """Revalidate current source absence and destination ownership before any delete.""" + service = candidate.get("service") + if service == "cosmos": + if _source_cosmos_identity_exists(candidate): + return False, "A Cosmos deletion candidate currently exists in the source." + container_definition = _get_migration_cosmos_container_definition( + candidate.get("target_type"), + candidate.get("container_name"), + ) + if not container_definition or target_database is None: + return False, "A Cosmos deletion candidate could not resolve its resource definition." + target_container = _get_reconciliation_target_cosmos_container( + target_database, + candidate.get("container_name"), + container_definition["partition_key_path"], + read_only=True, + ) + try: + current_target = target_container.read_item( + item=candidate.get("document_id"), + partition_key=candidate.get("partition_key"), + ) + except (CosmosResourceNotFoundError, ResourceNotFoundError): + current_target = None + if not current_target or not is_successful_migration_record( + get_cosmos_migration_provenance(current_target) + ): + return False, "A Cosmos deletion candidate lost migration ownership." + if not candidate.get("target_etag") or current_target.get("_etag") != candidate.get("target_etag"): + return False, "A Cosmos deletion candidate changed after reconciliation." + candidate["target_container"] = target_container + candidate["container_definition"] = container_definition + candidate["target_document"] = current_target + return True, "" + if service == "ai_search": + if _source_search_identity_exists(candidate): + return False, "An AI Search deletion candidate currently exists in the source." + artifact = _get_migration_search_artifact(candidate.get("target_type")) + if not artifact: + return False, "An AI Search deletion candidate could not resolve its index definition." + target_client = _get_target_search_client(settings, artifact["index_name"]) + target_documents = _get_target_search_documents_by_ids( + target_client, + [candidate.get("document_id")], + ) + current_target = target_documents.get(_safe_text(candidate.get("document_id"))) + if not current_target or not is_successful_migration_record( + get_search_migration_provenance(current_target) + ): + return False, "An AI Search deletion candidate lost migration ownership." + candidate["target_client"] = target_client + candidate["index_name"] = artifact["index_name"] + return True, "" + if service == "source_blobs": + if _source_blob_identity_exists(candidate): + return False, "A Blob deletion candidate currently exists in the source." + target_service = _get_target_enhanced_citations_blob_client(settings) + target_blob = target_service.get_blob_client( + container=candidate.get("container_name"), + blob=candidate.get("blob_name"), + ) + target_properties = _get_blob_properties_or_none(target_blob) + if not target_properties or not is_successful_migration_record( + get_blob_migration_provenance(getattr(target_properties, "metadata", None)) + ): + return False, "A Blob deletion candidate lost migration ownership." + current_etag = _safe_text(getattr(target_properties, "etag", None)) + if not candidate.get("target_etag") or current_etag != _safe_text(candidate.get("target_etag")): + return False, "A Blob deletion candidate changed after reconciliation." + candidate["target_blob"] = target_blob + candidate["target_etag"] = current_etag + return True, "" + return False, "An unsupported mirror deletion candidate was produced." + + +def _validate_mirror_deletion_candidates( + candidates, + heartbeat_callback=None, + settings=None, + target_database=None, +): + validated_count = 0 + blockers = set() + last_heartbeat_at = 0.0 + for index, candidate in enumerate(candidates, start=1): + now = time.monotonic() + if heartbeat_callback and ( + index == 1 or + index % 100 == 0 or + now - last_heartbeat_at >= DATA_MANAGEMENT_MIGRATION_HEARTBEAT_INTERVAL_SECONDS + ): + heartbeat_callback("Revalidating mirror deletion candidates", index) + last_heartbeat_at = now + is_valid, blocker = _validate_mirror_deletion_candidate( + candidate, + settings=settings, + target_database=target_database, + ) + if is_valid: + validated_count += 1 + elif len(blockers) < 20: + blockers.add(blocker) + return validated_count, sorted(blockers) + + +def _apply_mirror_deletion_candidates( + candidates, + heartbeat_callback=None, + settings=None, + target_database=None, + manifest_append=None, + manifest_flush=None, +): + deleted_counts = {"cosmos": 0, "ai_search": 0, "source_blobs": 0} + search_batch = [] + search_batch_key = None + + def record_deleted(candidate): + if not callable(manifest_append): + return + manifest_append({ + "service": candidate.get("service"), + "resource_name": "reconciliation:cutover", + "target_type": candidate.get("target_type"), + "source_identity": _hash_migration_manifest_identity( + "deleted-source", + candidate.get("service"), + candidate.get("container_name") or candidate.get("index_name"), + candidate.get("document_id") or candidate.get("blob_name"), + ), + "destination_identity": _hash_migration_manifest_identity( + "deleted-destination", + candidate.get("service"), + candidate.get("container_name") or candidate.get("index_name"), + candidate.get("document_id") or candidate.get("blob_name"), + ), + "status": "deleted", + "_locator": { + "service": candidate.get("service"), + "resource_name": "reconciliation:cutover", + "target_type": candidate.get("target_type"), + "document_id": candidate.get("document_id"), + "partition_key": candidate.get("partition_key"), + "container_name": candidate.get("container_name"), + "blob_name": candidate.get("blob_name"), + "index_name": candidate.get("index_name"), + }, + }) + if callable(manifest_flush): + manifest_flush() + + def flush_search_batch(): + nonlocal search_batch, search_batch_key + if not search_batch: + return + target_client = search_batch[0].get("target_client") + last_heartbeat_at = 0.0 + for index, candidate in enumerate(search_batch, start=1): + now = time.monotonic() + if heartbeat_callback and ( + index == 1 or + index % 100 == 0 or + now - last_heartbeat_at >= + DATA_MANAGEMENT_MIGRATION_HEARTBEAT_INTERVAL_SECONDS + ): + heartbeat_callback( + "Revalidating mirror AI Search deletions", + sum(deleted_counts.values()), + ) + last_heartbeat_at = now + is_valid, blocker = _validate_mirror_deletion_candidate( + candidate, + settings=settings, + target_database=target_database, + ) + if not is_valid: + raise DataManagementSettingsValidationError(blocker) + if heartbeat_callback: + heartbeat_callback( + "Applying mirror AI Search deletions", + sum(deleted_counts.values()), + ) + results = list(target_client.delete_documents( + documents=[{"id": candidate.get("document_id")} for candidate in search_batch], + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )) + successful_count = 0 + failure_count = 0 + for candidate, delete_result in zip(search_batch, results): + if _get_search_result_succeeded(delete_result): + successful_count += 1 + record_deleted(candidate) + else: + failure_count += 1 + failure_count += max(0, len(search_batch) - len(results)) + deleted_counts["ai_search"] += successful_count + if failure_count: + raise RuntimeError( + f"AI Search mirror reconciliation failed to delete {failure_count} document(s)." + ) + search_batch = [] + search_batch_key = None + + for candidate in candidates: + is_valid, blocker = _validate_mirror_deletion_candidate( + candidate, + settings=settings, + target_database=target_database, + ) + if not is_valid: + raise DataManagementSettingsValidationError(blocker) + if heartbeat_callback: + heartbeat_callback( + "Applying mirror deletions", + sum(deleted_counts.values()), + ) + if candidate.get("service") == "ai_search": + candidate_key = candidate.get("index_name") + if search_batch and ( + candidate_key != search_batch_key or + len(search_batch) >= DATA_MANAGEMENT_MIGRATION_BATCH_SIZE + ): + flush_search_batch() + search_batch_key = candidate_key + search_batch.append(candidate) + continue + flush_search_batch() + if candidate.get("service") == "cosmos": + _delete_target_cosmos_record( + candidate.get("target_container"), + candidate.get("target_document"), + (candidate.get("container_definition") or {}).get("partition_key_path"), + ) + deleted_counts["cosmos"] += 1 + record_deleted(candidate) + elif candidate.get("service") == "source_blobs": + candidate.get("target_blob").delete_blob( + etag=candidate.get("target_etag"), + match_condition=MatchConditions.IfNotModified, + ) + deleted_counts["source_blobs"] += 1 + record_deleted(candidate) + flush_search_batch() + return deleted_counts + + +def _public_reconciliation_report(report): + return { + key: value + for key, value in (report or {}).items() + if not str(key).startswith("_") + } + + +def _run_data_management_migration_reconciliation( + settings, + migration_plan, + job, + migration_state, + provenance_context, + target_database, + preview_snapshot=None, + migration_artifacts=None, +): + """Persist final source/destination reconciliation and mirror deletion outcomes.""" + resource_name = "reconciliation:cutover" + completed_resource = migration_state.get("resources", {}).get(resource_name) + if is_migration_resource_completed(migration_state, resource_name): + return completed_resource.get("result") if isinstance(completed_resource, dict) else {} + start_migration_resource(migration_state, resource_name) + deletion_plan_id = str(uuid.uuid4()) + append_deletion_candidate, flush_deletion_candidates = ( + _create_mirror_deletion_candidate_writer(job.get("id"), deletion_plan_id) + ) + append_reconciliation_manifest, flush_reconciliation_manifest = ( + _create_migration_manifest_writer(job.get("id"), resource_name) + ) + reconciliation_stop_event = Event() + reconciliation_failure_holder = {} + reconciliation_state_holder = {"state": migration_state} + reconciliation_persistence_lock = Lock() + reconciliation_heartbeat_thread = _start_reconciliation_heartbeat_thread( + job, + settings, + reconciliation_state_holder, + reconciliation_stop_event, + reconciliation_failure_holder, + reconciliation_persistence_lock, + ) + last_heartbeat_at = 0.0 + heartbeat_count = 0 + aggregate_fields = ( + "matched_count", + "missing_count", + "destination_only_owned_count", + "remaining_destination_only_owned_count", + "destination_only_unowned_count", + "conflict_count", + "deleted_count", + "unresolved_scope_count", + "create_count", + "update_count", + "unchanged_count", + "stale_count", + "delete_candidate_count", + "not_applicable_count", + "source_missing_count", + ) + service_reports = [] + + def stop_reconciliation_heartbeat(): + reconciliation_stop_event.set() + reconciliation_heartbeat_thread.join(timeout=5.0) + + def build_partial_result(error=None): + partial_result = { + field_name: sum( + _safe_int(report.get(field_name), default=0, minimum=0) + for report in service_reports + ) + for field_name in aggregate_fields + } + partial_result.update({ + "name": "migration_reconciliation", + "type": "migration_reconciliation", + "migration_mode": migration_plan.get("migration_mode"), + "readiness": "not_ready", + "services_completed": len(service_reports), + "services_total": 3, + "services": [ + _public_reconciliation_report(report) + for report in service_reports + ], + }) + if error: + partial_result["error"] = str(error)[:1000] + return partial_result + + def fail_reconciliation(error, failure_result=None): + nonlocal migration_state + stop_reconciliation_heartbeat() + migration_state = reconciliation_state_holder["state"] + persisted_result = copy.deepcopy( + failure_result if isinstance(failure_result, dict) else build_partial_result(error) + ) + resource = (migration_state.get("resources") or {}).get(resource_name) + if isinstance(resource, dict): + resource["result"] = persisted_result + fail_migration_resource( + migration_state, + resource_name, + str(error)[:1000], + ) + job["migration_state"] = migration_state + cleanup_errors = [] + for cleanup_name, cleanup in ( + ("deletion_plan", flush_deletion_candidates), + ("manifest", flush_reconciliation_manifest), + ): + try: + cleanup() + except Exception as cleanup_exc: + cleanup_errors.append(f"{cleanup_name}: {str(cleanup_exc)[:300]}") + if cleanup_errors: + persisted_result["persistence_warnings"] = cleanup_errors + resource = (migration_state.get("resources") or {}).get(resource_name) + if isinstance(resource, dict): + resource["result"] = persisted_result + try: + _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + str(error)[:1000], + ) + except Exception as checkpoint_exc: + persisted_result.setdefault("persistence_warnings", []).append( + f"checkpoint: {str(checkpoint_exc)[:300]}" + ) + resource = (migration_state.get("resources") or {}).get(resource_name) + if isinstance(resource, dict): + resource["result"] = persisted_result + resource["status"] = MIGRATION_RESOURCE_STATUS_FAILED + resource["last_error"] = str(error)[:1000] + job["migration_state"] = migration_state + + def persist_reconciliation_checkpoint(progress, message): + nonlocal migration_state + with reconciliation_persistence_lock: + migration_state = reconciliation_state_holder["state"] + migration_state = _persist_migration_checkpoint( + job, + migration_state, + settings, + resource_name, + progress, + message, + ) + reconciliation_state_holder["state"] = migration_state + return migration_state + + def heartbeat(message, completed_count=0, force=False): + nonlocal migration_state, last_heartbeat_at, heartbeat_count + _assert_migration_job_lease(job) + heartbeat_count += 1 + now = time.monotonic() + if not force and now - last_heartbeat_at < 2.0: + return + _raise_reconciliation_heartbeat_failure( + reconciliation_stop_event, + reconciliation_failure_holder, + ) + with reconciliation_persistence_lock: + migration_state = reconciliation_state_holder["state"] + migration_state = _persist_migration_checkpoint( + job, + migration_state, + settings, + resource_name, + { + "state": "reconciling", + "heartbeat_count": heartbeat_count, + "completed_count": completed_count, + }, + message, + ) + reconciliation_state_holder["state"] = migration_state + last_heartbeat_at = now + + try: + service_reports.append(_reconcile_cosmos_migration( + target_database, + migration_plan, + migration_state, + provenance_context, + apply_deletions=False, + heartbeat_callback=heartbeat, + deletion_candidate_callback=append_deletion_candidate, + stop_event=reconciliation_stop_event, + failure_holder=reconciliation_failure_holder, + )) + except Exception as exc: + fail_reconciliation(exc) + raise + try: + flush_deletion_candidates() + persist_reconciliation_checkpoint( + {"services_completed": 1, "services_total": 3}, + "Reconciling migrated Cosmos records", + ) + except Exception as exc: + fail_reconciliation(exc) + raise + try: + service_reports.append(_reconcile_ai_search_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=heartbeat, + deletion_candidate_callback=append_deletion_candidate, + stop_event=reconciliation_stop_event, + failure_holder=reconciliation_failure_holder, + )) + except Exception as exc: + fail_reconciliation(exc) + raise + try: + flush_deletion_candidates() + persist_reconciliation_checkpoint( + {"services_completed": 2, "services_total": 3}, + "Reconciling migrated AI Search documents", + ) + except Exception as exc: + fail_reconciliation(exc) + raise + try: + service_reports.append(_reconcile_blob_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=heartbeat, + deletion_candidate_callback=append_deletion_candidate, + stop_event=reconciliation_stop_event, + failure_holder=reconciliation_failure_holder, + )) + except Exception as exc: + fail_reconciliation(exc) + raise + try: + flush_deletion_candidates() + heartbeat("Completed read-only migration reconciliation", 3, force=True) + except Exception as exc: + fail_reconciliation(exc) + raise + result = { + field_name: sum(_safe_int(report.get(field_name), default=0, minimum=0) for report in service_reports) + for field_name in aggregate_fields + } + result.update({ + "name": "migration_reconciliation", + "type": "migration_reconciliation", + "migration_mode": migration_plan.get("migration_mode"), + "deletion_policy": ( + "migration_owned_only" + if migration_plan.get("migration_mode") == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR else + "retain_destination_only" + ), + "services": [_public_reconciliation_report(report) for report in service_reports], + }) + actual_outcomes = { + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "delete_count": result.get("delete_candidate_count", 0), + "not_applicable_count": 0, + "missing_count": 0, + "conflict_count": 0, + "failed_count": 0, + } + for artifact in migration_artifacts or []: + if artifact.get("type") not in { + "cosmos_container", + "ai_search_documents", + "source_blobs", + }: + continue + actual_outcomes["create_count"] += _safe_int( + artifact.get("created_count"), + default=0, + minimum=0, + ) + actual_outcomes["update_count"] += _safe_int( + artifact.get("updated_count"), + default=0, + minimum=0, + ) + actual_outcomes["unchanged_count"] += _safe_int( + artifact.get("unchanged_count"), + default=0, + minimum=0, + ) + actual_outcomes["not_applicable_count"] += _safe_int( + artifact.get("not_applicable_count"), + default=0, + minimum=0, + ) + actual_outcomes["missing_count"] += _safe_int( + artifact.get("missing_count"), + default=0, + minimum=0, + ) + actual_outcomes["conflict_count"] += _safe_int( + artifact.get("collision_count"), + default=0, + minimum=0, + ) + actual_outcomes["failed_count"] += _safe_int( + artifact.get("failed_count"), + default=0, + minimum=0, + ) + preview_counts = ( + preview_snapshot.get("estimated_outcomes") + if isinstance(preview_snapshot, dict) and + isinstance(preview_snapshot.get("estimated_outcomes"), dict) + else {} + ) + result["preview_snapshot"] = _sanitize_activity_value( + preview_snapshot if isinstance(preview_snapshot, dict) else {} + ) + result["actual_outcomes"] = actual_outcomes + result["preview_actual_divergence"] = { + field_name: actual_outcomes.get(field_name, 0) - _safe_int( + preview_counts.get(field_name), + default=0, + ) + for field_name in actual_outcomes + } + critical_mismatch_count = sum(( + result["missing_count"], + result["conflict_count"], + result["stale_count"], + result["source_missing_count"], + result["unresolved_scope_count"], + actual_outcomes["failed_count"], + )) + if critical_mismatch_count: + result["readiness"] = "not_ready" + elif ( + result["remaining_destination_only_owned_count"] or + result["destination_only_unowned_count"] or + result["unresolved_scope_count"] + ): + result["readiness"] = "ready_with_warnings" + else: + result["readiness"] = "ready" + + is_mirror = ( + migration_plan.get("migration_mode") == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR and + migration_plan.get("mirror_deletions_confirmed") is True + ) + if is_mirror: + preview_divergence = result["preview_actual_divergence"] + deletion_blockers = [] + if not preview_counts: + deletion_blockers.append("Mirror deletion requires a server-owned inventory preview.") + if any(value != 0 for value in preview_divergence.values()): + deletion_blockers.append("Migration outcomes diverged from the queued preview.") + if critical_mismatch_count: + deletion_blockers.append("Reconciliation found missing, stale, failed, or conflicting items.") + if result["destination_only_unowned_count"]: + deletion_blockers.append("Unowned destination-only items prevent an exact mirror.") + candidate_iterator = _iter_mirror_deletion_candidates( + job.get("id"), + deletion_plan_id, + ) + try: + validated_candidate_count, candidate_blockers = _validate_mirror_deletion_candidates( + candidate_iterator, + heartbeat_callback=heartbeat, + settings=settings, + target_database=target_database, + ) + except Exception as exc: + failure_result = copy.deepcopy(result) + failure_result.update({ + "readiness": "not_ready", + "error": str(exc)[:1000], + }) + fail_reconciliation(exc, failure_result=failure_result) + raise + deletion_blockers.extend(candidate_blockers) + if validated_candidate_count != result["delete_candidate_count"]: + deletion_blockers.append("Mirror deletion candidate count changed during validation.") + if deletion_blockers: + result["readiness"] = "not_ready" + result["deletion_status"] = "blocked" + result["deletion_blockers"] = sorted(set(deletion_blockers))[:20] + else: + target_search_write_gate_container = None + target_search_write_fence = None + requires_search_write_fence = any( + report.get("service") == "ai_search" and + _safe_int(report.get("delete_candidate_count"), default=0) > 0 + for report in service_reports + ) + + last_target_search_fence_renewal = 0.0 + + def mirror_deletion_heartbeat(message, completed_count=0): + nonlocal last_target_search_fence_renewal + heartbeat(message, completed_count) + if ( + target_search_write_gate_container is not None and + target_search_write_fence is not None and + time.monotonic() - last_target_search_fence_renewal >= + DATA_MANAGEMENT_MIGRATION_HEARTBEAT_INTERVAL_SECONDS + ): + renew_data_management_search_write_fence( + target_search_write_gate_container, + target_search_write_fence, + _get_target_search_write_fence_lease_seconds(settings), + ) + last_target_search_fence_renewal = time.monotonic() + + try: + if requires_search_write_fence: + target_search_write_gate_container = _get_target_data_management_search_write_gate_container( + target_database + ) + target_search_write_fence = acquire_data_management_search_write_fence( + target_search_write_gate_container, + job.get("id"), + _get_target_search_write_fence_lease_seconds(settings), + heartbeat_callback=lambda: heartbeat( + "Waiting for target AI Search writers to drain before mirror deletion", + result.get("deleted_count", 0), + ), + ) + deletion_counts = _apply_mirror_deletion_candidates( + _iter_mirror_deletion_candidates(job.get("id"), deletion_plan_id), + heartbeat_callback=mirror_deletion_heartbeat, + settings=settings, + target_database=target_database, + manifest_append=append_reconciliation_manifest, + manifest_flush=flush_reconciliation_manifest, + ) + if target_search_write_gate_container is not None: + release_data_management_search_write_fence( + target_search_write_gate_container, + target_search_write_fence, + ) + except DataManagementMigrationCanceledError: + if target_search_write_gate_container is not None: + release_data_management_search_write_fence( + target_search_write_gate_container, + target_search_write_fence, + ) + raise + except Exception as exc: + failure_result = copy.deepcopy(result) + failure_result.update({ + "readiness": "not_ready", + "error": str(exc)[:1000], + }) + fail_reconciliation(exc, failure_result=failure_result) + raise + flush_reconciliation_manifest() + result["deleted_count"] = sum(deletion_counts.values()) + result["remaining_destination_only_owned_count"] = max( + 0, + result["destination_only_owned_count"] - result["deleted_count"], + ) + result["actual_outcomes"]["delete_count"] = result["deleted_count"] + result["preview_actual_divergence"]["delete_count"] = ( + result["deleted_count"] - _safe_int(preview_counts.get("delete_count"), default=0) + ) + result["deletion_status"] = "completed" + result["deletion_counts"] = deletion_counts + result["readiness"] = ( + "ready" + if not result["preview_actual_divergence"]["delete_count"] else + "not_ready" + ) + heartbeat("Completed guarded mirror deletions", result["deleted_count"], force=True) + else: + result["actual_outcomes"]["delete_count"] = 0 + result["preview_actual_divergence"]["delete_count"] = ( + 0 - _safe_int(preview_counts.get("delete_count"), default=0) + ) + if result["readiness"] == "not_ready": + stop_reconciliation_heartbeat() + flush_reconciliation_manifest() + migration_state = reconciliation_state_holder["state"] + resource = (migration_state.get("resources") or {}).get(resource_name) + if isinstance(resource, dict): + resource["result"] = copy.deepcopy(result) + _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + "Migration reconciliation is not ready for cutover.", + ) + raise RuntimeError("Migration reconciliation is not ready for cutover.") + stop_reconciliation_heartbeat() + migration_state = reconciliation_state_holder["state"] + _complete_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + result, + "Completed migration source/destination reconciliation", + ) + return result + + +def _migration_selection_summary(migration_plan): + return { + target_type: { + "mode": (migration_plan.get(target_type) or {}).get("mode"), + "selected_count": len((migration_plan.get(target_type) or {}).get("ids") or []), + "include_documents": bool((migration_plan.get(target_type) or {}).get("include_documents")), + } + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES + } + + +def _job_lease_holder_id(): + return f"{socket.gethostname()}:{uuid.uuid4().hex}" + + +def _is_stale_job(job, stale_seconds=DATA_MANAGEMENT_DEFAULT_STALE_SECONDS): + last_heartbeat = _parse_iso_datetime((job or {}).get("last_heartbeat_at") or (job or {}).get("updated_at")) + if not last_heartbeat: + return True + return last_heartbeat <= _now_utc() - timedelta(seconds=stale_seconds) + + +def _read_job(job_id): + return cosmos_data_management_jobs_container.read_item(item=job_id, partition_key=job_id) + + +def _get_migration_destination_lock_id(settings, migration_plan): + """Use one global coordinator so partially overlapping migrations cannot race.""" + return f"{DATA_MANAGEMENT_MIGRATION_LOCK_PREFIX}_global" + + +def _get_migration_lock_lease_seconds(settings): + configured_lease_seconds = _safe_int( + (settings or {}).get("data_management_job_lease_seconds"), + default=DATA_MANAGEMENT_DEFAULT_LEASE_SECONDS, + minimum=60, + maximum=7200, + ) + minimum_quarantine_seconds = ( + DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS + + DATA_MANAGEMENT_MIGRATION_LOCK_RECOVERY_GRACE_SECONDS + ) + maximum_recovery_safe_seconds = ( + DATA_MANAGEMENT_DEFAULT_STALE_SECONDS - + DATA_MANAGEMENT_RECOVERY_RESUBMIT_DELAY_SECONDS + ) + return min( + max(configured_lease_seconds, minimum_quarantine_seconds), + maximum_recovery_safe_seconds, + ) + + +def _has_active_migration_destination_lock(job): + """Return whether a coordinator lock remains active and clear only conclusively stale metadata.""" + lock = job.get("migration_coordinator_lock") if isinstance(job, dict) else None + if not isinstance(lock, dict) or not lock.get("id") or not lock.get("lock_token"): + if isinstance(job, dict): + job.pop("migration_coordinator_lock", None) + return False + try: + current_lock = _read_job(lock["id"]) + except (CosmosResourceNotFoundError, ResourceNotFoundError, KeyError): + job.pop("migration_coordinator_lock", None) + return False + except Exception as exc: + raise DataManagementMigrationLeaseLostError( + "Migration destination coordinator lease could not be read while claiming a retry." + ) from exc + expires_at = _parse_iso_datetime(current_lock.get("expires_at")) + lock_is_active = bool( + current_lock.get("type") == DATA_MANAGEMENT_MIGRATION_LOCK_TYPE and + expires_at and + expires_at > _now_utc() + ) + if not lock_is_active: + job.pop("migration_coordinator_lock", None) + return lock_is_active + + +def _acquire_migration_destination_lock(job, settings, migration_plan): + """Acquire a destination-scoped lease so independent migration jobs cannot overlap.""" + if not isinstance(job, dict) or job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + return None + lock_id = _get_migration_destination_lock_id(settings, migration_plan) + now = _now_utc() + lease_seconds = _get_migration_lock_lease_seconds(settings) + lock_document = { + "id": lock_id, + "type": DATA_MANAGEMENT_MIGRATION_LOCK_TYPE, + "migration_job_id": job.get("id"), + "lock_token": str(uuid.uuid4()), + "acquired_at": now.isoformat(), + "updated_at": now.isoformat(), + "lease_seconds": lease_seconds, + "expires_at": (now + timedelta(seconds=lease_seconds)).isoformat(), + } + try: + cosmos_data_management_jobs_container.create_item(body=lock_document) + except Exception as exc: + if getattr(exc, "status_code", None) != 409: + raise + existing_lock = _read_job(lock_id) + expires_at = _parse_iso_datetime(existing_lock.get("expires_at")) + if expires_at and expires_at > now: + raise DataManagementSettingsValidationError( + "Another Data Management migration is active for this destination. Wait for it to complete or expire before retrying." + ) + replacement_lock = dict(existing_lock) + replacement_lock.update(lock_document) + try: + cosmos_data_management_jobs_container.replace_item( + item=lock_id, + body=replacement_lock, + etag=existing_lock.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception as replace_exc: + if getattr(replace_exc, "status_code", None) in {409, 412}: + raise DataManagementSettingsValidationError( + "Another Data Management migration acquired this destination before the retry could start." + ) from replace_exc + raise + + job["migration_coordinator_lock"] = { + "id": lock_id, + "lock_token": lock_document["lock_token"], + "lease_seconds": lease_seconds, + "expires_at": lock_document["expires_at"], + } + return job["migration_coordinator_lock"] + + +def _renew_migration_destination_lock(job, settings=None): + """Renew and verify the exclusive destination lease held by this worker.""" + lock = job.get("migration_coordinator_lock") if isinstance(job, dict) else None + if not isinstance(lock, dict) or not lock.get("id"): + return + now = _now_utc() + try: + current_lock = _read_job(lock["id"]) + except Exception as exc: + raise DataManagementMigrationLeaseLostError( + "Migration destination coordinator lease could not be read." + ) from exc + if ( + current_lock.get("type") != DATA_MANAGEMENT_MIGRATION_LOCK_TYPE or + current_lock.get("migration_job_id") != job.get("id") or + current_lock.get("lock_token") != lock.get("lock_token") + ): + raise DataManagementMigrationLeaseLostError( + "Migration destination coordinator lease was superseded." + ) + expires_at = _parse_iso_datetime(current_lock.get("expires_at")) + if expires_at is None or expires_at <= now: + raise DataManagementMigrationLeaseLostError( + "Migration destination coordinator lease expired." + ) + lease_seconds = _get_migration_lock_lease_seconds({ + "data_management_job_lease_seconds": ( + (settings or {}).get("data_management_job_lease_seconds") or + lock.get("lease_seconds") + ), + }) + current_lock["updated_at"] = now.isoformat() + current_lock["expires_at"] = ( + now + timedelta(seconds=lease_seconds) + ).isoformat() + try: + renewed_lock = cosmos_data_management_jobs_container.replace_item( + item=current_lock.get("id"), + body=current_lock, + etag=current_lock.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception as exc: + if getattr(exc, "status_code", None) in {409, 412}: + raise DataManagementMigrationLeaseLostError( + "Migration destination coordinator lease changed during renewal." + ) from exc + raise + lock["expires_at"] = renewed_lock.get("expires_at", current_lock["expires_at"]) + lock["lease_seconds"] = lease_seconds + + +def _release_migration_destination_lock(job): + """Release only the destination coordinator lease token owned by this job.""" + lock = job.get("migration_coordinator_lock") if isinstance(job, dict) else None + if not isinstance(lock, dict) or not lock.get("id"): + return + try: + current_lock = _read_job(lock["id"]) + except Exception: + return + if ( + current_lock.get("type") != DATA_MANAGEMENT_MIGRATION_LOCK_TYPE or + current_lock.get("migration_job_id") != job.get("id") or + current_lock.get("lock_token") != lock.get("lock_token") + ): + return + try: + cosmos_data_management_jobs_container.delete_item( + item=current_lock.get("id"), + partition_key=current_lock.get("id"), + etag=current_lock.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception: + return + + +def _assert_migration_job_lease(job, allow_cancel_requested=False): + """Stop a stale migration worker before it schedules another remote action.""" + if not isinstance(job, dict) or job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + return + lease_holder_id = _safe_text(job.get("lease_holder_id")) + if not lease_holder_id: + return + try: + persisted_job = _read_job(job.get("id")) + except Exception as exc: + raise DataManagementMigrationLeaseLostError( + "Migration worker could not verify its durable job lease." + ) from exc + + persisted_expiry = _parse_iso_datetime(persisted_job.get("lease_expires_at")) + if ( + persisted_job.get("status") != DATA_MANAGEMENT_STATUS_RUNNING or + _safe_text(persisted_job.get("lease_holder_id")) != lease_holder_id or + persisted_expiry is None or + persisted_expiry <= _now_utc() + ): + raise DataManagementMigrationLeaseLostError( + "Migration worker lease was superseded or expired." + ) + if persisted_job.get("cancel_requested_at") and not allow_cancel_requested: + raise DataManagementMigrationCanceledError( + "Migration cancellation was requested by an administrator." + ) + if persisted_job.get("_etag"): + job["_etag"] = persisted_job.get("_etag") + _renew_migration_destination_lock(job) + _renew_target_migration_coordinator(job, None) + + +def _replace_job(job): + return cosmos_data_management_jobs_container.replace_item( + item=job.get("id"), + body=job, + etag=job.get("_etag"), + match_condition=MatchConditions.IfNotModified, ) @@ -1196,6 +6535,10 @@ def queue_data_management_job(operation, backup_type=None, requested_by=None, re "last_heartbeat_at": None, "last_message": "Queued data management job", "last_error": None, + "cancel_requested_at": None, + "cancel_requested_by": None, + "cancel_requested_by_email": None, + "cancel_reason": None, "requested_by": _safe_text(requested_by), "requested_by_email": _safe_text(requested_by_email), "scheduled": bool(scheduled), @@ -1241,9 +6584,164 @@ def get_data_management_jobs(limit=DATA_MANAGEMENT_DEFAULT_JOB_LIMIT): ))[:safe_limit] +def get_recoverable_data_management_migration_jobs( + current_time=None, + limit=DATA_MANAGEMENT_DEFAULT_RECOVERY_JOB_LIMIT, +): + """Return queued or stale migration jobs that can safely be resubmitted.""" + safe_limit = _safe_int( + limit, + default=DATA_MANAGEMENT_DEFAULT_RECOVERY_JOB_LIMIT, + minimum=1, + maximum=100, + ) + now = current_time if isinstance(current_time, datetime) else _now_utc() + query = ( + "SELECT * FROM c WHERE c.type = @type AND c.operation = @operation " + "AND (c.status = @queued OR c.status = @running) ORDER BY c.updated_at ASC" + ) + parameters = [ + {"name": "@type", "value": DATA_MANAGEMENT_JOB_TYPE}, + {"name": "@operation", "value": DATA_MANAGEMENT_OPERATION_MIGRATION}, + {"name": "@queued", "value": DATA_MANAGEMENT_STATUS_QUEUED}, + {"name": "@running", "value": DATA_MANAGEMENT_STATUS_RUNNING}, + ] + candidates = list(cosmos_data_management_jobs_container.query_items( + query=query, + parameters=parameters, + enable_cross_partition_query=True, + max_item_count=safe_limit, + ))[:safe_limit] + recoverable_jobs = [] + for job in candidates: + status = job.get("status") + queued_since = _parse_iso_datetime(job.get("updated_at") or job.get("created_at")) + last_recovery = _parse_iso_datetime(job.get("last_recovery_submitted_at")) + resubmit_due = ( + last_recovery is None or + last_recovery <= now - timedelta(seconds=DATA_MANAGEMENT_RECOVERY_RESUBMIT_DELAY_SECONDS) + ) + if status == DATA_MANAGEMENT_STATUS_QUEUED: + is_queued_long_enough = ( + queued_since is None or + queued_since <= now - timedelta(seconds=DATA_MANAGEMENT_RECOVERY_QUEUE_DELAY_SECONDS) + ) + if is_queued_long_enough and resubmit_due: + recoverable_jobs.append({"job": job, "reason": "queued_recovery"}) + elif status == DATA_MANAGEMENT_STATUS_RUNNING and _is_stale_job(job) and resubmit_due: + recoverable_jobs.append({"job": job, "reason": "stale_recovery"}) + return recoverable_jobs + + +def recover_data_management_migration_jobs(app=None, settings=None, current_time=None): + """Resubmit recoverable migration jobs through the executor, never inline.""" + now = current_time if isinstance(current_time, datetime) else _now_utc() + recovery_results = [] + for recovery in get_recoverable_data_management_migration_jobs(current_time=now): + job = recovery["job"] + reason = recovery["reason"] + job.update({ + "last_recovery_submitted_at": now.isoformat(), + "recovery_attempt_count": _safe_int(job.get("recovery_attempt_count"), default=0, minimum=0) + 1, + "updated_at": now.isoformat(), + "last_message": ( + "Detected a stale migration worker; scheduling durable recovery." + if reason == "stale_recovery" else + "Detected a delayed queued migration; scheduling executor recovery." + ), + }) + saved_job = _save_data_management_job(job) + _record_data_management_job_event( + saved_job.get("id"), + "migration-stall-detected" if reason == "stale_recovery" else "migration-queued-recovery", + saved_job, + status=saved_job.get("status"), + message=saved_job.get("last_message"), + details={ + "recovery_attempt_count": saved_job.get("recovery_attempt_count"), + "cancel_requested": bool(saved_job.get("cancel_requested_at")), + }, + ) + submitted = submit_data_management_job(app, saved_job.get("id")) + if submitted: + _record_data_management_job_event( + saved_job.get("id"), + "migration-recovery-submitted", + saved_job, + status=saved_job.get("status"), + message="Submitted migration recovery to the executor.", + details={"reason": reason}, + ) + else: + log_event( + "[DataManagement] Migration recovery could not submit to the executor.", + {"job_id": saved_job.get("id"), "reason": reason}, + level=logging.WARNING, + ) + recovery_results.append({ + "job_id": saved_job.get("id"), + "reason": reason, + "submitted": submitted, + }) + return recovery_results + + +def _sanitize_data_management_migration_state_for_admin(state): + if not isinstance(state, dict): + return {} + resources = {} + for resource_name, resource in (state.get("resources") or {}).items(): + if not isinstance(resource, dict): + continue + resources[_safe_text(resource_name)] = { + "status": resource.get("status"), + "attempts": resource.get("attempts"), + "started_at": resource.get("started_at"), + "attempt_started_at": resource.get("attempt_started_at"), + "updated_at": resource.get("updated_at"), + "completed_at": resource.get("completed_at"), + "last_error": _safe_text(resource.get("last_error"))[:500], + "progress": _sanitize_activity_value(resource.get("progress") if isinstance(resource.get("progress"), dict) else {}), + "result": _sanitize_activity_value(resource.get("result") if isinstance(resource.get("result"), dict) else {}), + } + capacity = state.get("capacity") if isinstance(state.get("capacity"), dict) else {} + return { + "schema_version": state.get("schema_version"), + "migration_id": state.get("migration_id"), + "status": state.get("status"), + "created_at": state.get("created_at"), + "updated_at": state.get("updated_at"), + "completed_at": state.get("completed_at"), + "source_cutoff_at": state.get("source_cutoff_at"), + "last_progress_at": state.get("last_progress_at"), + "resume_count": state.get("resume_count"), + "last_error": _safe_text(state.get("last_error"))[:500], + "totals": _sanitize_activity_value(state.get("totals") if isinstance(state.get("totals"), dict) else {}), + "preflight": _sanitize_activity_value(state.get("preflight") if isinstance(state.get("preflight"), dict) else {}), + "capacity": { + "status": capacity.get("status"), + "restore_pending": bool(capacity.get("restore_pending")), + "target_ru": capacity.get("target_ru"), + "started_at": capacity.get("started_at"), + "completed_at": capacity.get("completed_at"), + "restored_at": capacity.get("restored_at"), + "restore_warnings": _sanitize_activity_value( + capacity.get("restore_warnings") if isinstance(capacity.get("restore_warnings"), list) else [] + ), + "targets": _sanitize_activity_value(capacity.get("targets") if isinstance(capacity.get("targets"), list) else []), + }, + "resources": resources, + } + + def sanitize_data_management_job_for_admin(job): if not isinstance(job, dict): return None + result = _sanitize_activity_value(job.get("result") if isinstance(job.get("result"), dict) else {}) + if isinstance(result, dict) and "migration_state" in result: + result["migration_state"] = _sanitize_data_management_migration_state_for_admin( + job.get("migration_state") + ) return { "id": job.get("id"), "operation": job.get("operation"), @@ -1254,13 +6752,28 @@ def sanitize_data_management_job_for_admin(job): "started_at": job.get("started_at"), "completed_at": job.get("completed_at"), "last_heartbeat_at": job.get("last_heartbeat_at"), + "last_progress_at": job.get("last_progress_at"), "last_message": job.get("last_message"), "last_error": job.get("last_error"), + "cancel_requested_at": job.get("cancel_requested_at"), "requested_by_email": job.get("requested_by_email"), "scheduled": bool(job.get("scheduled")), "progress": job.get("progress") if isinstance(job.get("progress"), dict) else {}, "warnings": job.get("warnings") if isinstance(job.get("warnings"), list) else [], - "result": job.get("result") if isinstance(job.get("result"), dict) else {}, + "result": result, + "migration_state": _sanitize_data_management_migration_state_for_admin(job.get("migration_state")), + "can_retry": bool( + job.get("operation") == DATA_MANAGEMENT_OPERATION_MIGRATION and + ( + job.get("status") in {DATA_MANAGEMENT_STATUS_FAILED, DATA_MANAGEMENT_STATUS_CANCELED} or + (job.get("status") == DATA_MANAGEMENT_STATUS_RUNNING and _is_stale_job(job)) + ) + ), + "can_cancel": bool( + job.get("operation") == DATA_MANAGEMENT_OPERATION_MIGRATION and + job.get("status") in {DATA_MANAGEMENT_STATUS_QUEUED, DATA_MANAGEMENT_STATUS_RUNNING} and + not job.get("cancel_requested_at") + ), } @@ -1284,9 +6797,173 @@ def get_data_management_job(job_id): if not safe_job_id: return None try: - return _read_job(safe_job_id) + job = _read_job(safe_job_id) except CosmosResourceNotFoundError: return None + return job if job.get("type") == DATA_MANAGEMENT_JOB_TYPE else None + + +def request_data_management_migration_cancellation( + job_id, + requested_by=None, + requested_by_email=None, + reason="", +): + """Request cooperative cancellation without invalidating durable checkpoints.""" + job = get_data_management_job(job_id) + if not job: + raise DataManagementSettingsValidationError("Data Management job was not found.") + if job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + raise DataManagementSettingsValidationError("Only migration jobs can be canceled from migration checkpoints.") + if job.get("status") == DATA_MANAGEMENT_STATUS_CANCELED: + return job + if job.get("status") in DATA_MANAGEMENT_TERMINAL_STATUSES: + raise DataManagementSettingsValidationError("Only queued or running migration jobs can be canceled.") + + now = _now_iso() + safe_reason = _safe_text(reason)[:500] + migration_state = job.get("migration_state") + if isinstance(migration_state, dict): + migration_state.update({ + "cancel_requested_at": now, + "cancel_requested_by": _safe_text(requested_by), + "cancel_reason": safe_reason, + "updated_at": now, + }) + + job.update({ + "cancel_requested_at": now, + "cancel_requested_by": _safe_text(requested_by), + "cancel_requested_by_email": _safe_text(requested_by_email), + "cancel_reason": safe_reason, + "updated_at": now, + "last_message": "Migration cancellation requested. The worker will stop at its next durable checkpoint.", + "migration_state": migration_state, + }) + if job.get("status") == DATA_MANAGEMENT_STATUS_QUEUED: + job.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "completed_at": now, + "last_heartbeat_at": now, + "last_message": "Queued migration canceled before execution started.", + "lease_holder_id": None, + "lease_expires_at": None, + }) + if isinstance(migration_state, dict): + migration_state.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "completed_at": now, + "updated_at": now, + }) + event_name = "migration-canceled" + event_message = "Queued migration canceled before execution started." + else: + event_name = "migration-cancel-requested" + event_message = "Migration cancellation requested; worker will stop at the next durable checkpoint." + + saved_job = _save_data_management_job(job) + _record_data_management_job_event( + saved_job.get("id"), + event_name, + saved_job, + status=saved_job.get("status"), + message=event_message, + details={ + "cancel_requested_at": now, + "has_reason": bool(safe_reason), + }, + ) + return saved_job + + +def retry_data_management_migration_job(job_id): + """Requeue a failed migration without changing its provenance identity or checkpoints.""" + job = get_data_management_job(job_id) + if not job: + raise DataManagementSettingsValidationError("Data Management job was not found.") + if job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + raise DataManagementSettingsValidationError("Only migration jobs can be retried from migration checkpoints.") + retryable_statuses = {DATA_MANAGEMENT_STATUS_FAILED, DATA_MANAGEMENT_STATUS_CANCELED} + if job.get("status") == DATA_MANAGEMENT_STATUS_RUNNING and _is_stale_job(job): + retryable_statuses.add(DATA_MANAGEMENT_STATUS_RUNNING) + if job.get("status") not in retryable_statuses: + raise DataManagementSettingsValidationError( + "Only failed, canceled, or stale migration jobs can be retried." + ) + migration_state = job.get("migration_state") + canceled_before_start = ( + job.get("status") == DATA_MANAGEMENT_STATUS_CANCELED and + not isinstance(migration_state, dict) + ) + if not canceled_before_start and ( + not isinstance(migration_state, dict) or not migration_state.get("migration_id") + ): + raise DataManagementSettingsValidationError("This migration does not have durable checkpoint state to retry.") + if ( + not canceled_before_start and + _safe_text(migration_state.get("migration_id")) != _safe_text(job.get("id")) + ): + raise DataManagementSettingsValidationError("Migration checkpoint identity does not match the job identity.") + + now = _now_iso() + if isinstance(migration_state, dict): + migration_state["last_cancellation"] = { + "requested_at": migration_state.get("cancel_requested_at"), + "requested_by": migration_state.get("cancel_requested_by"), + "reason": migration_state.get("cancel_reason"), + } + migration_state.update({ + "status": DATA_MANAGEMENT_STATUS_QUEUED, + "updated_at": now, + "last_error": None, + "retry_requested_at": now, + "cancel_requested_at": None, + "cancel_requested_by": None, + "cancel_reason": None, + }) + job.update({ + "status": DATA_MANAGEMENT_STATUS_QUEUED, + "updated_at": now, + "completed_at": None, + "last_heartbeat_at": None, + "last_message": "Migration retry queued from durable checkpoints", + "last_error": None, + "last_cancellation": { + "requested_at": job.get("cancel_requested_at"), + "requested_by": job.get("cancel_requested_by"), + "requested_by_email": job.get("cancel_requested_by_email"), + "reason": job.get("cancel_reason"), + }, + "cancel_requested_at": None, + "cancel_requested_by": None, + "cancel_requested_by_email": None, + "cancel_reason": None, + "lease_holder_id": None, + "lease_expires_at": None, + "migration_state": migration_state, + "result": {}, + }) + saved_job = _save_data_management_job(job) + _record_data_management_job_event( + saved_job.get("id"), + "migration-retry-queued", + saved_job, + status=DATA_MANAGEMENT_STATUS_QUEUED, + message="Migration retry queued from durable checkpoints", + details={ + "migration_id": ( + migration_state.get("migration_id") + if isinstance(migration_state, dict) else + job.get("id") + ), + "resume_count": ( + migration_state.get("resume_count") + if isinstance(migration_state, dict) else + 0 + ), + }, + ) + return saved_job def get_data_management_job_items(job_id, limit=200): @@ -1311,8 +6988,42 @@ def get_data_management_job_detail(job_id): job = get_data_management_job(job_id) if not job: return None + sanitized_job = sanitize_data_management_job_for_admin(job) + reconciliation_resource = ( + ((job.get("migration_state") or {}).get("resources") or {}).get( + "reconciliation:cutover" + ) + if isinstance(job.get("migration_state"), dict) else None + ) + reconciliation_result = ( + reconciliation_resource.get("result") + if isinstance(reconciliation_resource, dict) and + isinstance(reconciliation_resource.get("result"), dict) + else None + ) + if reconciliation_result and isinstance(sanitized_job, dict): + sanitized_result = ( + sanitized_job.get("result") + if isinstance(sanitized_job.get("result"), dict) + else {} + ) + artifacts = ( + list(sanitized_result.get("artifacts")) + if isinstance(sanitized_result.get("artifacts"), list) + else [] + ) + if not any( + isinstance(artifact, dict) and artifact.get("type") == "migration_reconciliation" + for artifact in artifacts + ): + reconciliation_artifact = _summarize_backup_artifact(reconciliation_result) + if reconciliation_artifact: + artifacts.append(reconciliation_artifact) + sanitized_result["artifacts"] = artifacts + sanitized_result["artifact_count"] = len(artifacts) + sanitized_job["result"] = sanitized_result return { - "job": sanitize_data_management_job_for_admin(job), + "job": sanitized_job, "items": [ sanitized_item for sanitized_item in ( @@ -1324,6 +7035,27 @@ def get_data_management_job_detail(job_id): } +def get_data_management_job_progress(job_id): + job = get_data_management_job(job_id) + if not job: + return None + public_job = sanitize_data_management_job_for_admin(job) + return { + "id": public_job.get("id"), + "operation": public_job.get("operation"), + "status": public_job.get("status"), + "updated_at": public_job.get("updated_at"), + "last_heartbeat_at": public_job.get("last_heartbeat_at"), + "last_progress_at": public_job.get("last_progress_at"), + "last_message": public_job.get("last_message"), + "last_error": public_job.get("last_error"), + "progress": public_job.get("progress"), + "migration_state": public_job.get("migration_state"), + "can_retry": public_job.get("can_retry"), + "can_cancel": public_job.get("can_cancel"), + } + + def _summarize_backup_artifact(artifact): if not isinstance(artifact, dict): return None @@ -1335,6 +7067,31 @@ def _summarize_backup_artifact(artifact): "path", "bytes", "item_count", + "copied_count", + "created_count", + "updated_count", + "unchanged_count", + "skipped_count", + "failed_count", + "collision_count", + "missing_count", + "not_applicable_count", + "processed_count", + "request_units", + "elapsed_seconds", + "items_per_second", + "bytes_per_second", + "request_units_per_second", + "parallel_operations", + "batch_size", + "retry_count", + "retry_attempt_count", + "prior_failed_count", + "checkpoint_count", + "source_read_count", + "destination_accepted_count", + "destination_failed_count", + "destination_provenance_skip_count", "blob_count", "encrypted", "container_name", @@ -1343,6 +7100,21 @@ def _summarize_backup_artifact(artifact): "partial_since_epoch", "partial_filter", "prefix", + "migration_mode", + "readiness", + "deletion_policy", + "deletion_status", + "deletion_blockers", + "deleted_count", + "delete_candidate_count", + "remaining_destination_only_owned_count", + "destination_only_unowned_count", + "unresolved_scope_count", + "stale_count", + "keyset_cursor", + "actual_outcomes", + "preview_actual_divergence", + "services", "warning", ] return { @@ -1372,7 +7144,11 @@ def _backup_artifact_totals(artifacts): continue totals["artifact_count"] += 1 totals["bytes"] += _safe_int(artifact.get("bytes"), default=0, minimum=0) - totals["record_count"] += _safe_int(artifact.get("item_count"), default=0, minimum=0) + totals["record_count"] += _safe_int( + artifact.get("item_count", artifact.get("copied_count")), + default=0, + minimum=0, + ) totals["blob_count"] += _safe_int(artifact.get("blob_count"), default=0, minimum=0) if artifact.get("warning") or artifact.get("status") == "warning": totals["warning_count"] += 1 @@ -1564,12 +7340,46 @@ def normalize_data_management_migration_plan(options): raw_plan = options.get("migration_plan") if isinstance(options, dict) else {} if not isinstance(raw_plan, dict): raw_plan = {} + migration_mode = _safe_text( + raw_plan.get("migration_mode"), + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ).lower() + if migration_mode not in DATA_MANAGEMENT_MIGRATION_MODES: + raise DataManagementSettingsValidationError( + "Migration mode must be new_only, delta_upsert, or mirror_with_deletions." + ) + baseline_job_id = _safe_text(raw_plan.get("baseline_job_id")) + if baseline_job_id: + try: + baseline_job_id = str(uuid.UUID(baseline_job_id)) + except (AttributeError, TypeError, ValueError) as exc: + raise DataManagementSettingsValidationError( + "Migration baseline job ID must be a valid GUID." + ) from exc + mirror_confirmation = _safe_text(raw_plan.get("mirror_confirmation")) + mirror_deletions_confirmed = ( + migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR and + mirror_confirmation == DATA_MANAGEMENT_MIRROR_CONFIRMATION + ) + if ( + migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR and + not mirror_deletions_confirmed + ): + raise DataManagementSettingsValidationError( + f"Mirror migration requires the exact confirmation phrase: {DATA_MANAGEMENT_MIRROR_CONFIRMATION}" + ) plan = { "users": _normalize_migration_selection(raw_plan.get("users")), "groups": _normalize_migration_selection(raw_plan.get("groups")), "public_workspaces": _normalize_migration_selection(raw_plan.get("public_workspaces")), "include_ai_search": raw_plan.get("include_ai_search") is not False, "include_source_blobs": bool(raw_plan.get("include_source_blobs")), + "target_ai_search_writes_frozen": ( + raw_plan.get("target_ai_search_writes_frozen") is True + ), + "migration_mode": migration_mode, + "baseline_job_id": baseline_job_id, + "mirror_deletions_confirmed": mirror_deletions_confirmed, } for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES: if plan[target_type].get("mode") == "none": @@ -1577,83 +7387,516 @@ def normalize_data_management_migration_plan(options): return plan -def _resolve_plan_scope_ids(target_type, plan_entry): - if plan_entry.get("mode") == "all": - catalog = get_data_management_migration_catalog(target_type, limit=1000) - return [item.get("id") for item in catalog.get("items", []) if item.get("id")] - return _dedupe_limited_strings(plan_entry.get("ids")) +def _resolve_plan_scope_ids(target_type, plan_entry): + if plan_entry.get("mode") == "all": + catalog = get_data_management_migration_catalog(target_type, limit=1000) + return [item.get("id") for item in catalog.get("items", []) if item.get("id")] + return _dedupe_limited_strings(plan_entry.get("ids")) + + +def summarize_data_management_migration_plan(options): + plan = normalize_data_management_migration_plan(options or {}) + summary = {} + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES: + ids = _resolve_plan_scope_ids(target_type, plan[target_type]) + summary[target_type] = { + "mode": plan[target_type].get("mode"), + "count": len(ids), + "include_documents": bool(plan[target_type].get("include_documents")), + "ids": ids[:50], + } + summary["include_ai_search"] = bool(plan.get("include_ai_search")) + summary["include_source_blobs"] = bool(plan.get("include_source_blobs")) + summary["target_ai_search_writes_frozen"] = bool( + plan.get("target_ai_search_writes_frozen") + ) + summary["migration_mode"] = plan.get("migration_mode") + summary["baseline_job_id"] = plan.get("baseline_job_id") + summary["mirror_deletions_confirmed"] = bool(plan.get("mirror_deletions_confirmed")) + return summary + + +def _count_preview_source_cosmos_records(migration_plan, source_cutoff_at): + source_cutoff = _parse_iso_datetime(source_cutoff_at) + source_cutoff_epoch = int(source_cutoff.timestamp()) if source_cutoff else None + record_count = 0 + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none": + continue + for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type]: + if container_definition.get("documents") and not selection.get("include_documents"): + continue + record_count += sum( + 1 + for _item in _iter_selected_cosmos_records( + container_definition, + selection, + source_cutoff_epoch=source_cutoff_epoch, + ) or [] + ) + return record_count + + +def preview_data_management_migration_plan( + settings, + options, + resolved_migration_plan=None, + source_cutoff_at=None, + migration_id=None, + heartbeat_callback=None, +): + """Build a read-only server-owned inventory preview for one migration plan.""" + if isinstance(resolved_migration_plan, dict): + migration_plan = copy.deepcopy(resolved_migration_plan) + else: + migration_plan = normalize_data_management_migration_plan(options or {}) + preview_job = {"id": str(uuid.uuid4())} + migration_plan = _resolve_data_management_migration_baseline( + preview_job, + settings, + migration_plan, + ) + source_cutoff_at = _safe_text(source_cutoff_at) or _now_iso() + preview_migration_id = _safe_text(migration_id) or str(uuid.uuid4()) + migration_state = {"source_cutoff_at": source_cutoff_at} + provenance_context = create_migration_provenance_context( + migration_id=preview_migration_id, + migrated_at_utc=source_cutoff_at, + ) + provenance_context.update({ + "migration_mode": migration_plan.get("migration_mode"), + "baseline_job_id": migration_plan.get("baseline_job_id"), + "baseline_source_cutoff_at": migration_plan.get("baseline_source_cutoff_at"), + }) + try: + target_database = _get_existing_target_cosmos_database(settings) + cosmos_report = _reconcile_cosmos_migration( + target_database, + migration_plan, + migration_state, + provenance_context, + apply_deletions=False, + heartbeat_callback=heartbeat_callback, + ) + except (CosmosResourceNotFoundError, ResourceNotFoundError): + source_record_count = _count_preview_source_cosmos_records( + migration_plan, + source_cutoff_at, + ) + cosmos_report = { + "service": "cosmos", + "create_count": source_record_count, + "update_count": 0, + "unchanged_count": 0, + "delete_candidate_count": 0, + "conflict_count": 0, + "missing_count": 0, + "warning": "Destination Cosmos database does not exist; selected source records are new.", + } + search_report = _reconcile_ai_search_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=heartbeat_callback, + ) + blob_report = _reconcile_blob_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=heartbeat_callback, + ) + service_reports = [cosmos_report, search_report, blob_report] + estimated_outcomes = { + "create_count": sum(_safe_int(report.get("create_count"), default=0) for report in service_reports), + "update_count": sum(_safe_int(report.get("update_count"), default=0) for report in service_reports), + "unchanged_count": sum(_safe_int(report.get("unchanged_count"), default=0) for report in service_reports), + "delete_count": sum(_safe_int(report.get("delete_candidate_count"), default=0) for report in service_reports), + "not_applicable_count": sum(_safe_int(report.get("not_applicable_count"), default=0) for report in service_reports), + "missing_count": sum( + _safe_int(report.get("source_missing_count"), default=0) + for report in service_reports + ), + "conflict_count": sum(_safe_int(report.get("conflict_count"), default=0) for report in service_reports), + "failed_count": 0, + } + normalized_preview_plan = copy.deepcopy(migration_plan) + normalized_preview_plan.pop("mirror_confirmation", None) + return { + "captured_at": source_cutoff_at, + "migration_mode": migration_plan.get("migration_mode"), + "baseline_job_id": migration_plan.get("baseline_job_id"), + "baseline_source_cutoff_at": migration_plan.get("baseline_source_cutoff_at"), + "plan_fingerprint": hashlib.sha256(json.dumps( + normalized_preview_plan, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + default=str, + ).encode("utf-8")).hexdigest(), + "estimated_outcomes": estimated_outcomes, + "services": service_reports, + } + + +def _try_claim_data_management_job(job_id, settings=None): + try: + job = _read_job(job_id) + except CosmosResourceNotFoundError: + return None + + status = _safe_text(job.get("status")) + if status in DATA_MANAGEMENT_TERMINAL_STATUSES: + return None + if status == DATA_MANAGEMENT_STATUS_RUNNING and not _is_stale_job(job): + return None + if job.get("operation") == DATA_MANAGEMENT_OPERATION_MIGRATION: + try: + if _has_active_migration_destination_lock(job): + return None + except DataManagementMigrationLeaseLostError as exc: + log_event( + "[DataManagement] Migration retry claim deferred because its coordinator lock could not be verified.", + {"job_id": job_id, "error": str(exc)}, + level=logging.WARNING, + ) + return None + + lease_seconds = _safe_int( + (settings or {}).get("data_management_job_lease_seconds"), + default=DATA_MANAGEMENT_DEFAULT_LEASE_SECONDS, + minimum=60, + maximum=7200, + ) + now = _now_utc() + job.update({ + "status": DATA_MANAGEMENT_STATUS_RUNNING, + "started_at": job.get("started_at") or now.isoformat(), + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "lease_holder_id": _job_lease_holder_id(), + "lease_expires_at": (now + timedelta(seconds=lease_seconds)).isoformat(), + "last_message": "Data management job claimed by a worker", + }) + if job.get("operation") == DATA_MANAGEMENT_OPERATION_MIGRATION: + job["migration_attempt_id"] = str(uuid.uuid4()) + try: + claimed_job = _replace_job(job) + _record_data_management_job_event( + claimed_job.get("id"), + "claimed", + claimed_job, + status=DATA_MANAGEMENT_STATUS_RUNNING, + message="Data management job claimed by a worker", + details={"lease_expires_at": claimed_job.get("lease_expires_at")}, + ) + return claimed_job + except Exception as exc: + status_code = getattr(exc, "status_code", None) + if status_code not in (409, 412): + log_event( + "[DataManagement] Job claim failed.", + {"job_id": job_id, "status_code": status_code, "error": str(exc)}, + level=logging.WARNING, + ) + return None + + +def _safe_job_item_id_part(value): + normalized = _safe_text(value, "event") + safe_value = "".join(character if character.isalnum() or character in {"-", "_"} else "-" for character in normalized) + return (safe_value.strip("-") or "event")[:80] + + +def _sanitize_migration_manifest_entry(entry): + allowed_fields = { + "service", + "resource_name", + "target_type", + "source_identity", + "destination_identity", + "item_ref", + "status", + "attempt", + "bytes", + "source_version", + "source_hash", + "error", + "recorded_at", + } + sanitized = { + field_name: _sanitize_activity_value((entry or {}).get(field_name)) + for field_name in allowed_fields + if (entry or {}).get(field_name) is not None + } + sanitized["error"] = _safe_text(sanitized.get("error"))[:500] + sanitized["recorded_at"] = _safe_text(sanitized.get("recorded_at")) or _now_iso() + return sanitized + + +def _hash_migration_manifest_identity(*identity_parts): + encoded = json.dumps( + [str(part) for part in identity_parts], + ensure_ascii=True, + separators=(",", ":"), + default=str, + ).encode("utf-8") + return f"sha256:{hashlib.sha256(encoded).hexdigest()}" + + +def _build_migration_manifest_item_ref( + job_id, + service="", + resource_name="", + target_type="", + document_id="", + partition_key=None, + container_name="", + blob_name="", + index_name="", +): + return str(uuid.uuid4()) + + +def resolve_data_management_migration_manifest_item(job_id, item_ref): + safe_job_id = _safe_text(job_id) + safe_item_ref = _safe_text(item_ref) + if not safe_job_id or not safe_item_ref: + raise DataManagementSettingsValidationError("Migration manifest item reference is required.") + job = get_data_management_job(safe_job_id) + if not job or job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + raise DataManagementSettingsValidationError("Data Management migration job was not found.") + query = ( + "SELECT * FROM c WHERE c.job_id = @job_id AND c.type = @type " + "ORDER BY c.created_at ASC" + ) + parameters = [ + {"name": "@job_id", "value": safe_job_id}, + {"name": "@type", "value": DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_TYPE}, + ] + batches = cosmos_data_management_job_items_container.query_items( + query=query, + parameters=parameters, + partition_key=safe_job_id, + max_item_count=100, + ) + for batch in batches: + for locator in batch.get("private_locators") or []: + if _safe_text(locator.get("item_ref")) != safe_item_ref: + continue + return { + "service": _safe_text(locator.get("service")), + "resource_name": _safe_text(locator.get("resource_name")), + "target_type": _safe_text(locator.get("target_type")), + "document_id": _safe_text(locator.get("document_id")), + "partition_key": locator.get("partition_key"), + "container_name": _safe_text(locator.get("container_name")), + "blob_name": _safe_text(locator.get("blob_name")), + "index_name": _safe_text(locator.get("index_name")), + } + raise DataManagementSettingsValidationError( + "Migration manifest item reference was not found." + ) + + +def _write_migration_manifest_batch(job_id, resource_name, entries): + normalized_entries = [] + private_locators = [] + for entry in entries or []: + if not isinstance(entry, dict): + continue + locator = entry.get("_locator") if isinstance(entry.get("_locator"), dict) else None + item_ref = _safe_text(entry.get("item_ref")) + if locator and not item_ref: + item_ref = _build_migration_manifest_item_ref(job_id, **locator) + public_entry = dict(entry) + public_entry.pop("_locator", None) + if item_ref: + public_entry["item_ref"] = item_ref + normalized_entries.append(_sanitize_migration_manifest_entry(public_entry)) + if locator and item_ref: + private_locators.append({ + "item_ref": item_ref, + "service": _safe_text(locator.get("service")), + "resource_name": _safe_text(locator.get("resource_name")), + "target_type": _safe_text(locator.get("target_type")), + "document_id": _safe_text(locator.get("document_id")), + "partition_key": copy.deepcopy(locator.get("partition_key")), + "container_name": _safe_text(locator.get("container_name")), + "blob_name": _safe_text(locator.get("blob_name")), + "index_name": _safe_text(locator.get("index_name")), + }) + if not normalized_entries: + return None + create_item = getattr(cosmos_data_management_job_items_container, "create_item", None) + if not callable(create_item): + raise RuntimeError("Migration manifest storage is not available.") + now = _now_iso() + body = { + "id": ( + f"{_safe_job_item_id_part(job_id)}:manifest:" + f"{_safe_job_item_id_part(resource_name)}:{uuid.uuid4().hex}" + ), + "job_id": _safe_text(job_id), + "type": DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_TYPE, + "resource_name": _safe_text(resource_name), + "created_at": now, + "updated_at": now, + "entry_count": len(normalized_entries), + "entries": normalized_entries, + "private_locators": private_locators, + } + return create_item(body=body) + + +def _create_migration_manifest_writer(job_id, resource_name): + buffer = [] + + def append(entry): + buffer.append(entry) + if len(buffer) >= DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_SIZE: + flush() + + def flush(): + if not buffer: + return + pending_entries = list(buffer) + _write_migration_manifest_batch(job_id, resource_name, pending_entries) + del buffer[:len(pending_entries)] + + return append, flush + + +def iter_data_management_migration_manifest_entries(job_id, statuses=None): + safe_job_id = _safe_text(job_id) + if not safe_job_id: + return + normalized_statuses = { + _safe_text(status).lower() + for status in (statuses or []) + if _safe_text(status) + } + query = ( + "SELECT * FROM c WHERE c.job_id = @job_id AND c.type = @type " + "ORDER BY c.created_at ASC" + ) + parameters = [ + {"name": "@job_id", "value": safe_job_id}, + {"name": "@type", "value": DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_TYPE}, + ] + batches = cosmos_data_management_job_items_container.query_items( + query=query, + parameters=parameters, + partition_key=safe_job_id, + max_item_count=100, + ) + for batch in batches: + for entry in batch.get("entries") or []: + sanitized = _sanitize_migration_manifest_entry(entry) + if normalized_statuses and _safe_text(sanitized.get("status")).lower() not in normalized_statuses: + continue + yield sanitized + + +def export_data_management_migration_manifest(job_id, statuses=None): + job = get_data_management_job(job_id) + if not job or job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + raise DataManagementSettingsValidationError("Data Management migration job was not found.") + entries = list(iter_data_management_migration_manifest_entries(job_id, statuses=statuses)) + lines = [ + json.dumps(entry, ensure_ascii=True, separators=(",", ":"), default=str) + for entry in entries + ] + return { + "content": "\n".join(lines) + ("\n" if lines else ""), + "entry_count": len(entries), + "statuses": sorted({ + _safe_text(status).lower() + for status in (statuses or []) + if _safe_text(status) + }), + } + + +def _sanitize_mirror_deletion_candidate(candidate): + allowed_fields = { + "service", + "target_type", + "container_name", + "document_id", + "partition_key", + "target_etag", + "index_name", + "blob_name", + } + return { + field_name: copy.deepcopy((candidate or {}).get(field_name)) + for field_name in allowed_fields + if (candidate or {}).get(field_name) is not None + } -def summarize_data_management_migration_plan(options): - plan = normalize_data_management_migration_plan(options or {}) - summary = {} - for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES: - ids = _resolve_plan_scope_ids(target_type, plan[target_type]) - summary[target_type] = { - "mode": plan[target_type].get("mode"), - "count": len(ids), - "include_documents": bool(plan[target_type].get("include_documents")), - "ids": ids[:50], - } - summary["include_ai_search"] = bool(plan.get("include_ai_search")) - summary["include_source_blobs"] = bool(plan.get("include_source_blobs")) - return summary +def _write_mirror_deletion_candidate_batch(job_id, plan_id, candidates): + normalized_candidates = [ + _sanitize_mirror_deletion_candidate(candidate) + for candidate in (candidates or []) + if isinstance(candidate, dict) + ] + if not normalized_candidates: + return None + now = _now_iso() + return cosmos_data_management_job_items_container.create_item(body={ + "id": ( + f"{_safe_job_item_id_part(job_id)}:mirror-plan:" + f"{_safe_job_item_id_part(plan_id)}:{uuid.uuid4().hex}" + ), + "job_id": _safe_text(job_id), + "type": DATA_MANAGEMENT_MIRROR_DELETION_BATCH_TYPE, + "plan_id": _safe_text(plan_id), + "created_at": now, + "updated_at": now, + "candidate_count": len(normalized_candidates), + "candidates": normalized_candidates, + }) -def _try_claim_data_management_job(job_id, settings=None): - try: - job = _read_job(job_id) - except CosmosResourceNotFoundError: - return None +def _create_mirror_deletion_candidate_writer(job_id, plan_id): + buffer = [] - status = _safe_text(job.get("status")) - if status in DATA_MANAGEMENT_TERMINAL_STATUSES: - return None - if status == DATA_MANAGEMENT_STATUS_RUNNING and not _is_stale_job(job): - return None + def append(candidate): + buffer.append(candidate) + if len(buffer) >= DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_SIZE: + flush() - lease_seconds = _safe_int( - (settings or {}).get("data_management_job_lease_seconds"), - default=DATA_MANAGEMENT_DEFAULT_LEASE_SECONDS, - minimum=60, - maximum=7200, - ) - now = _now_utc() - job.update({ - "status": DATA_MANAGEMENT_STATUS_RUNNING, - "started_at": job.get("started_at") or now.isoformat(), - "updated_at": now.isoformat(), - "last_heartbeat_at": now.isoformat(), - "lease_holder_id": _job_lease_holder_id(), - "lease_expires_at": (now + timedelta(seconds=lease_seconds)).isoformat(), - "last_message": "Data management job claimed by a worker", - }) - try: - claimed_job = _replace_job(job) - _record_data_management_job_event( - claimed_job.get("id"), - "claimed", - claimed_job, - status=DATA_MANAGEMENT_STATUS_RUNNING, - message="Data management job claimed by a worker", - details={"lease_expires_at": claimed_job.get("lease_expires_at")}, - ) - return claimed_job - except Exception as exc: - status_code = getattr(exc, "status_code", None) - if status_code not in (409, 412): - log_event( - "[DataManagement] Job claim failed.", - {"job_id": job_id, "status_code": status_code, "error": str(exc)}, - level=logging.WARNING, - ) - return None + def flush(): + if not buffer: + return + pending_candidates = list(buffer) + _write_mirror_deletion_candidate_batch(job_id, plan_id, pending_candidates) + del buffer[:len(pending_candidates)] + return append, flush -def _safe_job_item_id_part(value): - normalized = _safe_text(value, "event") - safe_value = "".join(character if character.isalnum() or character in {"-", "_"} else "-" for character in normalized) - return (safe_value.strip("-") or "event")[:80] + +def _iter_mirror_deletion_candidates(job_id, plan_id): + query = ( + "SELECT * FROM c WHERE c.job_id = @job_id AND c.type = @type " + "AND c.plan_id = @plan_id ORDER BY c.created_at ASC" + ) + parameters = [ + {"name": "@job_id", "value": _safe_text(job_id)}, + {"name": "@type", "value": DATA_MANAGEMENT_MIRROR_DELETION_BATCH_TYPE}, + {"name": "@plan_id", "value": _safe_text(plan_id)}, + ] + batches = cosmos_data_management_job_items_container.query_items( + query=query, + parameters=parameters, + partition_key=_safe_text(job_id), + max_item_count=100, + ) + for batch in batches: + for candidate in batch.get("candidates") or []: + yield _sanitize_mirror_deletion_candidate(candidate) def create_data_management_job_item(job_id, step_name, status=DATA_MANAGEMENT_STATUS_QUEUED, message=None, details=None): @@ -2116,20 +8359,50 @@ def save_data_management_cosmos_editor_document(container_name, document_id, par def _save_data_management_job(job): body = _strip_cosmos_system_fields(job) - saved = cosmos_data_management_jobs_container.upsert_item(body) + job_etag = job.get("_etag") if isinstance(job, dict) else None + if job_etag and hasattr(cosmos_data_management_jobs_container, "replace_item"): + try: + saved = cosmos_data_management_jobs_container.replace_item( + item=job.get("id"), + body=body, + etag=job_etag, + match_condition=MatchConditions.IfNotModified, + ) + except Exception as exc: + if getattr(exc, "status_code", None) in {409, 412}: + raise DataManagementMigrationLeaseLostError( + "Migration job state was changed by another worker." + ) from exc + raise + else: + saved = cosmos_data_management_jobs_container.upsert_item(body) job.clear() job.update(saved) return job -def _set_job_progress(job, message, completed_steps, total_steps, current_step=None, status=DATA_MANAGEMENT_STATUS_RUNNING): +def _set_job_progress( + job, + message, + completed_steps, + total_steps, + current_step=None, + status=DATA_MANAGEMENT_STATUS_RUNNING, + allow_cancel_requested=False, +): total_steps = max(1, total_steps) + _assert_migration_job_lease( + job, + allow_cancel_requested=allow_cancel_requested, + ) completed_steps = max(0, min(completed_steps, total_steps)) percent_complete = int((completed_steps / total_steps) * 100) + now = _now_iso() job.update({ "status": status, - "updated_at": _now_iso(), - "last_heartbeat_at": _now_iso(), + "updated_at": now, + "last_heartbeat_at": now, + "last_progress_at": now, "last_message": message, "progress": { "total_steps": total_steps, @@ -2573,73 +8846,396 @@ def execute_backup_job(job, settings): def execute_migration_job(job, settings): options = job.get("options") if isinstance(job.get("options"), dict) else {} migration_plan = normalize_data_management_migration_plan(options) + total_steps = 10 + _set_job_progress( + job, + "Resolving migration plan and baseline", + 0, + total_steps, + current_step="plan", + allow_cancel_requested=True, + ) + migration_plan = _resolve_data_management_migration_baseline( + job, + settings, + migration_plan, + ) dry_run = bool(options.get("dry_run")) warnings = [] artifacts = [] - total_steps = 5 plan_summary = summarize_data_management_migration_plan({"migration_plan": migration_plan}) + plan_summary["baseline_source_cutoff_at"] = migration_plan.get("baseline_source_cutoff_at") + selected_total = sum(plan_summary[target_type]["count"] for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES) + if selected_total == 0: + raise DataManagementSettingsValidationError("Choose at least one user, group, or public workspace before running migration.") + + migration_state, provenance_context = _initialize_data_management_migration_state( + job, + settings, + migration_plan, + ) + try: + migration_state = _persist_migration_state( + job, + migration_state, + settings, + "Initialized durable migration provenance and checkpoint state", + ) + except DataManagementMigrationCanceledError as exc: + migration_state.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "canceled_at": _now_iso(), + "last_error": None, + "cancel_message": str(exc), + }) + _persist_migration_state( + job, + migration_state, + settings, + "Migration cancellation acknowledged before execution started", + allow_cancel_requested=True, + ) + raise _set_job_progress(job, "Validated migration selection plan", 1, total_steps, current_step="plan") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state _record_data_management_job_event( job.get("id"), "migration-plan", job, status=DATA_MANAGEMENT_STATUS_RUNNING, message="Migration selection plan validated", - details={"migration_plan": plan_summary, "dry_run": dry_run}, + details={ + "migration_plan": plan_summary, + "dry_run": dry_run, + "migration_id": provenance_context["migration_id"], + "parallel_operations": _get_migration_parallel_operations(settings), + }, ) - selected_total = sum(plan_summary[target_type]["count"] for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES) - if selected_total == 0: - raise DataManagementSettingsValidationError("Choose at least one user, group, or public workspace before running migration.") - if dry_run: warning = "Migration dry run completed. No destination records were written." warnings.append(warning) + migration_state.update({ + "status": MIGRATION_RESOURCE_STATUS_COMPLETED, + "completed_at": _now_iso(), + }) + migration_state = _persist_migration_state(job, migration_state, settings, warning) _set_job_progress(job, warning, total_steps, total_steps, current_step="dry_run", status=DATA_MANAGEMENT_STATUS_COMPLETED_WITH_WARNINGS) return { "migration_plan": plan_summary, "dry_run": True, + "migration_id": provenance_context["migration_id"], + "migration_state": migration_state, "artifacts": [], "artifact_totals": _backup_artifact_totals([]), "warnings": warnings, } + _validate_target_ai_search_migration_write_safety(settings, migration_plan) + _acquire_migration_destination_lock(job, settings, migration_plan) + try: + _save_data_management_job(job) + except Exception: + _release_migration_destination_lock(job) + raise + + # The empty destination database must exist before it can host the target-wide coordinator. target_database = _get_target_cosmos_database(settings) - _set_job_progress(job, "Connected to target Cosmos database", 2, total_steps, current_step="target_cosmos") + try: + _acquire_target_migration_coordinator(job, target_database, settings) + _save_data_management_job(job) + except Exception: + _release_target_migration_coordinator(job) + _release_migration_destination_lock(job) + raise - for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES: - selection = migration_plan.get(target_type) or {} - if selection.get("mode") == "none": - continue - copied = _copy_cosmos_records_to_target(target_database, target_type, selection) - artifacts.extend(copied) + try: + _set_job_progress( + job, + "Building server-owned migration inventory", + 1, + total_steps, + current_step="inventory", + ) + preview_resource_name = "migration_preview:inventory" + preview_resource = migration_state.get("resources", {}).get(preview_resource_name) + migration_preview = ( + preview_resource.get("result") + if is_migration_resource_completed(migration_state, preview_resource_name) and + isinstance(preview_resource, dict) and + isinstance(preview_resource.get("result"), dict) + else None + ) + if not isinstance(migration_preview, dict): + start_migration_resource(migration_state, preview_resource_name) + migration_state = _persist_migration_checkpoint( + job, + migration_state, + settings, + preview_resource_name, + {"state": "inventory", "heartbeat_count": 0}, + "Building server-owned migration inventory preview", + ) + preview_heartbeat_count = 0 + preview_last_heartbeat_at = 0.0 + + def preview_heartbeat(message, completed_count=0): + nonlocal migration_state, preview_heartbeat_count, preview_last_heartbeat_at + preview_heartbeat_count += 1 + now = time.monotonic() + if now - preview_last_heartbeat_at < 2.0: + return + migration_state = _persist_migration_checkpoint( + job, + migration_state, + settings, + preview_resource_name, + { + "state": "inventory", + "heartbeat_count": preview_heartbeat_count, + "completed_count": completed_count, + }, + message, + ) + preview_last_heartbeat_at = now + + migration_preview = preview_data_management_migration_plan( + settings, + options, + resolved_migration_plan=migration_plan, + source_cutoff_at=migration_state.get("source_cutoff_at"), + migration_id=job.get("id"), + heartbeat_callback=preview_heartbeat, + ) + migration_state = _complete_migration_resource_checkpoint( + job, + migration_state, + settings, + preview_resource_name, + migration_preview, + "Completed server-owned migration inventory preview", + ) + options["migration_preview"] = migration_preview + job["options"] = options + migration_state = _persist_migration_state( + job, + migration_state, + settings, + "Pinned server-owned migration inventory preview", + ) + + _set_job_progress(job, "Migration inventory completed", 2, total_steps, current_step="inventory") + _set_job_progress(job, "Validating migration destinations", 2, total_steps, current_step="preflight") + migration_state = _run_data_management_migration_preflight( + job, + migration_state, + settings, + migration_plan, + ) + _set_job_progress(job, "Destination migration preflight completed", 3, total_steps, current_step="preflight") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state _record_data_management_job_event( job.get("id"), - f"migration-cosmos-{target_type}", + "migration-preflight", job, status=DATA_MANAGEMENT_STATUS_RUNNING, - message=f"Migrated {target_type.replace('_', ' ')} Cosmos records", - details={"target_type": target_type, "artifacts": copied}, + message="Verified source and destination migration access", + details=migration_state.get("preflight") if isinstance(migration_state.get("preflight"), dict) else {}, ) - _set_job_progress(job, "Cosmos migration completed", 3, total_steps, current_step="cosmos") - try: - search_artifacts = _copy_ai_search_to_target(settings, migration_plan) + _set_job_progress(job, "Preparing destination Cosmos capacity", 3, total_steps, current_step="capacity") + migration_state = _apply_temporary_destination_capacity( + job, + migration_state, + settings, + migration_plan, + ) + _set_job_progress(job, "Destination Cosmos capacity prepared", 4, total_steps, current_step="capacity") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state + + _set_job_progress(job, "Migrating Cosmos records", 4, total_steps, current_step="cosmos") + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none": + continue + copied = _copy_cosmos_records_to_target( + target_database, + target_type, + selection, + job, + migration_state, + provenance_context, + settings, + ) + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state + artifacts.extend(copied) + _record_data_management_job_event( + job.get("id"), + f"migration-cosmos-{target_type}", + job, + status=DATA_MANAGEMENT_STATUS_RUNNING, + message=f"Migrated {target_type.replace('_', ' ')} Cosmos records", + details={"target_type": target_type, "artifacts": copied}, + ) + _set_job_progress(job, "Cosmos migration completed", 5, total_steps, current_step="cosmos") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state + + search_artifacts = [] + if _migration_plan_includes_ai_search_documents(migration_plan): + _set_job_progress(job, "Freezing target AI Search writers", 5, total_steps, current_step="ai_search") + target_search_write_gate_container = _get_target_data_management_search_write_gate_container( + target_database + ) + target_search_write_fence = acquire_data_management_search_write_fence( + target_search_write_gate_container, + job.get("id"), + _get_target_search_write_fence_lease_seconds(settings), + heartbeat_callback=lambda: _persist_migration_heartbeat( + job, + settings, + "Waiting for target AI Search writers to drain", + ), + ) + release_target_search_write_fence = False + try: + _set_job_progress(job, "Migrating AI Search documents", 5, total_steps, current_step="ai_search") + search_artifacts = _copy_ai_search_to_target( + settings, + migration_plan, + job, + migration_state, + provenance_context, + target_search_write_fence=( + target_search_write_gate_container, + target_search_write_fence, + ), + ) + release_target_search_write_fence = True + except DataManagementMigrationCanceledError: + release_target_search_write_fence = True + raise + finally: + if release_target_search_write_fence: + release_data_management_search_write_fence( + target_search_write_gate_container, + target_search_write_fence, + ) + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state artifacts.extend(search_artifacts) - except Exception as exc: - warnings.append(f"AI Search migration warning: {exc}") - _set_job_progress(job, "AI Search migration step completed", 4, total_steps, current_step="ai_search") + _set_job_progress(job, "AI Search migration completed", 6, total_steps, current_step="ai_search") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state - try: - source_blob_artifacts = _copy_source_blobs_to_target(settings, migration_plan) + _set_job_progress(job, "Migrating source document blobs", 6, total_steps, current_step="source_blobs") + source_blob_artifacts = _copy_source_blobs_to_target( + settings, + migration_plan, + job, + migration_state, + provenance_context, + ) + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state artifacts.extend(source_blob_artifacts) + _set_job_progress(job, "Source blob migration completed", 7, total_steps, current_step="source_blobs") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state + + _set_job_progress(job, "Reconciling source and destination", 7, total_steps, current_step="reconciliation") + reconciliation_artifact = _run_data_management_migration_reconciliation( + settings, + migration_plan, + job, + migration_state, + provenance_context, + target_database, + preview_snapshot=migration_preview, + migration_artifacts=artifacts, + ) + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state + artifacts.append(reconciliation_artifact) + _set_job_progress(job, "Migration reconciliation completed", 8, total_steps, current_step="reconciliation") + _record_data_management_job_event( + job.get("id"), + "migration-reconciliation", + job, + status=DATA_MANAGEMENT_STATUS_RUNNING, + message="Reconciled migration source and destination identities", + details=reconciliation_artifact, + ) + except DataManagementMigrationCanceledError as exc: + migration_state.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "canceled_at": _now_iso(), + "last_error": None, + "cancel_message": str(exc), + }) + _persist_migration_state( + job, + migration_state, + settings, + "Migration cancellation acknowledged at a durable checkpoint", + allow_cancel_requested=True, + ) + raise except Exception as exc: - warnings.append(f"Enhanced Citations source blob migration warning: {exc}") + migration_state["status"] = MIGRATION_RESOURCE_STATUS_FAILED + migration_state["last_error"] = str(exc)[:1000] + _persist_migration_state(job, migration_state, settings, "Migration execution failed") + raise + finally: + try: + _set_job_progress( + job, + "Restoring destination Cosmos capacity", + 8, + total_steps, + current_step="capacity_restore", + allow_cancel_requested=True, + ) + restore_warnings, migration_state = _restore_temporary_destination_capacity( + job, + migration_state, + settings, + allow_cancel_requested=True, + ) + warnings.extend(restore_warnings) + except Exception as restore_exc: + warnings.append(f"Failed to restore temporary destination Cosmos capacity: {str(restore_exc)[:300]}") + + capacity = ( + migration_state.get("capacity") + if isinstance(migration_state.get("capacity"), dict) + else {} + ) + if capacity.get("restore_pending"): + migration_state.update({ + "status": MIGRATION_RESOURCE_STATUS_FAILED, + "last_error": "Destination Cosmos capacity restoration remains pending.", + }) + _persist_migration_state( + job, + migration_state, + settings, + "Destination Cosmos capacity restoration remains pending; migration can be retried", + ) + raise RuntimeError( + "Destination Cosmos capacity restoration remains pending. Retry the migration to restore the recorded capacity snapshot." + ) + for artifact in artifacts: if artifact.get("warning"): warnings.append(artifact.get("warning")) - _set_job_progress(job, "Source blob migration step completed", 5, total_steps, current_step="source_blobs") + migration_state.update({ + "status": MIGRATION_RESOURCE_STATUS_COMPLETED, + "completed_at": _now_iso(), + }) + _update_migration_state_totals(migration_state) + migration_state = _persist_migration_state( + job, + migration_state, + settings, + "Migration execution completed", + ) + _set_job_progress(job, "Migration execution completed", 10, total_steps, current_step="complete") artifact_summaries = summarize_backup_artifacts(artifacts) artifact_totals = _backup_artifact_totals(artifact_summaries) @@ -2651,6 +9247,12 @@ def execute_migration_job(job, settings): message="Migration execution completed", details={ "migration_plan": plan_summary, + "migration_id": provenance_context["migration_id"], + "migration_state": { + "resume_count": migration_state.get("resume_count"), + "totals": migration_state.get("totals"), + "capacity": migration_state.get("capacity"), + }, "artifact_count": len(artifacts), "artifact_totals": artifact_totals, "artifacts": artifact_summaries, @@ -2660,6 +9262,8 @@ def execute_migration_job(job, settings): return { "migration_plan": plan_summary, "dry_run": False, + "migration_id": provenance_context["migration_id"], + "migration_state": migration_state, "artifact_count": len(artifacts), "artifact_totals": artifact_totals, "artifacts": artifact_summaries, @@ -2672,6 +9276,7 @@ def process_data_management_job(job_id): job = _try_claim_data_management_job(job_id, settings=settings) if not job: return None + release_migration_lock = job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION try: if job.get("operation") == DATA_MANAGEMENT_OPERATION_BACKUP: @@ -2681,6 +9286,7 @@ def process_data_management_job(job_id): message = "Backup completed with warnings" if warnings else "Backup completed successfully" elif job.get("operation") == DATA_MANAGEMENT_OPERATION_MIGRATION: result = execute_migration_job(job, settings) + release_migration_lock = True warnings = list(job.get("warnings") or []) + result.get("warnings", []) status = DATA_MANAGEMENT_STATUS_COMPLETED_WITH_WARNINGS if warnings else DATA_MANAGEMENT_STATUS_COMPLETED message = "Migration completed with warnings" if warnings else "Migration completed successfully" @@ -2702,6 +9308,24 @@ def process_data_management_job(job_id): message = "Data Management job foundation completed with implementation warnings" now = _now_iso() + completed_progress = ( + copy.deepcopy(job.get("progress")) + if isinstance(job.get("progress"), dict) + else {} + ) + if not completed_progress: + completed_progress = { + "total_steps": 1, + "completed_steps": 1, + "current_step": "complete", + "percent_complete": 100, + } + else: + completed_progress.update({ + "completed_steps": completed_progress.get("total_steps", 1), + "current_step": "complete", + "percent_complete": 100, + }) job.update({ "status": status, "updated_at": now, @@ -2713,12 +9337,7 @@ def process_data_management_job(job_id): "lease_expires_at": None, "warnings": warnings, "result": result, - "progress": { - "total_steps": 1, - "completed_steps": 1, - "current_step": "complete", - "percent_complete": 100, - }, + "progress": completed_progress, }) saved_job = _save_data_management_job(job) _record_data_management_job_event( @@ -2730,6 +9349,54 @@ def process_data_management_job(job_id): details={"warnings": warnings, "result": result}, ) return saved_job + except DataManagementMigrationCanceledError as exc: + release_migration_lock = True + now = _now_iso() + migration_state = job.get("migration_state") + if isinstance(migration_state, dict): + migration_state.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "canceled_at": now, + "updated_at": now, + "last_error": None, + }) + job.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "updated_at": now, + "completed_at": now, + "last_heartbeat_at": now, + "last_message": "Migration canceled at a durable checkpoint.", + "last_error": None, + "lease_holder_id": None, + "lease_expires_at": None, + "migration_state": migration_state, + "progress": { + "total_steps": job.get("progress", {}).get("total_steps", 0), + "completed_steps": job.get("progress", {}).get("completed_steps", 0), + "current_step": "canceled", + "percent_complete": job.get("progress", {}).get("percent_complete", 0), + }, + }) + saved_job = _save_data_management_job(job) + _record_data_management_job_event( + saved_job.get("id"), + "migration-canceled", + saved_job, + status=DATA_MANAGEMENT_STATUS_CANCELED, + message="Migration canceled at a durable checkpoint.", + details={ + "cancel_requested_at": saved_job.get("cancel_requested_at"), + "reason_present": bool(saved_job.get("cancel_reason")), + }, + ) + return saved_job + except DataManagementMigrationLeaseLostError as exc: + log_event( + "[DataManagement] Migration worker stopped after losing its job lease.", + {"job_id": job_id, "operation": job.get("operation"), "error": str(exc)}, + level=logging.WARNING, + ) + return None except Exception as exc: now = _now_iso() job.update({ @@ -2758,12 +9425,14 @@ def process_data_management_job(job_id): details={"error": str(exc)}, ) return saved_job + finally: + if release_migration_lock: + _release_target_migration_coordinator(job) + _release_migration_destination_lock(job) def submit_data_management_job(app, job_id): - if not app: - return False - executor = app.extensions.get("executor") + executor = app.extensions.get("executor") if app else None if executor and hasattr(executor, "submit_stored"): executor.submit_stored( f"data_management_{job_id}", @@ -2774,7 +9443,14 @@ def submit_data_management_job(app, job_id): if executor and hasattr(executor, "submit"): executor.submit(process_data_management_job, job_id) return True - return False + worker_thread = Thread( + target=process_data_management_job, + kwargs={"job_id": job_id}, + daemon=True, + name=f"data-management-{job_id[:12]}", + ) + worker_thread.start() + return True def build_scheduled_occurrence_id(backup_type, run_at): @@ -2785,10 +9461,15 @@ def build_scheduled_occurrence_id(backup_type, run_at): def check_due_data_management_jobs_once(app=None): settings = get_data_management_settings() + current_time = _now_utc() + recovery_results = recover_data_management_migration_jobs( + app=app, + settings=settings, + current_time=current_time, + ) if not settings.get("enabled"): - return [] + return recovery_results - current_time = _now_utc() queued_jobs = [] for backup_type, next_key in ( (DATA_MANAGEMENT_BACKUP_FULL, "next_full_backup_run_at"), @@ -2823,4 +9504,4 @@ def check_due_data_management_jobs_once(app=None): if not submitted: process_data_management_job(job.get("id")) - return queued_jobs \ No newline at end of file + return queued_jobs + recovery_results \ No newline at end of file diff --git a/application/single_app/functions_data_management_migration_state.py b/application/single_app/functions_data_management_migration_state.py new file mode 100644 index 000000000..fd26a1ceb --- /dev/null +++ b/application/single_app/functions_data_management_migration_state.py @@ -0,0 +1,203 @@ +# functions_data_management_migration_state.py +"""Durable, secret-free checkpoint state for Data Management migrations.""" + +import copy +import hashlib +import json +from datetime import datetime, timezone + + +MIGRATION_STATE_SCHEMA_VERSION = 1 +MIGRATION_RESOURCE_STATUS_PENDING = "pending" +MIGRATION_RESOURCE_STATUS_IN_PROGRESS = "in_progress" +MIGRATION_RESOURCE_STATUS_COMPLETED = "completed" +MIGRATION_RESOURCE_STATUS_FAILED = "failed" + + +def _utc_now_iso(current_time=None): + """Return a normalized UTC timestamp without accepting local wall-clock time.""" + timestamp = current_time if isinstance(current_time, datetime) else datetime.now(timezone.utc) + if timestamp.tzinfo is None: + timestamp = timestamp.replace(tzinfo=timezone.utc) + return timestamp.astimezone(timezone.utc).isoformat() + + +def _parse_utc_datetime(value): + """Parse a durable checkpoint timestamp when it is valid.""" + try: + parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) + except (TypeError, ValueError): + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +def build_migration_configuration_fingerprint(configuration): + """Hash a caller-supplied, secret-free configuration snapshot.""" + encoded = json.dumps( + configuration if isinstance(configuration, dict) else {}, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + default=str, + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def initialize_migration_state(existing_state, migration_id, configuration, current_time=None): + """Create or resume an in-job migration state with fingerprint protection.""" + fingerprint = build_migration_configuration_fingerprint(configuration) + timestamp = _utc_now_iso(current_time) + state = copy.deepcopy(existing_state) if isinstance(existing_state, dict) else {} + + if state: + if state.get("schema_version") != MIGRATION_STATE_SCHEMA_VERSION: + raise ValueError("Migration state uses an unsupported schema version.") + if state.get("migration_id") != migration_id: + raise ValueError("Migration state belongs to a different migration ID.") + if state.get("configuration_fingerprint") != fingerprint: + raise ValueError("Migration settings changed after the job was checkpointed. Queue a new migration job.") + state["status"] = MIGRATION_RESOURCE_STATUS_IN_PROGRESS + state["updated_at"] = timestamp + state["resume_count"] = int(state.get("resume_count") or 0) + 1 + state.setdefault("resources", {}) + state.setdefault("preflight", {}) + state.setdefault("capacity", {}) + state.setdefault("totals", {}) + return state + + return { + "schema_version": MIGRATION_STATE_SCHEMA_VERSION, + "migration_id": migration_id, + "configuration": copy.deepcopy(configuration if isinstance(configuration, dict) else {}), + "configuration_fingerprint": fingerprint, + "status": MIGRATION_RESOURCE_STATUS_IN_PROGRESS, + "created_at": timestamp, + "updated_at": timestamp, + "completed_at": None, + "source_cutoff_at": timestamp, + "resume_count": 0, + "resources": {}, + "preflight": {}, + "capacity": {}, + "totals": {}, + } + + +def get_migration_resource(state, resource_name): + """Return a resource checkpoint without allocating a missing record.""" + resources = state.get("resources") if isinstance(state, dict) else None + return resources.get(resource_name) if isinstance(resources, dict) else None + + +def is_migration_resource_completed(state, resource_name): + """Return whether a resource completed successfully in this migration.""" + resource = get_migration_resource(state, resource_name) + return isinstance(resource, dict) and resource.get("status") == MIGRATION_RESOURCE_STATUS_COMPLETED + + +def start_migration_resource(state, resource_name, current_time=None): + """Start or resume one resource while preserving its last durable counters.""" + if not isinstance(state, dict): + raise ValueError("Migration state must be a dictionary.") + timestamp = _utc_now_iso(current_time) + resources = state.setdefault("resources", {}) + existing = resources.get(resource_name) if isinstance(resources, dict) else None + if isinstance(existing, dict) and existing.get("status") == MIGRATION_RESOURCE_STATUS_COMPLETED: + return existing + + progress = copy.deepcopy(existing.get("progress") if isinstance(existing, dict) else {}) + attempts = int(existing.get("attempts") or 0) + 1 if isinstance(existing, dict) else 1 + resource = { + "status": MIGRATION_RESOURCE_STATUS_IN_PROGRESS, + "attempts": attempts, + "started_at": existing.get("started_at") if isinstance(existing, dict) and existing.get("started_at") else timestamp, + "attempt_started_at": timestamp, + "updated_at": timestamp, + "completed_at": None, + "last_error": None, + "progress": progress, + "result": {}, + } + resources[resource_name] = resource + state["updated_at"] = timestamp + return resource + + +def update_migration_resource(state, resource_name, progress, current_time=None): + """Persist counters for an active resource checkpoint.""" + resource = get_migration_resource(state, resource_name) + if not isinstance(resource, dict): + raise ValueError(f"Migration resource '{resource_name}' was not started.") + if resource.get("status") == MIGRATION_RESOURCE_STATUS_COMPLETED: + raise ValueError(f"Migration resource '{resource_name}' is already completed.") + timestamp = _utc_now_iso(current_time) + resource["progress"] = copy.deepcopy(progress if isinstance(progress, dict) else {}) + resource["updated_at"] = timestamp + state["updated_at"] = timestamp + return resource + + +def complete_migration_resource(state, resource_name, result=None, current_time=None): + """Mark a resource completed only after its final target writes are known.""" + resource = get_migration_resource(state, resource_name) + if not isinstance(resource, dict): + raise ValueError(f"Migration resource '{resource_name}' was not started.") + timestamp = _utc_now_iso(current_time) + resource.update({ + "status": MIGRATION_RESOURCE_STATUS_COMPLETED, + "updated_at": timestamp, + "completed_at": timestamp, + "last_error": None, + "result": copy.deepcopy(result if isinstance(result, dict) else {}), + }) + state["updated_at"] = timestamp + return resource + + +def fail_migration_resource(state, resource_name, error_message, current_time=None): + """Persist an interrupted resource so the same job can resume it later.""" + resource = get_migration_resource(state, resource_name) + if not isinstance(resource, dict): + resource = start_migration_resource(state, resource_name, current_time=current_time) + timestamp = _utc_now_iso(current_time) + resource.update({ + "status": MIGRATION_RESOURCE_STATUS_FAILED, + "updated_at": timestamp, + "last_error": str(error_message or "Migration resource failed."), + }) + state["updated_at"] = timestamp + return resource + + +def build_transfer_metrics( + started_at, + copied_count=0, + skipped_count=0, + failed_count=0, + byte_count=0, + request_units=0.0, + current_time=None, +): + """Calculate bounded, JSON-safe counters and transfer rates for job details.""" + started = _parse_utc_datetime(started_at) + now = _parse_utc_datetime(current_time) if current_time is not None else datetime.now(timezone.utc) + elapsed_seconds = max(0.001, (now - started).total_seconds()) if started else 0.001 + copied = max(0, int(copied_count or 0)) + skipped = max(0, int(skipped_count or 0)) + failed = max(0, int(failed_count or 0)) + bytes_transferred = max(0, int(byte_count or 0)) + consumed_request_units = max(0.0, float(request_units or 0.0)) + return { + "copied_count": copied, + "skipped_count": skipped, + "failed_count": failed, + "processed_count": copied + skipped + failed, + "bytes": bytes_transferred, + "request_units": round(consumed_request_units, 3), + "elapsed_seconds": round(elapsed_seconds, 3), + "items_per_second": round((copied + skipped + failed) / elapsed_seconds, 3), + "bytes_per_second": round(bytes_transferred / elapsed_seconds, 3), + "request_units_per_second": round(consumed_request_units / elapsed_seconds, 3), + } \ No newline at end of file diff --git a/application/single_app/functions_data_management_search_write_fence.py b/application/single_app/functions_data_management_search_write_fence.py new file mode 100644 index 000000000..cb246230e --- /dev/null +++ b/application/single_app/functions_data_management_search_write_fence.py @@ -0,0 +1,644 @@ +# functions_data_management_search_write_fence.py +"""Cross-app Azure AI Search write fencing for Data Management migrations.""" + +import copy +import time +import uuid +from contextlib import contextmanager +from datetime import datetime, timedelta, timezone + +from azure.core import MatchConditions +from azure.core.exceptions import ResourceNotFoundError + + +DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID = "data_management_search_write_gate_global" +DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE = "data_management_search_write_gate" +DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_OPEN = "open" +DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_CLOSING = "closing" +DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_FROZEN = "frozen" +DATA_MANAGEMENT_SEARCH_WRITE_REQUEST_TIMEOUT_SECONDS = 30 +DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS = 150 +DATA_MANAGEMENT_SEARCH_WRITE_GATE_POLL_SECONDS = 0.1 +DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID = ( + "data_management_target_migration_coordinator_global" +) +DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE = ( + "data_management_target_migration_coordinator" +) + + +class DataManagementSearchWriteGateError(RuntimeError): + """Raised when the durable Search write gate cannot be coordinated safely.""" + + +class DataManagementSearchWritesFrozenError(DataManagementSearchWriteGateError): + """Raised when a target migration has frozen SimpleChat Search writes.""" + + +class DataManagementSearchWriteFenceLostError(DataManagementSearchWriteGateError): + """Raised when a migration no longer owns its target Search write fence.""" + + +class DataManagementTargetMigrationCoordinatorError(DataManagementSearchWriteGateError): + """Raised when a target migration coordinator cannot be acquired safely.""" + + +class DataManagementTargetMigrationCoordinatorLostError( + DataManagementTargetMigrationCoordinatorError +): + """Raised when a migration no longer owns its target coordinator lease.""" + + +def _now_utc(): + return datetime.now(timezone.utc) + + +def _parse_utc_datetime(value): + try: + parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) + except (TypeError, ValueError): + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +def _safe_int(value, default=0, minimum=0): + try: + normalized = int(value) + except (TypeError, ValueError): + return default + return max(minimum, normalized) + + +def _is_not_found_error(exc): + return ( + getattr(exc, "status_code", None) == 404 or + isinstance(exc, (KeyError, ResourceNotFoundError)) + ) + + +def _is_conflict_error(exc): + return getattr(exc, "status_code", None) in {409, 412} + + +def _read_gate(container): + try: + gate = container.read_item( + item=DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID, + partition_key=DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID, + ) + except Exception as exc: + if _is_not_found_error(exc): + return None + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate could not be read." + ) from exc + if not isinstance(gate, dict): + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate has an invalid record." + ) + return gate + + +def _new_open_gate(now=None): + timestamp = now or _now_utc() + return { + "id": DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID, + "type": DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE, + "state": DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_OPEN, + "writer_leases": [], + "active_writer_count": 0, + "updated_at": timestamp.isoformat(), + } + + +def _create_or_read_gate(container): + gate = _read_gate(container) + if gate is not None: + return gate + try: + return container.create_item(body=_new_open_gate()) + except Exception as exc: + if _is_conflict_error(exc): + gate = _read_gate(container) + if gate is not None: + return gate + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate could not be created." + ) from exc + + +def _replace_gate(container, existing_gate, replacement_gate): + try: + return container.replace_item( + item=DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID, + body=replacement_gate, + etag=existing_gate.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception as exc: + if _is_conflict_error(exc): + return None + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate could not be updated." + ) from exc + + +def _active_writer_leases(gate, now=None): + timestamp = now or _now_utc() + active_leases = [] + for lease in (gate or {}).get("writer_leases") or []: + if not isinstance(lease, dict) or not lease.get("token"): + continue + expires_at = _parse_utc_datetime(lease.get("expires_at")) + if expires_at and expires_at > timestamp: + active_leases.append({ + "token": str(lease.get("token")), + "expires_at": expires_at.isoformat(), + }) + return active_leases + + +def _is_active_migration_fence(gate, now=None): + timestamp = now or _now_utc() + expires_at = _parse_utc_datetime((gate or {}).get("expires_at")) + return bool( + (gate or {}).get("type") == DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE and + (gate or {}).get("state") in { + DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_CLOSING, + DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_FROZEN, + } and + expires_at and + expires_at > timestamp + ) + + +def _open_expired_gate(container, gate, now=None): + timestamp = now or _now_utc() + if _is_active_migration_fence(gate, timestamp): + return gate + replacement = _new_open_gate(timestamp) + replacement["writer_leases"] = _active_writer_leases(gate, timestamp) + replacement["active_writer_count"] = len(replacement["writer_leases"]) + return _replace_gate(container, gate, replacement) + + +def acquire_data_management_search_write_slot(container): + """Reserve one bounded target Search write before issuing the data-plane request.""" + for _attempt in range(12): + now = _now_utc() + gate = _create_or_read_gate(container) + if gate.get("type") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE: + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate has an unexpected record type." + ) + if _is_active_migration_fence(gate, now): + raise DataManagementSearchWritesFrozenError( + "AI Search writes are temporarily frozen while a Data Management migration is running." + ) + if gate.get("state") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_OPEN: + if _open_expired_gate(container, gate, now) is None: + continue + continue + lease_token = uuid.uuid4().hex + replacement = copy.deepcopy(gate) + active_leases = _active_writer_leases(gate, now) + active_leases.append({ + "token": lease_token, + "expires_at": ( + now + timedelta(seconds=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS) + ).isoformat(), + }) + replacement.update({ + "writer_leases": active_leases, + "active_writer_count": len(active_leases), + "updated_at": now.isoformat(), + }) + if _replace_gate(container, gate, replacement) is not None: + return lease_token + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate changed too often to reserve a write slot." + ) + + +def release_data_management_search_write_slot(container, lease_token): + """Release one target Search write slot without masking the caller's result.""" + for _attempt in range(12): + try: + gate = _read_gate(container) + except DataManagementSearchWriteGateError: + return False + if gate is None or gate.get("type") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE: + return False + now = _now_utc() + replacement = copy.deepcopy(gate) + active_leases = [ + lease + for lease in _active_writer_leases(gate, now) + if lease.get("token") != lease_token + ] + replacement.update({ + "writer_leases": active_leases, + "active_writer_count": len(active_leases), + "updated_at": now.isoformat(), + }) + if _replace_gate(container, gate, replacement) is not None: + return True + return False + + +@contextmanager +def hold_data_management_search_write_slot(container): + """Hold a target Search write slot until a response is known or its ambiguity lease expires.""" + lease_token = acquire_data_management_search_write_slot(container) + response_confirmed = False + try: + yield + response_confirmed = True + finally: + if response_confirmed: + release_data_management_search_write_slot(container, lease_token) + + +def acquire_data_management_search_write_fence( + container, + migration_id, + lease_seconds, + heartbeat_callback=None, +): + """Close target SimpleChat Search writes, drain active slots, and freeze the gate.""" + normalized_migration_id = str(migration_id or "").strip() + if not normalized_migration_id: + raise DataManagementSearchWriteGateError( + "A migration ID is required to acquire the target Search write fence." + ) + normalized_lease_seconds = _safe_int( + lease_seconds, + default=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + minimum=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + ) + fence_token = uuid.uuid4().hex + last_heartbeat = 0.0 + last_fence_renewal = 0.0 + while True: + now = _now_utc() + gate = _create_or_read_gate(container) + if gate.get("type") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE: + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate has an unexpected record type." + ) + state = gate.get("state") + same_fence = ( + gate.get("migration_id") == normalized_migration_id and + gate.get("fence_token") == fence_token + ) + if _is_active_migration_fence(gate, now) and not same_fence: + raise DataManagementSearchWritesFrozenError( + "Another migration is already freezing target AI Search writes." + ) + if state != DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_CLOSING or not same_fence: + replacement = copy.deepcopy(gate) + active_leases = _active_writer_leases(gate, now) + replacement.update({ + "state": DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_CLOSING, + "migration_id": normalized_migration_id, + "fence_token": fence_token, + "expires_at": ( + now + timedelta(seconds=normalized_lease_seconds) + ).isoformat(), + "writer_leases": active_leases, + "active_writer_count": len(active_leases), + "updated_at": now.isoformat(), + }) + if _replace_gate(container, gate, replacement) is None: + continue + gate = replacement + if ( + state == DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_CLOSING and + same_fence and + time.monotonic() - last_fence_renewal >= 2.0 + ): + replacement = copy.deepcopy(gate) + replacement.update({ + "expires_at": ( + now + timedelta(seconds=normalized_lease_seconds) + ).isoformat(), + "updated_at": now.isoformat(), + }) + persisted = _replace_gate(container, gate, replacement) + if persisted is None: + continue + gate = persisted + last_fence_renewal = time.monotonic() + active_leases = _active_writer_leases(gate, now) + if not active_leases: + replacement = copy.deepcopy(gate) + replacement.update({ + "state": DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_FROZEN, + "writer_leases": [], + "active_writer_count": 0, + "updated_at": now.isoformat(), + }) + if _replace_gate(container, gate, replacement) is not None: + return { + "id": DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID, + "migration_id": normalized_migration_id, + "fence_token": fence_token, + "lease_seconds": normalized_lease_seconds, + "expires_at": replacement.get("expires_at"), + } + continue + if callable(heartbeat_callback) and time.monotonic() - last_heartbeat >= 2.0: + heartbeat_callback() + last_heartbeat = time.monotonic() + time.sleep(DATA_MANAGEMENT_SEARCH_WRITE_GATE_POLL_SECONDS) + + +def renew_data_management_search_write_fence(container, fence, lease_seconds): + """Renew a migration-owned frozen target Search gate.""" + if not isinstance(fence, dict) or not fence.get("fence_token"): + return None + now = _now_utc() + gate = _read_gate(container) + if ( + not isinstance(gate, dict) or + gate.get("type") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE or + gate.get("migration_id") != fence.get("migration_id") or + gate.get("fence_token") != fence.get("fence_token") or + not _is_active_migration_fence(gate, now) + ): + raise DataManagementSearchWriteFenceLostError( + "The target AI Search write fence was lost or expired." + ) + normalized_lease_seconds = _safe_int( + lease_seconds, + default=fence.get("lease_seconds"), + minimum=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + ) + replacement = copy.deepcopy(gate) + replacement.update({ + "expires_at": (now + timedelta(seconds=normalized_lease_seconds)).isoformat(), + "updated_at": now.isoformat(), + }) + renewed = _replace_gate(container, gate, replacement) + if renewed is None: + raise DataManagementSearchWriteFenceLostError( + "The target AI Search write fence changed during renewal." + ) + fence["expires_at"] = renewed.get("expires_at", replacement.get("expires_at")) + fence["lease_seconds"] = normalized_lease_seconds + return fence + + +def release_data_management_search_write_fence(container, fence): + """Open only the exact target Search gate owned by the completed migration.""" + if not isinstance(fence, dict) or not fence.get("fence_token"): + return False + for _attempt in range(12): + gate = _read_gate(container) + if ( + not isinstance(gate, dict) or + gate.get("type") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE or + gate.get("migration_id") != fence.get("migration_id") or + gate.get("fence_token") != fence.get("fence_token") + ): + return False + now = _now_utc() + replacement = _new_open_gate(now) + replacement["writer_leases"] = _active_writer_leases(gate, now) + replacement["active_writer_count"] = len(replacement["writer_leases"]) + if _replace_gate(container, gate, replacement) is not None: + return True + return False + + +def _read_target_migration_coordinator(container): + try: + coordinator = container.read_item( + item=DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + partition_key=DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + ) + except Exception as exc: + if _is_not_found_error(exc): + return None + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator could not be read." + ) from exc + if not isinstance(coordinator, dict): + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator has an invalid record." + ) + return coordinator + + +def _replace_target_migration_coordinator(container, existing, replacement): + try: + return container.replace_item( + item=DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + body=replacement, + etag=existing.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception as exc: + if _is_conflict_error(exc): + return None + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator could not be updated." + ) from exc + + +def _target_migration_coordinator_is_active(coordinator, now=None): + timestamp = now or _now_utc() + expires_at = _parse_utc_datetime((coordinator or {}).get("expires_at")) + return bool( + (coordinator or {}).get("type") == + DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE and + expires_at and + expires_at > timestamp + ) + + +def _new_target_migration_coordinator(migration_id, lock_token, lease_seconds, now=None): + timestamp = now or _now_utc() + return { + "id": DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + "type": DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE, + "migration_id": str(migration_id), + "lock_token": str(lock_token), + "acquired_at": timestamp.isoformat(), + "updated_at": timestamp.isoformat(), + "lease_seconds": int(lease_seconds), + "expires_at": (timestamp + timedelta(seconds=int(lease_seconds))).isoformat(), + } + + +def acquire_data_management_target_migration_coordinator( + container, + migration_id, + lease_seconds, + existing_lock=None, +): + """Acquire one destination-wide coordinator before any migration mutation begins.""" + normalized_migration_id = str(migration_id or "").strip() + if not normalized_migration_id: + raise DataManagementTargetMigrationCoordinatorError( + "A migration ID is required to acquire the target migration coordinator." + ) + normalized_lease_seconds = _safe_int( + lease_seconds, + default=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + minimum=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + ) + existing_lock = existing_lock if isinstance(existing_lock, dict) else {} + existing_token = str(existing_lock.get("lock_token") or "").strip() + for _attempt in range(12): + now = _now_utc() + coordinator = _read_target_migration_coordinator(container) + if coordinator is None: + lock_token = existing_token or uuid.uuid4().hex + replacement = _new_target_migration_coordinator( + normalized_migration_id, + lock_token, + normalized_lease_seconds, + now, + ) + try: + persisted = container.create_item(body=replacement) + except Exception as exc: + if _is_conflict_error(exc): + continue + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator could not be created." + ) from exc + return { + "id": DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + "migration_id": normalized_migration_id, + "lock_token": lock_token, + "lease_seconds": normalized_lease_seconds, + "expires_at": persisted.get("expires_at", replacement["expires_at"]), + } + + if coordinator.get("type") != DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE: + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator has an unexpected record type." + ) + if _target_migration_coordinator_is_active(coordinator, now): + if ( + coordinator.get("migration_id") == normalized_migration_id and + existing_token and + coordinator.get("lock_token") == existing_token + ): + reused_lock = { + "id": DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + "migration_id": normalized_migration_id, + "lock_token": existing_token, + "lease_seconds": normalized_lease_seconds, + "expires_at": coordinator.get("expires_at"), + } + return renew_data_management_target_migration_coordinator( + container, + reused_lock, + normalized_lease_seconds, + ) + raise DataManagementTargetMigrationCoordinatorError( + "Another SimpleChat source is actively migrating to this destination." + ) + + replacement = _new_target_migration_coordinator( + normalized_migration_id, + existing_token or uuid.uuid4().hex, + normalized_lease_seconds, + now, + ) + persisted = _replace_target_migration_coordinator( + container, + coordinator, + replacement, + ) + if persisted is None: + continue + return { + "id": DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + "migration_id": normalized_migration_id, + "lock_token": replacement["lock_token"], + "lease_seconds": normalized_lease_seconds, + "expires_at": persisted.get("expires_at", replacement["expires_at"]), + } + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator changed too often to acquire safely." + ) + + +def renew_data_management_target_migration_coordinator(container, lock, lease_seconds): + """Renew an exact target-wide coordinator lease held by this migration.""" + if not isinstance(lock, dict) or not lock.get("lock_token"): + raise DataManagementTargetMigrationCoordinatorLostError( + "The target migration coordinator handle is missing." + ) + coordinator = _read_target_migration_coordinator(container) + now = _now_utc() + if ( + not isinstance(coordinator, dict) or + coordinator.get("type") != DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE or + coordinator.get("migration_id") != lock.get("migration_id") or + coordinator.get("lock_token") != lock.get("lock_token") or + not _target_migration_coordinator_is_active(coordinator, now) + ): + raise DataManagementTargetMigrationCoordinatorLostError( + "The target migration coordinator was lost, superseded, or expired." + ) + normalized_lease_seconds = _safe_int( + lease_seconds, + default=lock.get("lease_seconds"), + minimum=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + ) + replacement = copy.deepcopy(coordinator) + replacement.update({ + "lease_seconds": normalized_lease_seconds, + "updated_at": now.isoformat(), + "expires_at": ( + now + timedelta(seconds=normalized_lease_seconds) + ).isoformat(), + }) + persisted = _replace_target_migration_coordinator( + container, + coordinator, + replacement, + ) + if persisted is None: + raise DataManagementTargetMigrationCoordinatorLostError( + "The target migration coordinator changed during renewal." + ) + lock["lease_seconds"] = normalized_lease_seconds + lock["expires_at"] = persisted.get("expires_at", replacement["expires_at"]) + return lock + + +def release_data_management_target_migration_coordinator(container, lock): + """Release only the exact target coordinator held by a finished migration.""" + if not isinstance(lock, dict) or not lock.get("lock_token"): + return False + try: + coordinator = _read_target_migration_coordinator(container) + except DataManagementTargetMigrationCoordinatorError: + return False + if ( + not isinstance(coordinator, dict) or + coordinator.get("type") != DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE or + coordinator.get("migration_id") != lock.get("migration_id") or + coordinator.get("lock_token") != lock.get("lock_token") + ): + return False + try: + container.delete_item( + item=DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + partition_key=DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + etag=coordinator.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception as exc: + if _is_conflict_error(exc): + return False + return False + return True \ No newline at end of file diff --git a/application/single_app/functions_document_access_index.py b/application/single_app/functions_document_access_index.py index a34250e16..6cc690e7b 100644 --- a/application/single_app/functions_document_access_index.py +++ b/application/single_app/functions_document_access_index.py @@ -68,6 +68,7 @@ DOCUMENT_ACCESS_CACHE_VERSION_MAX_TTL_SECONDS = 86400 DOCUMENT_ACCESS_CACHE_VERSION_TTL_MULTIPLIER = 4 DOCUMENT_ACCESS_CACHE_VERSION_HYGIENE_BATCH_SIZE = 100 +DOCUMENT_ACCESS_BOUNDED_CATALOG_MAX_SCOPES = 25 DOCUMENT_ACCESS_CACHE_VERSION_HYGIENE_MAX_SCAN_ITERATIONS = 5 DOCUMENT_ACCESS_CACHE_KEY_PREFIX = 'DAI_LIST_CACHE' DOCUMENT_ACCESS_CACHE_VERSION_KEY_PREFIX = 'DAI_LIST_CACHE_VERSION' @@ -2127,7 +2128,7 @@ def _query_projection_rows_for_scope_with_diagnostics(scope_key, source_scope): 'WHERE c.type = @type ' 'AND c.source_scope = @source_scope ' 'AND c.scope_key = @scope_key ' - 'AND (c.access_granted = true OR c.approval_status = @approval_not_approved) ' + 'AND c.access_granted = true ' 'AND c.is_current_version = true ' 'AND c.projection_version = @projection_version' ) @@ -2139,7 +2140,6 @@ def _query_projection_rows_for_scope_with_diagnostics(scope_key, source_scope): {'name': '@type', 'value': DOCUMENT_ACCESS_INDEX_TYPE}, {'name': '@source_scope', 'value': source_scope}, {'name': '@scope_key', 'value': scope_key}, - {'name': '@approval_not_approved', 'value': DOCUMENT_ACCESS_APPROVAL_NOT_APPROVED}, {'name': '@projection_version', 'value': DOCUMENT_ACCESS_INDEX_SCHEMA_VERSION}, ], partition_key=scope_key, @@ -2162,7 +2162,7 @@ def _query_candidate_projection_rows_for_scope(scope_key, source_scope): 'WHERE c.type = @type ' 'AND c.source_scope = @source_scope ' 'AND c.scope_key = @scope_key ' - 'AND (c.access_granted = true OR c.approval_status = @approval_not_approved) ' + 'AND c.access_granted = true ' 'AND c.is_current_version = true ' 'AND c.projection_version = @projection_version' ) @@ -2174,13 +2174,130 @@ def _query_candidate_projection_rows_for_scope(scope_key, source_scope): {'name': '@type', 'value': DOCUMENT_ACCESS_INDEX_TYPE}, {'name': '@source_scope', 'value': source_scope}, {'name': '@scope_key', 'value': scope_key}, - {'name': '@approval_not_approved', 'value': DOCUMENT_ACCESS_APPROVAL_NOT_APPROVED}, {'name': '@projection_version', 'value': DOCUMENT_ACCESS_INDEX_SCHEMA_VERSION}, ], partition_key=scope_key, ) +def _query_bounded_projection_rows_for_scope(scope_key, source_scope, max_rows): + normalized_max_rows = max(1, min(_safe_int(max_rows), 1001)) + query = ( + f'SELECT TOP {normalized_max_rows} c.document_id, c.source_document_id, c.version, ' + 'c.revision_family_id, c.source_ts, c.file_name, ' + 'c.owner_user_id, c.owner_group_id, c.owner_public_workspace_id ' + 'FROM c ' + 'WHERE c.type = @type ' + 'AND c.source_scope = @source_scope ' + 'AND c.scope_key = @scope_key ' + 'AND c.access_granted = true ' + 'AND c.is_current_version = true ' + 'AND c.projection_version = @projection_version ' + 'ORDER BY c.source_ts DESC' + ) + return list(cosmos_document_access_index_container.query_items( + query=query, + parameters=[ + {'name': '@type', 'value': DOCUMENT_ACCESS_INDEX_TYPE}, + {'name': '@source_scope', 'value': source_scope}, + {'name': '@scope_key', 'value': scope_key}, + {'name': '@projection_version', 'value': DOCUMENT_ACCESS_INDEX_SCHEMA_VERSION}, + ], + partition_key=scope_key, + )) + + +def enumerate_bounded_document_access_index_ids( + source_scope, + max_documents, + user_id=None, + group_ids=None, + public_workspace_id=None, + public_workspace_ids=None, + settings=None, +): + """Return current candidate IDs only when the ready access-index catalog fits the bound.""" + source_scope = str(source_scope or '').strip().lower() + if source_scope not in DOCUMENT_ACCESS_SOURCE_SCOPES: + return {'success': False, 'status': 'invalid_source_scope', 'document_ids': []} + max_documents = _safe_int(max_documents) + if max_documents <= 0: + return {'success': False, 'status': 'invalid_document_limit', 'document_ids': []} + + readiness = _get_document_access_index_readiness(source_scope, settings=settings) + if not readiness.get('ready'): + return { + 'success': False, + 'status': readiness.get('status'), + 'document_ids': [], + 'readiness': readiness, + } + scope_keys = [ + scope_key + for scope_key in _build_shadow_scope( + source_scope, + user_id=user_id, + group_ids=group_ids, + public_workspace_id=public_workspace_id, + public_workspace_ids=public_workspace_ids, + ) + if scope_key + ] + if not scope_keys: + return {'success': False, 'status': 'missing_scope_keys', 'document_ids': []} + if len(scope_keys) > DOCUMENT_ACCESS_BOUNDED_CATALOG_MAX_SCOPES: + return { + 'success': False, + 'status': 'scope_limit_exceeded', + 'document_ids': [], + 'scope_count': len(scope_keys), + } + + rows_by_identity = {} + query_limit = max_documents + 1 + for scope_key in scope_keys: + for row in _query_bounded_projection_rows_for_scope( + scope_key, + source_scope, + query_limit, + ): + identity = _document_family_identity(row, source_scope) + if not identity: + continue + rows_by_identity[identity] = _prefer_projection_row( + rows_by_identity.get(identity), + row, + ) + if len(rows_by_identity) > max_documents: + return { + 'success': False, + 'status': 'document_limit_exceeded', + 'document_ids': [], + 'document_count_lower_bound': max_documents + 1, + } + + ordered_rows = sorted( + rows_by_identity.values(), + key=lambda row: ( + _safe_int(row.get('source_ts')), + str(row.get('source_document_id') or row.get('document_id') or ''), + ), + reverse=True, + ) + document_ids = [ + str(row.get('source_document_id') or row.get('document_id') or '').strip() + for row in ordered_rows + if str(row.get('source_document_id') or row.get('document_id') or '').strip() + ] + return { + 'success': True, + 'status': 'bounded_catalog_ready', + 'document_ids': document_ids, + 'document_count': len(document_ids), + 'scope_count': len(scope_keys), + } + + def _query_tag_projection_rows_for_scope(scope_key, source_scope): query = ( 'SELECT c.document_id, c.source_document_id, c.version, c.scope_key, ' diff --git a/application/single_app/functions_document_actions.py b/application/single_app/functions_document_actions.py index d7a6b3823..d2432b95a 100644 --- a/application/single_app/functions_document_actions.py +++ b/application/single_app/functions_document_actions.py @@ -17,6 +17,7 @@ DOCUMENT_ACTION_TYPE_COMPARISON = 'comparison' DOCUMENT_ACTION_ANALYSIS_MODE_COMBINED = 'combined' DOCUMENT_ACTION_ANALYSIS_MODE_PER_DOCUMENT = 'per_document' +DOCUMENT_ACTION_TARGET_MODE_ALL = 'all' DOCUMENT_ACTION_TARGET_MODE_SELECTED = 'selected' DOCUMENT_ACTION_TARGET_MODE_RECENT = 'recent' DEFAULT_RECENT_DOCUMENT_WINDOW_MINUTES = 10 @@ -25,6 +26,7 @@ DOCUMENT_ACTION_ANALYSIS_MODE_PER_DOCUMENT, } VALID_DOCUMENT_ACTION_TARGET_MODES = { + DOCUMENT_ACTION_TARGET_MODE_ALL, DOCUMENT_ACTION_TARGET_MODE_SELECTED, DOCUMENT_ACTION_TARGET_MODE_RECENT, } @@ -270,6 +272,8 @@ def normalize_document_action_config( if action_type == DOCUMENT_ACTION_TYPE_SEARCH: target_mode = normalize_document_action_target_mode(source_action.get('target_mode')) + if target_mode == DOCUMENT_ACTION_TARGET_MODE_ALL: + raise ValueError('All Documents is exhaustive only for Analyze.') document_ids = normalize_search_id_list(source_action.get('document_ids')) if resolved_max_documents is not None and len(document_ids) > resolved_max_documents: raise ValueError(f'Document search supports up to {resolved_max_documents} documents at a time.') @@ -297,6 +301,19 @@ def normalize_document_action_config( if action_type == DOCUMENT_ACTION_TYPE_ANALYZE: target_mode = normalize_document_action_target_mode(source_action.get('target_mode')) + if target_mode == DOCUMENT_ACTION_TARGET_MODE_ALL: + normalized_action.update({ + 'doc_scope': source_action.get('doc_scope', 'all'), + 'active_group_ids': normalize_search_id_list(source_action.get('active_group_ids')), + 'active_public_workspace_id': normalize_search_id_list(source_action.get('active_public_workspace_id')), + 'window_unit': source_action.get('window_unit') or 'pages', + 'window_size': source_action.get('window_size'), + 'window_percent': source_action.get('window_percent'), + 'max_retries_per_window': source_action.get('max_retries_per_window', 1), + 'analysis_mode': normalize_document_action_analysis_mode(source_action.get('analysis_mode')), + 'target_mode': target_mode, + }) + return normalized_action recent_targets_resolved = bool(source_action.get('recent_targets_resolved')) if target_mode == DOCUMENT_ACTION_TARGET_MODE_RECENT and not normalize_search_id_list(source_action.get('document_ids')): source_action = dict(source_action) @@ -323,6 +340,8 @@ def normalize_document_action_config( return normalized_action target_mode = normalize_document_action_target_mode(source_action.get('target_mode')) + if target_mode == DOCUMENT_ACTION_TARGET_MODE_ALL: + raise ValueError('All Documents is supported only for Analyze.') recent_targets_resolved = bool(source_action.get('recent_targets_resolved')) if target_mode == DOCUMENT_ACTION_TARGET_MODE_RECENT and not recent_targets_resolved: normalized_action.update({ diff --git a/application/single_app/functions_document_analysis.py b/application/single_app/functions_document_analysis.py index d82d8de04..df7844c2c 100644 --- a/application/single_app/functions_document_analysis.py +++ b/application/single_app/functions_document_analysis.py @@ -8,6 +8,10 @@ from functions_appinsights import log_event from functions_debug import debug_print +from functions_mixed_source_orchestration import ( + MixedSourceCancellationError, + raise_if_mixed_source_cancelled, +) from functions_search import normalize_search_id_list, normalize_search_scope @@ -679,6 +683,8 @@ def _reduce_document_analysis_items( failed_range_labels, reduction_batch_size, max_reduction_rounds, + cancel_requested=None, + request_correlation_id=None, ): current_items = list(items or []) reduction_round = 1 @@ -687,6 +693,11 @@ def _reduce_document_analysis_items( next_items = [] batches = _build_reduction_batches(current_items, reduction_batch_size) for batch_index, batch_items in enumerate(batches, start=1): + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative_reduction', + request_correlation_id=request_correlation_id, + ) reduction_prompt = _build_document_reduction_prompt( analysis_prompt, document_name, @@ -705,6 +716,11 @@ def _reduce_document_analysis_items( 'item_count': len(batch_items), }, ) or '').strip() + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative_reduction', + request_correlation_id=request_correlation_id, + ) if not reduced_text: raise RuntimeError( f'Document analysis document reduction returned an empty response for {document_name} ' @@ -784,12 +800,19 @@ def run_document_analysis( activity_callback=None, max_documents=None, include_coverage_summary=True, + cancel_requested=None, + request_correlation_id=None, ): normalized_analysis_prompt = str(analysis_prompt or '').strip() if not normalized_analysis_prompt: raise ValueError('An analysis prompt is required for document analysis.') if not callable(invoke_prompt): raise ValueError('A callable invoke_prompt handler is required for document analysis.') + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative_manifest', + request_correlation_id=request_correlation_id, + ) build_document_chunk_windows, get_document_chunks_payload = _get_search_service_helpers() @@ -861,6 +884,11 @@ def run_document_analysis( json_code_block_requested = analysis_intent.get('json_code_block_requested') for document_index, document_id in enumerate(targets.get('document_ids', []), start=1): + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative_manifest', + request_correlation_id=request_correlation_id, + ) document_payload = get_document_chunks_payload( document_id=document_id, user_id=user_id, @@ -872,6 +900,11 @@ def run_document_analysis( window_size=targets.get('window_size'), window_percent=targets.get('window_percent'), ) + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative_manifest', + request_correlation_id=request_correlation_id, + ) windows = build_document_chunk_windows( document_payload.get('chunks', []), window_unit=targets.get('window_unit'), @@ -919,6 +952,11 @@ def run_document_analysis( }) for document_run in document_runs: + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative', + request_correlation_id=request_correlation_id, + ) document_id = document_run.get('document_id') document_payload = document_run.get('document_payload') or {} document_metadata = document_payload.get('document') if isinstance(document_payload.get('document'), dict) else {} @@ -963,6 +1001,11 @@ def run_document_analysis( }) for window_payload in windows: + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative', + request_correlation_id=request_correlation_id, + ) window_range = _serialize_window_range(window_payload) document_summary['ranges'].append(window_range) window_label = _build_window_label(document_name, window_range) @@ -1004,6 +1047,11 @@ def run_document_analysis( last_error = '' max_attempts = targets.get('max_retries_per_window', DEFAULT_MAX_RETRIES_PER_WINDOW) + 1 for attempt_number in range(1, max_attempts + 1): + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative', + request_correlation_id=request_correlation_id, + ) if attempt_number > 1: coverage['retries'] += 1 @@ -1024,9 +1072,16 @@ def run_document_analysis( 'attempt_number': attempt_number, }, ) or '').strip() + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative', + request_correlation_id=request_correlation_id, + ) if not analysis_text: raise ValueError('The analysis runner returned an empty response.') break + except MixedSourceCancellationError: + raise except Exception as exc: last_error = str(exc) debug_print( @@ -1165,6 +1220,8 @@ def run_document_analysis( document_summary.get('failed_ranges', []), reduction_batch_size, max_reduction_rounds, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, ) document_result_text = str(document_result.get('text', '') or '').strip() @@ -1264,6 +1321,11 @@ def run_document_analysis( next_items = [] batches = _build_reduction_batches(current_items, reduction_batch_size) for batch_index, batch_items in enumerate(batches, start=1): + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative_reduction', + request_correlation_id=request_correlation_id, + ) reduction_step_index = completed_reduction_steps + 1 reduction_progress_percent = 90 if reduction_step_total > 0: @@ -1316,6 +1378,11 @@ def run_document_analysis( 'item_count': len(batch_items), }, ) or '').strip() + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative_reduction', + request_correlation_id=request_correlation_id, + ) if not reduced_text: debug_print( '[DocumentAnalysis] Reduction failed | ' @@ -1350,6 +1417,11 @@ def run_document_analysis( final_analysis_reply = current_items[0].get('text', '').strip() + raise_if_mixed_source_cancelled( + cancel_requested, + 'narrative_finalization', + request_correlation_id=request_correlation_id, + ) _set_progress_meta( coverage, phase='completed', diff --git a/application/single_app/functions_document_comparison.py b/application/single_app/functions_document_comparison.py index 565f45eb3..0643ea00b 100644 --- a/application/single_app/functions_document_comparison.py +++ b/application/single_app/functions_document_comparison.py @@ -5,6 +5,10 @@ from functions_appinsights import log_event from functions_debug import debug_print +from functions_mixed_source_orchestration import ( + MixedSourceCancellationError, + raise_if_mixed_source_cancelled, +) from functions_document_actions import DOCUMENT_ACTION_TYPE_COMPARISON from functions_document_analysis import ( build_document_analysis_progress_snapshot, @@ -224,6 +228,158 @@ def _format_comparison_coverage_summary(coverage, left_document_name, right_docu return '\n'.join(lines) +def run_evidence_document_comparison( + comparison_prompt, + left_source, + right_sources, + invoke_prompt, + activity_callback=None, + cancel_requested=None, + request_correlation_id=None, +): + """Run the established one-left-to-many comparison over native evidence envelopes.""" + normalized_prompt = str(comparison_prompt or '').strip() + if not normalized_prompt or not callable(invoke_prompt): + raise ValueError('A comparison prompt and callable invoke_prompt handler are required.') + + left_source = left_source if isinstance(left_source, dict) else {} + right_sources = [source for source in list(right_sources or []) if isinstance(source, dict)] + left_name = str(left_source.get('document_name') or 'Source').strip() or 'Source' + left_summary = str(left_source.get('summary') or '').strip() + left_status = str(left_source.get('status') or '').strip().lower() + if left_status not in {'completed', 'partial'} or not left_summary: + raise RuntimeError('The comparison Source could not be prepared from native evidence.') + if cancel_requested is not None: + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison', + request_correlation_id=request_correlation_id, + ) + comparison_items = [] + compared_targets = [] + failed_targets = [] + + for comparison_index, right_source in enumerate(right_sources, start=1): + if cancel_requested is not None: + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison', + request_correlation_id=request_correlation_id, + ) + right_name = str(right_source.get('document_name') or f'Target {comparison_index}').strip() or f'Target {comparison_index}' + right_status = str(right_source.get('status') or '').strip().lower() + if right_status not in {'completed', 'partial'} or not left_summary: + failed_targets.append(right_name) + continue + + if callable(activity_callback): + activity_callback({ + 'type': 'comparison_started', + 'left_document_id': left_source.get('document_id'), + 'left_document_name': left_name, + 'right_document_id': right_source.get('document_id'), + 'right_document_name': right_name, + 'comparison_index': comparison_index, + 'comparison_count': len(right_sources), + }) + try: + pairwise_text = str(invoke_prompt( + _build_pairwise_comparison_prompt( + normalized_prompt, + left_name, + right_name, + left_summary, + str(right_source.get('summary') or ''), + ), + stage='comparison', + metadata={ + 'comparison_index': comparison_index, + 'comparison_count': len(right_sources), + 'left_document_id': left_source.get('document_id'), + 'right_document_id': right_source.get('document_id'), + }, + ) or '').strip() + if cancel_requested is not None: + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison', + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + raise + except Exception: + failed_targets.append(right_name) + continue + if not pairwise_text: + failed_targets.append(right_name) + continue + compared_targets.append(right_name) + comparison_items.append({ + 'right_document_id': right_source.get('document_id'), + 'right_document_name': right_name, + 'text': pairwise_text, + }) + if callable(activity_callback): + activity_callback({ + 'type': 'comparison_completed', + 'left_document_id': left_source.get('document_id'), + 'left_document_name': left_name, + 'right_document_id': right_source.get('document_id'), + 'right_document_name': right_name, + 'comparison_index': comparison_index, + 'comparison_count': len(right_sources), + }) + + if not comparison_items: + final_reply = 'No target comparison could be completed from the available source evidence.' + elif len(comparison_items) == 1: + final_reply = comparison_items[0]['text'] + else: + if cancel_requested is not None: + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison_reduction', + request_correlation_id=request_correlation_id, + ) + final_reply = str(invoke_prompt( + _build_comparison_reduction_prompt(normalized_prompt, left_name, comparison_items), + stage='comparison_reduction', + metadata={'comparison_count': len(comparison_items), 'left_document_id': left_source.get('document_id')}, + ) or '').strip() or 'The completed target comparisons could not be reduced into a final response.' + if cancel_requested is not None: + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison_reduction', + request_correlation_id=request_correlation_id, + ) + + evidence_engines = sorted({ + str(source.get('engine') or 'unknown') + for source in [left_source, *right_sources] + }) + conclusion_level = 'aggregate or narrative' + if all(str(source.get('source_kind') or '') == 'tabular' for source in [left_source, *right_sources]): + conclusion_level = 'aggregate; row-level conclusions require validated structured table operations' + coverage_note = ( + f'\n\n## Comparison Coverage\n- Targets compared: {len(compared_targets)}\n' + f'- Failed or partial targets: {len(failed_targets)}\n' + f'- Evidence engines: {", ".join(evidence_engines)}\n' + f'- Conclusion level: {conclusion_level}' + ) + return { + 'reply': f'{final_reply}{coverage_note}', + 'analysis_reply': f'{final_reply}{coverage_note}', + 'coverage': {'document_count': 1 + len(right_sources), 'partial_coverage': bool(failed_targets), 'failed_targets': failed_targets}, + 'documents': [left_source, *right_sources], + 'left_document': {'document_id': left_source.get('document_id'), 'document_name': left_name}, + 'right_documents': [ + {'document_id': source.get('document_id'), 'document_name': source.get('document_name')} + for source in right_sources + ], + 'comparison_items': comparison_items, + } + + def run_document_comparison( user_id, comparison_prompt, @@ -231,6 +387,8 @@ def run_document_comparison( invoke_prompt, activity_callback=None, conversation_id=None, + cancel_requested=None, + request_correlation_id=None, ): normalized_prompt = str(comparison_prompt or '').strip() if not normalized_prompt: @@ -245,6 +403,11 @@ def run_document_comparison( right_document_ids = list(action_config.get('right_document_ids') or []) if not left_document_id or not right_document_ids: raise ValueError('Document comparison requires one Source document and at least one Target document.') + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison_manifest', + request_correlation_id=request_correlation_id, + ) debug_print( '[DocumentComparison] Starting comparison | ' @@ -288,6 +451,11 @@ def run_document_comparison( document_summaries = {} for document_index, document_id in enumerate(document_order, start=1): + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison', + request_correlation_id=request_correlation_id, + ) document_state = document_states[document_id] role_label = document_state.get('role_label', 'right') debug_print( @@ -347,6 +515,8 @@ def summary_activity_callback(event, current_document_id=document_id, current_do activity_callback=summary_activity_callback, max_documents=1, include_coverage_summary=False, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, ) document_summaries[document_id] = summary_result document_state['document_name'] = ( @@ -370,6 +540,11 @@ def summary_activity_callback(event, current_document_id=document_id, current_do left_document_name = document_states[left_document_id].get('document_name') or left_document_id comparison_items = [] for comparison_index, right_document_id in enumerate(right_document_ids, start=1): + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison', + request_correlation_id=request_correlation_id, + ) right_document_name = document_states[right_document_id].get('document_name') or right_document_id comparison_progress_percent = ((comparison_index - 1) / len(right_document_ids)) * 100 if right_document_ids else 0 _set_comparison_progress_meta( @@ -417,6 +592,11 @@ def summary_activity_callback(event, current_document_id=document_id, current_do 'right_document_id': right_document_id, }, ) or '').strip() + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison', + request_correlation_id=request_correlation_id, + ) if not pairwise_text: debug_print( '[DocumentComparison] Pairwise comparison failed | ' @@ -466,6 +646,11 @@ def summary_activity_callback(event, current_document_id=document_id, current_do if len(comparison_items) == 1: final_reply = comparison_items[0].get('text', '').strip() else: + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison_reduction', + request_correlation_id=request_correlation_id, + ) _set_comparison_progress_meta( coverage, phase='reducing', @@ -501,6 +686,11 @@ def summary_activity_callback(event, current_document_id=document_id, current_do 'left_document_id': left_document_id, }, ) or '').strip() + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison_reduction', + request_correlation_id=request_correlation_id, + ) if not final_reply: debug_print( '[DocumentComparison] Comparison reduction failed | ' @@ -513,6 +703,11 @@ def summary_activity_callback(event, current_document_id=document_id, current_do f'comparison_count={len(comparison_items)}' ) + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison_finalization', + request_correlation_id=request_correlation_id, + ) _set_comparison_progress_meta( coverage, phase='completed', diff --git a/application/single_app/functions_documents.py b/application/single_app/functions_documents.py index df5d1da88..824c69b73 100644 --- a/application/single_app/functions_documents.py +++ b/application/single_app/functions_documents.py @@ -19,6 +19,10 @@ sync_document_access_index_for_document_fail_open, validate_document_access_index_shadow, ) +from functions_data_management_search_write_fence import ( + DataManagementSearchWritesFrozenError, + hold_data_management_search_write_slot, +) from functions_visio import build_visio_page_markdown, parse_vsdx_pages from functions_content import * from functions_settings import * @@ -32,6 +36,39 @@ _AUDIO_RUNTIME_CAPABILITIES_CACHE = None + +class DocumentSearchAclProjectionDeferredError(RuntimeError): + """Raised when an authorization-reducing Search ACL update must be retried safely.""" + + +def _search_indexing_results_succeeded(results): + """Return whether Azure AI Search acknowledged every requested document mutation.""" + normalized_results = list(results or []) + return bool(normalized_results) and all( + bool( + result.get("succeeded", result.get("status", False)) + if isinstance(result, dict) else + getattr(result, "succeeded", getattr(result, "status", False)) + ) + for result in normalized_results + ) + + +def _execute_document_search_write(search_client, operation_name, *args, **kwargs): + """Serialize bounded target Search writes with an active Data Management migration fence.""" + kwargs.update({ + "connection_timeout": 30, + "read_timeout": 30, + "retry_total": 0, + }) + with hold_data_management_search_write_slot(cosmos_data_management_jobs_container): + results = getattr(search_client, operation_name)(*args, **kwargs) + if not _search_indexing_results_succeeded(results): + raise RuntimeError( + f"Azure AI Search did not acknowledge every {operation_name} document mutation." + ) + return results + def allowed_file(filename, allowed_extensions=None): if not allowed_extensions: allowed_extensions = ALLOWED_EXTENSIONS @@ -786,7 +823,11 @@ def set_document_chunk_visibility(document_item, active=True): documents_to_update.append(chunk_item) - search_client.upload_documents(documents=documents_to_update) + _execute_document_search_write( + search_client, + "upload_documents", + documents=documents_to_update, + ) return len(documents_to_update) @@ -1333,7 +1374,11 @@ def save_video_chunk( # 3) upload to search index try: debug_print(f"[VIDEO CHUNK] Uploading chunk to search index") - client.upload_documents(documents=[chunk]) + _execute_document_search_write( + client, + "upload_documents", + documents=[chunk], + ) debug_print(f"[VIDEO CHUNK] Upload successful for chunk: {chunk_id}") print(f"[VideoChunk] UPLOAD OK for {chunk_id}", flush=True) except Exception as e: @@ -2675,7 +2720,11 @@ def save_chunks(page_text_content, page_number, file_name, user_id, document_id, else: search_client = CLIENTS["search_client_user"] # Upload as a single-document list - search_client.upload_documents(documents=[chunk_document]) + _execute_document_search_write( + search_client, + "upload_documents", + documents=[chunk_document], + ) except Exception as e: print(f"Error uploading chunk document for document {document_id}: {e}") @@ -2861,7 +2910,11 @@ def save_chunks_batch(chunks_data, user_id, document_id, group_id=None, public_w upload_batch_size = 32 for i in range(0, len(chunk_documents), upload_batch_size): sub_batch = chunk_documents[i:i + upload_batch_size] - search_client.upload_documents(documents=sub_batch) + _execute_document_search_write( + search_client, + "upload_documents", + documents=sub_batch, + ) except Exception as e: log_event(f"[save_chunks_batch] Error uploading batch to AI Search for document {document_id}: {e}", level=logging.ERROR) @@ -3030,7 +3083,11 @@ def update_chunk_metadata(chunk_id, user_id, group_id=None, public_workspace_id= else: chunk_item[field] = kwargs[field] - search_client.upload_documents(documents=[chunk_item]) + _execute_document_search_write( + search_client, + "upload_documents", + documents=[chunk_item], + ) except Exception as e: print(f"Error updating chunk metadata for chunk {chunk_id}: {e}") @@ -3944,7 +4001,11 @@ def delete_document_chunks(document_id, group_id=None, public_workspace_id=None) documents_to_delete = [{"id": doc_id} for doc_id in ids_to_delete] batch = IndexDocumentsBatch() batch.add_delete_actions(documents_to_delete) - result = search_client.index_documents(batch) + result = _execute_document_search_write( + search_client, + "index_documents", + batch, + ) except Exception as e: raise @@ -3955,7 +4016,9 @@ def delete_document_version_chunks(document_id, version, group_id=None, public_w search_client = CLIENTS["search_client_public"] if is_public_workspace else CLIENTS["search_client_group"] if is_group else CLIENTS["search_client_user"] - search_client.delete_documents( + _execute_document_search_write( + search_client, + "delete_documents", actions=[ {"@search.action": "delete", "id": chunk['id']} for chunk in search_client.search( @@ -4395,7 +4458,7 @@ def extract_document_metadata(document_id, user_id, group_id=None, public_worksp json.dumps(meta_data), user_id, document_id=document_id, - top_n=12, + top_n=50, doc_scope=document_scope ) elif document_scope == "group": @@ -4403,7 +4466,7 @@ def extract_document_metadata(document_id, user_id, group_id=None, public_worksp json.dumps(meta_data), user_id, document_id=document_id, - top_n=12, + top_n=50, doc_scope=document_scope, active_group_id=scope_id ) @@ -4412,7 +4475,7 @@ def extract_document_metadata(document_id, user_id, group_id=None, public_worksp json.dumps(meta_data), user_id, document_id=document_id, - top_n=12, + top_n=50, doc_scope=document_scope, active_public_workspace_id=scope_id ) @@ -4424,7 +4487,7 @@ def extract_document_metadata(document_id, user_id, group_id=None, public_worksp json.dumps(meta_data), user_id, document_id=document_id, - top_n=12, + top_n=50, doc_scope="public", active_public_workspace_id=public_workspace_id ) @@ -9110,41 +9173,41 @@ def unshare_document_from_user(document_id, owner_user_id, target_user_id): # Remove all entries for the target user (by oid prefix) new_shared_user_ids = [entry for entry in shared_user_ids if not entry.startswith(f"{target_user_id},")] if len(new_shared_user_ids) != len(shared_user_ids): + # A revoked user must never retain search access after the API reports success. + # Update every chunk projection before committing the authoritative Cosmos ACL. + chunks = get_all_chunks(document_id, actual_owner_id) + try: + for chunk in chunks: + chunk_id = chunk.get('id') + if not chunk_id: + continue + update_chunk_metadata( + chunk_id=chunk_id, + user_id=actual_owner_id, + group_id=None, + public_workspace_id=None, + document_id=document_id, + shared_user_ids=new_shared_user_ids + ) + except DataManagementSearchWritesFrozenError as exc: + raise DocumentSearchAclProjectionDeferredError( + "Document access was not changed because target Search writes are temporarily frozen. Retry after the migration finishes." + ) from exc + document_item['shared_user_ids'] = new_shared_user_ids document_item['last_updated'] = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ') - # Update the document _upsert_document_and_sync_access_index( cosmos_user_documents_container, document_item, operation='document_unshared_from_user', ) - # Update all chunks with the new shared_user_ids - try: - chunks = get_all_chunks(document_id, actual_owner_id) - for chunk in chunks: - chunk_id = chunk.get('id') - if chunk_id: - try: - update_chunk_metadata( - chunk_id=chunk_id, - user_id=actual_owner_id, - group_id=None, - public_workspace_id=None, - document_id=document_id, - shared_user_ids=new_shared_user_ids - ) - except Exception as chunk_e: - print(f"Warning: Failed to update chunk {chunk_id}: {chunk_e}") - # Continue with other chunks - except Exception as e: - print(f"Warning: Failed to update chunks for document {document_id}: {e}") - # Don't fail the whole operation if chunk update fails - return True except CosmosResourceNotFoundError: return False + except DocumentSearchAclProjectionDeferredError: + raise except Exception as e: print(f"Error unsharing document {document_id}: {e}") return False diff --git a/application/single_app/functions_generated_file_exports.py b/application/single_app/functions_generated_file_exports.py new file mode 100644 index 000000000..623a90c67 --- /dev/null +++ b/application/single_app/functions_generated_file_exports.py @@ -0,0 +1,609 @@ +# functions_generated_file_exports.py +"""Format-neutral planning and rendering for generated chat file exports.""" + +import html +import io +import json +import os +import re +import tempfile +from datetime import datetime +from typing import Any, Dict, List, Optional, Sequence, Tuple + +from functions_assistant_table_exports import ( + assistant_table_export_requested, + build_assistant_table_csv, + build_csv_output_clarification_guidance, + extract_assistant_table_entries, +) + + +GENERATED_FILE_FORMAT_CSV = 'csv' +GENERATED_FILE_FORMAT_DOCX = 'docx' +GENERATED_FILE_FORMAT_PDF = 'pdf' +GENERATED_FILE_FORMATS = { + GENERATED_FILE_FORMAT_CSV, + GENERATED_FILE_FORMAT_DOCX, + GENERATED_FILE_FORMAT_PDF, +} +GENERATED_FILE_PREVIEW_ROWS = 3 +FUNCTION_RESULT_ROW_KEYS = ( + 'rows', + 'data', + 'items', + 'results', + 'records', + 'value', + 'values', + 'result', + 'body', + 'output', + 'payload', +) +FUNCTION_RESULT_CONTROL_KEYS = { + 'count', + 'detail', + 'error', + 'errormessage', + 'hasmore', + 'message', + 'metadata', + 'meta', + 'nextlink', + 'nextpage', + 'pagination', + 'returnedrows', + 'status', + 'statuscode', + 'success', + 'summary', + 'total', + 'totalcount', + 'totalmatches', +} +FUNCTION_RESULT_SENSITIVE_KEY_FRAGMENTS = ( + 'accesstoken', + 'apikey', + 'authorization', + 'clientsecret', + 'connectionstring', + 'credential', + 'password', + 'privatekey', + 'secret', + 'sharedaccesssignature', + 'subscriptionkey', + 'token', +) +TABULAR_FUNCTION_RESULT_PLUGIN_NAMES = {'tabularprocessingplugin'} +DOCX_OUTPUT_REQUEST_PATTERNS = ( + re.compile( + r'\b(?:build|create|download|export|generate|make|prepare|save|turn|convert)\b' + r'.{0,120}\b(?:a\s+)?(?:word|docx)(?:\s+(?:document|file|output|report))?\b' + ), + re.compile(r'\b(?:word|docx)\s+(?:document|file|output|report|version)\b'), + re.compile(r'\b(?:get|give)\s+(?:me\s+)?(?:a|the|one)\s+(?:word|docx)\b'), + re.compile(r'\b(?:need|want)\s+(?:(?:a|the|one)\s+)?(?:word|docx)\b'), + re.compile(r'\b(?:in|as)\s+(?:a\s+)?(?:word|docx)(?:\s+(?:document|file|report))?\b'), +) +PDF_OUTPUT_REQUEST_PATTERNS = ( + re.compile( + r'\b(?:build|create|download|export|generate|make|prepare|save|turn|convert)\b' + r'.{0,120}\b(?:a\s+)?pdf(?:\s+(?:document|file|output|report))?\b' + ), + re.compile(r'\bpdf\s+(?:document|file|output|report|version)\b'), + re.compile(r'\b(?:get|give)\s+(?:me\s+)?(?:a|the|one)\s+pdf\b'), + re.compile(r'\b(?:need|want)\s+(?:(?:a|the|one)\s+)?pdf\b'), + re.compile(r'\b(?:in|as)\s+(?:a\s+)?pdf(?:\s+(?:document|file|report))?\b'), +) +PDF_EXPORT_CSS = """ +body { font-family: sans-serif; font-size: 10pt; color: #172033; } +h1 { font-size: 20pt; color: #173b5f; margin-bottom: 10pt; } +h2 { font-size: 14pt; color: #173b5f; margin-top: 16pt; } +p { line-height: 1.35; margin-bottom: 8pt; } +table { border-collapse: collapse; width: 100%; margin-top: 8pt; } +th { background-color: #e8eef5; font-weight: bold; } +th, td { border: 0.6pt solid #aab7c4; padding: 4pt; vertical-align: top; } +""" + + +def get_requested_generated_file_format(user_question: str) -> Optional[str]: + """Return the requested generated file format, if any.""" + if assistant_table_export_requested(user_question): + return GENERATED_FILE_FORMAT_CSV + + normalized_question = re.sub(r'\s+', ' ', str(user_question or '').strip().casefold()) + if not normalized_question: + return None + if any(pattern.search(normalized_question) for pattern in DOCX_OUTPUT_REQUEST_PATTERNS): + return GENERATED_FILE_FORMAT_DOCX + if any(pattern.search(normalized_question) for pattern in PDF_OUTPUT_REQUEST_PATTERNS): + return GENERATED_FILE_FORMAT_PDF + return None + + +def generated_file_export_requested(user_question: str) -> bool: + """Return whether the user asked for a supported generated file artifact.""" + return get_requested_generated_file_format(user_question) is not None + + +def build_generated_file_output_guidance(user_question: str) -> str: + """Return shared model guidance for a requested generated output format.""" + output_format = get_requested_generated_file_format(user_question) + if output_format == GENERATED_FILE_FORMAT_CSV: + return build_csv_output_clarification_guidance(user_question) + if output_format in {GENERATED_FILE_FORMAT_DOCX, GENERATED_FILE_FORMAT_PDF}: + return ( + f'The user requested a downloadable {output_format.upper()} artifact. Provide a clear final ' + 'response grounded in the available evidence. Structured function results from this turn may ' + 'be included as labeled tables in the generated file; do not invent rows or claim an attachment ' + 'exists before the file-output finalizer publishes it.' + ) + return '' + + +def get_generated_file_export_content(assistant_result: Any) -> str: + """Return the structured document-action reply when it supersedes a concise artifact reply.""" + if not isinstance(assistant_result, dict): + return str(assistant_result or '') + + analysis_result = assistant_result.get('analysis_result') + if isinstance(analysis_result, dict): + analysis_reply = str(analysis_result.get('analysis_reply') or '').strip() + if analysis_reply: + return analysis_reply + + return str(assistant_result.get('reply') or '') + + +def build_generated_file_export( + user_question: str, + assistant_content: str, + function_results: Optional[List[Dict[str, Any]]] = None, +) -> Optional[Dict[str, Any]]: + """Build a generated file payload from final assistant content and function-result evidence.""" + output_format = get_requested_generated_file_format(user_question) + if output_format not in GENERATED_FILE_FORMATS: + return None + + assistant_text = str(assistant_content or '').strip() + assistant_rows = extract_assistant_table_entries(assistant_text) + function_rows = extract_authorized_function_result_rows(function_results) + + if output_format == GENERATED_FILE_FORMAT_CSV: + rows = assistant_rows or function_rows + if not rows: + return None + row_source = 'assistant response' if assistant_rows else 'structured function result' + return _build_generated_file_payload( + output_format=output_format, + file_content=build_assistant_table_csv(rows), + rows=rows, + row_source=row_source, + assistant_content=assistant_text, + ) + + if not assistant_text and not assistant_rows and not function_rows: + return None + rows = function_rows or assistant_rows + row_source = 'structured function result' if function_rows else 'assistant response' + title = _build_generated_file_title(output_format) + if output_format == GENERATED_FILE_FORMAT_DOCX: + file_content = _render_docx_file_export(title, assistant_text, rows, row_source) + else: + file_content = _render_pdf_file_export(title, assistant_text, rows, row_source) + return _build_generated_file_payload( + output_format=output_format, + file_content=file_content, + rows=rows, + row_source=row_source, + assistant_content=assistant_text, + title=title, + ) + + +def extract_authorized_function_result_rows(function_results: Optional[List[Dict[str, Any]]]) -> List[Dict[str, Any]]: + """Return structured rows from successful current-turn non-tabular function results.""" + function_row_groups: List[Tuple[str, List[Dict[str, Any]]]] = [] + for function_result in function_results or []: + if not isinstance(function_result, dict): + continue + if function_result.get('success') is False or _is_tabular_function_result(function_result): + continue + + structured_rows = _extract_function_result_rows( + _parse_function_result_payload(function_result.get('function_result')), + ) + if not structured_rows: + continue + function_row_groups.append(( + _get_function_result_label(function_result), + structured_rows, + )) + + if not function_row_groups: + return [] + if len(function_row_groups) == 1: + return function_row_groups[0][1] + + source_column = _get_function_result_source_column(function_row_groups) + combined_rows = [] + for function_label, rows in function_row_groups: + for row in rows: + normalized_row = dict(row) + normalized_row[source_column] = function_label + combined_rows.append(normalized_row) + return combined_rows + + +def has_generated_file_output(existing_outputs: Optional[List[Dict[str, Any]]], output_format: str) -> bool: + """Return whether an existing generated artifact already covers an output format.""" + normalized_output_format = str(output_format or '').strip().lower() + if not normalized_output_format: + return False + for output in existing_outputs or []: + if not isinstance(output, dict): + continue + existing_output_format = str(output.get('output_format') or '').strip().lower() + existing_file_name = str(output.get('file_name') or '').strip().lower() + if existing_output_format == normalized_output_format or existing_file_name.endswith(f'.{normalized_output_format}'): + return True + return False + + +def build_generated_file_artifact_metadata( + export_payload: Dict[str, Any], + upload_result: Dict[str, Any], + conversation_id: str, +) -> Optional[Dict[str, Any]]: + """Build public artifact metadata after an authorized generated-file upload.""" + uploaded_message = upload_result.get('message') if isinstance(upload_result, dict) else {} + uploaded_message = uploaded_message if isinstance(uploaded_message, dict) else {} + artifact_message_id = str(uploaded_message.get('id') or '').strip() + if not artifact_message_id: + return None + + generated_file_name = str(export_payload.get('file_name') or '').strip() + artifact_metadata = { + 'capability': str(export_payload.get('capability') or 'file_export').strip().lower() or 'file_export', + 'artifact_message_id': artifact_message_id, + 'conversation_id': str(conversation_id or '').strip(), + 'storage_scope': 'chat', + 'file_name': uploaded_message.get('file_name') or generated_file_name, + 'output_format': str(export_payload.get('output_format') or '').strip().lower(), + 'summary': str(export_payload.get('summary') or '').strip(), + } + row_count = export_payload.get('row_count') + if isinstance(row_count, int) and row_count > 0: + artifact_metadata['row_count'] = row_count + preview_rows = export_payload.get('preview_rows') + if isinstance(preview_rows, list) and preview_rows: + artifact_metadata['preview_rows'] = preview_rows + preview_lines = export_payload.get('preview_lines') + if isinstance(preview_lines, list) and preview_lines: + artifact_metadata['preview_lines'] = preview_lines + row_source = str(export_payload.get('row_source') or '').strip() + if row_source: + artifact_metadata['row_source'] = row_source + return artifact_metadata + + +def _build_generated_file_payload( + output_format: str, + file_content: Any, + rows: Sequence[Dict[str, Any]], + row_source: str, + assistant_content: str, + title: str = '', +) -> Dict[str, Any]: + normalized_output_format = str(output_format or '').strip().lower() + row_count = len(rows or []) + normalized_title = str(title or _build_generated_file_title(normalized_output_format)).strip() + return { + 'capability': 'file_export', + 'file_name': _build_generated_file_name(normalized_output_format), + 'file_content': file_content, + 'output_format': normalized_output_format, + 'row_count': row_count, + 'preview_rows': list(rows or [])[:GENERATED_FILE_PREVIEW_ROWS], + 'preview_lines': _build_preview_lines(assistant_content), + 'row_source': row_source, + '_structured_rows': list(rows or []), + 'summary': _build_generated_file_summary( + normalized_output_format, + row_count, + row_source, + normalized_title, + ), + } + + +def _build_generated_file_name(output_format: str) -> str: + timestamp_suffix = datetime.utcnow().strftime('%Y%m%d_%H%M%S') + return f'generated_output_{timestamp_suffix}.{output_format}' + + +def _build_generated_file_title(output_format: str) -> str: + return f'Generated {str(output_format or "file").upper()} export' + + +def _build_generated_file_summary( + output_format: str, + row_count: int, + row_source: str, + title: str, +) -> str: + row_detail = f' with {row_count} structured row(s)' if row_count else '' + return f'Prepared {title}{row_detail} from the {row_source}.' + + +def _build_preview_lines(assistant_content: str) -> List[str]: + normalized_lines = [ + line.strip() + for line in str(assistant_content or '').splitlines() + if line.strip() + ] + return normalized_lines[:3] + + +def _parse_function_result_payload(value: Any) -> Any: + if not isinstance(value, str): + return value + + normalized_value = value.strip() + if not normalized_value or normalized_value[0] not in '[{': + return value + try: + return json.loads(normalized_value) + except (TypeError, ValueError, json.JSONDecodeError): + return value + + +def _extract_function_result_rows( + payload: Any, + depth: int = 0, + is_data_row: bool = False, +) -> List[Dict[str, Any]]: + if depth > 4: + return [] + if isinstance(payload, list): + rows = [] + for item in payload: + if isinstance(item, dict): + normalized_row = _normalize_function_result_row(item, is_data_row=True) + if normalized_row: + rows.append(normalized_row) + elif item not in (None, ''): + rows.append({'value': _sanitize_function_result_value(item)}) + return rows + if not isinstance(payload, dict): + return [] + + normalized_keys = { + _normalize_function_result_key(key): key + for key in payload + } + for row_key in FUNCTION_RESULT_ROW_KEYS: + matching_key = normalized_keys.get(_normalize_function_result_key(row_key)) + if matching_key is None: + continue + rows = _extract_function_result_rows( + payload.get(matching_key), + depth + 1, + is_data_row=True, + ) + if rows: + return rows + + normalized_row = _normalize_function_result_row(payload, is_data_row=is_data_row) + return [normalized_row] if normalized_row else [] + + +def _normalize_function_result_row(row: Dict[str, Any], is_data_row: bool) -> Dict[str, Any]: + normalized_row = {} + for raw_key, raw_value in row.items(): + key = str(raw_key or '').strip() + normalized_key = _normalize_function_result_key(key) + if not key or _is_sensitive_function_result_key(key): + continue + if not is_data_row and normalized_key in FUNCTION_RESULT_CONTROL_KEYS: + continue + + value = _sanitize_function_result_value(raw_value) + if value in (None, '', [], {}) or value == '***REDACTED***': + continue + normalized_row[key] = value + return normalized_row + + +def _sanitize_function_result_value(value: Any, depth: int = 0) -> Any: + if depth > 4: + return '[truncated]' + if isinstance(value, dict): + return { + str(key): _sanitize_function_result_value(item, depth + 1) + for key, item in value.items() + if not _is_sensitive_function_result_key(key) + } + if isinstance(value, (list, tuple, set)): + return [ + _sanitize_function_result_value(item, depth + 1) + for item in value + ] + return value + + +def _normalize_function_result_key(key: Any) -> str: + return re.sub(r'[^a-z0-9]', '', str(key or '').casefold()) + + +def _is_sensitive_function_result_key(key: Any) -> bool: + normalized_key = _normalize_function_result_key(key) + if not normalized_key: + return False + return any(fragment in normalized_key for fragment in FUNCTION_RESULT_SENSITIVE_KEY_FRAGMENTS) + + +def _is_tabular_function_result(function_result: Dict[str, Any]) -> bool: + plugin_name = str(function_result.get('plugin_name') or '').strip().casefold() + return plugin_name in TABULAR_FUNCTION_RESULT_PLUGIN_NAMES + + +def _get_function_result_label(function_result: Dict[str, Any]) -> str: + return ( + str(function_result.get('function_name') or '').strip() + or str(function_result.get('plugin_name') or '').strip() + or 'function result' + ) + + +def _get_function_result_source_column( + function_row_groups: Sequence[Tuple[str, Sequence[Dict[str, Any]]]], +) -> str: + existing_columns = { + str(column_name).casefold() + for _, rows in function_row_groups + for row in rows + for column_name in row + } + source_column = 'Source action' + suffix = 2 + while source_column.casefold() in existing_columns: + source_column = f'Source action {suffix}' + suffix += 1 + return source_column + + +def _render_docx_file_export( + title: str, + assistant_content: str, + rows: Sequence[Dict[str, Any]], + row_source: str, +) -> bytes: + from docx import Document as DocxDocument + + document = DocxDocument() + document.add_heading(title, level=1) + _append_docx_text(document, assistant_content) + if rows: + document.add_heading(_build_structured_rows_heading(row_source), level=2) + _append_docx_table(document, rows) + + output_buffer = io.BytesIO() + document.save(output_buffer) + return output_buffer.getvalue() + + +def _append_docx_text(document: Any, assistant_content: str) -> None: + normalized_content = str(assistant_content or '').strip() + if not normalized_content: + return + for paragraph_text in re.split(r'\n\s*\n', normalized_content): + cleaned_paragraph = paragraph_text.strip() + if cleaned_paragraph: + document.add_paragraph(cleaned_paragraph) + + +def _append_docx_table(document: Any, rows: Sequence[Dict[str, Any]]) -> None: + columns = _collect_structured_row_columns(rows) + if not columns: + return + table = document.add_table(rows=1, cols=len(columns)) + table.style = 'Table Grid' + for index, column_name in enumerate(columns): + table.rows[0].cells[index].text = str(column_name) + for row in rows: + cells = table.add_row().cells + for index, column_name in enumerate(columns): + cells[index].text = _format_structured_cell(row.get(column_name)) + + +def _render_pdf_file_export( + title: str, + assistant_content: str, + rows: Sequence[Dict[str, Any]], + row_source: str, +) -> bytes: + import fitz + + html_parts = [f'

{html.escape(title)}

'] + normalized_content = str(assistant_content or '').strip() + if normalized_content: + html_parts.append('

Response

') + for paragraph_text in re.split(r'\n\s*\n', normalized_content): + cleaned_paragraph = paragraph_text.strip() + if cleaned_paragraph: + html_parts.append(f'

{html.escape(cleaned_paragraph).replace(chr(10), "
")}

') + if rows: + html_parts.append(f'

{html.escape(_build_structured_rows_heading(row_source))}

') + html_parts.append(_build_structured_rows_html(rows)) + + media_box = fitz.paper_rect('letter') + content_box = media_box + (36, 36, -36, -36) + story = fitz.Story(html='\n'.join(html_parts), user_css=PDF_EXPORT_CSS) + temporary_path = None + try: + with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as temporary_file: + temporary_path = temporary_file.name + + writer = fitz.DocumentWriter(temporary_path) + has_more = True + while has_more: + device = writer.begin_page(media_box) + has_more, _ = story.place(content_box) + story.draw(device) + writer.end_page() + writer.close() + with open(temporary_path, 'rb') as generated_file: + return generated_file.read() + finally: + if temporary_path: + try: + os.unlink(temporary_path) + except OSError: + pass + + +def _build_structured_rows_html(rows: Sequence[Dict[str, Any]]) -> str: + columns = _collect_structured_row_columns(rows) + if not columns: + return '

No structured rows were available.

' + table_parts = [''] + table_parts.extend(f'' for column_name in columns) + table_parts.append('') + for row in rows: + table_parts.append('') + for column_name in columns: + table_parts.append(f'') + table_parts.append('') + table_parts.append('
{html.escape(str(column_name))}
{html.escape(_format_structured_cell(row.get(column_name))).replace(chr(10), "
")}
') + return ''.join(table_parts) + + +def _collect_structured_row_columns(rows: Sequence[Dict[str, Any]]) -> List[str]: + columns = [] + seen_columns = set() + for row in rows or []: + if not isinstance(row, dict): + continue + for raw_column_name in row: + column_name = str(raw_column_name or '').strip() + if not column_name or column_name.casefold() in seen_columns: + continue + seen_columns.add(column_name.casefold()) + columns.append(column_name) + return columns + + +def _format_structured_cell(value: Any) -> str: + if value is None: + return '' + if isinstance(value, (dict, list, tuple, set)): + return json.dumps(value, default=str, ensure_ascii=False) + return str(value) + + +def _build_structured_rows_heading(row_source: str) -> str: + if row_source == 'structured function result': + return 'Structured function results' + return 'Structured response rows' diff --git a/application/single_app/functions_group.py b/application/single_app/functions_group.py index 94ce8a075..df59c1c01 100644 --- a/application/single_app/functions_group.py +++ b/application/single_app/functions_group.py @@ -101,7 +101,7 @@ def search_all_groups(search_query, limit=10): parameters=params, enable_cross_partition_query=True )) - return results[:max(1, min(int(limit or 10), 25))] + return results[:max(1, min(int(limit or 10), 50))] def get_user_groups(user_id): """ diff --git a/application/single_app/functions_migration_provenance.py b/application/single_app/functions_migration_provenance.py new file mode 100644 index 000000000..92134623d --- /dev/null +++ b/application/single_app/functions_migration_provenance.py @@ -0,0 +1,251 @@ +# functions_migration_provenance.py +"""Durable provenance metadata shared by in-app migration destinations.""" + +from datetime import datetime, timedelta, timezone +import uuid + + +MIGRATION_PROVENANCE_STATUS_SUCCEEDED = "succeeded" +COSMOS_MIGRATION_PROVENANCE_FIELD = "simplechatMigration" +SEARCH_MIGRATION_ID_FIELD = "simplechatMigrationId" +SEARCH_MIGRATED_AT_FIELD = "simplechatMigratedAtUtc" +SEARCH_MIGRATION_STATUS_FIELD = "simplechatMigrationStatus" +SEARCH_MIGRATION_SOURCE_HASH_FIELD = "simplechatMigrationSourceHash" +SEARCH_MIGRATION_SOURCE_VERSION_FIELD = "simplechatMigrationSourceVersion" +BLOB_MIGRATION_ID_METADATA_KEY = "simplechatMigrationId" +BLOB_MIGRATED_AT_METADATA_KEY = "simplechatMigratedAtUtc" +BLOB_MIGRATION_STATUS_METADATA_KEY = "simplechatMigrationStatus" +BLOB_MIGRATION_SOURCE_HASH_METADATA_KEY = "simplechatMigrationSourceHash" +BLOB_MIGRATION_SOURCE_VERSION_METADATA_KEY = "simplechatMigrationSourceVersion" +BLOB_MIGRATION_SCOPE_HASH_METADATA_KEY = "simplechatMigrationScopeHash" + + +def _parse_utc_datetime(value): + """Return an aware UTC timestamp when the supplied value is valid.""" + if isinstance(value, datetime): + parsed = value + else: + try: + parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) + except (TypeError, ValueError): + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +def _utc_iso(value): + """Format a UTC datetime using the durable migration metadata format.""" + parsed = _parse_utc_datetime(value) + if parsed is None: + raise ValueError("Migration timestamp must be a valid UTC datetime.") + return parsed.isoformat() + + +def _record_value(record, *names): + """Read the first available provenance value without changing stored data.""" + if not isinstance(record, dict): + return None + for name in names: + if name in record: + return record.get(name) + return None + + +def create_migration_provenance_context(migration_id=None, migrated_at_utc=None, skip_within_hours=0): + """Create a validated provenance context for one durable migration job.""" + try: + normalized_migration_id = str(uuid.UUID(str(migration_id or uuid.uuid4()))) + except (AttributeError, TypeError, ValueError) as exc: + raise ValueError("Migration ID must be a valid GUID.") from exc + + try: + normalized_skip_hours = int(skip_within_hours) + except (TypeError, ValueError) as exc: + raise ValueError("Migration provenance skip window must be an integer number of hours.") from exc + if normalized_skip_hours < 0 or normalized_skip_hours > 8760: + raise ValueError("Migration provenance skip window must be between 0 and 8760 hours.") + + return { + "migration_id": normalized_migration_id, + "migrated_at_utc": _utc_iso(migrated_at_utc or datetime.now(timezone.utc)), + "skip_within_hours": normalized_skip_hours, + } + + +def build_migration_provenance_record( + context, + source_hash="", + source_version="", + scope_hash="", +): + """Return the neutral, destination-independent success record.""" + record = { + "migrationId": str((context or {}).get("migration_id") or ""), + "migratedAtUtc": str((context or {}).get("migrated_at_utc") or ""), + "status": MIGRATION_PROVENANCE_STATUS_SUCCEEDED, + } + if source_hash: + record["sourceHash"] = str(source_hash) + if source_version: + record["sourceVersion"] = str(source_version) + if scope_hash: + record["scopeHash"] = str(scope_hash) + return record + + +def is_successful_migration_record(record): + """Return whether a destination item has durable successful migration ownership.""" + if not isinstance(record, dict): + return False + status = str(_record_value(record, "status", "migrationStatus") or "").strip().lower() + migration_id = str(_record_value(record, "migrationId", "migration_id") or "").strip() + return status == MIGRATION_PROVENANCE_STATUS_SUCCEEDED and bool(migration_id) + + +def migration_record_matches_source(record, source_hash="", source_version=""): + """Compare an owned destination marker with the current source fingerprint.""" + if not is_successful_migration_record(record): + return False + normalized_hash = str(source_hash or "").strip() + normalized_version = str(source_version or "").strip() + if not normalized_hash and not normalized_version: + return False + if normalized_hash and ( + str(_record_value(record, "sourceHash", "source_hash") or "").strip() != normalized_hash + ): + return False + if normalized_version and ( + str(_record_value(record, "sourceVersion", "source_version") or "").strip() != + normalized_version + ): + return False + return True + + +def should_skip_migration_record(record, context, current_time=None): + """Return whether a successful current or recent migration owns an item.""" + status = str(_record_value(record, "status", "migrationStatus") or "").strip().lower() + if status != MIGRATION_PROVENANCE_STATUS_SUCCEEDED: + return False + + migration_id = str(_record_value(record, "migrationId", "migration_id") or "").strip() + context_migration_id = str((context or {}).get("migration_id") or "").strip() + if migration_id and migration_id.lower() == context_migration_id.lower(): + return True + + try: + skip_within_hours = int((context or {}).get("skip_within_hours") or 0) + except (TypeError, ValueError): + return False + if skip_within_hours <= 0: + return False + + migrated_at = _parse_utc_datetime(_record_value(record, "migratedAtUtc", "migrated_at_utc")) + now = _parse_utc_datetime(current_time or datetime.now(timezone.utc)) + if migrated_at is None or now is None: + return False + return migrated_at >= now - timedelta(hours=skip_within_hours) + + +def add_cosmos_migration_provenance(document, context, source_hash="", source_version=""): + """Add migration provenance to a writable Cosmos document.""" + if not isinstance(document, dict): + raise ValueError("Cosmos migration documents must be dictionaries.") + document[COSMOS_MIGRATION_PROVENANCE_FIELD] = build_migration_provenance_record( + context, + source_hash=source_hash, + source_version=source_version, + ) + return document + + +def get_cosmos_migration_provenance(document): + """Get a Cosmos document's migration record when present.""" + if not isinstance(document, dict): + return None + value = document.get(COSMOS_MIGRATION_PROVENANCE_FIELD) + return value if isinstance(value, dict) else None + + +def add_search_migration_provenance(document, context, source_hash="", source_version=""): + """Add migration provenance to a writable Azure AI Search document.""" + if not isinstance(document, dict): + raise ValueError("AI Search migration documents must be dictionaries.") + record = build_migration_provenance_record( + context, + source_hash=source_hash, + source_version=source_version, + ) + document[SEARCH_MIGRATION_ID_FIELD] = record["migrationId"] + document[SEARCH_MIGRATED_AT_FIELD] = record["migratedAtUtc"] + document[SEARCH_MIGRATION_STATUS_FIELD] = record["status"] + if record.get("sourceHash"): + document[SEARCH_MIGRATION_SOURCE_HASH_FIELD] = record["sourceHash"] + if record.get("sourceVersion"): + document[SEARCH_MIGRATION_SOURCE_VERSION_FIELD] = record["sourceVersion"] + return document + + +def get_search_migration_provenance(document): + """Translate Azure AI Search provenance fields into the neutral record.""" + if not isinstance(document, dict): + return None + record = { + "migrationId": document.get(SEARCH_MIGRATION_ID_FIELD), + "migratedAtUtc": document.get(SEARCH_MIGRATED_AT_FIELD), + "status": document.get(SEARCH_MIGRATION_STATUS_FIELD), + } + if document.get(SEARCH_MIGRATION_SOURCE_HASH_FIELD): + record["sourceHash"] = document.get(SEARCH_MIGRATION_SOURCE_HASH_FIELD) + if document.get(SEARCH_MIGRATION_SOURCE_VERSION_FIELD): + record["sourceVersion"] = document.get(SEARCH_MIGRATION_SOURCE_VERSION_FIELD) + return record + + +def merge_blob_migration_metadata( + metadata, + context, + source_hash="", + source_version="", + scope_hash="", +): + """Preserve existing Blob metadata while appending migration provenance.""" + merged = { + str(key): str(value) + for key, value in (metadata or {}).items() + if key is not None and value is not None + } + record = build_migration_provenance_record( + context, + source_hash=source_hash, + source_version=source_version, + scope_hash=scope_hash, + ) + merged[BLOB_MIGRATION_ID_METADATA_KEY] = record["migrationId"] + merged[BLOB_MIGRATED_AT_METADATA_KEY] = record["migratedAtUtc"] + merged[BLOB_MIGRATION_STATUS_METADATA_KEY] = record["status"] + if record.get("sourceHash"): + merged[BLOB_MIGRATION_SOURCE_HASH_METADATA_KEY] = record["sourceHash"] + if record.get("sourceVersion"): + merged[BLOB_MIGRATION_SOURCE_VERSION_METADATA_KEY] = record["sourceVersion"] + if record.get("scopeHash"): + merged[BLOB_MIGRATION_SCOPE_HASH_METADATA_KEY] = record["scopeHash"] + return merged + + +def get_blob_migration_provenance(metadata): + """Translate Blob metadata provenance fields into the neutral record.""" + values = metadata or {} + record = { + "migrationId": values.get(BLOB_MIGRATION_ID_METADATA_KEY), + "migratedAtUtc": values.get(BLOB_MIGRATED_AT_METADATA_KEY), + "status": values.get(BLOB_MIGRATION_STATUS_METADATA_KEY), + } + if values.get(BLOB_MIGRATION_SOURCE_HASH_METADATA_KEY): + record["sourceHash"] = values.get(BLOB_MIGRATION_SOURCE_HASH_METADATA_KEY) + if values.get(BLOB_MIGRATION_SOURCE_VERSION_METADATA_KEY): + record["sourceVersion"] = values.get(BLOB_MIGRATION_SOURCE_VERSION_METADATA_KEY) + if values.get(BLOB_MIGRATION_SCOPE_HASH_METADATA_KEY): + record["scopeHash"] = values.get(BLOB_MIGRATION_SCOPE_HASH_METADATA_KEY) + return record \ No newline at end of file diff --git a/application/single_app/functions_mixed_source_orchestration.py b/application/single_app/functions_mixed_source_orchestration.py new file mode 100644 index 000000000..f7fb226e8 --- /dev/null +++ b/application/single_app/functions_mixed_source_orchestration.py @@ -0,0 +1,1824 @@ +# functions_mixed_source_orchestration.py +"""Authorization-safe source manifest and evidence contracts for mixed sources.""" + +import json +import logging +import math +import os +import time +import uuid + +from functions_appinsights import log_event + + +SOURCE_KIND_TABULAR = "tabular" +SOURCE_KIND_NARRATIVE = "narrative" +SOURCE_KIND_UNSUPPORTED = "unsupported" +SOURCE_KIND_UNRESOLVED = "unresolved" +SOURCE_KINDS = frozenset({ + SOURCE_KIND_TABULAR, + SOURCE_KIND_NARRATIVE, + SOURCE_KIND_UNSUPPORTED, + SOURCE_KIND_UNRESOLVED, +}) + +SOURCE_SCOPE_PERSONAL = "personal" +SOURCE_SCOPE_GROUP = "group" +SOURCE_SCOPE_PUBLIC = "public" +SOURCE_SCOPE_CHAT = "chat" +SOURCE_SCOPES = frozenset({ + SOURCE_SCOPE_PERSONAL, + SOURCE_SCOPE_GROUP, + SOURCE_SCOPE_PUBLIC, + SOURCE_SCOPE_CHAT, +}) + +AUTHORIZATION_STATUS_AUTHORIZED = "authorized" +AUTHORIZATION_STATUS_UNRESOLVED = "unresolved" +SOURCE_MANIFEST_MAX_SOURCES = 100 + +SELECTION_MODE_SELECTED = "selected" +SELECTION_MODE_ALL = "all" +SELECTION_MODE_HISTORY = "history" +SELECTION_MODE_RELEVANCE = "relevance" +SELECTION_MODES = frozenset({ + SELECTION_MODE_SELECTED, + SELECTION_MODE_ALL, + SELECTION_MODE_HISTORY, + SELECTION_MODE_RELEVANCE, +}) + +TABULAR_SOURCE_EXTENSIONS = frozenset({".csv", ".xls", ".xlsx", ".xlsm"}) +NARRATIVE_SOURCE_EXTENSIONS = frozenset({ + ".txt", ".doc", ".docm", ".docx", ".html", ".htm", ".md", ".markdown", + ".json", ".xml", ".yaml", ".yml", ".log", ".pdf", ".ppt", ".pptx", + ".msg", ".vsdx", ".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", + ".heif", ".heic", ".3ga", ".aac", ".ac3", ".aif", ".aifc", ".aiff", + ".amr", ".ape", ".au", ".caf", ".dts", ".f4a", ".flac", ".m4a", + ".m4b", ".m4r", ".mka", ".mp2", ".mp3", ".mpa", ".oga", ".ogg", + ".opus", ".spx", ".wav", ".weba", ".wma", ".wv", ".mp4", ".mov", + ".avi", ".mkv", ".flv", ".mxf", ".gxf", ".ts", ".ps", ".3gp", + ".3gpp", ".mpg", ".wmv", ".asf", ".m4v", ".isma", ".ismv", + ".dvr-ms", ".webm", ".mpeg", +}) + +EVIDENCE_ENGINE_TABULAR_TOOLS = "tabular_tools" +EVIDENCE_ENGINE_DOCUMENT_ANALYSIS = "document_analysis" +EVIDENCE_ENGINE_HYBRID_SEARCH = "hybrid_search" +EVIDENCE_ENGINES = frozenset({ + EVIDENCE_ENGINE_TABULAR_TOOLS, + EVIDENCE_ENGINE_DOCUMENT_ANALYSIS, + EVIDENCE_ENGINE_HYBRID_SEARCH, +}) + +EVIDENCE_STATUS_COMPLETED = "completed" +EVIDENCE_STATUS_PARTIAL = "partial" +EVIDENCE_STATUS_FAILED = "failed" +EVIDENCE_STATUS_SKIPPED = "skipped" +EVIDENCE_STATUSES = frozenset({ + EVIDENCE_STATUS_COMPLETED, + EVIDENCE_STATUS_PARTIAL, + EVIDENCE_STATUS_FAILED, + EVIDENCE_STATUS_SKIPPED, +}) + +EVIDENCE_ENVELOPE_MAX_BYTES = 65536 +EVIDENCE_SUMMARY_MAX_BYTES = 4096 +EVIDENCE_ERROR_MAX_BYTES = 1024 +EVIDENCE_LIST_MAX_ITEMS = 10 +EVIDENCE_ITEM_MAX_BYTES = 1536 +EVIDENCE_COVERAGE_MAX_BYTES = 4096 +EVIDENCE_JSON_MAX_DEPTH = 4 +EVIDENCE_JSON_MAX_COLLECTION_ITEMS = 20 +EVIDENCE_JSON_MAX_STRING_BYTES = 1024 +MIXED_SOURCE_HANDOFF_MAX_BYTES = 49152 +MIXED_SOURCE_HANDOFF_MAX_ENVELOPES = 20 +MIXED_SOURCE_MODES = frozenset({"chat", "search", "analyze", "compare"}) +MIXED_SOURCE_TERMINAL_REASON_MAX_BYTES = 128 +MIXED_SOURCE_TELEMETRY_EVENTS = frozenset({ + "authorization_failure", + "background_export", + "cancellation", + "continuity", + "native_execution", + "reduction", + "terminal_coverage", +}) +MIXED_SOURCE_TELEMETRY_METRICS = frozenset({ + "artifact_count", + "authorization_failure_count", + "background_export_count", + "cancellation_count", + "citation_count", + "completed_source_count", + "duplicate_evidence_count", + "engine_call_count", + "evidence_omitted_count", + "failed_source_count", + "history_rerun_count", + "history_reuse_count", + "history_source_count", + "latency_ms", + "missing_coverage_violation_count", + "model_request_count", + "narrative_source_count", + "partial_failure_count", + "partial_source_count", + "prompt_tokens", + "request_count", + "skipped_source_count", + "successful_source_count", + "tabular_source_count", + "token_request_count", + "total_source_count", + "total_tokens", + "unexpected_evidence_count", + "unsupported_source_count", + "unresolved_source_count", +}) +MIXED_SOURCE_TELEMETRY_DIMENSIONS = frozenset({ + "cancellation_phase", + "continuity_decision", + "outcome_status", + "selection_mode", +}) + + +class MixedSourceCancellationError(RuntimeError): + """Stop mixed-source work without converting cancellation into source failure.""" + + def __init__(self, phase="unknown"): + self.phase = str(phase or "unknown").strip().lower() or "unknown" + super().__init__(f"Mixed-source execution canceled during {self.phase}.") + + +class MixedSourceFinalizationError(RuntimeError): + """Prevent publication when a fresh manifest no longer matches execution evidence.""" + + def __init__(self, reason): + self.reason = str(reason or "finalization_failed").strip().lower() + super().__init__("Mixed-source evidence changed or became unavailable before publication.") + + +def normalize_mixed_source_correlation_id(request_correlation_id=None): + """Return an internal UUID correlation value without trusting caller-shaped text.""" + try: + return str(uuid.UUID(str(request_correlation_id or "").strip())) + except (TypeError, ValueError, AttributeError): + return str(uuid.uuid4()) + + +def raise_if_mixed_source_cancelled( + cancel_requested, + phase, + request_correlation_id=None, +): + """Raise the shared cancellation signal when an optional predicate is set.""" + if cancel_requested is None: + return + if not callable(cancel_requested): + raise TypeError("cancel_requested must be callable") + if not cancel_requested(): + return + + normalized_phase = str(phase or "unknown").strip().lower() or "unknown" + log_event( + "[MixedSourceLifecycle] Execution canceled.", + extra={ + "request_correlation_id": normalize_mixed_source_correlation_id( + request_correlation_id + ), + "cancellation_phase": normalized_phase, + }, + level=logging.INFO, + ) + raise MixedSourceCancellationError(normalized_phase) + + +def emit_mixed_source_telemetry( + settings, + event_name, + mode, + request_correlation_id=None, + metrics=None, + dimensions=None, +): + """Emit only allowlisted aggregate lifecycle telemetry when explicitly enabled.""" + if not bool((settings or {}).get("enable_mixed_source_development_telemetry", False)): + return False + + normalized_event_name = str(event_name or "").strip().lower() + normalized_mode = str(mode or "").strip().lower() + if normalized_event_name not in MIXED_SOURCE_TELEMETRY_EVENTS: + raise ValueError(f"Unsupported mixed-source telemetry event: {normalized_event_name}") + if normalized_mode not in MIXED_SOURCE_MODES: + raise ValueError(f"Unsupported mixed-source telemetry mode: {normalized_mode}") + + metrics = metrics if isinstance(metrics, dict) else {} + dimensions = dimensions if isinstance(dimensions, dict) else {} + unknown_metrics = set(metrics) - MIXED_SOURCE_TELEMETRY_METRICS + unknown_dimensions = set(dimensions) - MIXED_SOURCE_TELEMETRY_DIMENSIONS + if unknown_metrics or unknown_dimensions: + raise ValueError("Mixed-source telemetry contains non-allowlisted fields") + + extra = { + "request_correlation_id": normalize_mixed_source_correlation_id( + request_correlation_id + ), + "event_name": normalized_event_name, + "mode": normalized_mode, + } + for field_name, raw_value in metrics.items(): + if isinstance(raw_value, bool): + metric_value = int(raw_value) + elif isinstance(raw_value, (int, float)) and math.isfinite(raw_value): + metric_value = max(0, raw_value) + else: + raise ValueError("Mixed-source telemetry metrics must be finite numbers") + extra[field_name] = metric_value + for field_name, raw_value in dimensions.items(): + extra[field_name] = _truncate_utf8( + str(raw_value or "").strip().lower(), + 64, + ) + + log_event( + "[MixedSourceTelemetry] Aggregate lifecycle metrics.", + extra=extra, + level=logging.INFO, + ) + return True + + +def emit_mixed_source_coverage_telemetry( + settings, + mode, + coverage, + request_correlation_id=None, +): + """Emit terminal status and source-kind counts derived from the bounded ledger.""" + coverage = coverage if isinstance(coverage, dict) else {} + terminal_ledger = [ + entry + for entry in list(coverage.get("terminal_ledger") or []) + if isinstance(entry, dict) + ] + source_kind_counts = {source_kind: 0 for source_kind in SOURCE_KINDS} + for entry in terminal_ledger: + source_kind = str(entry.get("source_kind") or "").strip().lower() + if source_kind in source_kind_counts: + source_kind_counts[source_kind] += 1 + + outcome_status = ( + EVIDENCE_STATUS_PARTIAL + if coverage.get("partial_coverage") + else EVIDENCE_STATUS_COMPLETED + ) + if coverage.get("successful_source_count", 0) == 0 and coverage.get( + "requested_source_count", + 0, + ): + outcome_status = EVIDENCE_STATUS_FAILED + return emit_mixed_source_telemetry( + settings, + "terminal_coverage", + mode, + request_correlation_id=request_correlation_id, + metrics={ + "total_source_count": coverage.get("requested_source_count", 0), + "completed_source_count": coverage.get("completed_source_count", 0), + "partial_source_count": coverage.get("partial_source_count", 0), + "failed_source_count": coverage.get("failed_source_count", 0), + "skipped_source_count": coverage.get("skipped_source_count", 0), + "successful_source_count": coverage.get("successful_source_count", 0), + "tabular_source_count": source_kind_counts[SOURCE_KIND_TABULAR], + "narrative_source_count": source_kind_counts[SOURCE_KIND_NARRATIVE], + "unsupported_source_count": source_kind_counts[SOURCE_KIND_UNSUPPORTED], + "unresolved_source_count": source_kind_counts[SOURCE_KIND_UNRESOLVED], + "missing_coverage_violation_count": coverage.get( + "missing_coverage_violation_count", + 0, + ), + "duplicate_evidence_count": coverage.get("duplicate_evidence_count", 0), + "unexpected_evidence_count": coverage.get("unexpected_evidence_count", 0), + "evidence_omitted_count": coverage.get("evidence_omitted_count", 0), + "partial_failure_count": int(bool(coverage.get("partial_coverage"))), + }, + dimensions={ + "selection_mode": coverage.get("selection_mode") or "selected", + "outcome_status": outcome_status, + }, + ) + + +def normalize_selection_mode(selection_mode, default=SELECTION_MODE_SELECTED): + """Return a supported selection mode or raise for an invalid explicit value.""" + normalized_default = str(default or "").strip().lower() + if normalized_default not in SELECTION_MODES: + raise ValueError("Invalid default selection_mode") + + normalized_mode = str(selection_mode or "").strip().lower() + if not normalized_mode: + return normalized_default + if normalized_mode not in SELECTION_MODES: + raise ValueError( + f"selection_mode must be one of: {', '.join(sorted(SELECTION_MODES))}" + ) + return normalized_mode + + +def normalize_document_context_request( + selection_mode=None, + selected_document_ids=None, + document_context_requested=None, + hybrid_search=False, +): + """Validate the Phase 2 request contract and derive effective context intent.""" + normalized_document_ids = [] + seen_document_ids = set() + for document_id in _normalize_identifier_list(selected_document_ids): + if document_id == "all" or document_id in seen_document_ids: + continue + seen_document_ids.add(document_id) + normalized_document_ids.append(document_id) + + normalized_hybrid_search = _normalize_boolean( + hybrid_search, + field_name="hybrid_search", + ) + normalized_context_requested = _normalize_optional_boolean( + document_context_requested, + field_name="document_context_requested", + ) + has_explicit_selection_mode = str(selection_mode or "").strip() != "" + + if normalized_document_ids: + normalized_selection_mode = normalize_selection_mode( + selection_mode, + default=SELECTION_MODE_SELECTED, + ) + if normalized_selection_mode != SELECTION_MODE_SELECTED: + raise ValueError( + "selection_mode must be selected when selected_document_ids are provided" + ) + normalized_context_requested = True + elif has_explicit_selection_mode: + normalized_selection_mode = normalize_selection_mode(selection_mode) + if normalized_selection_mode == SELECTION_MODE_SELECTED: + raise ValueError( + "selection_mode selected requires at least one selected_document_id" + ) + if ( + normalized_selection_mode in {SELECTION_MODE_ALL, SELECTION_MODE_RELEVANCE} + and normalized_context_requested is None + ): + normalized_context_requested = True + elif normalized_context_requested is True or normalized_hybrid_search: + normalized_selection_mode = SELECTION_MODE_RELEVANCE + normalized_context_requested = True + else: + normalized_selection_mode = None + normalized_context_requested = False + + return { + "selection_mode": normalized_selection_mode, + "selected_document_ids": normalized_document_ids, + "document_context_requested": bool(normalized_context_requested), + "hybrid_search": normalized_hybrid_search, + "explicit_selection": bool(normalized_document_ids), + } + + +def should_run_tabular_evidence(user_question, has_narrative_sources=False): + """Return whether a mixed-source question needs tabular data or schema evidence.""" + normalized_question = " ".join(str(user_question or "").strip().lower().split()) + if not normalized_question: + return True + + tabular_markers = ( + "calculate", "calculation", "count", "average", "mean", "median", + "minimum", "maximum", "total", "sum", "percentage", "percent", + "rows", "columns", "spreadsheet", "workbook", "worksheet", "sheet", + "csv", "xlsx", "xls", "table", "tabular", "data set", "dataset", + "trend", "group by", "how many", "highest", "lowest", + ) + collective_markers = ( + "both files", "both documents", "all files", "all documents", + "all selected", "each file", "each document", "each source", + "across the files", "across the documents", "across the sources", + "mixed sources", + ) + narrative_markers = ( + "pdf", "docx", "word document", "presentation", "powerpoint", + "paragraph", "section", "policy", "procedure", "contract", + "agreement", "memo", "letter", "narrative", "prose", "report", + ) + + if any(marker in normalized_question for marker in tabular_markers): + return True + if any(marker in normalized_question for marker in collective_markers): + return True + if has_narrative_sources and any( + marker in normalized_question for marker in narrative_markers + ): + return False + if normalized_question in {"summarize", "summary", "summarize the selected sources"}: + return True + if has_narrative_sources: + return False + return True + + +def build_tabular_file_contexts_from_manifest(tabular_sources): + """Build canonical per-file contexts for the existing tabular runner.""" + contexts = [] + seen_source_identities = set() + for raw_source in list(tabular_sources or []): + source = raw_source if isinstance(raw_source, dict) else {} + if ( + source.get("authorization_status") != AUTHORIZATION_STATUS_AUTHORIZED + or source.get("source_kind") != SOURCE_KIND_TABULAR + ): + continue + + document_id = str(source.get("document_id") or "").strip() + file_name = str(source.get("file_name") or "").strip() + scope = str(source.get("scope") or "").strip().lower() + if not document_id or not file_name or scope not in SOURCE_SCOPES: + continue + source_hint = "workspace" if scope == SOURCE_SCOPE_PERSONAL else scope + source_identity = ( + document_id, + source_hint, + str(source.get("scope_id") or "").strip(), + ) + if source_identity in seen_source_identities: + continue + seen_source_identities.add(source_identity) + contexts.append({ + "document_id": document_id, + "file_name": file_name, + "source_hint": source_hint, + "group_id": source.get("group_id"), + "public_workspace_id": source.get("public_workspace_id"), + "conversation_id": source.get("conversation_id"), + "storage_locator": dict(source.get("storage_locator") or {}), + }) + return contexts + + +def classify_source_kind(file_name, document_item=None): + """Classify a resolved source by native capability without reading its content.""" + normalized_file_name = str(file_name or "").strip() + extension = os.path.splitext(normalized_file_name)[1].lower() + if extension in TABULAR_SOURCE_EXTENSIONS: + return SOURCE_KIND_TABULAR + if extension in NARRATIVE_SOURCE_EXTENSIONS: + return SOURCE_KIND_NARRATIVE + + document_item = document_item if isinstance(document_item, dict) else {} + if any( + document_item.get(field_name) + for field_name in ( + "num_file_chunks", + "comparison_text", + "extracted_text", + "vision_analysis", + ) + ): + return SOURCE_KIND_NARRATIVE + return SOURCE_KIND_UNSUPPORTED + + +def _normalize_document_id(requested_source): + if isinstance(requested_source, dict): + requested_source = requested_source.get("document_id") or requested_source.get("id") + return str(requested_source or "").strip() + + +def _normalize_identifier_list(values): + if values is None: + return [] + if isinstance(values, (str, int)): + values = [values] + return [ + normalized_value + for normalized_value in ( + str(value or "").strip() + for value in list(values) + ) + if normalized_value + ] + + +def _normalize_boolean(value, field_name): + if isinstance(value, bool): + return value + if value in (None, ""): + return False + if isinstance(value, int) and value in (0, 1): + return bool(value) + if isinstance(value, str): + normalized_value = value.strip().lower() + if normalized_value in {"true", "1"}: + return True + if normalized_value in {"false", "0"}: + return False + raise ValueError(f"{field_name} must be a boolean") + + +def _normalize_optional_boolean(value, field_name): + if value is None or value == "": + return None + return _normalize_boolean(value, field_name=field_name) + + +def _safe_file_name(file_name): + return str(file_name or "").replace("\\", "/").split("/")[-1].strip() + + +def _unresolved_manifest_entry(document_id): + return { + "document_id": document_id, + "display_name": None, + "file_name": None, + "extension": None, + "source_kind": SOURCE_KIND_UNRESOLVED, + "scope": None, + "scope_id": None, + "group_id": None, + "public_workspace_id": None, + "conversation_id": None, + "source_version": None, + "storage_locator": None, + "authorization_status": AUTHORIZATION_STATUS_UNRESOLVED, + } + + +def _build_authorized_manifest_entry(document_id, user_id, document_context): + if not isinstance(document_context, dict): + return _unresolved_manifest_entry(document_id) + + document_item = document_context.get("document") + if not isinstance(document_item, dict): + return _unresolved_manifest_entry(document_id) + + resolved_document_id = str(document_item.get("id") or "").strip() + if resolved_document_id != document_id: + return _unresolved_manifest_entry(document_id) + + scope = str(document_context.get("scope") or "").strip().lower() + if scope not in SOURCE_SCOPES: + return _unresolved_manifest_entry(document_id) + + group_id = None + public_workspace_id = None + conversation_id = str( + document_context.get("conversation_id") + or document_item.get("conversation_id") + or "" + ).strip() or None + + if scope == SOURCE_SCOPE_PERSONAL: + scope_id = str(document_item.get("user_id") or user_id or "").strip() + if scope_id != str(user_id or "").strip() and not any( + str(shared_entry or "").strip() == f"{user_id},approved" + for shared_entry in document_item.get("shared_user_ids", []) or [] + ): + return _unresolved_manifest_entry(document_id) + elif scope == SOURCE_SCOPE_GROUP: + group_id = str(document_context.get("group_id") or "").strip() or None + scope_id = group_id + document_group_id = str(document_item.get("group_id") or "").strip() + if document_group_id != group_id and not any( + str(shared_entry or "").strip() == f"{group_id},approved" + for shared_entry in document_item.get("shared_group_ids", []) or [] + ): + return _unresolved_manifest_entry(document_id) + elif scope == SOURCE_SCOPE_PUBLIC: + public_workspace_id = str( + document_context.get("public_workspace_id") or "" + ).strip() or None + scope_id = public_workspace_id + else: + scope_id = conversation_id + + if not scope_id: + return _unresolved_manifest_entry(document_id) + + file_name = _safe_file_name( + document_item.get("file_name") + or document_item.get("filename") + or document_item.get("title") + ) + display_name = str(document_item.get("title") or file_name or document_id).strip() + extension = os.path.splitext(file_name)[1].lower() or None + source_version = document_item.get("version") + if source_version is None: + source_version = document_item.get("source_version") + if source_version is not None and not isinstance(source_version, (str, int, float)): + source_version = str(source_version) + + storage_locator = None + if scope != SOURCE_SCOPE_CHAT: + try: + from functions_documents import get_document_blob_storage_info + + blob_container, blob_path = get_document_blob_storage_info( + document_item, + user_id=( + document_item.get("user_id") + if scope == SOURCE_SCOPE_PERSONAL + else None + ), + group_id=( + document_item.get("group_id") + if scope == SOURCE_SCOPE_GROUP + else None + ), + public_workspace_id=( + document_item.get("public_workspace_id") + if scope == SOURCE_SCOPE_PUBLIC + else None + ), + ) + if blob_container and blob_path: + storage_locator = { + "container": str(blob_container), + "blob_path": str(blob_path), + } + except Exception: + storage_locator = None + + return { + "document_id": document_id, + "display_name": display_name, + "file_name": file_name or None, + "extension": extension, + "source_kind": classify_source_kind(file_name, document_item=document_item), + "scope": scope, + "scope_id": scope_id, + "group_id": group_id, + "public_workspace_id": public_workspace_id, + "conversation_id": conversation_id, + "source_version": source_version, + "storage_locator": storage_locator, + "authorization_status": AUTHORIZATION_STATUS_AUTHORIZED, + } + + +def _default_document_context_batch_resolver(**resolver_arguments): + # Imported lazily so this contract module remains usable by startup code and isolated tests. + from functions_search_service import resolve_document_contexts + + resolver_arguments["include_content"] = False + return resolve_document_contexts(**resolver_arguments) + + +def resolve_authorized_source_manifest( + requested_sources, + user_id, + selection_mode=SELECTION_MODE_SELECTED, + conversation_id=None, + active_group_ids=None, + active_public_workspace_ids=None, + doc_scope="all", + context_resolver=None, + cancel_requested=None, + request_correlation_id=None, +): + """Resolve each unique requested ID once into an ordered, authorized manifest.""" + request_correlation_id = normalize_mixed_source_correlation_id( + request_correlation_id + ) + raise_if_mixed_source_cancelled( + cancel_requested, + "manifest", + request_correlation_id=request_correlation_id, + ) + normalized_selection_mode = normalize_selection_mode(selection_mode) + normalized_user_id = str(user_id or "").strip() + if not normalized_user_id: + raise ValueError("user_id is required") + + if isinstance(requested_sources, (str, int, dict)): + requested_source_list = [requested_sources] + else: + requested_source_list = list(requested_sources or []) + if len(requested_source_list) > SOURCE_MANIFEST_MAX_SOURCES: + log_event( + "[MixedSourceManifest] Rejected over-limit source manifest request.", + extra={ + "selection_mode": normalized_selection_mode, + "requested_source_count": len(requested_source_list), + "source_limit": SOURCE_MANIFEST_MAX_SOURCES, + }, + level=logging.WARNING, + ) + raise ValueError( + f"A source manifest supports at most {SOURCE_MANIFEST_MAX_SOURCES} requested sources" + ) + + unique_document_ids = [] + seen_document_ids = set() + duplicate_ids_removed = 0 + for requested_source in requested_source_list: + document_id = _normalize_document_id(requested_source) + if not document_id: + continue + if document_id in seen_document_ids: + duplicate_ids_removed += 1 + continue + seen_document_ids.add(document_id) + unique_document_ids.append(document_id) + + if context_resolver is not None and not callable(context_resolver): + raise TypeError("context_resolver must be callable") + + started_at = time.perf_counter() + manifest = [] + resolution_error_count = 0 + normalized_active_group_ids = _normalize_identifier_list(active_group_ids) + normalized_public_workspace_ids = _normalize_identifier_list( + active_public_workspace_ids + ) + normalized_conversation_id = str(conversation_id or "").strip() or None + normalized_doc_scope = str(doc_scope or "all").strip().lower() + if normalized_doc_scope not in {"all", "personal", "group", "public"}: + raise ValueError("doc_scope must be all, personal, group, or public") + + resolved_contexts = None + if context_resolver is None: + try: + raise_if_mixed_source_cancelled( + cancel_requested, + "manifest", + request_correlation_id=request_correlation_id, + ) + resolved_contexts = _default_document_context_batch_resolver( + document_ids=unique_document_ids, + user_id=normalized_user_id, + doc_scope=normalized_doc_scope, + active_group_ids=normalized_active_group_ids, + active_public_workspace_id=normalized_public_workspace_ids, + conversation_id=normalized_conversation_id, + ) + raise_if_mixed_source_cancelled( + cancel_requested, + "manifest", + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + raise + except Exception: + resolved_contexts = [None] * len(unique_document_ids) + resolution_error_count = len(unique_document_ids) + if ( + not isinstance(resolved_contexts, list) + or len(resolved_contexts) != len(unique_document_ids) + ): + resolved_contexts = [None] * len(unique_document_ids) + resolution_error_count = len(unique_document_ids) + + for document_index, document_id in enumerate(unique_document_ids): + raise_if_mixed_source_cancelled( + cancel_requested, + "manifest", + request_correlation_id=request_correlation_id, + ) + document_context = None + if resolved_contexts is not None: + document_context = resolved_contexts[document_index] + else: + try: + document_context = context_resolver( + document_id=document_id, + user_id=normalized_user_id, + doc_scope=normalized_doc_scope, + active_group_ids=normalized_active_group_ids, + active_public_workspace_id=normalized_public_workspace_ids, + conversation_id=normalized_conversation_id, + ) + raise_if_mixed_source_cancelled( + cancel_requested, + "manifest", + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + raise + except Exception: + resolution_error_count += 1 + if ( + isinstance(document_context, dict) + and normalized_doc_scope != "all" + and str(document_context.get("scope") or "").strip().lower() + != normalized_doc_scope + ): + document_context = None + manifest.append( + _build_authorized_manifest_entry( + document_id, + normalized_user_id, + document_context, + ) + ) + + source_kind_counts = {source_kind: 0 for source_kind in SOURCE_KINDS} + scope_distribution = {scope: 0 for scope in SOURCE_SCOPES} + for entry in manifest: + source_kind_counts[entry["source_kind"]] += 1 + if entry["scope"] in scope_distribution: + scope_distribution[entry["scope"]] += 1 + + duration_ms = round((time.perf_counter() - started_at) * 1000, 3) + resolved_source_count = len(manifest) - source_kind_counts[SOURCE_KIND_UNRESOLVED] + log_event( + "[MixedSourceManifest] Resolved authorized source manifest.", + extra={ + "selection_mode": normalized_selection_mode, + "requested_source_count": len(requested_source_list), + "unique_source_count": len(unique_document_ids), + "resolved_source_count": resolved_source_count, + "tabular_source_count": source_kind_counts[SOURCE_KIND_TABULAR], + "narrative_source_count": source_kind_counts[SOURCE_KIND_NARRATIVE], + "unsupported_source_count": source_kind_counts[SOURCE_KIND_UNSUPPORTED], + "unresolved_or_unauthorized_count": source_kind_counts[SOURCE_KIND_UNRESOLVED], + "duplicate_ids_removed": duplicate_ids_removed, + "resolution_error_count": resolution_error_count, + "scope_distribution": scope_distribution, + "manifest_resolution_duration_ms": duration_ms, + "request_correlation_id": request_correlation_id, + }, + level=logging.INFO, + ) + return manifest + + +def partition_source_manifest(manifest): + """Partition a manifest by capability while preserving order within each cohort.""" + partitions = { + "tabular_sources": [], + "narrative_sources": [], + "unsupported_sources": [], + "unresolved_sources": [], + } + partition_key_by_source_kind = { + SOURCE_KIND_TABULAR: "tabular_sources", + SOURCE_KIND_NARRATIVE: "narrative_sources", + SOURCE_KIND_UNSUPPORTED: "unsupported_sources", + SOURCE_KIND_UNRESOLVED: "unresolved_sources", + } + + for raw_entry in list(manifest or []): + entry = raw_entry if isinstance(raw_entry, dict) else {} + document_id = _normalize_document_id(entry) + if entry.get("authorization_status") != AUTHORIZATION_STATUS_AUTHORIZED: + partitions["unresolved_sources"].append( + _unresolved_manifest_entry(document_id) + ) + continue + + partition_key = partition_key_by_source_kind.get( + entry.get("source_kind"), + "unsupported_sources", + ) + partitions[partition_key].append(entry) + + return partitions + + +def _truncate_utf8(value, max_bytes): + normalized_value = str(value or "") + encoded_value = normalized_value.encode("utf-8") + if len(encoded_value) <= max_bytes: + return normalized_value + if max_bytes <= 3: + return encoded_value[:max_bytes].decode("utf-8", errors="ignore") + return ( + encoded_value[:max_bytes - 3].decode("utf-8", errors="ignore").rstrip() + + "..." + ) + + +def _make_json_safe(value, depth=0): + if value is None or isinstance(value, (bool, int)): + return value, False + if isinstance(value, float): + return (value, False) if math.isfinite(value) else (None, True) + if isinstance(value, str): + bounded_value = _truncate_utf8(value, EVIDENCE_JSON_MAX_STRING_BYTES) + return bounded_value, bounded_value != value + if depth >= EVIDENCE_JSON_MAX_DEPTH: + return _truncate_utf8(str(value), EVIDENCE_JSON_MAX_STRING_BYTES), True + if isinstance(value, dict): + source_items = list(value.items()) + bounded_value = {} + was_truncated = len(source_items) > EVIDENCE_JSON_MAX_COLLECTION_ITEMS + for key, item_value in source_items[:EVIDENCE_JSON_MAX_COLLECTION_ITEMS]: + normalized_key = _truncate_utf8(key, 128) + bounded_item, item_was_truncated = _make_json_safe( + item_value, + depth + 1, + ) + if ( + not isinstance(key, str) + or normalized_key != key + or normalized_key in bounded_value + ): + was_truncated = True + bounded_value[normalized_key] = bounded_item + was_truncated = was_truncated or item_was_truncated + return bounded_value, was_truncated + if isinstance(value, (list, tuple, set)): + source_items = list(value) + bounded_value = [] + was_truncated = len(source_items) > EVIDENCE_JSON_MAX_COLLECTION_ITEMS + for item in source_items[:EVIDENCE_JSON_MAX_COLLECTION_ITEMS]: + bounded_item, item_was_truncated = _make_json_safe(item, depth + 1) + bounded_value.append(bounded_item) + was_truncated = was_truncated or item_was_truncated + return bounded_value, was_truncated + return _truncate_utf8(str(value), EVIDENCE_JSON_MAX_STRING_BYTES), True + + +def _json_size_bytes(value): + return len( + json.dumps( + value, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + ) + + +def _bound_json_value(value, max_bytes): + safe_value, safe_value_was_truncated = _make_json_safe(value) + if _json_size_bytes(safe_value) <= max_bytes: + return safe_value, safe_value_was_truncated + + serialized_preview = json.dumps( + safe_value, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + ) + preview_max_bytes = max(16, max_bytes // 3) + while preview_max_bytes > 0: + bounded_value = { + "truncated": True, + "preview": _truncate_utf8(serialized_preview, preview_max_bytes), + } + if _json_size_bytes(bounded_value) <= max_bytes: + return bounded_value, True + preview_max_bytes //= 2 + return {"truncated": True}, True + + +def _bound_json_list(values): + if values is None: + return [], False + if not isinstance(values, (list, tuple)): + raise ValueError("Evidence collection values must be lists") + + source_values = list(values) + bounded_values = [] + was_truncated = len(source_values) > EVIDENCE_LIST_MAX_ITEMS + for value in source_values[:EVIDENCE_LIST_MAX_ITEMS]: + bounded_value, value_was_truncated = _bound_json_value( + value, + EVIDENCE_ITEM_MAX_BYTES, + ) + bounded_values.append(bounded_value) + was_truncated = was_truncated or value_was_truncated + return bounded_values, was_truncated + + +def deduplicate_mixed_source_references(references, reference_type="citation"): + """Deduplicate structured citations or artifacts while preserving first payloads.""" + normalized_reference_type = str(reference_type or "citation").strip().lower() + if normalized_reference_type not in {"citation", "artifact"}: + raise ValueError("reference_type must be citation or artifact") + + deduplicated = [] + seen_keys = set() + for reference in list(references or []): + if not isinstance(reference, dict): + dedupe_key = ("scalar", str(reference)) + elif normalized_reference_type == "artifact": + identity_value = ( + reference.get("artifact_message_id") + or reference.get("document_id") + or reference.get("export_run_id") + ) + dedupe_key = ( + ("artifact_id", str(identity_value).strip()) + if identity_value + else ( + "artifact_location", + str(reference.get("file_name") or "").strip(), + str(reference.get("output_format") or "").strip().lower(), + str(reference.get("capability") or "").strip().lower(), + ) + ) + else: + identity_value = ( + reference.get("artifact_id") + or reference.get("citation_id") + or reference.get("chunk_id") + ) + plugin_name = str(reference.get("plugin_name") or "").strip() + function_name = str(reference.get("function_name") or "").strip() + tool_name = str(reference.get("tool_name") or "").strip() + if identity_value: + dedupe_key = ("citation_id", str(identity_value).strip()) + elif plugin_name or function_name or tool_name: + tool_arguments = ( + reference.get("function_arguments") + if reference.get("function_arguments") is not None + else reference.get("parameters") + ) + dedupe_key = ( + "tool_citation", + plugin_name, + function_name, + tool_name, + json.dumps( + tool_arguments, + allow_nan=False, + default=str, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ), + ) + else: + dedupe_key = ( + "citation_location", + str(reference.get("document_id") or "").strip(), + str(reference.get("file_name") or "").strip(), + str(reference.get("page_number") or "").strip(), + str(reference.get("chunk_sequence") or "").strip(), + str(reference.get("sheet_name") or "").strip(), + ) + if dedupe_key in seen_keys: + continue + seen_keys.add(dedupe_key) + deduplicated.append(reference) + return deduplicated + + +def _build_truncated_coverage(coverage, coverage_was_truncated=False): + normalized_coverage = dict(coverage or {}) + normalized_coverage["evidence_envelope_truncated"] = True + if coverage_was_truncated: + normalized_coverage["coverage_truncated"] = True + + bounded_coverage, additional_truncation = _bound_json_value( + normalized_coverage, + EVIDENCE_COVERAGE_MAX_BYTES, + ) + if additional_truncation: + return { + "evidence_envelope_truncated": True, + "coverage_truncated": True, + } + return bounded_coverage + + +def build_evidence_envelope( + document_id, + source_kind, + engine, + status, + summary="", + evidence=None, + citations=None, + generated_artifacts=None, + coverage=None, + error=None, +): + """Build a bounded, JSON-safe evidence envelope for later synthesis phases.""" + normalized_document_id = str(document_id or "").strip() + if not normalized_document_id: + raise ValueError("document_id is required") + + normalized_source_kind = str(source_kind or "").strip().lower() + if normalized_source_kind not in {SOURCE_KIND_TABULAR, SOURCE_KIND_NARRATIVE}: + raise ValueError("Evidence source_kind must be tabular or narrative") + + normalized_engine = str(engine or "").strip().lower() + if normalized_engine not in EVIDENCE_ENGINES: + raise ValueError(f"Unsupported evidence engine: {normalized_engine}") + + normalized_status = str(status or "").strip().lower() + if normalized_status not in EVIDENCE_STATUSES: + raise ValueError(f"Unsupported evidence status: {normalized_status}") + + if coverage is not None and not isinstance(coverage, dict): + raise ValueError("coverage must be a dictionary") + + bounded_evidence, evidence_was_truncated = _bound_json_list(evidence) + bounded_citations, citations_were_truncated = _bound_json_list( + deduplicate_mixed_source_references(citations, reference_type="citation") + ) + bounded_artifacts, artifacts_were_truncated = _bound_json_list( + deduplicate_mixed_source_references(generated_artifacts, reference_type="artifact") + ) + normalized_summary = _truncate_utf8(summary, EVIDENCE_SUMMARY_MAX_BYTES) + normalized_error = ( + _truncate_utf8(error, EVIDENCE_ERROR_MAX_BYTES) + if error is not None + else None + ) + bounds_applied = bool( + evidence_was_truncated + or citations_were_truncated + or artifacts_were_truncated + or len(str(summary or "").encode("utf-8")) > EVIDENCE_SUMMARY_MAX_BYTES + or ( + error is not None + and len(str(error).encode("utf-8")) > EVIDENCE_ERROR_MAX_BYTES + ) + ) + normalized_coverage = dict(coverage or {}) + if bounds_applied: + normalized_coverage["evidence_envelope_truncated"] = True + bounded_coverage, coverage_was_truncated = _bound_json_value( + normalized_coverage, + EVIDENCE_COVERAGE_MAX_BYTES, + ) + if coverage_was_truncated: + bounded_coverage = _build_truncated_coverage( + {}, + coverage_was_truncated=True, + ) + + envelope = { + "document_id": normalized_document_id, + "source_kind": normalized_source_kind, + "engine": normalized_engine, + "status": normalized_status, + "summary": normalized_summary, + "evidence": bounded_evidence, + "citations": bounded_citations, + "generated_artifacts": bounded_artifacts, + "coverage": bounded_coverage, + "error": normalized_error, + } + + while _json_size_bytes(envelope) > EVIDENCE_ENVELOPE_MAX_BYTES: + candidate_field = max( + ("evidence", "citations", "generated_artifacts"), + key=lambda field_name: len(envelope[field_name]), + ) + if envelope[candidate_field]: + envelope[candidate_field].pop() + envelope["coverage"] = _build_truncated_coverage( + envelope["coverage"], + ) + continue + envelope["summary"] = _truncate_utf8( + envelope["summary"], + max(128, len(envelope["summary"].encode("utf-8")) // 2), + ) + if len(envelope["summary"].encode("utf-8")) <= 128: + raise ValueError("Unable to bound evidence envelope") + + return envelope + + +def serialize_evidence_envelope(envelope): + """Validate and serialize a bounded evidence envelope.""" + if not isinstance(envelope, dict): + raise ValueError("Evidence envelope must be a dictionary") + + required_fields = { + "document_id", + "source_kind", + "engine", + "status", + "summary", + "evidence", + "citations", + "generated_artifacts", + "coverage", + "error", + } + if set(envelope) != required_fields: + raise ValueError("Evidence envelope fields do not match the contract") + + bounded_envelope = build_evidence_envelope( + document_id=envelope["document_id"], + source_kind=envelope["source_kind"], + engine=envelope["engine"], + status=envelope["status"], + summary=envelope["summary"], + evidence=envelope["evidence"], + citations=envelope["citations"], + generated_artifacts=envelope["generated_artifacts"], + coverage=envelope["coverage"], + error=envelope["error"], + ) + + serialized_envelope = json.dumps( + bounded_envelope, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + if len(serialized_envelope.encode("utf-8")) > EVIDENCE_ENVELOPE_MAX_BYTES: + raise ValueError("Evidence envelope exceeds its serialized size bound") + return serialized_envelope + + +def build_narrative_evidence_envelopes( + narrative_sources, + search_results, + selection_mode, +): + """Normalize bounded narrative search results into one envelope per source.""" + normalized_selection_mode = normalize_selection_mode( + selection_mode, + default=SELECTION_MODE_RELEVANCE, + ) + results_by_document_id = {} + for raw_result in list(search_results or []): + result = raw_result if isinstance(raw_result, dict) else {} + document_id = str(result.get("document_id") or "").strip() + if not document_id: + continue + results_by_document_id.setdefault(document_id, []).append(result) + + envelopes = [] + for source in list(narrative_sources or []): + source = source if isinstance(source, dict) else {} + document_id = str(source.get("document_id") or "").strip() + if not document_id: + continue + source_results = results_by_document_id.get(document_id, []) + evidence = [] + citations = [] + for result in source_results: + evidence.append({ + "chunk_text": result.get("chunk_text"), + "page_number": result.get("page_number"), + "chunk_sequence": result.get("chunk_sequence"), + "score": result.get("score"), + }) + citations.append({ + "citation_id": result.get("id") or result.get("chunk_id"), + "page_number": result.get("page_number"), + "chunk_sequence": result.get("chunk_sequence"), + }) + + result_count = len(source_results) + envelopes.append(build_evidence_envelope( + document_id=document_id, + source_kind=SOURCE_KIND_NARRATIVE, + engine=EVIDENCE_ENGINE_HYBRID_SEARCH, + status=( + EVIDENCE_STATUS_COMPLETED + if result_count + else EVIDENCE_STATUS_PARTIAL + ), + summary=( + f"Retrieved {result_count} bounded narrative excerpt(s)." + if result_count + else "No relevant narrative excerpts were returned." + ), + evidence=evidence, + citations=citations, + coverage={ + "selection_mode": normalized_selection_mode, + "terminal": True, + "result_count": result_count, + }, + error=( + None + if result_count + else "Narrative retrieval returned no relevant excerpts." + ), + )) + return envelopes + + +def build_failed_narrative_evidence_envelopes( + narrative_sources, + selection_mode, + reason="narrative_retrieval_failed", +): + """Build one scrubbed terminal failure for every narrative source in a failed cohort.""" + normalized_selection_mode = normalize_selection_mode( + selection_mode, + default=SELECTION_MODE_RELEVANCE, + ) + normalized_reason = _truncate_utf8( + str(reason or "narrative_retrieval_failed").strip().lower(), + MIXED_SOURCE_TERMINAL_REASON_MAX_BYTES, + ) + envelopes = [] + for source in list(narrative_sources or []): + source = source if isinstance(source, dict) else {} + document_id = str(source.get("document_id") or "").strip() + if not document_id: + continue + envelopes.append(build_evidence_envelope( + document_id=document_id, + source_kind=SOURCE_KIND_NARRATIVE, + engine=EVIDENCE_ENGINE_HYBRID_SEARCH, + status=EVIDENCE_STATUS_FAILED, + summary="Narrative evidence could not be retrieved for this source.", + coverage={ + "selection_mode": normalized_selection_mode, + "terminal": True, + "reason": normalized_reason, + }, + error="Narrative retrieval could not be completed.", + )) + return envelopes + + +def execute_tabular_evidence_sources( + tabular_sources, + execute_source, + selection_mode, + execute=True, + cancel_requested=None, + request_correlation_id=None, +): + """Execute the existing tabular runner once per source and require terminal coverage.""" + normalized_selection_mode = normalize_selection_mode( + selection_mode, + default=SELECTION_MODE_RELEVANCE, + ) + request_correlation_id = normalize_mixed_source_correlation_id( + request_correlation_id + ) + if execute and not callable(execute_source): + raise TypeError("execute_source must be callable") + + envelopes = [] + completed_count = 0 + failed_count = 0 + skipped_count = 0 + for raw_source in list(tabular_sources or []): + raise_if_mixed_source_cancelled( + cancel_requested, + "tabular", + request_correlation_id=request_correlation_id, + ) + source = raw_source if isinstance(raw_source, dict) else {} + document_id = str(source.get("document_id") or "").strip() + if not document_id: + continue + + if not execute: + skipped_count += 1 + envelopes.append(build_evidence_envelope( + document_id=document_id, + source_kind=SOURCE_KIND_TABULAR, + engine=EVIDENCE_ENGINE_TABULAR_TOOLS, + status=EVIDENCE_STATUS_SKIPPED, + summary="Tabular processing was not needed for this narrative-only request.", + coverage={ + "selection_mode": normalized_selection_mode, + "terminal": True, + "reason": "narrative_only_request", + }, + )) + continue + + try: + raw_result = execute_source(source) + raise_if_mixed_source_cancelled( + cancel_requested, + "tabular", + request_correlation_id=request_correlation_id, + ) + result = raw_result if isinstance(raw_result, dict) else {} + summary = str(result.get("summary") or "").strip() + if not summary: + raise ValueError("Tabular execution returned no bounded summary") + completed_count += 1 + envelopes.append(build_evidence_envelope( + document_id=document_id, + source_kind=SOURCE_KIND_TABULAR, + engine=EVIDENCE_ENGINE_TABULAR_TOOLS, + status=EVIDENCE_STATUS_COMPLETED, + summary=summary, + evidence=result.get("evidence"), + citations=result.get("citations"), + generated_artifacts=result.get("generated_artifacts"), + coverage={ + "selection_mode": normalized_selection_mode, + "terminal": True, + **dict(result.get("coverage") or {}), + }, + )) + except MixedSourceCancellationError: + raise + except Exception: + failed_count += 1 + envelopes.append(build_evidence_envelope( + document_id=document_id, + source_kind=SOURCE_KIND_TABULAR, + engine=EVIDENCE_ENGINE_TABULAR_TOOLS, + status=EVIDENCE_STATUS_FAILED, + summary="Tabular evidence could not be completed for this source.", + coverage={ + "selection_mode": normalized_selection_mode, + "terminal": True, + }, + error="Tabular evidence could not be completed.", + )) + + log_event( + "[MixedSourceChatSearch] Tabular source execution reached terminal coverage.", + extra={ + "selection_mode": normalized_selection_mode, + "tabular_candidate_count": len(list(tabular_sources or [])), + "tabular_completed_count": completed_count, + "tabular_failed_count": failed_count, + "tabular_skipped_count": skipped_count, + "request_correlation_id": request_correlation_id, + }, + level=logging.INFO, + ) + return envelopes + + +def _get_terminal_reason(status, source, envelope=None): + """Return one bounded, non-sensitive reason for a terminal non-success state.""" + envelope = envelope if isinstance(envelope, dict) else {} + coverage = envelope.get("coverage") if isinstance(envelope.get("coverage"), dict) else {} + explicit_reason = str(coverage.get("reason") or "").strip().lower() + if explicit_reason: + return _truncate_utf8(explicit_reason, MIXED_SOURCE_TERMINAL_REASON_MAX_BYTES) + if source.get("authorization_status") != AUTHORIZATION_STATUS_AUTHORIZED: + return "source_unavailable" + if source.get("source_kind") == SOURCE_KIND_UNSUPPORTED: + return "unsupported_source" + if status == EVIDENCE_STATUS_SKIPPED: + return "bounded_policy_skip" + if status == EVIDENCE_STATUS_PARTIAL: + return "incomplete_native_coverage" + if status == EVIDENCE_STATUS_FAILED: + return "native_execution_failed" + return None + + +def build_terminal_coverage_ledger( + manifest, + evidence_envelopes, + max_handoff_envelopes=MIXED_SOURCE_HANDOFF_MAX_ENVELOPES, +): + """Align exactly one terminal state and bounded evidence item to each manifest source.""" + manifest_entries = [entry for entry in list(manifest or []) if isinstance(entry, dict)] + raw_envelopes = [ + envelope + for envelope in list(evidence_envelopes or []) + if isinstance(envelope, dict) + ] + manifest_document_ids = { + str(entry.get("document_id") or "").strip() + for entry in manifest_entries + if str(entry.get("document_id") or "").strip() + } + envelopes_by_document_id = {} + unexpected_evidence_count = 0 + for envelope in raw_envelopes: + document_id = str(envelope.get("document_id") or "").strip() + if not document_id or document_id not in manifest_document_ids: + unexpected_evidence_count += 1 + continue + envelopes_by_document_id.setdefault(document_id, []).append(envelope) + + ledger_entries = [] + aligned_envelopes = [] + status_counts = {status: 0 for status in EVIDENCE_STATUSES} + missing_coverage_violation_count = 0 + duplicate_evidence_count = 0 + evidence_omitted_count = 0 + + for request_order, source in enumerate(manifest_entries): + document_id = str(source.get("document_id") or "").strip() + source_kind = str(source.get("source_kind") or SOURCE_KIND_UNRESOLVED).strip().lower() + source_envelopes = envelopes_by_document_id.get(document_id, []) + envelope = source_envelopes[0] if len(source_envelopes) == 1 else None + reason = None + + if len(source_envelopes) > 1: + status = EVIDENCE_STATUS_FAILED + reason = "duplicate_terminal_evidence" + duplicate_evidence_count += len(source_envelopes) - 1 + missing_coverage_violation_count += 1 + elif source.get("authorization_status") != AUTHORIZATION_STATUS_AUTHORIZED: + status = EVIDENCE_STATUS_FAILED + reason = "source_unavailable" + elif source_kind == SOURCE_KIND_UNSUPPORTED: + status = EVIDENCE_STATUS_SKIPPED + reason = "unsupported_source" + elif envelope is None: + status = EVIDENCE_STATUS_FAILED + reason = "missing_terminal_evidence" + missing_coverage_violation_count += 1 + elif str(envelope.get("source_kind") or "").strip().lower() != source_kind: + status = EVIDENCE_STATUS_FAILED + reason = "evidence_identity_mismatch" + missing_coverage_violation_count += 1 + envelope = None + else: + status = str(envelope.get("status") or "").strip().lower() + if status not in EVIDENCE_STATUSES: + status = EVIDENCE_STATUS_FAILED + reason = "invalid_terminal_status" + missing_coverage_violation_count += 1 + envelope = None + else: + reason = _get_terminal_reason(status, source, envelope=envelope) + + status_counts[status] += 1 + handoff_included = False + if envelope is not None: + if len(aligned_envelopes) < max(0, int(max_handoff_envelopes)): + aligned_envelopes.append(envelope) + handoff_included = True + else: + evidence_omitted_count += 1 + + source_role = str( + source.get("comparison_role") + or source.get("source_role") + or source.get("role") + or "selected" + ).strip().lower() or "selected" + ledger_entry = { + "document_id": document_id, + "scope": source.get("scope"), + "scope_id": source.get("scope_id"), + "source_version": source.get("source_version"), + "source_kind": source_kind, + "role": source_role, + "request_order": request_order, + "status": status, + "reason": reason, + "handoff_included": handoff_included, + } + ledger_entries.append(ledger_entry) + + partial_coverage = bool( + status_counts[EVIDENCE_STATUS_PARTIAL] + or status_counts[EVIDENCE_STATUS_FAILED] + or status_counts[EVIDENCE_STATUS_SKIPPED] + or missing_coverage_violation_count + or duplicate_evidence_count + or unexpected_evidence_count + or evidence_omitted_count + ) + return { + "entries": ledger_entries, + "evidence_envelopes": aligned_envelopes, + "requested_source_count": len(manifest_entries), + "completed_source_count": status_counts[EVIDENCE_STATUS_COMPLETED], + "partial_source_count": status_counts[EVIDENCE_STATUS_PARTIAL], + "failed_source_count": status_counts[EVIDENCE_STATUS_FAILED], + "skipped_source_count": status_counts[EVIDENCE_STATUS_SKIPPED], + "successful_source_count": ( + status_counts[EVIDENCE_STATUS_COMPLETED] + + status_counts[EVIDENCE_STATUS_PARTIAL] + ), + "missing_coverage_violation_count": missing_coverage_violation_count, + "duplicate_evidence_count": duplicate_evidence_count, + "unexpected_evidence_count": unexpected_evidence_count, + "evidence_omitted_count": evidence_omitted_count, + "partial_coverage": partial_coverage, + } + + +def evaluate_mixed_source_mode_outcome(mode, coverage_ledger): + """Apply the Phase 6 terminal failure policy to one aggregate-only ledger.""" + normalized_mode = str(mode or "").strip().lower() + if normalized_mode not in MIXED_SOURCE_MODES: + raise ValueError(f"Unsupported mixed-source mode: {normalized_mode}") + + coverage_ledger = coverage_ledger if isinstance(coverage_ledger, dict) else {} + entries = [ + entry + for entry in list(coverage_ledger.get("entries") or []) + if isinstance(entry, dict) + ] + successful_statuses = {EVIDENCE_STATUS_COMPLETED, EVIDENCE_STATUS_PARTIAL} + successful_source_count = sum( + str(entry.get("status") or "").strip().lower() in successful_statuses + for entry in entries + ) + partial_coverage = bool(coverage_ledger.get("partial_coverage")) + should_reduce = successful_source_count > 0 + reason = None + + if normalized_mode == "compare": + source_entry = next( + ( + entry + for entry in entries + if str(entry.get("role") or "").strip().lower() in {"left", "source"} + ), + entries[0] if entries else None, + ) + source_succeeded = bool( + source_entry + and str(source_entry.get("status") or "").strip().lower() in successful_statuses + ) + target_entries = [entry for entry in entries if entry is not source_entry] + successful_target_count = sum( + str(entry.get("status") or "").strip().lower() in successful_statuses + for entry in target_entries + ) + should_reduce = source_succeeded and successful_target_count > 0 + if not source_succeeded: + reason = "source_preparation_failed" + elif not successful_target_count: + reason = "no_target_prepared" + partial_coverage = partial_coverage or successful_target_count < len(target_entries) + + if not should_reduce: + status = EVIDENCE_STATUS_FAILED + reason = reason or "no_successful_source" + elif partial_coverage: + status = EVIDENCE_STATUS_PARTIAL + else: + status = EVIDENCE_STATUS_COMPLETED + + return { + "mode": normalized_mode, + "status": status, + "should_reduce": should_reduce, + "successful_source_count": successful_source_count, + "partial_coverage": partial_coverage, + "reason": reason, + } + + +def compare_reauthorized_source_manifests(execution_manifest, fresh_manifest): + """Return aggregate canonical-identity and version differences for authorized sources.""" + fresh_by_document_id = { + str(source.get("document_id") or "").strip(): source + for source in list(fresh_manifest or []) + if isinstance(source, dict) and str(source.get("document_id") or "").strip() + } + authorization_failure_count = 0 + source_version_changed_count = 0 + for source in list(execution_manifest or []): + if ( + not isinstance(source, dict) + or source.get("authorization_status") != AUTHORIZATION_STATUS_AUTHORIZED + ): + continue + document_id = str(source.get("document_id") or "").strip() + fresh_source = fresh_by_document_id.get(document_id) or {} + same_canonical_identity = ( + fresh_source.get("authorization_status") == AUTHORIZATION_STATUS_AUTHORIZED + and str(fresh_source.get("scope") or "").strip().lower() + == str(source.get("scope") or "").strip().lower() + and str(fresh_source.get("scope_id") or "").strip() + == str(source.get("scope_id") or "").strip() + ) + if not same_canonical_identity: + authorization_failure_count += 1 + continue + prior_version = source.get("source_version") + fresh_version = fresh_source.get("source_version") + if prior_version != fresh_version: + source_version_changed_count += 1 + return { + "authorization_failure_count": authorization_failure_count, + "source_version_changed_count": source_version_changed_count, + } + + +def build_mixed_source_evidence_handoff( + manifest, + evidence_envelopes, + selection_mode, + mode=None, + telemetry_settings=None, + request_correlation_id=None, +): + """Build one bounded synthesis handoff from Phase 1 evidence envelopes.""" + normalized_selection_mode = normalize_selection_mode( + selection_mode, + default=SELECTION_MODE_RELEVANCE, + ) + manifest_entries = [entry for entry in list(manifest or []) if isinstance(entry, dict)] + ledger = build_terminal_coverage_ledger(manifest_entries, evidence_envelopes) + envelopes = list(ledger["evidence_envelopes"]) + source_coverage = [] + for source_index, (entry, ledger_entry) in enumerate( + zip(manifest_entries, ledger["entries"]), + start=1, + ): + if entry.get("authorization_status") != AUTHORIZATION_STATUS_AUTHORIZED: + source_label = f"Unavailable selected source {source_index}" + elif entry.get("source_kind") == SOURCE_KIND_UNSUPPORTED: + source_label = str(entry.get("display_name") or f"Unsupported source {source_index}") + else: + source_label = str(entry.get("display_name") or f"Source {source_index}") + source_coverage.append({ + "source": _truncate_utf8(source_label, 128), + "source_kind": ledger_entry.get("source_kind"), + "status": ledger_entry.get("status"), + "reason": ledger_entry.get("reason"), + }) + + coverage = { + "selection_mode": normalized_selection_mode, + **{ + key: value + for key, value in ledger.items() + if key not in {"entries", "evidence_envelopes"} + }, + "sources": source_coverage, + "terminal_ledger": ledger["entries"], + } + prompt_terminal_ledger = [] + for entry, ledger_entry in zip(manifest_entries, ledger["entries"]): + prompt_entry = dict(ledger_entry) + if entry.get("authorization_status") != AUTHORIZATION_STATUS_AUTHORIZED: + prompt_entry.update({ + "document_id": None, + "scope": None, + "scope_id": None, + "source_version": None, + }) + prompt_terminal_ledger.append(prompt_entry) + prompt_coverage = dict(coverage) + prompt_coverage["terminal_ledger"] = prompt_terminal_ledger + payload = { + "coverage": prompt_coverage, + "evidence_envelopes": envelopes, + } + serialized_payload = json.dumps( + payload, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + ) + if len(serialized_payload.encode("utf-8")) > MIXED_SOURCE_HANDOFF_MAX_BYTES: + payload["evidence_envelopes"] = [ + { + "document_id": envelope.get("document_id"), + "source_kind": envelope.get("source_kind"), + "engine": envelope.get("engine"), + "status": envelope.get("status"), + "summary": _truncate_utf8(envelope.get("summary"), 512), + "coverage": { + "selection_mode": (envelope.get("coverage") or {}).get("selection_mode"), + "terminal": bool((envelope.get("coverage") or {}).get("terminal")), + "result_count": (envelope.get("coverage") or {}).get("result_count"), + "tool_call_count": (envelope.get("coverage") or {}).get("tool_call_count"), + }, + } + for envelope in envelopes + ] + payload["coverage"]["handoff_compacted"] = True + payload["coverage"]["partial_coverage"] = True + payload["coverage"]["evidence_compacted_count"] = len(envelopes) + coverage["handoff_compacted"] = True + coverage["partial_coverage"] = True + coverage["evidence_compacted_count"] = len(envelopes) + serialized_payload = json.dumps( + payload, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + ) + if len(serialized_payload.encode("utf-8")) > MIXED_SOURCE_HANDOFF_MAX_BYTES: + raise ValueError("Mixed-source evidence handoff exceeds its size bound") + + partial_coverage_instruction = ( + "State clearly that source coverage was partial and identify unavailable authorized source labels." + if coverage["partial_coverage"] + else "Do not claim that selected sources were omitted." + ) + if mode: + emit_mixed_source_coverage_telemetry( + telemetry_settings, + mode, + coverage, + request_correlation_id=request_correlation_id, + ) + return { + "role": "system", + "content": ( + "Use the mixed-source evidence handoff below with the other bounded narrative excerpts " + "and computed tabular results. Synthesize one answer. Preserve narrative source citations " + "and tabular tool citations; do not convert computed table facts into unsupported narrative claims. " + "When selection_mode is selected, current selected-source evidence supersedes prior document " + "grounding; do not use prior source claims to fill missing current coverage. " + f"{partial_coverage_instruction}\n\n{serialized_payload}" + ), + "mixed_source_coverage": coverage, + "evidence_envelopes": list(payload["evidence_envelopes"]), + } \ No newline at end of file diff --git a/application/single_app/functions_search.py b/application/single_app/functions_search.py index 011e2c01a..8b8bc38da 100644 --- a/application/single_app/functions_search.py +++ b/application/single_app/functions_search.py @@ -22,7 +22,7 @@ logger = logging.getLogger(__name__) -SEARCH_DEFAULT_TOP_N = 12 +SEARCH_DEFAULT_TOP_N = 50 SEARCH_MAX_TOP_N = 500 VALID_SEARCH_SCOPES = {"all", "personal", "group", "public"} BASE_SEARCH_SELECT_FIELDS = [ @@ -259,7 +259,7 @@ def _build_odata_any_eq(collection_field: str, iterator_name: str, value: Any) - escaped_value = _escape_odata_literal(value) return f"{collection_field}/any({iterator_name}: {iterator_name} eq '{escaped_value}')" -def hybrid_search(query, user_id, document_id=None, document_ids=None, top_n=12, doc_scope="all", active_group_id=None, active_group_ids=None, active_public_workspace_id=None, enable_file_sharing=True, tags_filter=None, document_filter_mode="intersection", enforce_public_workspace_visibility=True): +def hybrid_search(query, user_id, document_id=None, document_ids=None, top_n=50, doc_scope="all", active_group_id=None, active_group_ids=None, active_public_workspace_id=None, enable_file_sharing=True, tags_filter=None, document_filter_mode="intersection", enforce_public_workspace_visibility=True): """ Hybrid search that queries the user doc index, group doc index, or public doc index depending on doc type. diff --git a/application/single_app/functions_search_service.py b/application/single_app/functions_search_service.py index 58b9793b0..2dc18b832 100644 --- a/application/single_app/functions_search_service.py +++ b/application/single_app/functions_search_service.py @@ -12,7 +12,12 @@ from azure.cosmos.exceptions import CosmosResourceNotFoundError from openai import AzureOpenAI -from config import CLIENTS, cognitive_services_scope, cosmos_messages_container +from config import ( + CLIENTS, + cognitive_services_scope, + cosmos_conversations_container, + cosmos_messages_container, +) from functions_appinsights import log_event from functions_debug import debug_print from functions_documents import get_document_record, get_ordered_document_chunks @@ -40,6 +45,9 @@ SUMMARY_MAX_WINDOW_SIZE = 50 CHAT_UPLOAD_CHUNK_WORD_SIZE = 400 CHAT_UPLOAD_CHUNK_WORD_OVERLAP = 40 +MIXED_SOURCE_TABULAR_CANDIDATE_TOP_N = 50 +MIXED_SOURCE_TABULAR_CANDIDATE_LIMIT = 100 +MIXED_SOURCE_TABULAR_EXTENSIONS = frozenset({".csv", ".xls", ".xlsx", ".xlsm"}) def _coerce_positive_int(value, default_value, min_value=1, max_value=None): @@ -238,34 +246,101 @@ def _build_chat_upload_chunks(text_content, max_chunks=None): return chunks -def _resolve_chat_upload_context(document_id, conversation_id=None): +def _authorize_chat_upload_conversation(user_id, conversation_id): + normalized_user_id = str(user_id or "").strip() + normalized_conversation_id = str(conversation_id or "").strip() + if not normalized_user_id or not normalized_conversation_id: + return False + + try: + conversation_item = cosmos_conversations_container.read_item( + item=normalized_conversation_id, + partition_key=normalized_conversation_id, + ) + except CosmosResourceNotFoundError: + return False + except Exception as exc: + log_event( + "[SearchService] Failed to authorize chat upload conversation.", + extra={"exception_type": type(exc).__name__}, + level=logging.WARNING, + exceptionTraceback=True, + debug_only=True, + ) + return False + + return str(conversation_item.get("user_id") or "").strip() == normalized_user_id + + +def _resolve_chat_upload_context( + document_id, + user_id=None, + conversation_id=None, + include_content=True, + authorization_prechecked=False, +): normalized_conversation_id = str(conversation_id or "").strip() normalized_document_id = str(document_id or "").strip() if not normalized_conversation_id or not normalized_document_id: return None + if ( + not authorization_prechecked + and not _authorize_chat_upload_conversation(user_id, normalized_conversation_id) + ): + return None try: - message_item = cosmos_messages_container.read_item( - item=normalized_document_id, - partition_key=normalized_conversation_id, - ) + if include_content: + message_item = cosmos_messages_container.read_item( + item=normalized_document_id, + partition_key=normalized_conversation_id, + ) + else: + metadata_items = list(cosmos_messages_container.query_items( + query=""" + SELECT TOP 1 + c.id, + c.role, + c.filename, + c.title, + c.version, + c.metadata.is_user_upload AS is_user_upload, + c.metadata.is_generated_chat_artifact AS is_generated_chat_artifact, + c.metadata.generated_artifact_capability AS generated_artifact_capability, + c.metadata.generated_artifact_output_format AS generated_artifact_output_format + FROM c + WHERE c.id = @document_id + """, + parameters=[ + {"name": "@document_id", "value": normalized_document_id}, + ], + partition_key=normalized_conversation_id, + )) + if not metadata_items: + return None + message_item = metadata_items[0] except CosmosResourceNotFoundError: return None except Exception as exc: - debug_print( - "[SearchService] Failed to resolve chat upload context | " - f"document_id={normalized_document_id} | conversation_id={normalized_conversation_id} | error={exc}" + log_event( + "[SearchService] Failed to resolve authorized chat upload context.", + extra={"exception_type": type(exc).__name__}, + level=logging.WARNING, + exceptionTraceback=True, + debug_only=True, ) return None role_name = str(message_item.get("role") or "").strip().lower() metadata = message_item.get("metadata", {}) or {} - is_uploaded_image = role_name == "image" and bool((message_item.get("metadata") or {}).get("is_user_upload")) + is_uploaded_image = role_name == "image" and bool( + metadata.get("is_user_upload") or message_item.get("is_user_upload") + ) if role_name not in {"file", "image"} or (role_name == "image" and not is_uploaded_image): return None - comparison_text = _coerce_chat_upload_text(message_item) - if not comparison_text: + comparison_text = _coerce_chat_upload_text(message_item) if include_content else "" + if include_content and not comparison_text: return None message_title = str(message_item.get("filename") or message_item.get("title") or normalized_document_id).strip() or normalized_document_id @@ -275,11 +350,24 @@ def _resolve_chat_upload_context(document_id, conversation_id=None): "title": message_title, "conversation_id": normalized_conversation_id, "source_type": "chat_upload", - "source_subtype": "generated_chat_artifact" if metadata.get("is_generated_chat_artifact") else "chat_upload", - "artifact_capability": str(metadata.get("generated_artifact_capability") or "").strip().lower() or None, - "artifact_output_format": str(metadata.get("generated_artifact_output_format") or "").strip().lower() or None, - "comparison_text": comparison_text, + "source_subtype": "generated_chat_artifact" if ( + metadata.get("is_generated_chat_artifact") + or message_item.get("is_generated_chat_artifact") + ) else "chat_upload", + "artifact_capability": str( + metadata.get("generated_artifact_capability") + or message_item.get("generated_artifact_capability") + or "" + ).strip().lower() or None, + "artifact_output_format": str( + metadata.get("generated_artifact_output_format") + or message_item.get("generated_artifact_output_format") + or "" + ).strip().lower() or None, + "version": message_item.get("version"), } + if include_content: + resolved_document["comparison_text"] = comparison_text return { "scope": "chat", "group_id": None, @@ -289,6 +377,59 @@ def _resolve_chat_upload_context(document_id, conversation_id=None): } +def _resolve_personal_document_context(document_id, user_id): + personal_document = get_document_record( + user_id=user_id, + document_id=document_id, + ) + if not personal_document: + return None + return { + "scope": "personal", + "group_id": None, + "public_workspace_id": None, + "document": personal_document, + } + + +def _resolve_group_document_context(document_id, user_id, authorized_group_ids): + for group_id in authorized_group_ids or []: + group_document = get_document_record( + user_id=user_id, + document_id=document_id, + group_id=group_id, + ) + if group_document: + return { + "scope": "group", + "group_id": group_id, + "public_workspace_id": None, + "document": group_document, + } + return None + + +def _resolve_public_document_context( + document_id, + user_id, + authorized_public_workspace_ids, +): + for public_workspace_id in authorized_public_workspace_ids or []: + public_document = get_document_record( + user_id=user_id, + document_id=document_id, + public_workspace_id=public_workspace_id, + ) + if public_document: + return { + "scope": "public", + "group_id": None, + "public_workspace_id": public_workspace_id, + "document": public_document, + } + return None + + def resolve_document_context( document_id, user_id, @@ -296,59 +437,45 @@ def resolve_document_context( active_group_ids=None, active_public_workspace_id=None, conversation_id=None, + include_content=True, ): normalized_scope = normalize_search_scope(doc_scope) if normalized_scope in ("all", "personal"): - personal_document = get_document_record(user_id=user_id, document_id=document_id) - if personal_document: - return { - "scope": "personal", - "group_id": None, - "public_workspace_id": None, - "document": personal_document, - } + personal_context = _resolve_personal_document_context(document_id, user_id) + if personal_context: + return personal_context if normalized_scope in ("all", "group"): - for group_id in _resolve_active_group_ids( + group_context = _resolve_group_document_context( + document_id, user_id, - active_group_ids=active_group_ids, - fallback_to_memberships=True, - ): - group_document = get_document_record( - user_id=user_id, - document_id=document_id, - group_id=group_id, - ) - if group_document: - return { - "scope": "group", - "group_id": group_id, - "public_workspace_id": None, - "document": group_document, - } + _resolve_active_group_ids( + user_id, + active_group_ids=active_group_ids, + fallback_to_memberships=True, + ), + ) + if group_context: + return group_context if normalized_scope in ("all", "public"): - for public_workspace_id in _resolve_public_workspace_ids( + public_context = _resolve_public_document_context( + document_id, user_id, - active_public_workspace_id=active_public_workspace_id, - ): - public_document = get_document_record( - user_id=user_id, - document_id=document_id, - public_workspace_id=public_workspace_id, - ) - if public_document: - return { - "scope": "public", - "group_id": None, - "public_workspace_id": public_workspace_id, - "document": public_document, - } + _resolve_public_workspace_ids( + user_id, + active_public_workspace_id=active_public_workspace_id, + ), + ) + if public_context: + return public_context chat_upload_context = _resolve_chat_upload_context( document_id=document_id, + user_id=user_id, conversation_id=conversation_id, + include_content=include_content, ) if chat_upload_context: return chat_upload_context @@ -356,6 +483,68 @@ def resolve_document_context( return None +def resolve_document_contexts( + document_ids, + user_id, + doc_scope="all", + active_group_ids=None, + active_public_workspace_id=None, + conversation_id=None, + include_content=True, +): + """Resolve ordered document contexts using one current authorization snapshot.""" + normalized_scope = normalize_search_scope(doc_scope) + normalized_document_ids = normalize_search_id_list(document_ids) + authorized_group_ids = [] + if normalized_scope in ("all", "group"): + authorized_group_ids = _resolve_active_group_ids( + user_id, + active_group_ids=active_group_ids, + fallback_to_memberships=True, + ) + authorized_public_workspace_ids = [] + if normalized_scope in ("all", "public"): + authorized_public_workspace_ids = _resolve_public_workspace_ids( + user_id, + active_public_workspace_id=active_public_workspace_id, + ) + + normalized_conversation_id = str(conversation_id or "").strip() + chat_conversation_authorized = bool( + normalized_conversation_id + and _authorize_chat_upload_conversation(user_id, normalized_conversation_id) + ) + + resolved_contexts = [] + for document_id in normalized_document_ids: + document_context = None + if normalized_scope in ("all", "personal"): + document_context = _resolve_personal_document_context(document_id, user_id) + if not document_context and normalized_scope in ("all", "group"): + document_context = _resolve_group_document_context( + document_id, + user_id, + authorized_group_ids, + ) + if not document_context and normalized_scope in ("all", "public"): + document_context = _resolve_public_document_context( + document_id, + user_id, + authorized_public_workspace_ids, + ) + if not document_context and chat_conversation_authorized: + document_context = _resolve_chat_upload_context( + document_id=document_id, + user_id=user_id, + conversation_id=normalized_conversation_id, + include_content=include_content, + authorization_prechecked=True, + ) + resolved_contexts.append(document_context) + + return resolved_contexts + + def build_search_request( query, user_id, @@ -367,6 +556,7 @@ def build_search_request( active_group_ids=None, active_public_workspace_id=None, enable_file_sharing=True, + include_all_public_workspaces=False, ): normalized_query = str(query or "").strip() if not normalized_query: @@ -406,7 +596,11 @@ def build_search_request( active_public_workspace_id=active_public_workspace_id, ) if resolved_public_workspace_ids and normalized_scope in ("all", "public"): - search_request["active_public_workspace_id"] = resolved_public_workspace_ids[0] + search_request["active_public_workspace_id"] = ( + resolved_public_workspace_ids + if include_all_public_workspaces + else resolved_public_workspace_ids[0] + ) return search_request @@ -422,6 +616,7 @@ def search_documents( active_group_ids=None, active_public_workspace_id=None, enable_file_sharing=True, + include_all_public_workspaces=False, ): search_request = build_search_request( query=query, @@ -434,6 +629,7 @@ def search_documents( active_group_ids=active_group_ids, active_public_workspace_id=active_public_workspace_id, enable_file_sharing=enable_file_sharing, + include_all_public_workspaces=include_all_public_workspaces, ) results = hybrid_search(**search_request) or [] unique_document_ids = { @@ -456,6 +652,70 @@ def search_documents( } +def search_relevant_tabular_candidates( + query, + user_id, + doc_scope="all", + document_ids=None, + tags_filter=None, + active_group_ids=None, + active_public_workspace_id=None, + max_candidates=MIXED_SOURCE_TABULAR_CANDIDATE_LIMIT, +): + """Find a bounded set of authorized table candidates from indexed schema chunks.""" + normalized_limit = _coerce_positive_int( + max_candidates, + MIXED_SOURCE_TABULAR_CANDIDATE_LIMIT, + min_value=1, + max_value=MIXED_SOURCE_TABULAR_CANDIDATE_LIMIT, + ) + candidate_query = ( + f"{str(query or '').strip()}\n" + "Relevant spreadsheet, workbook, worksheet, CSV, table schema, columns, and data fields." + ).strip() + search_result = search_documents( + query=candidate_query, + user_id=user_id, + top_n=MIXED_SOURCE_TABULAR_CANDIDATE_TOP_N, + doc_scope=doc_scope, + document_ids=document_ids, + tags_filter=tags_filter, + active_group_ids=active_group_ids, + active_public_workspace_id=active_public_workspace_id, + include_all_public_workspaces=True, + ) + + candidate_document_ids = [] + seen_document_ids = set() + for result in search_result.get("results") or []: + file_name = str(result.get("file_name") or "").strip() + if os.path.splitext(file_name)[1].lower() not in MIXED_SOURCE_TABULAR_EXTENSIONS: + continue + document_id = str(result.get("document_id") or "").strip() + if not document_id or document_id in seen_document_ids: + continue + seen_document_ids.add(document_id) + candidate_document_ids.append(document_id) + if len(candidate_document_ids) >= normalized_limit: + break + + log_event( + "[MixedSourceChatSearch] Completed bounded authorized tabular candidate search.", + extra={ + "candidate_search_result_count": search_result.get("result_count", 0), + "tabular_candidate_count": len(candidate_document_ids), + "candidate_limit": normalized_limit, + }, + level=logging.INFO, + ) + return { + "document_ids": candidate_document_ids, + "candidate_count": len(candidate_document_ids), + "search_result_count": search_result.get("result_count", 0), + "query": search_result.get("query"), + } + + def _derive_window_size(chunks, window_unit, window_size=None, window_percent=None): if not chunks: return 0 diff --git a/application/single_app/functions_settings.py b/application/single_app/functions_settings.py index d54d3824f..536edc010 100644 --- a/application/single_app/functions_settings.py +++ b/application/single_app/functions_settings.py @@ -210,6 +210,50 @@ def is_tabular_processing_enabled(settings): return bool((settings or {}).get('enable_enhanced_citations', False)) +def is_mixed_source_manifest_enabled(settings): + """Return whether Phase 1 mixed-source manifest diagnostics are enabled.""" + return bool((settings or {}).get('enable_mixed_source_manifest', False)) + + +def is_mixed_source_development_telemetry_enabled(settings): + """Return whether aggregate-only mixed-source development telemetry is enabled.""" + return bool((settings or {}).get('enable_mixed_source_development_telemetry', False)) + + +def is_mixed_source_chat_search_enabled(settings): + """Return whether Phase 2 mixed-source Chat and Search behavior is enabled.""" + return bool((settings or {}).get('enable_mixed_source_chat_search', False)) + + +def is_mixed_source_conversation_continuity_enabled(settings): + """Return whether Phase 5 reauthorized source-continuity metadata is enabled.""" + return bool( + (settings or {}).get('enable_mixed_source_chat_search', False) + and (settings or {}).get('enable_mixed_source_conversation_continuity', False) + ) + + +def is_cross_format_compare_enabled(settings): + """Return whether Phase 4 native mixed-source Compare behavior is enabled.""" + return bool((settings or {}).get('enable_cross_format_compare', False)) + + +def is_cross_format_compare_one_to_many_enabled(settings): + """Return whether the separately staged one-to-many mixed-target rollout is enabled.""" + return bool( + (settings or {}).get('enable_cross_format_compare', False) + and (settings or {}).get('enable_cross_format_compare_one_to_many', False) + ) + + +def is_mixed_source_relevance_candidates_enabled(settings): + """Return whether Phase 2 relevance-mode table candidates are enabled.""" + return bool( + (settings or {}).get('enable_mixed_source_chat_search', False) + and (settings or {}).get('enable_mixed_source_relevance_candidates', False) + ) + + CHAT_FILE_UPLOAD_APP_ROLE = "ChatFileUploadUser" WORKFLOW_USER_APP_ROLE = "WorkflowUser" DOCUMENT_INTELLIGENCE_PDF_IMAGE_EXTRACTION_MODES = {"read", "layout", "auto"} @@ -759,6 +803,13 @@ def get_settings(use_cosmos=False, include_source=False): 'enable_fact_memory_plugin': True, 'enable_tabular_processing_plugin': False, 'enable_multi_agent_orchestration': False, + 'enable_mixed_source_development_telemetry': False, + 'enable_mixed_source_manifest': False, + 'enable_mixed_source_chat_search': False, + 'enable_mixed_source_relevance_candidates': False, + 'enable_mixed_source_conversation_continuity': False, + 'enable_cross_format_compare': False, + 'enable_cross_format_compare_one_to_many': False, 'max_rounds_per_agent': 1, 'workflow_max_auto_invoke_attempts': 60, 'enable_semantic_kernel': False, diff --git a/application/single_app/functions_simplechat_operations.py b/application/single_app/functions_simplechat_operations.py index 10c17d70b..e89a24e6f 100644 --- a/application/single_app/functions_simplechat_operations.py +++ b/application/single_app/functions_simplechat_operations.py @@ -610,6 +610,70 @@ def upload_generated_analysis_artifact_for_current_user( ) +def delete_generated_chat_artifact_for_current_user( + conversation_id: str, + artifact_message_id: str, +) -> bool: + """Delete one generated artifact after reauthorizing its conversation and stored identity.""" + current_user_info = _require_current_user_info() + current_user_id = str(current_user_info.get("userId") or "").strip() + return delete_generated_chat_artifact_for_user( + current_user_id, + conversation_id, + artifact_message_id, + ) + + +def delete_generated_chat_artifact_for_user( + current_user_id: str, + conversation_id: str, + artifact_message_id: str, +) -> bool: + """Delete one generated artifact for a known user after object-level authorization.""" + current_user_id = str(current_user_id or "").strip() + normalized_conversation_id = str(conversation_id or "").strip() + normalized_message_id = str(artifact_message_id or "").strip() + if not current_user_id or not normalized_conversation_id or not normalized_message_id: + return False + + try: + conversation_item = cosmos_conversations_container.read_item( + item=normalized_conversation_id, + partition_key=normalized_conversation_id, + ) + message_item = cosmos_messages_container.read_item( + item=normalized_message_id, + partition_key=normalized_conversation_id, + ) + except CosmosResourceNotFoundError: + return False + + if str(conversation_item.get("user_id") or "").strip() != current_user_id: + raise PermissionError("Forbidden") + message_metadata = message_item.get("metadata") if isinstance(message_item.get("metadata"), dict) else {} + if ( + str(message_item.get("conversation_id") or "").strip() != normalized_conversation_id + or message_item.get("role") != "file" + or not message_metadata.get("is_generated_chat_artifact") + ): + raise PermissionError("Forbidden") + + delete_blob_backed_chat_message_files([message_item]) + cosmos_messages_container.delete_item( + item=normalized_message_id, + partition_key=normalized_conversation_id, + ) + log_event( + "[SimpleChat] Generated chat artifact rolled back after cancellation", + { + "artifact_count": 1, + "rollback_reason": "cancellation", + }, + debug_only=True, + ) + return True + + def upload_generated_analysis_artifact_for_user( current_user_id: str, conversation_id: str, @@ -667,6 +731,66 @@ def upload_generated_analysis_artifact_for_user( ) +def upload_generated_analysis_artifact_stream_for_user( + current_user_id: str, + conversation_id: str, + file_name: str, + file_stream: Any, + file_size: int, + capability: str = "analysis", + output_format: str = "", + summary: str = "", + artifact_idempotency_key: str = "", +) -> Dict[str, Any]: + """Upload a bounded-memory generated artifact stream for an authorized user.""" + normalized_user_id = str(current_user_id or "").strip() + normalized_conversation_id = str(conversation_id or "").strip() + normalized_file_name = _normalize_generated_document_file_name(file_name) + normalized_capability = str(capability or "analysis").strip().lower() or "analysis" + normalized_output_format = str(output_format or "").strip().lower() or os.path.splitext(normalized_file_name)[1].lower().lstrip(".") + normalized_summary = str(summary or "").strip() + + if not normalized_user_id: + raise ValueError("current_user_id is required") + if not normalized_conversation_id: + raise ValueError("conversation_id is required") + if not hasattr(file_stream, "read") or not hasattr(file_stream, "seek"): + raise ValueError("file_stream must be seekable and readable") + if not allowed_file(normalized_file_name): + raise ValueError("Generated file type is not supported") + + normalized_file_size = max(0, int(file_size or 0)) + if normalized_file_size <= 0: + raise ValueError("Generated artifact is empty") + + settings = get_settings() + max_artifact_size_mb = settings.get("max_generated_chat_artifact_size_mb", 500) + try: + max_artifact_size_mb = max(1, int(max_artifact_size_mb)) + except (TypeError, ValueError): + max_artifact_size_mb = 500 + + max_artifact_size_bytes = max_artifact_size_mb * 1024 * 1024 + if normalized_file_size > max_artifact_size_bytes: + raise ValueError( + f"Generated artifact exceeds the {max_artifact_size_mb} MB size limit" + ) + + file_stream.seek(0) + return _upload_generated_chat_artifact_for_current_user( + current_user_id=normalized_user_id, + conversation_id=normalized_conversation_id, + normalized_file_name=normalized_file_name, + file_content_bytes=file_stream, + artifact_metadata={ + "capability": normalized_capability, + "output_format": normalized_output_format, + "summary": normalized_summary, + }, + artifact_idempotency_key=artifact_idempotency_key, + ) + + def delete_blob_backed_chat_message_files(messages: Iterable[Dict[str, Any]]) -> int: """Delete blob-backed chat files referenced by the provided message documents.""" blob_service_client = CLIENTS.get("storage_account_office_docs_client") @@ -1219,7 +1343,7 @@ def search_directory_users(query: str, limit: int = 10) -> List[Dict[str, str]]: f"or startswith(mail, '{escaped_query}') " f"or startswith(userPrincipalName, '{escaped_query}')" ), - "$top": max(1, min(int(limit or 10), 25)), + "$top": max(1, min(int(limit or 10), 50)), "$select": "id,displayName,mail,userPrincipalName", }, ) @@ -1743,6 +1867,7 @@ def _upload_generated_chat_artifact_for_current_user( normalized_file_name: str, file_content_bytes: bytes, artifact_metadata: Optional[Dict[str, Any]] = None, + artifact_idempotency_key: str = "", ) -> Dict[str, Any]: try: conversation_item = cosmos_conversations_container.read_item( @@ -1759,7 +1884,15 @@ def _upload_generated_chat_artifact_for_current_user( if not blob_service_client: raise RuntimeError("Blob storage client not available") - artifact_message_id = f"{conversation_id}_generated_file_{uuid.uuid4().hex}" + normalized_idempotency_key = str(artifact_idempotency_key or "").strip() + if normalized_idempotency_key: + artifact_suffix = uuid.uuid5( + uuid.NAMESPACE_URL, + f"simplechat-generated-artifact:{conversation_id}:{normalized_idempotency_key}", + ).hex + else: + artifact_suffix = uuid.uuid4().hex + artifact_message_id = f"{conversation_id}_generated_file_{artifact_suffix}" blob_path = ( f"{current_user_id}/{conversation_id}/generated/" f"{artifact_message_id}/{normalized_file_name}" @@ -1768,6 +1901,33 @@ def _upload_generated_chat_artifact_for_current_user( container=storage_account_personal_chat_container_name, blob=blob_path, ) + if normalized_idempotency_key: + try: + existing_message = cosmos_messages_container.read_item( + item=artifact_message_id, + partition_key=conversation_id, + ) + except CosmosResourceNotFoundError: + existing_message = None + if ( + isinstance(existing_message, dict) + and existing_message.get("role") == "file" + and existing_message.get("filename") == normalized_file_name + and existing_message.get("blob_path") == blob_path + and blob_client.exists() + ): + existing_metadata = existing_message.get("metadata") or {} + return { + "message": { + "id": artifact_message_id, + "file_name": normalized_file_name, + "blob_container": storage_account_personal_chat_container_name, + "blob_path": blob_path, + "capability": existing_metadata.get("generated_artifact_capability") or "analysis", + "output_format": existing_metadata.get("generated_artifact_output_format") or "", + }, + "conversation_id": conversation_id, + } blob_client.upload_blob( file_content_bytes, overwrite=True, @@ -1775,6 +1935,7 @@ def _upload_generated_chat_artifact_for_current_user( "conversation_id": conversation_id, "user_id": current_user_id, "generated_artifact": "true", + "idempotent_artifact": str(bool(normalized_idempotency_key)).lower(), }, ) @@ -1804,6 +1965,7 @@ def _upload_generated_chat_artifact_for_current_user( "generated_artifact_capability": artifact_capability, "generated_artifact_output_format": artifact_output_format, "generated_artifact_summary": artifact_summary, + "generated_artifact_idempotency_key": normalized_idempotency_key or None, "thread_info": { "thread_id": current_thread_id, "previous_thread_id": previous_thread_id, diff --git a/application/single_app/functions_tabular_analysis.py b/application/single_app/functions_tabular_analysis.py index ae6421c4e..f18cfac58 100644 --- a/application/single_app/functions_tabular_analysis.py +++ b/application/single_app/functions_tabular_analysis.py @@ -2,35 +2,39 @@ """ Reusable tabular analysis helpers for chat and workflow execution. -This module is the non-route import surface for tabular analysis behavior that -is still implemented in route_backend_chats.py. Keeping workflow code pointed at -this module lets the implementation move out of the chat route incrementally -without changing workflow callers again. +This module owns extracted tabular coordination helpers and provides the +non-route import surface for behavior still implemented in +route_backend_chats.py. Keeping workflow code pointed here lets the remaining +implementation move incrementally without changing workflow callers again. """ def _load_chat_helper(helper_name): # Import lazily because route_backend_chats imports functions_workflow_runner during app startup. from route_backend_chats import ( + _execute_mixed_source_tabular_evidence, augment_tabular_invocations_with_related_document_evidence, build_tabular_computed_results_system_message, build_tabular_related_document_evidence_summary, - get_new_plugin_invocations, maybe_create_tabular_generated_output, run_tabular_analysis_with_thought_tracking, ) helpers = { + 'execute_mixed_source_tabular_evidence': _execute_mixed_source_tabular_evidence, 'augment_tabular_invocations_with_related_document_evidence': augment_tabular_invocations_with_related_document_evidence, 'build_tabular_computed_results_system_message': build_tabular_computed_results_system_message, 'build_tabular_related_document_evidence_summary': build_tabular_related_document_evidence_summary, - 'get_new_plugin_invocations': get_new_plugin_invocations, 'maybe_create_tabular_generated_output': maybe_create_tabular_generated_output, 'run_tabular_analysis_with_thought_tracking': run_tabular_analysis_with_thought_tracking, } return helpers[helper_name] +def execute_mixed_source_tabular_evidence(*args, **kwargs): + return _load_chat_helper('execute_mixed_source_tabular_evidence')(*args, **kwargs) + + def augment_tabular_invocations_with_related_document_evidence(*args, **kwargs): return _load_chat_helper('augment_tabular_invocations_with_related_document_evidence')(*args, **kwargs) @@ -43,8 +47,18 @@ def build_tabular_related_document_evidence_summary(*args, **kwargs): return _load_chat_helper('build_tabular_related_document_evidence_summary')(*args, **kwargs) -def get_new_plugin_invocations(*args, **kwargs): - return _load_chat_helper('get_new_plugin_invocations')(*args, **kwargs) +def get_new_plugin_invocations(invocations, baseline_count): + """Return only the plugin invocations created after the baseline count.""" + if not invocations: + return [] + + if baseline_count <= 0: + return list(invocations) + + if baseline_count >= len(invocations): + return [] + + return list(invocations[baseline_count:]) async def maybe_create_tabular_generated_output(*args, **kwargs): diff --git a/application/single_app/functions_tabular_csv_query.py b/application/single_app/functions_tabular_csv_query.py new file mode 100644 index 000000000..a8915fb4d --- /dev/null +++ b/application/single_app/functions_tabular_csv_query.py @@ -0,0 +1,204 @@ +# functions_tabular_csv_query.py +"""Shared bounded CSV query evaluation for foreground tools and durable exports.""" + +import ast + +import pandas + + +TABULAR_ROW_LOCAL_QUERY_AST_NODES = ( + ast.Expression, + ast.BoolOp, + ast.BinOp, + ast.UnaryOp, + ast.Compare, + ast.Name, + ast.Load, + ast.Constant, + ast.List, + ast.Tuple, + ast.Set, + ast.And, + ast.Or, + ast.Not, + ast.Eq, + ast.NotEq, + ast.Lt, + ast.LtE, + ast.Gt, + ast.GtE, + ast.In, + ast.NotIn, + ast.Is, + ast.IsNot, + ast.Add, + ast.Sub, + ast.Mult, + ast.Div, + ast.FloorDiv, + ast.Mod, + ast.Pow, + ast.BitAnd, + ast.BitOr, + ast.BitXor, + ast.USub, + ast.UAdd, + ast.Invert, +) + + +def _replace_backtick_column_references(expression): + """Replace pandas backtick column labels with parseable placeholder names.""" + output = [] + placeholder_index = 0 + character_index = 0 + active_quote = None + escaped = False + while character_index < len(expression): + character = expression[character_index] + if escaped: + output.append(character) + escaped = False + character_index += 1 + continue + if character == '\\' and active_quote: + output.append(character) + escaped = True + character_index += 1 + continue + if active_quote: + output.append(character) + if character == active_quote: + active_quote = None + character_index += 1 + continue + if character in {'\'', '"'}: + active_quote = character + output.append(character) + character_index += 1 + continue + if character != '`': + output.append(character) + character_index += 1 + continue + + closing_index = expression.find('`', character_index + 1) + if closing_index < 0 or not expression[character_index + 1:closing_index].strip(): + raise ValueError('Source query contains an invalid backtick column reference') + output.append(f'__simplechat_column_{placeholder_index}') + placeholder_index += 1 + character_index = closing_index + 1 + + return ''.join(output) + + +def validate_tabular_csv_query_expression(query_expression): + """Return a row-local expression or reject operations that change across chunks.""" + normalized_expression = str(query_expression or '').strip() + if not normalized_expression: + raise ValueError('Source query expression is required') + if '@' in normalized_expression: + raise ValueError( + 'Source query variables cannot be replayed in bounded chunks' + ) + + parseable_expression = _replace_backtick_column_references(normalized_expression) + try: + parsed_expression = ast.parse(parseable_expression, mode='eval') + except SyntaxError as exc: + raise ValueError('Source query is not a valid row-local expression') from exc + + unsupported_nodes = [ + node.__class__.__name__ + for node in ast.walk(parsed_expression) + if not isinstance(node, TABULAR_ROW_LOCAL_QUERY_AST_NODES) + ] + if unsupported_nodes: + raise ValueError( + 'Source query uses an operation that cannot be replayed equivalently in bounded chunks: ' + f'{unsupported_nodes[0]}' + ) + for node in ast.walk(parsed_expression): + if isinstance(node, ast.Name) and node.id.startswith('__') and not node.id.startswith('__simplechat_column_'): + raise ValueError('Source query contains an unsupported private identifier') + return normalized_expression + + +def detect_tabular_csv_numeric_columns(csv_stream, source_chunk_rows, tabular_plugin): + """Find columns that pandas can convert to numeric across every bounded chunk.""" + numeric_columns = None + csv_stream.seek(0) + for source_chunk in pandas.read_csv( + csv_stream, + keep_default_na=False, + dtype=str, + chunksize=max(1, int(source_chunk_rows or 1)), + ): + source_chunk = tabular_plugin._normalize_dataframe_columns(source_chunk) + if numeric_columns is None: + numeric_columns = set(source_chunk.columns) + for column_name in list(numeric_columns): + try: + pandas.to_numeric(source_chunk[column_name]) + except (TypeError, ValueError): + numeric_columns.discard(column_name) + csv_stream.seek(0) + return numeric_columns or set() + + +def iter_tabular_csv_query_rows( + csv_stream, + query_expression, + return_columns, + source_chunk_rows, + tabular_plugin, + start_source_row=0, + replay_stats=None, +): + """Yield physical source row numbers and query-matched records in bounded chunks.""" + source_chunk_rows = max(1, int(source_chunk_rows or 1)) + start_source_row = max(0, int(start_source_row or 0)) + normalized_query_expression = validate_tabular_csv_query_expression(query_expression) + numeric_columns = detect_tabular_csv_numeric_columns( + csv_stream, + source_chunk_rows, + tabular_plugin, + ) + parsed_return_columns = tabular_plugin._parse_optional_column_list_argument(return_columns) + + csv_stream.seek(0) + read_options = { + 'keep_default_na': False, + 'dtype': str, + 'chunksize': source_chunk_rows, + } + if start_source_row: + read_options['skiprows'] = lambda row_index: 0 < row_index <= start_source_row + + source_row_offset = start_source_row + for source_chunk in pandas.read_csv(csv_stream, **read_options): + source_chunk = tabular_plugin._normalize_dataframe_columns(source_chunk) + source_chunk.index = range(source_row_offset, source_row_offset + len(source_chunk)) + source_row_offset += len(source_chunk) + for column_name in numeric_columns: + if column_name in source_chunk.columns: + source_chunk[column_name] = pandas.to_numeric(source_chunk[column_name]) + + filtered_chunk, used_reviewer_style_fallback = tabular_plugin._apply_query_expression_with_fallback( + source_chunk, + query_expression=normalized_query_expression, + normalize_match=False, + ) + if isinstance(replay_stats, dict) and used_reviewer_style_fallback: + replay_stats['used_reviewer_style_fallback'] = True + selected_columns = [ + column_name + for column_name in (parsed_return_columns or list(filtered_chunk.columns)) + if column_name in filtered_chunk.columns + ] + output_records = tabular_plugin._build_row_output_records( + filtered_chunk, + selected_columns, + ) + for source_row_index, output_record in zip(filtered_chunk.index, output_records): + yield int(source_row_index) + 1, output_record diff --git a/application/single_app/functions_tabular_generated_exports.py b/application/single_app/functions_tabular_generated_exports.py index 6d27ef76d..822d65aa5 100644 --- a/application/single_app/functions_tabular_generated_exports.py +++ b/application/single_app/functions_tabular_generated_exports.py @@ -2,35 +2,52 @@ """Durable background runs for large tabular generated exports.""" import asyncio +from collections import Counter import csv import io import json import logging +import math import os import re import socket +import tempfile import time import uuid from datetime import datetime, timedelta, timezone from azure.core import MatchConditions +from azure.core.exceptions import ResourceExistsError from azure.cosmos.exceptions import CosmosResourceNotFoundError from flask import current_app, has_app_context from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.azure_chat_prompt_execution_settings import AzureChatPromptExecutionSettings from semantic_kernel.contents.chat_history import ChatHistory as SKChatHistory +from semantic_kernel_plugins.tabular_processing_plugin import TabularProcessingPlugin from config import ( CLIENTS, + cosmos_conversations_container, cosmos_tabular_export_runs_container, + storage_account_group_documents_container_name, storage_account_personal_chat_container_name, + storage_account_public_documents_container_name, + storage_account_user_documents_container_name, ) from functions_appinsights import log_event +from functions_assistant_table_exports import build_safe_csv_headers, neutralize_csv_spreadsheet_formula +from functions_tabular_csv_query import ( + iter_tabular_csv_query_rows, + validate_tabular_csv_query_expression, +) +from functions_group import assert_group_role from functions_model_endpoint_runtime import build_semantic_kernel_chat_service_for_model +from functions_public_workspaces import get_user_visible_public_workspace_ids_from_settings from functions_settings import get_settings -from functions_simplechat_operations import upload_generated_analysis_artifact_for_user +from functions_simplechat_operations import upload_generated_analysis_artifact_stream_for_user TABULAR_EXPORT_RUN_TYPE = 'tabular_generated_output_run' +TABULAR_EXPORT_CONTRACT_VERSION = 2 TABULAR_EXPORT_STATUS_QUEUED = 'queued' TABULAR_EXPORT_STATUS_RUNNING = 'running' TABULAR_EXPORT_STATUS_COMPLETED = 'completed' @@ -49,8 +66,16 @@ TABULAR_EXPORT_DEFAULT_STALE_SECONDS = 420 TABULAR_EXPORT_DEFAULT_SCAN_LIMIT = 5 TABULAR_EXPORT_DEFAULT_MAX_TRANSIENT_FAILURES = 20 -TABULAR_EXPORT_DEFAULT_BATCH_CONCURRENCY = 2 +TABULAR_EXPORT_DEFAULT_BATCH_CONCURRENCY = 3 TABULAR_EXPORT_MAX_BATCH_CONCURRENCY = 5 +TABULAR_EXPORT_DEFAULT_BATCH_TIMEOUT_SECONDS = 300 +TABULAR_EXPORT_FINAL_SPOOL_MAX_MEMORY_BYTES = 1024 * 1024 +TABULAR_EXPORT_DEFAULT_SOURCE_CHUNK_ROWS = 1000 +TABULAR_EXPORT_DEFAULT_SOURCE_BATCH_ROWS = 50 +TABULAR_EXPORT_DEFAULT_SOURCE_BATCH_CHARS = 60000 +TABULAR_EXPORT_SUMMARY_MAX_FIELDS = 25 +TABULAR_EXPORT_SUMMARY_MAX_VALUES_PER_FIELD = 5 +TABULAR_EXPORT_SUMMARY_AGGREGATE_MAX_VALUES = 25 TABULAR_EXPORT_PROGRESS_LOG_INTERVAL_SECONDS = 30 TABULAR_EXPORT_SCHEDULER_STATUSES = ( TABULAR_EXPORT_STATUS_QUEUED, @@ -87,6 +112,19 @@ 'worker exiting', 'worker restart', ) +TABULAR_EXPORT_INPUT_ROW_NUMBER_FIELD = '__simplechat_source_row_number' +TABULAR_EXPORT_INPUT_ROW_IDENTITY_FIELD = '__simplechat_source_row_identity' +TABULAR_EXPORT_INPUT_ROW_TOKEN_FIELD = '__simplechat_source_row_token' +TABULAR_EXPORT_OUTPUT_ROW_NUMBER_FIELD = 'source_row_number' +TABULAR_EXPORT_OUTPUT_ROW_IDENTITY_FIELD = 'source_row_identity' + + +class TabularExportCanceledError(RuntimeError): + """Raised when a durable export is canceled between checkpoints.""" + + +class TabularExportLeaseLostError(RuntimeError): + """Raised when a stale worker no longer owns the durable run claim.""" def _now_utc(): @@ -196,13 +234,269 @@ def _serialize_generated_output_value(value): if value is None: return '' if isinstance(value, (dict, list)): - return json.dumps(value, default=str, ensure_ascii=False) + return neutralize_csv_spreadsheet_formula(json.dumps(value, default=str, ensure_ascii=False)) if hasattr(value, 'isoformat') and not isinstance(value, str): try: - return value.isoformat() + return neutralize_csv_spreadsheet_formula(value.isoformat()) except TypeError: pass - return str(value) + return neutralize_csv_spreadsheet_formula(value) + + +def _normalize_source_identity_label(value): + return re.sub(r'[^a-z0-9]+', '', str(value or '').strip().casefold()) + + +def _select_source_row_identity(row, source_row_number): + if not isinstance(row, dict): + return str(source_row_number) + + identity_priorities = ( + 'sourceidentity', + 'sourceid', + 'caseid', + 'recordid', + 'rowid', + 'commentid', + 'submissionid', + 'id', + ) + normalized_values = {} + for field_name, field_value in row.items(): + if field_name in { + TABULAR_EXPORT_INPUT_ROW_NUMBER_FIELD, + TABULAR_EXPORT_INPUT_ROW_IDENTITY_FIELD, + TABULAR_EXPORT_INPUT_ROW_TOKEN_FIELD, + }: + continue + normalized_label = _normalize_source_identity_label(field_name) + if normalized_label and field_value not in (None, '', [], {}): + normalized_values.setdefault(normalized_label, field_value) + + for identity_label in identity_priorities: + identity_value = normalized_values.get(identity_label) + if identity_value not in (None, '', [], {}): + return str(identity_value) + + for identity_label, identity_value in normalized_values.items(): + if identity_label.endswith('id') and not isinstance(identity_value, (dict, list, tuple, set)): + return str(identity_value) + + return str(source_row_number) + + +def _prepare_tabular_source_rows(rows, start_row=0, token_namespace=''): + try: + normalized_start_row = max(0, int(start_row or 0)) + except (TypeError, ValueError): + normalized_start_row = 0 + + prepared_rows = [] + for row_offset, row in enumerate(rows or []): + source_row_number = normalized_start_row + row_offset + 1 + prepared_row = dict(row) if isinstance(row, dict) else {'value': row} + prepared_row[TABULAR_EXPORT_INPUT_ROW_NUMBER_FIELD] = source_row_number + prepared_row[TABULAR_EXPORT_INPUT_ROW_IDENTITY_FIELD] = _select_source_row_identity( + prepared_row, + source_row_number, + ) + token_seed = ( + f'simplechat-tabular-row:{token_namespace}:{source_row_number}:' + f'{prepared_row[TABULAR_EXPORT_INPUT_ROW_IDENTITY_FIELD]}' + ) + prepared_row[TABULAR_EXPORT_INPUT_ROW_TOKEN_FIELD] = uuid.uuid5( + uuid.NAMESPACE_URL, + token_seed, + ).hex + prepared_rows.append(prepared_row) + return prepared_rows + + +def _normalize_generated_batch_entries( + source_rows, + generated_entries, + expected_output_schema=None, + require_source_token=True, +): + source_rows = list(source_rows or []) + generated_entries = list(generated_entries or []) + if len(source_rows) != len(generated_entries): + raise ValueError( + f'Generated row count {len(generated_entries)} does not match source row count {len(source_rows)}' + ) + + normalized_entries = [] + for row_index, (source_row, generated_entry) in enumerate(zip(source_rows, generated_entries), start=1): + if not isinstance(source_row, dict): + raise ValueError(f'Source row {row_index} is not an object') + if not isinstance(generated_entry, dict): + raise ValueError(f'Generated row {row_index} is not an object') + + source_row_number = source_row.get(TABULAR_EXPORT_INPUT_ROW_NUMBER_FIELD) + source_row_identity = source_row.get(TABULAR_EXPORT_INPUT_ROW_IDENTITY_FIELD) + source_row_token = str(source_row.get(TABULAR_EXPORT_INPUT_ROW_TOKEN_FIELD) or '').strip() + if source_row_number in (None, '') or source_row_identity in (None, ''): + raise ValueError(f'Source identity is missing for generated row {row_index}') + generated_row_token = str( + generated_entry.get(TABULAR_EXPORT_INPUT_ROW_TOKEN_FIELD) or '' + ).strip() + if require_source_token and ( + not source_row_token + or generated_row_token != source_row_token + ): + raise ValueError( + f'Generated source row token mismatch at row {row_index}' + ) + + normalized_entry = { + TABULAR_EXPORT_OUTPUT_ROW_NUMBER_FIELD: source_row_number, + TABULAR_EXPORT_OUTPUT_ROW_IDENTITY_FIELD: str(source_row_identity), + } + normalized_entry.update({ + str(field_name): field_value + for field_name, field_value in generated_entry.items() + if str(field_name) not in { + TABULAR_EXPORT_INPUT_ROW_NUMBER_FIELD, + TABULAR_EXPORT_INPUT_ROW_IDENTITY_FIELD, + TABULAR_EXPORT_INPUT_ROW_TOKEN_FIELD, + TABULAR_EXPORT_OUTPUT_ROW_NUMBER_FIELD, + TABULAR_EXPORT_OUTPUT_ROW_IDENTITY_FIELD, + } + }) + normalized_entries.append(normalized_entry) + + output_schema = list(expected_output_schema or []) + if not output_schema and normalized_entries: + output_schema = list(normalized_entries[0]) + + expected_fields = set(output_schema) + for row_index, normalized_entry in enumerate(normalized_entries, start=1): + actual_fields = set(normalized_entry) + if actual_fields != expected_fields: + missing_fields = sorted(expected_fields - actual_fields) + unexpected_fields = sorted(actual_fields - expected_fields) + raise ValueError( + f'Generated output schema mismatch at row {row_index}; ' + f'missing={missing_fields}; unexpected={unexpected_fields}' + ) + + ordered_entries = [ + {field_name: entry.get(field_name) for field_name in output_schema} + for entry in normalized_entries + ] + return ordered_entries, output_schema + + +def _build_generated_batch_summary(entries): + entries = [entry for entry in (entries or []) if isinstance(entry, dict)] + summary = { + 'row_count': len(entries), + 'source_row_start': None, + 'source_row_end': None, + 'fields': {}, + } + if not entries: + return summary + + summary['source_row_start'] = _safe_int(entries[0].get(TABULAR_EXPORT_OUTPUT_ROW_NUMBER_FIELD)) or None + summary['source_row_end'] = _safe_int(entries[-1].get(TABULAR_EXPORT_OUTPUT_ROW_NUMBER_FIELD)) or None + field_names = [ + field_name + for field_name in entries[0] + if field_name not in { + TABULAR_EXPORT_OUTPUT_ROW_NUMBER_FIELD, + TABULAR_EXPORT_OUTPUT_ROW_IDENTITY_FIELD, + } + ][:TABULAR_EXPORT_SUMMARY_MAX_FIELDS] + for field_name in field_names: + populated_count = 0 + value_counts = Counter() + for entry in entries: + field_value = entry.get(field_name) + if field_value in (None, '', [], {}): + continue + populated_count += 1 + if isinstance(field_value, (str, int, float, bool)): + normalized_value = str(field_value).strip() + if normalized_value and len(normalized_value) <= 100: + value_counts[normalized_value] += 1 + + summary['fields'][field_name] = { + 'populated_count': populated_count, + 'empty_count': len(entries) - populated_count, + 'top_values': [ + {'value': value, 'count': count} + for value, count in value_counts.most_common(TABULAR_EXPORT_SUMMARY_MAX_VALUES_PER_FIELD) + ], + } + return summary + + +def _build_compact_post_run_summary(run): + batch_count = _safe_int(run.get('batch_count')) + row_count = _safe_int(run.get('row_count')) + aggregate_fields = {} + summarized_batch_count = 0 + for batch_number in range(1, batch_count + 1): + summary_blob_path = _output_summary_blob_path( + run.get('user_id'), + run.get('conversation_id'), + run.get('id'), + batch_number, + ) + if not _blob_exists(summary_blob_path): + continue + batch_summary = _download_json_blob(summary_blob_path) + if not isinstance(batch_summary, dict): + continue + summarized_batch_count += 1 + for field_name, field_summary in list((batch_summary.get('fields') or {}).items())[ + :TABULAR_EXPORT_SUMMARY_MAX_FIELDS + ]: + aggregate_field = aggregate_fields.setdefault(field_name, { + 'populated_count': 0, + 'empty_count': 0, + 'value_counts': Counter(), + }) + aggregate_field['populated_count'] += _safe_int(field_summary.get('populated_count')) + aggregate_field['empty_count'] += _safe_int(field_summary.get('empty_count')) + for value_summary in field_summary.get('top_values') or []: + value = str(value_summary.get('value') or '').strip() + if value: + aggregate_field['value_counts'][value] += _safe_int(value_summary.get('count')) + if len(aggregate_field['value_counts']) > TABULAR_EXPORT_SUMMARY_AGGREGATE_MAX_VALUES * 2: + aggregate_field['value_counts'] = Counter(dict( + aggregate_field['value_counts'].most_common(TABULAR_EXPORT_SUMMARY_AGGREGATE_MAX_VALUES) + )) + + summary_parts = [ + f'Processed {row_count:,} ordered row(s) across {batch_count:,} checkpointed batch(es).' + ] + if aggregate_fields: + field_names = list(aggregate_fields)[:10] + summary_parts.append(f"Output fields: {', '.join(field_names)}.") + completeness_parts = [] + for field_name in field_names[:5]: + populated_count = aggregate_fields[field_name]['populated_count'] + completeness_percent = round((populated_count / row_count) * 100) if row_count else 0 + completeness_parts.append(f'{field_name} {completeness_percent}% populated') + if completeness_parts: + summary_parts.append(f"Completeness: {', '.join(completeness_parts)}.") + + common_value_parts = [] + for field_name in field_names[:5]: + top_values = aggregate_fields[field_name]['value_counts'].most_common(3) + if 1 < len(top_values) <= 3: + rendered_values = ', '.join(f'{value} ({count:,})' for value, count in top_values) + common_value_parts.append(f'{field_name}: {rendered_values}') + if common_value_parts: + summary_parts.append(f"Common values: {'; '.join(common_value_parts)}.") + + if summarized_batch_count != batch_count: + summary_parts.append( + f'Batch summaries available for {summarized_batch_count:,} of {batch_count:,} batch(es).' + ) + return ' '.join(summary_parts)[:2000] def _build_generated_output_csv(entries): @@ -221,14 +515,15 @@ def _build_generated_output_csv(entries): if not ordered_columns: ordered_columns = ['value'] + safe_ordered_columns = build_safe_csv_headers(ordered_columns) output_buffer = io.StringIO() - writer = csv.DictWriter(output_buffer, fieldnames=ordered_columns) + writer = csv.DictWriter(output_buffer, fieldnames=safe_ordered_columns) writer.writeheader() for entry in entries or []: serialized_row = {} if isinstance(entry, dict): - for field_name in ordered_columns: - serialized_row[field_name] = _serialize_generated_output_value(entry.get(field_name)) + for field_name, safe_field_name in zip(ordered_columns, safe_ordered_columns): + serialized_row[safe_field_name] = _serialize_generated_output_value(entry.get(field_name)) writer.writerow(serialized_row) return output_buffer.getvalue() @@ -252,14 +547,18 @@ def _output_blob_path(user_id, conversation_id, run_id, batch_number): return f"{user_id}/{conversation_id}/generated/tabular_runs/{run_id}/output/batch_{batch_number:06d}.json" -def _upload_json_blob(blob_path, payload, metadata=None): +def _output_summary_blob_path(user_id, conversation_id, run_id, batch_number): + return f"{user_id}/{conversation_id}/generated/tabular_runs/{run_id}/summary/batch_{batch_number:06d}.json" + + +def _upload_json_blob(blob_path, payload, metadata=None, overwrite=True): blob_client = _get_blob_service_client().get_blob_client( container=storage_account_personal_chat_container_name, blob=blob_path, ) blob_client.upload_blob( json.dumps(payload, default=str, ensure_ascii=False).encode('utf-8'), - overwrite=True, + overwrite=overwrite, metadata={str(key): str(value) for key, value in (metadata or {}).items()}, ) @@ -283,6 +582,109 @@ def _blob_exists(blob_path): return bool(blob_client.exists()) +def _delete_blob_if_exists(blob_path): + blob_client = _get_blob_service_client().get_blob_client( + container=storage_account_personal_chat_container_name, + blob=blob_path, + ) + if blob_client.exists(): + blob_client.delete_blob() + + +def _authorize_tabular_export_run_execution(run): + user_id = str((run or {}).get('user_id') or '').strip() + conversation_id = str((run or {}).get('conversation_id') or '').strip() + if not user_id or not conversation_id: + raise PermissionError('Export run identity is incomplete') + + try: + conversation = cosmos_conversations_container.read_item( + item=conversation_id, + partition_key=conversation_id, + ) + except CosmosResourceNotFoundError as exc: + raise PermissionError('Export conversation no longer exists') from exc + if str(conversation.get('user_id') or '').strip() != user_id: + raise PermissionError('Export conversation ownership changed') + + source_authorization = ( + (run or {}).get('source_descriptor') + or (run or {}).get('source_authorization') + or {} + ) + if not source_authorization: + return conversation + + source = str(source_authorization.get('source') or '').strip().lower() + container_name = str(source_authorization.get('container') or '').strip() + blob_path = str(source_authorization.get('blob_path') or '').strip() + scope_id = str(source_authorization.get('scope_id') or '').strip() + if not source: + raise PermissionError('Export source authorization is incomplete') + + if source == 'chat': + expected_prefix = f'{user_id}/{conversation_id}/' + authorized = not container_name and not blob_path or ( + container_name == storage_account_personal_chat_container_name + and blob_path.startswith(expected_prefix) + ) + elif source == 'workspace': + expected_prefix = f'{user_id}/' + authorized = not container_name and not blob_path or ( + container_name == storage_account_user_documents_container_name + and blob_path.startswith(expected_prefix) + ) + elif source == 'group': + if not scope_id: + raise PermissionError('Export group scope is missing') + assert_group_role( + user_id, + scope_id, + allowed_roles=('Owner', 'Admin', 'DocumentManager', 'User'), + ) + authorized = not container_name and not blob_path or ( + container_name == storage_account_group_documents_container_name + and blob_path.startswith(f'{scope_id}/') + ) + elif source == 'public': + if not scope_id: + raise PermissionError('Export public workspace scope is missing') + visible_workspace_ids = set(get_user_visible_public_workspace_ids_from_settings(user_id) or []) + authorized = scope_id in visible_workspace_ids and ( + not container_name and not blob_path + or ( + container_name == storage_account_public_documents_container_name + and blob_path.startswith(f'{scope_id}/') + ) + ) + else: + raise PermissionError('Export source type is not supported') + + if not authorized: + raise PermissionError('Export source is no longer authorized') + return conversation + + +def _get_versioned_source_blob_client(source_descriptor): + container_name = str((source_descriptor or {}).get('container') or '').strip() + blob_path = str((source_descriptor or {}).get('blob_path') or '').strip() + expected_etag = str((source_descriptor or {}).get('blob_etag') or '').strip() + if not container_name or not blob_path or not expected_etag: + raise ValueError('Source query descriptor is incomplete') + + blob_client = _get_blob_service_client().get_blob_client( + container=container_name, + blob=blob_path, + ) + blob_properties = blob_client.get_blob_properties() + current_etag = str(getattr(blob_properties, 'etag', '') or '').strip() + if isinstance(blob_properties, dict): + current_etag = current_etag or str(blob_properties.get('etag') or '').strip() + if current_etag != expected_etag: + raise ValueError('Source CSV changed after the export was queued') + return blob_client + + def _clean_generated_json_code_fence(response_content): cleaned = str(response_content or '').strip() if not cleaned: @@ -336,17 +738,286 @@ def _dump_generated_output_json(value): return json.dumps(value, default=str, ensure_ascii=False, separators=(',', ':')) -def _build_batch_prompt(user_question, batch_rows, batch_index, total_batches, source_file_name, selected_sheet=''): +def _checkpoint_source_input_batch(run, batch_rows, source_scan_row_count): + _raise_if_tabular_export_canceled(run) + staged_row_count = _safe_int(run.get('source_staged_rows')) + batch_number = _safe_int(run.get('source_staged_batches')) + 1 + prepared_rows = _prepare_tabular_source_rows( + batch_rows, + start_row=staged_row_count, + token_namespace=run.get('id'), + ) + input_blob_path = _input_blob_path( + run.get('user_id'), + run.get('conversation_id'), + run.get('id'), + batch_number, + ) + _upload_json_blob( + input_blob_path, + prepared_rows, + metadata={ + 'run_id': run.get('id'), + 'conversation_id': run.get('conversation_id'), + 'generated_output_input': 'true', + 'batch_number': batch_number, + 'source_query_checkpoint': 'true', + }, + ) + _raise_if_tabular_export_canceled(run) + + now = _now_iso() + run.update({ + 'source_staged_rows': staged_row_count + len(prepared_rows), + 'source_staged_batches': batch_number, + 'source_scan_row_count': _safe_int(source_scan_row_count), + 'updated_at': now, + 'last_heartbeat_at': now, + 'last_message': ( + f'Staged source query batch {batch_number} ' + f'with {staged_row_count + len(prepared_rows)} row(s) ready' + ), + }) + return _replace_claimed_run(run) + + +def _stage_tabular_generated_output_source(run, settings): + source_descriptor = run.get('source_descriptor') or {} + if source_descriptor.get('kind') != 'query_tabular_data': + raise ValueError('Unsupported generated export source descriptor') + if not str(source_descriptor.get('blob_path') or '').lower().endswith('.csv'): + raise ValueError('Source-backed generated exports require a CSV source') + + expected_row_count = _safe_int(source_descriptor.get('expected_row_count')) + if expected_row_count <= 0: + raise ValueError('Source query descriptor has no expected rows') + source_chunk_rows = _settings_int( + settings, + 'tabular_generated_output_source_chunk_rows', + TABULAR_EXPORT_DEFAULT_SOURCE_CHUNK_ROWS, + minimum=100, + maximum=10000, + ) + max_batch_rows = _safe_int( + source_descriptor.get('batch_max_rows'), + default=TABULAR_EXPORT_DEFAULT_SOURCE_BATCH_ROWS, + minimum=1, + maximum=100, + ) + max_batch_chars = _safe_int( + source_descriptor.get('batch_max_chars'), + default=TABULAR_EXPORT_DEFAULT_SOURCE_BATCH_CHARS, + minimum=6000, + maximum=120000, + ) + source_blob_client = _get_versioned_source_blob_client(source_descriptor) + resume_source_row = _safe_int(run.get('source_scan_row_count')) + pending_rows = [] + pending_chars = 0 + last_source_row_number = resume_source_row + + with tempfile.SpooledTemporaryFile( + max_size=TABULAR_EXPORT_FINAL_SPOOL_MAX_MEMORY_BYTES, + mode='w+b', + ) as source_stream: + source_blob_client.download_blob( + etag=source_descriptor.get('blob_etag'), + match_condition=MatchConditions.IfNotModified, + ).readinto(source_stream) + source_stream.seek(0) + + tabular_plugin = TabularProcessingPlugin() + for source_row_number, source_row in iter_tabular_csv_query_rows( + csv_stream=source_stream, + query_expression=source_descriptor.get('query_expression'), + return_columns=source_descriptor.get('return_columns'), + source_chunk_rows=source_chunk_rows, + tabular_plugin=tabular_plugin, + start_source_row=resume_source_row, + ): + source_row_text = _dump_generated_output_json(source_row) + if pending_rows and ( + len(pending_rows) >= max_batch_rows + or pending_chars + len(source_row_text) > max_batch_chars + ): + run = _checkpoint_source_input_batch( + run, + pending_rows, + source_scan_row_count=source_row_number - 1, + ) + pending_rows = [] + pending_chars = 0 + + pending_rows.append(source_row) + pending_chars += len(source_row_text) + last_source_row_number = source_row_number + if len(pending_rows) >= max_batch_rows: + run = _checkpoint_source_input_batch( + run, + pending_rows, + source_scan_row_count=source_row_number, + ) + pending_rows = [] + pending_chars = 0 + + if pending_rows: + run = _checkpoint_source_input_batch( + run, + pending_rows, + source_scan_row_count=last_source_row_number, + ) + + staged_row_count = _safe_int(run.get('source_staged_rows')) + staged_batch_count = _safe_int(run.get('source_staged_batches')) + if staged_row_count != expected_row_count: + raise ValueError( + f'Source query returned {staged_row_count} row(s); expected {expected_row_count}' + ) + if staged_batch_count <= 0: + raise ValueError('Source query produced no input checkpoints') + + now = _now_iso() + run.update({ + 'source_staging_complete': True, + 'row_count': staged_row_count, + 'batch_count': staged_batch_count, + 'updated_at': now, + 'last_heartbeat_at': now, + 'last_message': ( + f'Source query staging complete: {staged_row_count} row(s) ' + f'across {staged_batch_count} batch(es)' + ), + }) + return _replace_claimed_run(run) + + +def _migrate_legacy_tabular_export_run(run): + if _safe_int(run.get('contract_version')) >= TABULAR_EXPORT_CONTRACT_VERSION: + return run + + batch_count = _safe_int(run.get('batch_count')) + expected_row_count = _safe_int(run.get('row_count')) + if batch_count <= 0 or expected_row_count <= 0: + raise ValueError('Legacy export run has invalid input counts') + + aggregate_input_batches = None + aggregate_input_blob_path = str(run.get('input_blob_path') or '').strip() + if aggregate_input_blob_path: + aggregate_input_batches = _download_json_blob(aggregate_input_blob_path) + if not isinstance(aggregate_input_batches, list): + raise ValueError('Legacy input batches blob was not a JSON array') + + migrated_row_count = 0 + for batch_number in range(1, batch_count + 1): + _raise_if_tabular_export_canceled(run) + if aggregate_input_batches is not None: + try: + batch_rows = aggregate_input_batches[batch_number - 1] + except IndexError as exc: + raise ValueError(f'Legacy input batch {batch_number}/{batch_count} is missing') from exc + else: + batch_rows = _download_json_blob(_input_blob_path( + run.get('user_id'), + run.get('conversation_id'), + run.get('id'), + batch_number, + )) + if not isinstance(batch_rows, list): + raise ValueError(f'Legacy input batch {batch_number}/{batch_count} was not a JSON array') + + prepared_rows = _prepare_tabular_source_rows( + batch_rows, + start_row=migrated_row_count, + token_namespace=run.get('id'), + ) + _upload_json_blob( + _input_blob_path( + run.get('user_id'), + run.get('conversation_id'), + run.get('id'), + batch_number, + ), + prepared_rows, + metadata={ + 'run_id': run.get('id'), + 'conversation_id': run.get('conversation_id'), + 'generated_output_input': 'true', + 'batch_number': batch_number, + 'contract_version': TABULAR_EXPORT_CONTRACT_VERSION, + }, + ) + _delete_blob_if_exists(_output_blob_path( + run.get('user_id'), + run.get('conversation_id'), + run.get('id'), + batch_number, + )) + _delete_blob_if_exists(_output_summary_blob_path( + run.get('user_id'), + run.get('conversation_id'), + run.get('id'), + batch_number, + )) + migrated_row_count += len(prepared_rows) + + if migrated_row_count != expected_row_count: + raise ValueError( + f'Legacy input migration found {migrated_row_count} row(s); expected {expected_row_count}' + ) + + now = _now_iso() + run.update({ + 'contract_version': TABULAR_EXPORT_CONTRACT_VERSION, + 'input_blob_path': None, + 'completed_batches': 0, + 'processed_rows': 0, + 'output_schema': None, + 'regenerate_legacy_output_checkpoints': False, + 'updated_at': now, + 'last_heartbeat_at': now, + 'last_message': 'Legacy export inputs migrated; output checkpoints will be regenerated', + }) + _raise_if_tabular_export_canceled(run) + return _replace_claimed_run(run) + + +def _build_batch_prompt( + user_question, + batch_rows, + batch_index, + total_batches, + source_file_name, + selected_sheet='', + output_schema=None, +): source_file_name = str(source_file_name or 'unknown file').strip() or 'unknown file' selected_sheet = str(selected_sheet or '').strip() batch_rows_json = _dump_generated_output_json(batch_rows) selected_sheet_line = f"Worksheet: {selected_sheet}\n" if selected_sheet else '' + model_output_schema = [ + field_name + for field_name in (output_schema or []) + if field_name not in { + TABULAR_EXPORT_OUTPUT_ROW_NUMBER_FIELD, + TABULAR_EXPORT_OUTPUT_ROW_IDENTITY_FIELD, + } + ] + output_schema_line = ( + f'Use exactly these output fields for every object, in this order: ' + f'{json.dumps(model_output_schema, ensure_ascii=False)}.\n' + if model_output_schema + else '' + ) return ( 'Transform the tabular input rows below into structured output for the user.\n\n' f'User instructions:\n{user_question}\n\n' 'Return ONLY a valid JSON array.\n' f'Return exactly {len(batch_rows)} JSON object(s), one per input row, in the same order.\n' + f'{output_schema_line}' + f'Copy {TABULAR_EXPORT_INPUT_ROW_TOKEN_FIELD} exactly from each input row into its matching output object. ' + f'The {TABULAR_EXPORT_INPUT_ROW_NUMBER_FIELD} and {TABULAR_EXPORT_INPUT_ROW_IDENTITY_FIELD} fields are internal; ' + 'do not include those two fields in generated objects.\n' 'Do not drop, merge, summarize, or cap rows.\n' 'Input rows may include normalized helper fields such as comment_id, body_text, source_file, attachment_present, attachment_names, and attachment_text. Use those normalized fields when they are present.\n' 'Input rows may include a referenced_documents array containing row-linked evidence from explicitly referenced non-tabular documents. Use that evidence as part of the source row context when it is relevant to the requested output.\n' @@ -380,6 +1051,8 @@ async def _generate_batch_entries( selected_sheet, retry_attempts, run_id, + expected_output_schema=None, + batch_timeout_seconds=TABULAR_EXPORT_DEFAULT_BATCH_TIMEOUT_SECONDS, ): batch_number = batch_index + 1 batch_prompt = _build_batch_prompt( @@ -389,11 +1062,20 @@ async def _generate_batch_entries( total_batches, source_file_name, selected_sheet=selected_sheet, + output_schema=expected_output_schema, ) parsed_entries = None raw_response_content = '' mismatch_count = 0 + last_validation_error = None + timeout_seconds = max( + _safe_float( + batch_timeout_seconds, + default=TABULAR_EXPORT_DEFAULT_BATCH_TIMEOUT_SECONDS, + ), + 0.001, + ) for attempt_number in range(1, retry_attempts + 1): chat_history = SKChatHistory() chat_history.add_system_message( @@ -409,12 +1091,29 @@ async def _generate_batch_entries( chat_history.add_user_message(batch_prompt) execution_settings = AzureChatPromptExecutionSettings(service_id='tabular-generated-output-background') - result = await chat_service.get_chat_message_contents(chat_history, execution_settings) + try: + result = await asyncio.wait_for( + chat_service.get_chat_message_contents(chat_history, execution_settings), + timeout=timeout_seconds, + ) + except asyncio.TimeoutError as exc: + raise TimeoutError( + f'Background structured export batch {batch_number}/{total_batches} ' + f'timed out after {timeout_seconds:g} seconds.' + ) from exc raw_response_content = result[0].content if result and result[0].content else '' parsed_entries = _parse_generated_json_entries(raw_response_content) if raw_response_content else None parsed_entry_count = len(parsed_entries) if parsed_entries is not None else 0 if parsed_entries is not None and parsed_entry_count == len(batch_rows): - return parsed_entries, mismatch_count + try: + normalized_entries, output_schema = _normalize_generated_batch_entries( + batch_rows, + parsed_entries, + expected_output_schema=expected_output_schema, + ) + return normalized_entries, mismatch_count, output_schema + except ValueError as exc: + last_validation_error = str(exc) mismatch_count += 1 log_event( @@ -426,15 +1125,19 @@ async def _generate_batch_entries( 'attempt_number': attempt_number, 'expected_row_count': len(batch_rows), 'parsed_row_count': parsed_entry_count, + 'validation_error': last_validation_error, 'response_char_count': len(raw_response_content), 'response_preview': _truncate_response_preview(raw_response_content), }, debug_only=True, ) + failure_detail = last_validation_error or ( + f'returned {len(parsed_entries) if parsed_entries is not None else 0} object(s) ' + f'for {len(batch_rows)} input row(s)' + ) raise ValueError( - f'Background structured export batch {batch_number}/{total_batches} returned ' - f'{len(parsed_entries) if parsed_entries is not None else 0} object(s) for {len(batch_rows)} input row(s).' + f'Background structured export batch {batch_number}/{total_batches} failed validation: {failure_detail}.' ) @@ -448,10 +1151,12 @@ async def _generate_batch_entries_for_window( selected_sheet, retry_attempts, run_id, + expected_output_schema, + batch_timeout_seconds, ): async with semaphore: batch_started_at = time.monotonic() - batch_entries, mismatch_count = await _generate_batch_entries( + batch_entries, mismatch_count, output_schema = await _generate_batch_entries( chat_service, user_question, batch_request['rows'], @@ -461,13 +1166,17 @@ async def _generate_batch_entries_for_window( selected_sheet, retry_attempts, run_id, + expected_output_schema=expected_output_schema, + batch_timeout_seconds=batch_timeout_seconds, ) return { 'batch_number': batch_request['batch_number'], 'batch_entries': batch_entries, + 'batch_summary': _build_generated_batch_summary(batch_entries), 'batch_row_count': len(batch_entries), 'elapsed_seconds': time.monotonic() - batch_started_at, 'mismatch_count': mismatch_count, + 'output_schema': output_schema, } @@ -481,6 +1190,8 @@ async def _generate_batch_window_entries( retry_attempts, run_id, batch_concurrency, + expected_output_schema=None, + batch_timeout_seconds=TABULAR_EXPORT_DEFAULT_BATCH_TIMEOUT_SECONDS, ): semaphore = asyncio.Semaphore(max(1, batch_concurrency)) tasks = [ @@ -494,6 +1205,8 @@ async def _generate_batch_window_entries( selected_sheet, retry_attempts, run_id, + expected_output_schema, + batch_timeout_seconds, ) for batch_request in batch_requests ] @@ -530,6 +1243,46 @@ def should_queue_tabular_generated_output_background(row_count, batch_count, set return _safe_int(row_count) > inline_max_rows or _safe_int(batch_count) > inline_max_batches +def build_tabular_generated_output_row_batches(rows, settings=None): + """Split structured rows using the shared generated-export size budget.""" + settings = settings or {} + max_batch_rows = _settings_int( + settings, + 'tabular_generated_output_max_batch_rows', + TABULAR_EXPORT_DEFAULT_SOURCE_BATCH_ROWS, + minimum=1, + maximum=100, + ) + max_batch_chars = _settings_int( + settings, + 'tabular_generated_output_max_batch_chars', + TABULAR_EXPORT_DEFAULT_SOURCE_BATCH_CHARS, + minimum=6000, + maximum=120000, + ) + batches = [] + current_batch = [] + current_batch_chars = 0 + + for row in rows or []: + row_text = json.dumps(row, default=str, ensure_ascii=False, separators=(',', ':')) + if current_batch and ( + len(current_batch) >= max_batch_rows + or current_batch_chars + len(row_text) > max_batch_chars + ): + batches.append(current_batch) + current_batch = [] + current_batch_chars = 0 + + current_batch.append(row) + current_batch_chars += len(row_text) + + if current_batch: + batches.append(current_batch) + + return batches + + def _parse_iso_datetime(value): normalized_value = str(value or '').strip() if not normalized_value: @@ -628,23 +1381,32 @@ def _scheduler_candidate_sort_key(run): ) -def _query_scheduler_candidates_by_status(status, scan_limit): +def _query_scheduler_candidates_by_status(status, scan_limit, settings): per_status_limit = _safe_int(scan_limit, default=TABULAR_EXPORT_DEFAULT_SCAN_LIMIT, minimum=1, maximum=10) query = ( - f"SELECT TOP {per_status_limit} " + "SELECT " "c.id, c.user_id, c.status, c.created_at, c.updated_at, c.last_heartbeat_at, " "c.next_attempt_at, c.last_error, c.transient_failure_count " - "FROM c WHERE c.type = @type AND c.status = @status" + "FROM c WHERE c.type = @type AND c.status = @status " + "ORDER BY c.updated_at ASC" ) try: - return list(cosmos_tabular_export_runs_container.query_items( + query_results = cosmos_tabular_export_runs_container.query_items( query=query, parameters=[ {'name': '@type', 'value': TABULAR_EXPORT_RUN_TYPE}, {'name': '@status', 'value': status}, ], enable_cross_partition_query=True, - )) + ) + eligible_candidates = [] + for run in query_results: + if not _scheduler_candidate_reason(run, settings): + continue + eligible_candidates.append(run) + if len(eligible_candidates) >= per_status_limit: + break + return eligible_candidates except Exception as exc: log_event( '[Tabular Generated Output] Scheduler candidate query failed', @@ -679,6 +1441,14 @@ def _can_resume_run(run, settings=None): return False +def _can_cancel_run(run): + status = str((run or {}).get('status') or '').strip().lower() + return not run.get('publishing_started_at') and status not in { + TABULAR_EXPORT_STATUS_COMPLETED, + TABULAR_EXPORT_STATUS_CANCELED, + } + + def _build_checkpoint_summary(completed_batches, batch_count, processed_rows, row_count): checkpoint_parts = [] if batch_count: @@ -727,6 +1497,16 @@ def _build_run_status_detail(run, settings, retryable_failure, can_resume): 'retry_delay_seconds': None, } if status == TABULAR_EXPORT_STATUS_RUNNING: + if run.get('publishing_started_at'): + return { + 'status_label': 'Finalizing', + 'status_tone': 'info', + 'status_detail': 'Export validation passed and the final artifact is being published.', + 'is_stale': False, + 'waiting_for_retry': False, + 'retry_due': False, + 'retry_delay_seconds': None, + } return { 'status_label': 'Running', 'status_tone': 'info', @@ -828,6 +1608,8 @@ def _build_run_public_status(run, settings=None): 'storage_scope': 'chat', 'source_file_name': run.get('source_file_name'), 'selected_sheet': run.get('selected_sheet'), + 'summary': run.get('post_run_summary'), + 'suppress_assistant_table_export': True, } return { @@ -865,11 +1647,13 @@ def _build_run_public_status(run, settings=None): 'manual_resume_count': _safe_int(run.get('manual_resume_count')), 'next_attempt_at': run.get('next_attempt_at'), 'can_resume': can_resume, + 'can_cancel': _can_cancel_run(run), 'retryable_failure': retryable_failure, 'artifact_message_id': final_artifact.get('artifact_message_id'), 'file_name': final_artifact.get('file_name') or run.get('generated_file_name'), 'generated_artifact': generated_artifact, 'capability': 'tabular', + 'suppress_assistant_table_export': True, 'background_export': not ( str(run.get('status') or '').strip().lower() == TABULAR_EXPORT_STATUS_COMPLETED and generated_artifact @@ -884,6 +1668,7 @@ def build_background_tabular_generated_output_metadata(run): 'export_run_id': public_status.get('run_id'), 'background_export': True, 'capability': 'tabular', + 'suppress_assistant_table_export': True, 'summary': ( f"Queued structured {str(public_status.get('output_format') or 'json').upper()} export " f"for {public_status.get('row_count', 0)} row(s) across {public_status.get('batch_count', 0)} batch(es)." @@ -922,6 +1707,18 @@ def resume_tabular_generated_output_run(user_id, run_id): except CosmosResourceNotFoundError: return None + try: + _authorize_tabular_export_run_execution(run) + except (LookupError, PermissionError, ValueError): + return { + 'success': False, + 'resumed': False, + 'submitted': False, + 'authorization_failed': True, + 'message': 'Background export access is no longer authorized.', + 'run': _build_run_public_status(run, settings=settings), + } + status = str(run.get('status') or '').strip().lower() if status == TABULAR_EXPORT_STATUS_COMPLETED: return { @@ -939,6 +1736,14 @@ def resume_tabular_generated_output_run(user_id, run_id): 'message': 'Canceled background exports cannot be continued.', 'run': _build_run_public_status(run, settings=settings), } + if not _can_cancel_run(run): + return { + 'success': False, + 'resumed': False, + 'submitted': False, + 'message': 'Background export is already publishing and cannot be resumed.', + 'run': _build_run_public_status(run, settings=settings), + } if status == TABULAR_EXPORT_STATUS_RUNNING and not _is_stale_running_run(run, settings): return { 'success': True, @@ -970,10 +1775,21 @@ def resume_tabular_generated_output_run(user_id, run_id): 'manual_resume_count': _safe_int(run.get('manual_resume_count')) + 1, 'last_manual_resume_at': now, }) - run = _upsert_run(run) + try: + run = _replace_run(run) + except Exception as exc: + if getattr(exc, 'status_code', None) not in (409, 412): + raise + current_run = _read_run(normalized_user_id, normalized_run_id) + return { + 'success': False, + 'resumed': False, + 'submitted': False, + 'message': 'Background export state changed. Refresh its status before continuing.', + 'run': _build_run_public_status(current_run, settings=settings), + } submitted = submit_tabular_generated_output_run(normalized_run_id, normalized_user_id) run['submitted_to_executor'] = submitted - run = _upsert_run(run) log_event( '[Tabular Generated Output] Background export manually resumed', { @@ -998,6 +1814,86 @@ def resume_tabular_generated_output_run(user_id, run_id): } +def cancel_tabular_generated_output_run(user_id, run_id): + """Cancel a queued, running, retryable, or failed generated-output run.""" + normalized_user_id = str(user_id or '').strip() + normalized_run_id = str(run_id or '').strip() + if not normalized_user_id or not normalized_run_id: + return None + + settings = get_settings() + try: + run = _read_run(normalized_user_id, normalized_run_id) + except CosmosResourceNotFoundError: + return None + + status = str(run.get('status') or '').strip().lower() + if status == TABULAR_EXPORT_STATUS_COMPLETED: + return { + 'success': False, + 'canceled': False, + 'message': 'Completed background exports cannot be canceled.', + 'run': _build_run_public_status(run, settings=settings), + } + if status == TABULAR_EXPORT_STATUS_CANCELED: + return { + 'success': True, + 'canceled': True, + 'message': 'Background export is already canceled.', + 'run': _build_run_public_status(run, settings=settings), + } + if not _can_cancel_run(run): + return { + 'success': False, + 'canceled': False, + 'message': 'Background export is already publishing and can no longer be canceled.', + 'run': _build_run_public_status(run, settings=settings), + } + + now = _now_iso() + run.update({ + 'status': TABULAR_EXPORT_STATUS_CANCELED, + 'updated_at': now, + 'completed_at': now, + 'last_heartbeat_at': now, + 'lease_holder_id': None, + 'lease_expires_at': None, + 'next_attempt_at': None, + 'last_message': 'Background structured export canceled by the user', + 'last_error': None, + 'canceled_at': now, + }) + try: + run = _replace_run(run) + except Exception as exc: + if getattr(exc, 'status_code', None) not in (409, 412): + raise + current_run = _read_run(normalized_user_id, normalized_run_id) + return { + 'success': False, + 'canceled': False, + 'message': 'Background export state changed. Refresh its status before canceling.', + 'run': _build_run_public_status(current_run, settings=settings), + } + log_event( + '[Tabular Generated Output] Background export canceled', + { + 'run_id': normalized_run_id, + 'conversation_id': run.get('conversation_id'), + 'user_id': normalized_user_id, + 'completed_batches': run.get('completed_batches'), + 'processed_rows': run.get('processed_rows'), + }, + level=logging.INFO, + ) + return { + 'success': True, + 'canceled': True, + 'message': 'Background export canceled.', + 'run': _build_run_public_status(run, settings=settings), + } + + def _read_run(user_id, run_id): return cosmos_tabular_export_runs_container.read_item( item=run_id, @@ -1005,6 +1901,37 @@ def _read_run(user_id, run_id): ) +def _raise_if_tabular_export_canceled(run): + current_run = _read_run(run.get('user_id'), run.get('id')) + current_status = str(current_run.get('status') or '').strip().lower() + if current_status == TABULAR_EXPORT_STATUS_CANCELED: + run.clear() + run.update(current_run) + raise TabularExportCanceledError('Background structured export was canceled') + + claim_matches = ( + current_status == TABULAR_EXPORT_STATUS_RUNNING + and str(current_run.get('lease_holder_id') or '') == str(run.get('lease_holder_id') or '') + and _safe_int(current_run.get('lease_generation')) == _safe_int(run.get('lease_generation')) + ) + if not claim_matches: + raise TabularExportLeaseLostError('Background structured export worker lost its claim') + + run['_etag'] = current_run.get('_etag') + return current_run + + +def _replace_claimed_run(run): + try: + return _replace_run(run) + except Exception as exc: + if getattr(exc, 'status_code', None) in (409, 412): + raise TabularExportLeaseLostError( + 'Background structured export worker lost its claim' + ) from exc + raise + + def _replace_run(run): return cosmos_tabular_export_runs_container.replace_item( item=run.get('id'), @@ -1014,10 +1941,6 @@ def _replace_run(run): ) -def _upsert_run(run): - return cosmos_tabular_export_runs_container.upsert_item(run) - - def _lease_holder_id(): return f"{socket.gethostname()}:{os.getpid()}:{uuid.uuid4().hex}" @@ -1080,6 +2003,7 @@ def _try_claim_run(user_id, run_id, settings): 'completed_at': None, 'last_heartbeat_at': now.isoformat(), 'lease_holder_id': _lease_holder_id(), + 'lease_generation': _safe_int(run.get('lease_generation')) + 1, 'lease_expires_at': (now + timedelta(seconds=lease_seconds)).isoformat(), 'next_attempt_at': None, 'last_message': 'Background structured export is running', @@ -1107,7 +2031,10 @@ def _mark_run_failed(run, error_message): 'last_error': str(error_message or 'Unknown error')[:1000], 'last_message': 'Background structured export failed', }) - _upsert_run(run) + try: + run = _replace_claimed_run(run) + except TabularExportLeaseLostError: + return _read_run(run.get('user_id'), run.get('id')) log_event( '[Tabular Generated Output] Background export run failed', { @@ -1156,7 +2083,10 @@ def _mark_run_retryable(run, error_message, settings): 'transient_failure_count': transient_failure_count, 'next_attempt_at': next_attempt_at, }) - _upsert_run(run) + try: + run = _replace_claimed_run(run) + except TabularExportLeaseLostError: + return _read_run(run.get('user_id'), run.get('id')) log_event( '[Tabular Generated Output] Background export run requeued after transient failure', { @@ -1231,7 +2161,7 @@ def _update_run_progress(run, completed_batches, processed_rows, batch_rows, bat }) run['recent_batches'] = recent_batches - return _upsert_run(run) + return _replace_claimed_run(run) def _log_progress_if_due(run, last_logged_at): @@ -1263,44 +2193,124 @@ def _log_progress_if_due(run, last_logged_at): return now_monotonic -def _assemble_output_entries(run): - output_entries = [] +def _write_ordered_output_stream(run, output_stream): user_id = run.get('user_id') conversation_id = run.get('conversation_id') run_id = run.get('id') batch_count = _safe_int(run.get('batch_count')) + expected_row_count = _safe_int(run.get('row_count')) + output_format = str(run.get('output_format') or 'json').strip().lower() or 'json' + output_schema = list(run.get('output_schema') or []) + if not output_schema: + raise ValueError('Generated output schema is missing') + if TABULAR_EXPORT_OUTPUT_ROW_NUMBER_FIELD not in output_schema: + raise ValueError('Generated output schema is missing source row order') + + csv_writer = None + safe_output_schema = None + if output_format == 'csv': + safe_output_schema = build_safe_csv_headers(output_schema) + csv_writer = csv.DictWriter(output_stream, fieldnames=safe_output_schema, lineterminator='\n') + csv_writer.writeheader() + else: + output_stream.write('[\n') + + written_row_count = 0 + expected_source_row_number = 1 for batch_number in range(1, batch_count + 1): batch_blob_path = _output_blob_path(user_id, conversation_id, run_id, batch_number) batch_entries = _download_json_blob(batch_blob_path) - if isinstance(batch_entries, list): - output_entries.extend(batch_entries) - return output_entries + if not isinstance(batch_entries, list): + raise ValueError(f'Output checkpoint {batch_number}/{batch_count} was not a JSON array') + + for batch_row_index, entry in enumerate(batch_entries, start=1): + if not isinstance(entry, dict): + raise ValueError( + f'Output checkpoint {batch_number}/{batch_count} row {batch_row_index} was not an object' + ) + if set(entry) != set(output_schema): + raise ValueError( + f'Output checkpoint {batch_number}/{batch_count} row {batch_row_index} has schema drift' + ) + + source_row_number = _safe_int(entry.get(TABULAR_EXPORT_OUTPUT_ROW_NUMBER_FIELD)) + if source_row_number != expected_source_row_number: + raise ValueError( + f'Source row order gap or overlap: expected {expected_source_row_number}, ' + f'found {source_row_number}' + ) + ordered_entry = { + field_name: entry.get(field_name) + for field_name in output_schema + } + if csv_writer: + csv_writer.writerow({ + safe_field_name: _serialize_generated_output_value(ordered_entry.get(field_name)) + for field_name, safe_field_name in zip(output_schema, safe_output_schema) + }) + else: + if written_row_count: + output_stream.write(',\n') + output_stream.write(json.dumps(ordered_entry, default=str, ensure_ascii=False)) + + written_row_count += 1 + expected_source_row_number += 1 + + if output_format != 'csv': + output_stream.write('\n]\n') + if written_row_count != expected_row_count: + raise ValueError( + f'Final output row count {written_row_count} does not match expected count {expected_row_count}' + ) + return written_row_count def _complete_run(run): - output_entries = _assemble_output_entries(run) output_format = str(run.get('output_format') or 'json').strip().lower() or 'json' - if output_format == 'csv': - serialized_output = _build_generated_output_csv(output_entries) - else: - serialized_output = json.dumps(output_entries, indent=2, default=str, ensure_ascii=False) - generated_file_name = run.get('generated_file_name') or _build_generated_file_name( run.get('source_file_name'), output_format, ) - upload_result = upload_generated_analysis_artifact_for_user( - current_user_id=run.get('user_id'), - conversation_id=run.get('conversation_id'), - file_name=generated_file_name, - file_content=serialized_output, - capability='tabular', - output_format=output_format, - summary=( - f"Saved {len(output_entries)} row(s) to {generated_file_name} " - 'from a durable background tabular export.' - ), - ) + with tempfile.SpooledTemporaryFile( + max_size=TABULAR_EXPORT_FINAL_SPOOL_MAX_MEMORY_BYTES, + mode='w+b', + ) as binary_output_stream: + text_output_stream = io.TextIOWrapper( + binary_output_stream, + encoding='utf-8', + newline='', + write_through=True, + ) + try: + output_entry_count = _write_ordered_output_stream(run, text_output_stream) + post_run_summary = _build_compact_post_run_summary(run) + _authorize_tabular_export_run_execution(run) + _raise_if_tabular_export_canceled(run) + if not run.get('publishing_started_at'): + run.update({ + 'publishing_started_at': _now_iso(), + 'last_message': 'Final validation passed; publishing the generated artifact', + }) + run = _replace_claimed_run(run) + _authorize_tabular_export_run_execution(run) + text_output_stream.flush() + output_size = binary_output_stream.tell() + binary_output_stream.seek(0) + upload_result = upload_generated_analysis_artifact_stream_for_user( + current_user_id=run.get('user_id'), + conversation_id=run.get('conversation_id'), + file_name=generated_file_name, + file_stream=binary_output_stream, + file_size=output_size, + capability='tabular', + output_format=output_format, + summary=post_run_summary, + artifact_idempotency_key=f"tabular-generated-output:{run.get('id')}", + ) + _raise_if_tabular_export_canceled(run) + finally: + text_output_stream.detach() + uploaded_message = upload_result.get('message') or {} now = _now_iso() run.update({ @@ -1308,9 +2318,10 @@ def _complete_run(run): 'updated_at': now, 'completed_at': now, 'last_heartbeat_at': now, - 'processed_rows': len(output_entries), + 'processed_rows': output_entry_count, 'completed_batches': _safe_int(run.get('batch_count')), 'last_message': 'Background structured export completed', + 'post_run_summary': post_run_summary, 'generated_file_name': uploaded_message.get('file_name') or generated_file_name, 'final_artifact': { 'artifact_message_id': uploaded_message.get('id'), @@ -1322,7 +2333,7 @@ def _complete_run(run): }, 'estimated_remaining_seconds': 0, }) - _upsert_run(run) + run = _replace_claimed_run(run) log_event( '[Tabular Generated Output] Background export completed', { @@ -1331,7 +2342,7 @@ def _complete_run(run): 'user_id': run.get('user_id'), 'source_file_name': run.get('source_file_name'), 'output_format': output_format, - 'row_count': len(output_entries), + 'row_count': output_entry_count, 'batch_count': run.get('batch_count'), 'artifact_message_id': uploaded_message.get('id'), 'generated_file_name': uploaded_message.get('file_name') or generated_file_name, @@ -1372,8 +2383,37 @@ def _build_batch_window(run, input_batches, user_id, run_id, window_start, windo batch_number, ) - if _blob_exists(output_blob_path): + if _blob_exists(output_blob_path) and not run.get('regenerate_legacy_output_checkpoints'): batch_entries = _download_json_blob(output_blob_path) + expected_output_schema = set(run.get('output_schema') or []) + if not isinstance(batch_entries, list) or not batch_entries: + raise ValueError(f'Output checkpoint {batch_number}/{batch_count} is empty or malformed') + for checkpoint_row_index, checkpoint_entry in enumerate(batch_entries, start=1): + if not isinstance(checkpoint_entry, dict): + raise ValueError( + f'Output checkpoint {batch_number}/{batch_count} row {checkpoint_row_index} is not an object' + ) + if set(checkpoint_entry) != expected_output_schema: + raise ValueError( + f'Output checkpoint {batch_number}/{batch_count} row {checkpoint_row_index} has schema drift' + ) + summary_blob_path = _output_summary_blob_path( + user_id, + run.get('conversation_id'), + run_id, + batch_number, + ) + if isinstance(batch_entries, list) and not _blob_exists(summary_blob_path): + _upload_json_blob( + summary_blob_path, + _build_generated_batch_summary(batch_entries), + metadata={ + 'run_id': run_id, + 'conversation_id': run.get('conversation_id'), + 'batch_number': batch_number, + 'generated_output_summary': 'true', + }, + ) batch_results[batch_number] = { 'batch_number': batch_number, 'batch_row_count': len(batch_entries) if isinstance(batch_entries, list) else 0, @@ -1404,8 +2444,29 @@ def _build_batch_window(run, input_batches, user_id, run_id, window_start, windo def _checkpoint_generated_batch_results(run, generated_results): + _raise_if_tabular_export_canceled(run) batch_results = {} - for generated_result in generated_results: + expected_output_schema = list(run.get('output_schema') or []) + ordered_results = sorted(generated_results, key=lambda result: result['batch_number']) + for generated_result in ordered_results: + generated_output_schema = list(generated_result.get('output_schema') or []) + if not expected_output_schema: + expected_output_schema = generated_output_schema + if generated_output_schema != expected_output_schema: + raise ValueError( + f"Generated output schema drifted in batch {generated_result['batch_number']}" + ) + + if not expected_output_schema: + raise ValueError('Generated output schema could not be established') + if list(run.get('output_schema') or []) != expected_output_schema: + run['output_schema'] = expected_output_schema + persisted_run = _replace_claimed_run(run) + run.clear() + run.update(persisted_run) + + for generated_result in ordered_results: + _raise_if_tabular_export_canceled(run) batch_number = generated_result['batch_number'] output_blob_path = _output_blob_path( run.get('user_id'), @@ -1413,14 +2474,49 @@ def _checkpoint_generated_batch_results(run, generated_results): run.get('id'), batch_number, ) + try: + _upload_json_blob( + output_blob_path, + generated_result['batch_entries'], + metadata={ + 'run_id': run.get('id'), + 'conversation_id': run.get('conversation_id'), + 'batch_number': batch_number, + 'generated_output': 'true', + 'lease_generation': run.get('lease_generation'), + }, + overwrite=bool(run.get('regenerate_legacy_output_checkpoints')), + ) + except ResourceExistsError: + checkpoint_entries = _download_json_blob(output_blob_path) + if ( + not isinstance(checkpoint_entries, list) + or len(checkpoint_entries) != generated_result['batch_row_count'] + or any( + not isinstance(entry, dict) or set(entry) != set(expected_output_schema) + for entry in checkpoint_entries + ) + ): + raise ValueError( + f'Concurrent output checkpoint {batch_number} failed schema or row-count validation' + ) + generated_result['batch_entries'] = checkpoint_entries + generated_result['batch_summary'] = _build_generated_batch_summary(checkpoint_entries) _upload_json_blob( - output_blob_path, - generated_result['batch_entries'], + _output_summary_blob_path( + run.get('user_id'), + run.get('conversation_id'), + run.get('id'), + batch_number, + ), + generated_result.get('batch_summary') or _build_generated_batch_summary( + generated_result['batch_entries'] + ), metadata={ 'run_id': run.get('id'), 'conversation_id': run.get('conversation_id'), 'batch_number': batch_number, - 'generated_output': 'true', + 'generated_output_summary': 'true', }, ) batch_results[batch_number] = { @@ -1433,6 +2529,31 @@ def _checkpoint_generated_batch_results(run, generated_results): return batch_results +def _build_passthrough_batch_results(run, batch_requests): + """Create checkpoint entries directly when rows are already final export output.""" + expected_output_schema = list(run.get('output_schema') or []) + generated_results = [] + for batch_request in batch_requests: + batch_started_at = time.monotonic() + batch_entries, output_schema = _normalize_generated_batch_entries( + batch_request['rows'], + batch_request['rows'], + expected_output_schema=expected_output_schema, + ) + if not expected_output_schema: + expected_output_schema = output_schema + generated_results.append({ + 'batch_number': batch_request['batch_number'], + 'batch_entries': batch_entries, + 'batch_summary': _build_generated_batch_summary(batch_entries), + 'batch_row_count': len(batch_entries), + 'elapsed_seconds': time.monotonic() - batch_started_at, + 'mismatch_count': 0, + 'output_schema': output_schema, + }) + return generated_results + + def _advance_run_progress_for_window(run, batch_results, completed_batches, processed_rows, window_start, window_end): for batch_number in range(window_start, window_end + 1): batch_result = batch_results.get(batch_number) @@ -1467,6 +2588,11 @@ def process_tabular_generated_output_run(run_id, user_id): return None try: + _authorize_tabular_export_run_execution(run) + run = _migrate_legacy_tabular_export_run(run) + if run.get('source_descriptor') and not run.get('source_staging_complete'): + run = _stage_tabular_generated_output_source(run, settings) + retry_attempts = _settings_int( settings, 'tabular_generated_output_batch_retry_attempts', @@ -1481,11 +2607,29 @@ def process_tabular_generated_output_run(run_id, user_id): minimum=1, maximum=TABULAR_EXPORT_MAX_BATCH_CONCURRENCY, ) - chat_service = _build_chat_service( - run.get('gpt_model'), + stale_seconds = _settings_int( settings, - model_context=run.get('model_context'), + 'tabular_generated_output_stale_seconds', + TABULAR_EXPORT_DEFAULT_STALE_SECONDS, + minimum=60, + ) + batch_timeout_seconds = min( + _settings_int( + settings, + 'tabular_generated_output_batch_timeout_seconds', + TABULAR_EXPORT_DEFAULT_BATCH_TIMEOUT_SECONDS, + minimum=30, + maximum=900, + ), + max(30, stale_seconds - 30), ) + chat_service = None + if not run.get('passthrough_input_rows'): + chat_service = _build_chat_service( + run.get('gpt_model'), + settings, + model_context=run.get('model_context'), + ) completed_batches = _safe_int(run.get('completed_batches')) processed_rows = _safe_int(run.get('processed_rows')) batch_count = _safe_int(run.get('batch_count')) @@ -1509,13 +2653,17 @@ def process_tabular_generated_output_run(run_id, user_id): 'batch_count': batch_count, 'resume_completed_batches': completed_batches, 'batch_concurrency': batch_concurrency, + 'batch_timeout_seconds': batch_timeout_seconds, }, level=logging.INFO, ) while completed_batches < batch_count: + _raise_if_tabular_export_canceled(run) window_start = completed_batches + 1 window_end = min(batch_count, window_start + batch_concurrency - 1) + if not run.get('output_schema'): + window_end = window_start batch_results, batch_requests = _build_batch_window( run, input_batches, @@ -1538,23 +2686,30 @@ def process_tabular_generated_output_run(run_id, user_id): 'window_end': window_end, 'batch_count': batch_count, 'batch_concurrency': batch_concurrency, + 'batch_timeout_seconds': batch_timeout_seconds, 'generation_request_count': len(batch_requests), }, debug_only=True, ) - generated_results, generation_error = asyncio.run( - _generate_batch_window_entries( - chat_service, - run.get('user_question'), - batch_requests, - batch_count, - run.get('source_file_name'), - run.get('selected_sheet'), - retry_attempts, - normalized_run_id, - batch_concurrency, + if run.get('passthrough_input_rows'): + generated_results = _build_passthrough_batch_results(run, batch_requests) + else: + generated_results, generation_error = asyncio.run( + _generate_batch_window_entries( + chat_service, + run.get('user_question'), + batch_requests, + batch_count, + run.get('source_file_name'), + run.get('selected_sheet'), + retry_attempts, + normalized_run_id, + batch_concurrency, + expected_output_schema=run.get('output_schema'), + batch_timeout_seconds=batch_timeout_seconds, + ) ) - ) + _raise_if_tabular_export_canceled(run) batch_results.update(_checkpoint_generated_batch_results(run, generated_results)) previous_completed_batches = completed_batches @@ -1572,7 +2727,12 @@ def process_tabular_generated_output_run(run_id, user_id): if completed_batches == previous_completed_batches: raise RuntimeError(f'No progress was made for batch window {window_start}-{window_end}') + _raise_if_tabular_export_canceled(run) return _complete_run(run) + except TabularExportCanceledError: + return _read_run(normalized_user_id, normalized_run_id) + except TabularExportLeaseLostError: + return _read_run(normalized_user_id, normalized_run_id) except Exception as exc: if _is_retryable_export_error(exc): return _mark_run_retryable(run, exc, settings) @@ -1609,6 +2769,8 @@ def queue_tabular_generated_output_run( gpt_model, settings=None, model_context=None, + source_descriptor=None, + passthrough_input_rows=False, ): """Stage batch input blobs, create a run record, and submit background processing.""" normalized_user_id = str(user_id or '').strip() @@ -1622,34 +2784,80 @@ def queue_tabular_generated_output_run( selected_sheet = str(source_candidate.get('selected_sheet') or '').strip() normalized_output_format = str(output_format or 'json').strip().lower() or 'json' generated_file_name = _build_generated_file_name(source_file_name, normalized_output_format) - row_batches = list(row_batches or []) + settings = settings or {} + source_descriptor = dict(source_descriptor or {}) + source_authorization = dict(source_candidate.get('source_authorization') or {}) staged_row_count = 0 staged_char_count = 0 - normalized_row_batches = [] - - for index, batch_rows in enumerate(row_batches, start=1): - if not isinstance(batch_rows, list): - batch_rows = list(batch_rows or []) - normalized_row_batches.append(batch_rows) - staged_row_count += len(batch_rows) - staged_char_count += len(json.dumps(batch_rows, default=str, ensure_ascii=False)) + staged_batch_count = 0 + + if source_descriptor: + staged_row_count = _safe_int(source_descriptor.get('expected_row_count')) + if staged_row_count <= 0: + raise ValueError('Source query descriptor must include the expected row count') + source_descriptor['batch_max_rows'] = _safe_int( + source_descriptor.get('batch_max_rows'), + default=_settings_int( + settings, + 'tabular_generated_output_max_batch_rows', + TABULAR_EXPORT_DEFAULT_SOURCE_BATCH_ROWS, + minimum=1, + maximum=100, + ), + minimum=1, + maximum=100, + ) + source_descriptor['batch_max_chars'] = _safe_int( + source_descriptor.get('batch_max_chars'), + default=_settings_int( + settings, + 'tabular_generated_output_max_batch_chars', + TABULAR_EXPORT_DEFAULT_SOURCE_BATCH_CHARS, + minimum=6000, + maximum=120000, + ), + minimum=6000, + maximum=120000, + ) + staged_batch_count = max( + 1, + math.ceil(staged_row_count / source_descriptor['batch_max_rows']), + ) + source_authorization = { + field_name: source_descriptor.get(field_name) + for field_name in ('source', 'scope_id', 'container', 'blob_path') + } + else: + for index, batch_rows in enumerate(row_batches or [], start=1): + if not isinstance(batch_rows, list): + batch_rows = list(batch_rows or []) + prepared_batch_rows = _prepare_tabular_source_rows( + batch_rows, + start_row=staged_row_count, + token_namespace=run_id, + ) + _upload_json_blob( + _input_blob_path(normalized_user_id, normalized_conversation_id, run_id, index), + prepared_batch_rows, + metadata={ + 'run_id': run_id, + 'conversation_id': normalized_conversation_id, + 'generated_output_input': 'true', + 'batch_number': index, + }, + ) + staged_row_count += len(prepared_batch_rows) + staged_char_count += len(json.dumps(prepared_batch_rows, default=str, ensure_ascii=False)) + staged_batch_count = index - input_blob_path = _input_batches_blob_path(normalized_user_id, normalized_conversation_id, run_id) - _upload_json_blob( - input_blob_path, - normalized_row_batches, - metadata={ - 'run_id': run_id, - 'conversation_id': normalized_conversation_id, - 'generated_output_input': 'true', - 'batch_count': len(normalized_row_batches), - }, - ) + if not staged_batch_count or not staged_row_count: + raise ValueError('At least one source row is required for a background tabular export') now = _now_iso() run = { 'id': run_id, 'type': TABULAR_EXPORT_RUN_TYPE, + 'contract_version': TABULAR_EXPORT_CONTRACT_VERSION, 'user_id': normalized_user_id, 'conversation_id': normalized_conversation_id, 'status': TABULAR_EXPORT_STATUS_QUEUED, @@ -1664,13 +2872,21 @@ def queue_tabular_generated_output_run( 'output_format': normalized_output_format, 'gpt_model': str(gpt_model or '').strip(), 'model_context': model_context if isinstance(model_context, dict) else {}, + 'passthrough_input_rows': bool(passthrough_input_rows), 'generated_file_name': generated_file_name, 'row_count': staged_row_count, - 'batch_count': len(row_batches), + 'batch_count': staged_batch_count, 'completed_batches': 0, 'processed_rows': 0, + 'output_schema': None, + 'source_descriptor': source_descriptor or None, + 'source_authorization': source_authorization or None, + 'source_staging_complete': not bool(source_descriptor), + 'source_staged_rows': 0 if source_descriptor else staged_row_count, + 'source_staged_batches': 0 if source_descriptor else staged_batch_count, + 'source_scan_row_count': 0, 'input_blob_container': storage_account_personal_chat_container_name, - 'input_blob_path': input_blob_path, + 'input_blob_path': None, 'input_blob_prefix': f'{normalized_user_id}/{normalized_conversation_id}/generated/tabular_runs/{run_id}/input/', 'output_blob_container': storage_account_personal_chat_container_name, 'output_blob_prefix': f'{normalized_user_id}/{normalized_conversation_id}/generated/tabular_runs/{run_id}/output/', @@ -1697,8 +2913,9 @@ def queue_tabular_generated_output_run( 'selected_sheet': selected_sheet, 'output_format': normalized_output_format, 'row_count': staged_row_count, - 'batch_count': len(row_batches), + 'batch_count': staged_batch_count, 'staged_input_char_count': staged_char_count, + 'source_backed': bool(source_descriptor), 'submitted_to_executor': submitted, }, level=logging.INFO, @@ -1725,7 +2942,7 @@ def check_due_tabular_generated_output_runs_once(limit=None): scanned_candidates = [] status_counts = {} for status in TABULAR_EXPORT_SCHEDULER_STATUSES: - status_candidates = _query_scheduler_candidates_by_status(status, scan_limit) + status_candidates = _query_scheduler_candidates_by_status(status, scan_limit, settings) status_counts[status] = len(status_candidates) scanned_candidates.extend(status_candidates) @@ -1754,15 +2971,19 @@ def check_due_tabular_generated_output_runs_once(limit=None): for candidate in candidates: run = candidate.get('run') or {} status = str(run.get('status') or '').strip().lower() - processed_run = process_tabular_generated_output_run(run.get('id'), run.get('user_id')) - if processed_run: - processed.append(processed_run.get('id')) + submitted = submit_tabular_generated_output_run(run.get('id'), run.get('user_id')) + if submitted: + processed.append(run.get('id')) else: - skipped.append({ - 'run_id': run.get('id'), - 'status': status, - 'reason': f"{candidate.get('reason')}; claim or processing did not start", - }) + processed_run = process_tabular_generated_output_run(run.get('id'), run.get('user_id')) + if processed_run: + processed.append(processed_run.get('id')) + else: + skipped.append({ + 'run_id': run.get('id'), + 'status': status, + 'reason': f"{candidate.get('reason')}; claim or processing did not start", + }) if scanned_candidates or candidates: log_event( diff --git a/application/single_app/functions_workflow_runner.py b/application/single_app/functions_workflow_runner.py index aff2f3848..701309d8b 100644 --- a/application/single_app/functions_workflow_runner.py +++ b/application/single_app/functions_workflow_runner.py @@ -10,6 +10,7 @@ import logging import os import re +import time import uuid from contextlib import contextmanager from datetime import datetime, timezone @@ -40,9 +41,23 @@ cosmos_messages_container, cosmos_public_documents_container, cosmos_user_documents_container, + storage_account_personal_chat_container_name, ) from functions_activity_logging import log_conversation_creation, log_token_usage, log_workflow_run from functions_appinsights import log_event +from functions_assistant_table_exports import ( + build_safe_csv_headers, + has_generated_tabular_csv_output, + neutralize_csv_spreadsheet_formula, +) +from functions_generated_file_exports import ( + build_generated_file_artifact_metadata, + build_generated_file_export, + build_generated_file_output_guidance, + get_generated_file_export_content, + get_requested_generated_file_format, + has_generated_file_output, +) from functions_chart_operations import append_proactive_chart_guidance from functions_collaboration import ( create_collaboration_message_notifications, @@ -52,6 +67,7 @@ from functions_document_actions import ( DOCUMENT_ACTION_ANALYSIS_MODE_PER_DOCUMENT, DOCUMENT_ACTION_CONTEXT_WORKFLOW, + DOCUMENT_ACTION_TARGET_MODE_ALL, DOCUMENT_ACTION_TARGET_MODE_RECENT, DOCUMENT_ACTION_TYPE_COMPARISON, DOCUMENT_ACTION_TYPE_ANALYZE, @@ -66,7 +82,13 @@ normalize_document_action_analysis_mode, ) from functions_documents import select_current_documents, sort_documents -from functions_document_comparison import run_document_comparison +from functions_document_access_index import ( + DOCUMENT_ACCESS_SCOPE_GROUP, + DOCUMENT_ACCESS_SCOPE_PERSONAL, + DOCUMENT_ACCESS_SCOPE_PUBLIC, + enumerate_bounded_document_access_index_ids, +) +from functions_document_comparison import run_document_comparison, run_evidence_document_comparison from functions_debug import debug_print from functions_document_analysis import run_document_analysis from functions_file_sync import get_authorized_sync_source, queue_file_sync_source_run @@ -78,6 +100,27 @@ build_agent_citation_artifact_documents, make_json_serializable, ) +from functions_mixed_source_orchestration import ( + EVIDENCE_ENGINE_DOCUMENT_ANALYSIS, + EVIDENCE_ENGINE_TABULAR_TOOLS, + EVIDENCE_STATUS_COMPLETED, + EVIDENCE_STATUS_FAILED, + MixedSourceCancellationError, + MixedSourceFinalizationError, + SELECTION_MODE_SELECTED, + build_evidence_envelope, + build_failed_narrative_evidence_envelopes, + build_mixed_source_evidence_handoff, + compare_reauthorized_source_manifests, + evaluate_mixed_source_mode_outcome, + build_narrative_evidence_envelopes, + deduplicate_mixed_source_references, + emit_mixed_source_telemetry, + partition_source_manifest, + normalize_mixed_source_correlation_id, + resolve_authorized_source_manifest, + raise_if_mixed_source_cancelled, +) from model_endpoint_clients import ( MODEL_ENDPOINT_PROTOCOL_AZURE_OPENAI, ) @@ -85,10 +128,35 @@ from functions_notifications import create_workflow_priority_notification from functions_personal_workflows import save_personal_workflow_run, save_personal_workflow_run_item from functions_public_workspaces import get_user_visible_public_workspace_ids_from_settings -from functions_search_service import resolve_document_context, search_documents +from functions_search_service import ( + resolve_document_context, + search_documents, + search_relevant_tabular_candidates, +) from functions_search import normalize_search_id_list, normalize_search_scope, normalize_search_top_n -from functions_simplechat_operations import upload_generated_analysis_artifact_for_current_user -from functions_settings import get_settings, get_user_settings, is_tabular_processing_enabled, normalize_model_endpoints, resolve_model_endpoint_foundry_scope +from functions_simplechat_operations import ( + delete_generated_chat_artifact_for_current_user, + delete_generated_chat_artifact_for_user, + upload_generated_analysis_artifact_for_current_user, + upload_generated_analysis_artifact_for_user, +) +from functions_settings import ( + get_settings, + get_user_settings, + is_mixed_source_chat_search_enabled, + is_mixed_source_manifest_enabled, + is_cross_format_compare_enabled, + is_cross_format_compare_one_to_many_enabled, + is_tabular_processing_enabled, + normalize_model_endpoints, + resolve_model_endpoint_foundry_scope, +) +from functions_tabular_generated_exports import ( + build_background_tabular_generated_output_metadata, + build_tabular_generated_output_row_batches, + queue_tabular_generated_output_run, + should_queue_tabular_generated_output_background, +) from functions_source_review import ( URL_ACCESS_CONTEXT_WORKFLOW, compact_source_review_result_for_metadata, @@ -96,6 +164,7 @@ validate_url_access_request, ) from functions_thoughts import ThoughtTracker +from functions_tabular_generated_exports import cancel_tabular_generated_output_run from semantic_kernel_loader import load_user_semantic_kernel from semantic_kernel_plugins.plugin_invocation_logger import get_plugin_logger, sanitize_plugin_invocation_value from semantic_kernel_plugins.plugin_invocation_thoughts import register_plugin_invocation_thought_callback @@ -785,13 +854,13 @@ def _serialize_document_analysis_csv_value(value): if value is None: return '' if isinstance(value, (dict, list)): - return json.dumps(value, ensure_ascii=False, default=str) + return neutralize_csv_spreadsheet_formula(json.dumps(value, ensure_ascii=False, default=str)) if hasattr(value, 'isoformat') and not isinstance(value, str): try: - return value.isoformat() + return neutralize_csv_spreadsheet_formula(value.isoformat()) except TypeError: pass - return str(value) + return neutralize_csv_spreadsheet_formula(value) def _add_document_analysis_source_context(row, source_context): @@ -888,13 +957,14 @@ def add_column(column_name): for column_name in row.keys(): add_column(column_name) + safe_ordered_columns = build_safe_csv_headers(ordered_columns) output_buffer = io.StringIO() - writer = csv.DictWriter(output_buffer, fieldnames=ordered_columns, lineterminator='\n') + writer = csv.DictWriter(output_buffer, fieldnames=safe_ordered_columns, lineterminator='\n') writer.writeheader() for row in rows: writer.writerow({ - column_name: _serialize_document_analysis_csv_value(row.get(column_name)) - for column_name in ordered_columns + safe_column_name: _serialize_document_analysis_csv_value(row.get(column_name)) + for column_name, safe_column_name in zip(ordered_columns, safe_ordered_columns) }) return output_buffer.getvalue() @@ -995,6 +1065,8 @@ def _upload_document_analysis_generated_artifact( preview_rows=None, preview_items=None, preview_lines=None, + cancel_requested=None, + request_correlation_id=None, ): try: upload_result = upload_generated_analysis_artifact_for_current_user( @@ -1005,6 +1077,20 @@ def _upload_document_analysis_generated_artifact( output_format=output_format, summary=summary, ) + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + delete_generated_chat_artifact_for_current_user( + normalized_conversation_id, + (upload_result.get('message') or {}).get('id'), + ) + raise + except MixedSourceCancellationError: + raise except Exception as exc: debug_print( '[WorkflowDocumentAnalysis] Generated artifact upload skipped | ' @@ -1083,7 +1169,14 @@ def _maybe_create_document_analysis_generated_artifacts( analysis_prompt, conversation_id='', primary_generated_outputs=None, + cancel_requested=None, + request_correlation_id=None, ): + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) normalized_conversation_id = str(conversation_id or '').strip() if not normalized_conversation_id or not has_request_context(): return {'artifacts': [], 'assistant_reply': None} @@ -1109,72 +1202,85 @@ def _maybe_create_document_analysis_generated_artifacts( if create_lossless_artifacts: artifacts = [] structured_rows = _build_document_analysis_structured_rows(analysis_result) - - if artifact_intent.get('csv_artifact_recommended') and structured_rows: - csv_output = _build_document_analysis_rows_csv(structured_rows) - csv_file_name = _build_document_analysis_artifact_file_name(analysis_result, 'csv') - csv_summary = ( - f'Saved {len(structured_rows)} extracted analysis row(s) for {document_count} ' - 'source document(s) as a downloadable CSV artifact.' - ) - csv_artifact = _upload_document_analysis_generated_artifact( - normalized_conversation_id, - csv_file_name, - csv_output, - 'csv', - csv_summary, - preview_rows=structured_rows[:DOCUMENT_ANALYSIS_ARTIFACT_PREVIEW_ROW_COUNT], - ) - if csv_artifact: - artifacts.append(csv_artifact) - - markdown_output = _build_document_analysis_markdown_artifact(analysis_result) - should_create_markdown_artifact = bool( - ( - artifact_intent.get('markdown_analysis_artifact_recommended') - or (json_payload is not None and not json_artifact_requested) - ) - and markdown_output - and ( - not primary_tabular_outputs - or _prompt_explicitly_requests_markdown_artifact(analysis_prompt) - ) - ) - if should_create_markdown_artifact: - markdown_file_name = _build_document_analysis_artifact_file_name(analysis_result, 'md') - markdown_summary = ( - f'Saved the final analysis plus retained raw analysis notes for {document_count} ' - 'source document(s) as a downloadable Markdown artifact.' - ) - markdown_artifact = _upload_document_analysis_generated_artifact( - normalized_conversation_id, - markdown_file_name, - markdown_output, - 'md', - markdown_summary, - preview_lines=_build_document_analysis_preview_lines(analysis_reply), - ) - if markdown_artifact: - artifacts.append(markdown_artifact) - - if json_payload is not None and json_artifact_requested and not primary_tabular_outputs: - json_file_name = _build_document_analysis_artifact_file_name(analysis_result, 'json') - json_summary = _build_document_analysis_artifact_summary(document_count, 'json') - json_preview_items = [] - if isinstance(json_payload, list): - json_preview_items = json_payload[:DOCUMENT_ANALYSIS_ARTIFACT_PREVIEW_ITEM_COUNT] - elif isinstance(json_payload, dict): - json_preview_items = [json_payload] - json_artifact = _upload_document_analysis_generated_artifact( - normalized_conversation_id, - json_file_name, - json.dumps(json_payload, indent=2, ensure_ascii=False), - 'json', - json_summary, - preview_items=json_preview_items, + try: + if artifact_intent.get('csv_artifact_recommended') and structured_rows: + csv_output = _build_document_analysis_rows_csv(structured_rows) + csv_file_name = _build_document_analysis_artifact_file_name(analysis_result, 'csv') + csv_summary = ( + f'Saved {len(structured_rows)} extracted analysis row(s) for {document_count} ' + 'source document(s) as a downloadable CSV artifact.' + ) + csv_artifact = _upload_document_analysis_generated_artifact( + normalized_conversation_id, + csv_file_name, + csv_output, + 'csv', + csv_summary, + preview_rows=structured_rows[:DOCUMENT_ANALYSIS_ARTIFACT_PREVIEW_ROW_COUNT], + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + if csv_artifact: + artifacts.append(csv_artifact) + + markdown_output = _build_document_analysis_markdown_artifact(analysis_result) + should_create_markdown_artifact = bool( + ( + artifact_intent.get('markdown_analysis_artifact_recommended') + or (json_payload is not None and not json_artifact_requested) + ) + and markdown_output + and ( + not primary_tabular_outputs + or _prompt_explicitly_requests_markdown_artifact(analysis_prompt) + ) ) - if json_artifact: - artifacts.append(json_artifact) + if should_create_markdown_artifact: + markdown_file_name = _build_document_analysis_artifact_file_name(analysis_result, 'md') + markdown_summary = ( + f'Saved the final analysis plus retained raw analysis notes for {document_count} ' + 'source document(s) as a downloadable Markdown artifact.' + ) + markdown_artifact = _upload_document_analysis_generated_artifact( + normalized_conversation_id, + markdown_file_name, + markdown_output, + 'md', + markdown_summary, + preview_lines=_build_document_analysis_preview_lines(analysis_reply), + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + if markdown_artifact: + artifacts.append(markdown_artifact) + + if json_payload is not None and json_artifact_requested and not primary_tabular_outputs: + json_file_name = _build_document_analysis_artifact_file_name(analysis_result, 'json') + json_summary = _build_document_analysis_artifact_summary(document_count, 'json') + json_preview_items = [] + if isinstance(json_payload, list): + json_preview_items = json_payload[:DOCUMENT_ANALYSIS_ARTIFACT_PREVIEW_ITEM_COUNT] + elif isinstance(json_payload, dict): + json_preview_items = [json_payload] + json_artifact = _upload_document_analysis_generated_artifact( + normalized_conversation_id, + json_file_name, + json.dumps(json_payload, indent=2, ensure_ascii=False), + 'json', + json_summary, + preview_items=json_preview_items, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + if json_artifact: + artifacts.append(json_artifact) + except MixedSourceCancellationError: + for artifact in reversed(artifacts): + delete_generated_chat_artifact_for_current_user( + normalized_conversation_id, + artifact.get('artifact_message_id'), + ) + raise if artifacts or primary_tabular_outputs: assistant_reply = _build_document_analysis_multi_artifact_reply( @@ -1247,6 +1353,13 @@ def _maybe_create_document_analysis_generated_artifacts( summary, preview_items=preview_items, preview_lines=preview_lines, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, ) if not artifact_payload: return {'artifacts': [], 'assistant_reply': None} @@ -1289,7 +1402,18 @@ def _build_comparison_artifact_reply(left_document_name, right_count, output_for ) -def _maybe_create_comparison_generated_artifacts(comparison_result, comparison_prompt, conversation_id=''): +def _maybe_create_comparison_generated_artifacts( + comparison_result, + comparison_prompt, + conversation_id='', + cancel_requested=None, + request_correlation_id=None, +): + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) normalized_conversation_id = str(conversation_id or '').strip() if not normalized_conversation_id or not has_request_context(): return {'artifacts': [], 'assistant_reply': None} @@ -1335,6 +1459,11 @@ def _maybe_create_comparison_generated_artifacts(comparison_result, comparison_p summary = _build_comparison_artifact_summary(left_document_name, len(right_documents), output_format) try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) upload_result = upload_generated_analysis_artifact_for_current_user( conversation_id=normalized_conversation_id, file_name=file_name, @@ -1343,6 +1472,20 @@ def _maybe_create_comparison_generated_artifacts(comparison_result, comparison_p output_format=output_format, summary=summary, ) + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + delete_generated_chat_artifact_for_current_user( + normalized_conversation_id, + (upload_result.get('message') or {}).get('id'), + ) + raise + except MixedSourceCancellationError: + raise except Exception as exc: debug_print( '[WorkflowDocumentComparison] Generated artifact upload skipped | ' @@ -1447,6 +1590,191 @@ def _finalize_token_usage(aggregate): return token_usage +def _accumulate_token_usage_summary(aggregate, token_usage): + """Merge an already-normalized token summary into the existing aggregate.""" + if not isinstance(aggregate, dict) or not isinstance(token_usage, dict): + return + for key in ('prompt_tokens', 'completion_tokens', 'total_tokens'): + value = _coerce_token_count(token_usage.get(key)) + if value is not None: + aggregate[key] = int(aggregate.get(key, 0) or 0) + value + request_count = _coerce_token_count(token_usage.get('request_count')) + if request_count is None and any( + token_usage.get(key) not in (None, 0, '') + for key in ('prompt_tokens', 'completion_tokens', 'total_tokens') + ): + request_count = 1 + if request_count: + aggregate['request_count'] = int(aggregate.get('request_count', 0) or 0) + request_count + + +def _emit_mixed_source_token_telemetry( + settings, + mode, + token_usage, + request_correlation_id=None, +): + """Link existing aggregate token accounting to mixed-source request telemetry.""" + token_usage = token_usage if isinstance(token_usage, dict) else {} + return emit_mixed_source_telemetry( + settings, + 'native_execution', + mode, + request_correlation_id=request_correlation_id, + metrics={ + 'prompt_tokens': token_usage.get('prompt_tokens', 0), + 'completion_tokens': token_usage.get('completion_tokens', 0), + 'total_tokens': token_usage.get('total_tokens', 0), + 'token_request_count': token_usage.get('request_count', 0), + 'request_count': token_usage.get('request_count', 0), + }, + ) + + +def _rollback_mixed_source_generated_outputs( + user_id, + conversation_id, + generated_outputs, + reason='finalization', +): + """Cancel queued exports and delete exact generated artifacts by stored identity.""" + normalized_user_id = str(user_id or '').strip() + normalized_conversation_id = str(conversation_id or '').strip() + if not normalized_user_id or not normalized_conversation_id: + return {'canceled_export_count': 0, 'deleted_artifact_count': 0} + + export_run_ids = [] + artifact_message_ids = [] + for output in list(generated_outputs or []): + if not isinstance(output, dict): + continue + export_run_id = str(output.get('export_run_id') or '').strip() + artifact_message_id = str(output.get('artifact_message_id') or '').strip() + if export_run_id and export_run_id not in export_run_ids: + export_run_ids.append(export_run_id) + if artifact_message_id and artifact_message_id not in artifact_message_ids: + artifact_message_ids.append(artifact_message_id) + + canceled_export_count = 0 + deleted_artifact_count = 0 + rollback_failure_count = 0 + for export_run_id in export_run_ids: + try: + cancel_result = cancel_tabular_generated_output_run( + normalized_user_id, + export_run_id, + ) + if isinstance(cancel_result, dict) and cancel_result.get('canceled'): + canceled_export_count += 1 + except Exception: + rollback_failure_count += 1 + for artifact_message_id in artifact_message_ids: + try: + if delete_generated_chat_artifact_for_user( + normalized_user_id, + normalized_conversation_id, + artifact_message_id, + ): + deleted_artifact_count += 1 + except Exception: + rollback_failure_count += 1 + + log_event( + '[MixedSourceLifecycle] Generated output rollback completed.', + extra={ + 'canceled_export_count': canceled_export_count, + 'deleted_artifact_count': deleted_artifact_count, + 'rollback_failure_count': rollback_failure_count, + 'rollback_reason': str(reason or 'finalization')[:64], + }, + level=logging.INFO if not rollback_failure_count else logging.WARNING, + ) + return { + 'canceled_export_count': canceled_export_count, + 'deleted_artifact_count': deleted_artifact_count, + 'rollback_failure_count': rollback_failure_count, + } + + +def _reauthorize_mixed_source_workflow_result( + workflow, + action_config, + result, + settings, + conversation_id, + cancel_requested=None, + request_correlation_id=None, + additional_generated_outputs=None, +): + """Reauthorize a mixed result and roll back outputs before propagating failure.""" + result = result if isinstance(result, dict) else {} + execution_manifest = result.get('mixed_source_manifest') + if not isinstance(execution_manifest, list): + return + + user_id = str((workflow or {}).get('user_id') or '').strip() + requested_ids = [ + str(source.get('document_id') or '').strip() + for source in execution_manifest + if isinstance(source, dict) and str(source.get('document_id') or '').strip() + ] + selection_mode = str( + (action_config or {}).get('target_mode') or SELECTION_MODE_SELECTED + ).strip().lower() + if selection_mode not in {'selected', 'all', 'history', 'relevance'}: + selection_mode = SELECTION_MODE_SELECTED + + generated_outputs = list(result.get('generated_tabular_outputs') or []) + generated_outputs.extend(list(additional_generated_outputs or [])) + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'finalization', + request_correlation_id=request_correlation_id, + ) + fresh_manifest = resolve_authorized_source_manifest( + requested_ids, + user_id=user_id, + selection_mode=selection_mode, + conversation_id=conversation_id, + active_group_ids=(action_config or {}).get('active_group_ids'), + active_public_workspace_ids=(action_config or {}).get('active_public_workspace_id'), + doc_scope=(action_config or {}).get('doc_scope', 'all'), + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + finalization_result = compare_reauthorized_source_manifests( + execution_manifest, + fresh_manifest, + ) + if finalization_result['authorization_failure_count']: + emit_mixed_source_telemetry( + settings, + 'authorization_failure', + ( + 'compare' + if (action_config or {}).get('type') == DOCUMENT_ACTION_TYPE_COMPARISON + else 'analyze' + ), + request_correlation_id=request_correlation_id, + metrics={ + 'authorization_failure_count': finalization_result['authorization_failure_count'], + }, + dimensions={'outcome_status': 'failed'}, + ) + raise MixedSourceFinalizationError('authorization_lost') + if finalization_result['source_version_changed_count']: + raise MixedSourceFinalizationError('source_version_changed') + except (MixedSourceCancellationError, MixedSourceFinalizationError): + _rollback_mixed_source_generated_outputs( + user_id, + conversation_id, + generated_outputs, + reason='finalization', + ) + raise + + def _strip_agent_citation_artifact_refs(agent_citations): compact_citations = [] for citation in agent_citations or []: @@ -1562,7 +1890,7 @@ def _normalize_tabular_source_hint(scope): return 'workspace' -def _resolve_tabular_document_action_documents(action_config, user_id, conversation_id=''): +def _get_document_action_source_ids(action_config): action_config = action_config if isinstance(action_config, dict) else {} action_type = str(action_config.get('type') or '').strip().lower() @@ -1588,6 +1916,16 @@ def _resolve_tabular_document_action_documents(action_config, user_id, conversat document_ids.append(document_id) role_by_document_id[document_id] = 'right' + return document_ids, role_by_document_id + + +def _resolve_tabular_document_action_documents( + action_config, + user_id, + conversation_id='', +): + action_config = action_config if isinstance(action_config, dict) else {} + document_ids, role_by_document_id = _get_document_action_source_ids(action_config) if not document_ids: return [] @@ -1776,6 +2114,501 @@ def _build_tabular_comparison_action_prompt(comparison_prompt, left_document, ri ) +def _build_mixed_source_analyze_reduction_prompt(analysis_prompt, handoff): + """Build the one bounded collective reduction prompt for combined Analyze.""" + return ( + 'Answer the exact Analyze request using only the bounded evidence handoff below. ' + 'Treat computed tabular facts as tool-backed calculations and narrative facts as document excerpts. ' + 'Identify cross-source relationships only where the evidence supports them. Keep facts source-separated when needed. ' + 'Explicitly state missing, failed, unsupported, unresolved, or unprocessed evidence and never claim coverage for it.\n\n' + f'Analyze request:\n{str(analysis_prompt or "").strip()}\n\n' + f'Bounded evidence handoff:\n{json.dumps(handoff, ensure_ascii=False, separators=(",", ":"))}' + ) + + +def _build_mixed_source_analysis_coverage(handoff): + """Expose terminal coverage for every manifest entry with engine/status totals.""" + coverage = dict((handoff or {}).get('mixed_source_coverage') or {}) + engine_status_totals = {} + for envelope in list((handoff or {}).get('evidence_envelopes') or []): + engine = str(envelope.get('engine') or 'unknown') + status = str(envelope.get('status') or 'failed') + engine_status_totals.setdefault(engine, {})[status] = ( + engine_status_totals.setdefault(engine, {}).get(status, 0) + 1 + ) + coverage['engine_status_totals'] = engine_status_totals + coverage['document_count'] = coverage.get('requested_source_count', 0) + coverage['progress_meta'] = { + 'phase': 'complete', + 'phase_label': 'Complete' if not coverage.get('partial_coverage') else 'Partial', + 'phase_detail': 'Mixed-source Analyze completed with terminal source coverage', + 'status': 'partial' if coverage.get('partial_coverage') else 'completed', + 'percent_override': 100, + } + return coverage + + +def _execute_mixed_source_analyze_workflow( + workflow, + analysis_config, + settings, + invoke_prompt, + conversation_id='', + activity_callback=None, + thought_tracker=None, + live_thought_callback=None, + max_documents=None, + cancel_requested=None, + request_correlation_id=None, + token_usage_callback=None, +): + """Run combined Analyze cohorts natively, then reduce bounded evidence once.""" + started_at = time.perf_counter() + request_correlation_id = normalize_mixed_source_correlation_id( + request_correlation_id + ) + user_id = str(workflow.get('user_id') or '').strip() + requested_ids, _ = _get_document_action_source_ids(analysis_config) + requested_selection_mode = str( + analysis_config.get('selection_mode') + or analysis_config.get('target_mode') + or SELECTION_MODE_SELECTED + ).strip().lower() + if requested_selection_mode == 'all': + if not bool(settings.get('enable_mixed_source_analyze_all', False)): + raise ValueError( + 'Analyze All Documents is temporarily unavailable while its rollout flag is disabled.' + ) + if max_documents is None: + raise ValueError('Analyze All Documents requires a configured document limit.') + all_targets = _resolve_analyze_all_document_ids( + workflow, + analysis_config, + settings, + int(max_documents), + ) + analysis_config = {**analysis_config, **all_targets} + requested_ids = list(all_targets['document_ids']) + if max_documents is not None and len(requested_ids) > int(max_documents): + raise ValueError( + f'Analyze supports up to {int(max_documents)} authorized documents at a time.' + ) + if callable(activity_callback): + activity_callback({'type': 'mixed_source_progress', 'phase': 'resolving_sources', 'label': 'Resolving sources'}) + manifest = resolve_authorized_source_manifest( + requested_ids, + user_id=user_id, + selection_mode=requested_selection_mode, + conversation_id=conversation_id, + active_group_ids=analysis_config.get('active_group_ids'), + active_public_workspace_ids=analysis_config.get('active_public_workspace_id'), + doc_scope=analysis_config.get('doc_scope', 'all'), + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + partitions = partition_source_manifest(manifest) + evidence_envelopes = [] + generated_tabular_outputs = [] + tabular_agent_citations = [] + + narrative_sources = partitions['narrative_sources'] + if narrative_sources: + if callable(activity_callback): + activity_callback({'type': 'mixed_source_progress', 'phase': 'analyzing_narrative', 'label': 'Analyzing narrative documents'}) + try: + narrative_result = run_document_analysis( + user_id=user_id, + analysis_prompt=workflow.get('task_prompt', ''), + document_ids=[source.get('document_id') for source in narrative_sources], + invoke_prompt=invoke_prompt, + doc_scope=analysis_config.get('doc_scope'), + active_group_ids=analysis_config.get('active_group_ids'), + active_public_workspace_id=analysis_config.get('active_public_workspace_id'), + conversation_id=conversation_id, + window_unit=analysis_config.get('window_unit'), + window_size=analysis_config.get('window_size'), + window_percent=analysis_config.get('window_percent'), + max_retries_per_window=analysis_config.get('max_retries_per_window'), + activity_callback=activity_callback, + max_documents=max_documents, + include_coverage_summary=False, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + documents_by_id = { + str(document.get('document_id') or ''): document + for document in list((narrative_result.get('coverage') or {}).get('documents') or []) + } + narrative_items_by_id = { + str(item.get('document_id') or ''): item + for item in list(narrative_result.get('document_analysis_items') or []) + if str(item.get('document_id') or '').strip() + } + for source in narrative_sources: + document_id = source.get('document_id') + document_coverage = documents_by_id.get(str(document_id), {}) + narrative_item = narrative_items_by_id.get(str(document_id), {}) + failed_windows = int(document_coverage.get('failed_windows') or 0) + total_windows = int(document_coverage.get('total_windows') or 0) + processed_windows = int(document_coverage.get('processed_windows') or 0) + status = ( + EVIDENCE_STATUS_COMPLETED + if total_windows and processed_windows == total_windows and not failed_windows + else ('partial' if processed_windows else EVIDENCE_STATUS_FAILED) + ) + evidence_envelopes.append(build_evidence_envelope( + document_id=document_id, + source_kind='narrative', + engine=EVIDENCE_ENGINE_DOCUMENT_ANALYSIS, + status=status, + summary=str(narrative_item.get('text') or ''), + citations=[], + generated_artifacts=[], + coverage={'terminal': True, 'processed_windows': processed_windows, 'total_windows': total_windows, 'failed_windows': failed_windows}, + error='Narrative analysis could not be completed.' if status == EVIDENCE_STATUS_FAILED else None, + )) + except MixedSourceCancellationError: + raise + except Exception: + for source in narrative_sources: + evidence_envelopes.append(build_evidence_envelope( + document_id=source.get('document_id'), source_kind='narrative', + engine=EVIDENCE_ENGINE_DOCUMENT_ANALYSIS, status=EVIDENCE_STATUS_FAILED, + summary='Narrative evidence could not be completed for this source.', + coverage={'terminal': True}, error='Narrative analysis could not be completed.', + )) + + tabular_sources = partitions['tabular_sources'] + if tabular_sources: + if callable(activity_callback): + activity_callback({'type': 'mixed_source_progress', 'phase': 'analyzing_tabular', 'label': 'Analyzing tabular documents'}) + for source in tabular_sources: + tabular_config = dict(analysis_config) + tabular_config['document_ids'] = [source.get('document_id')] + tabular_payload = _maybe_execute_tabular_document_action( + DOCUMENT_ACTION_TYPE_ANALYZE, workflow, tabular_config, settings, + conversation_id=conversation_id, invoke_prompt=invoke_prompt, + thought_tracker=thought_tracker, live_thought_callback=live_thought_callback, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + token_usage_callback=token_usage_callback, + ) + if not tabular_payload: + evidence_envelopes.append(build_evidence_envelope( + document_id=source.get('document_id'), source_kind='tabular', + engine=EVIDENCE_ENGINE_TABULAR_TOOLS, status=EVIDENCE_STATUS_FAILED, + summary='Tabular evidence could not be completed for this source.', + coverage={'terminal': True}, error='Tabular analysis could not be completed.', + )) + continue + tabular_result = tabular_payload.get('result') or {} + evidence_envelopes.append(build_evidence_envelope( + document_id=source.get('document_id'), source_kind='tabular', + engine=EVIDENCE_ENGINE_TABULAR_TOOLS, status=EVIDENCE_STATUS_COMPLETED, + summary=str(tabular_result.get('analysis_reply') or tabular_result.get('reply') or ''), + citations=list(tabular_payload.get('agent_citations') or []), + generated_artifacts=list(tabular_payload.get('generated_tabular_outputs') or []), + coverage={'terminal': True, 'tool_call_count': 1}, + )) + generated_tabular_outputs.extend(tabular_payload.get('generated_tabular_outputs') or []) + tabular_agent_citations.extend(tabular_payload.get('agent_citations') or []) + + handoff = build_mixed_source_evidence_handoff( + manifest, + evidence_envelopes, + requested_selection_mode, + mode='analyze', + telemetry_settings=settings, + request_correlation_id=request_correlation_id, + ) + mode_outcome = evaluate_mixed_source_mode_outcome( + 'analyze', + { + 'entries': (handoff.get('mixed_source_coverage') or {}).get('terminal_ledger') or [], + 'partial_coverage': bool((handoff.get('mixed_source_coverage') or {}).get('partial_coverage')), + }, + ) + if not mode_outcome['should_reduce']: + raise RuntimeError( + 'Mixed-source Analyze could not prepare evidence from any selected source.' + ) + raise_if_mixed_source_cancelled( + cancel_requested, + 'reduction', + request_correlation_id=request_correlation_id, + ) + if callable(activity_callback): + activity_callback({'type': 'mixed_source_progress', 'phase': 'combining_findings', 'label': 'Combining findings'}) + collective_reply = str(invoke_prompt( + _build_mixed_source_analyze_reduction_prompt(workflow.get('task_prompt', ''), handoff), + stage='mixed_source_reduction', metadata={'requested_source_count': len(manifest)}, + ) or '').strip() + raise_if_mixed_source_cancelled( + cancel_requested, + 'reduction', + request_correlation_id=request_correlation_id, + ) + if not collective_reply: + collective_reply = 'The selected sources could not be combined into a final analysis.' + emit_mixed_source_telemetry( + settings, + 'reduction', + 'analyze', + request_correlation_id=request_correlation_id, + metrics={ + 'engine_call_count': len(evidence_envelopes), + 'model_request_count': 1, + 'latency_ms': round((time.perf_counter() - started_at) * 1000, 3), + }, + dimensions={ + 'selection_mode': requested_selection_mode, + 'outcome_status': mode_outcome['status'], + }, + ) + coverage = _build_mixed_source_analysis_coverage(handoff) + if callable(activity_callback): + activity_callback({ + 'type': 'mixed_source_progress', + 'phase': 'complete', + 'label': coverage['progress_meta']['phase_label'], + 'status': coverage['progress_meta']['status'], + }) + return { + 'reply': collective_reply, + 'analysis_reply': collective_reply, + 'coverage': coverage, + 'documents': list(coverage.get('sources') or []), + 'document_ids': [source.get('document_id') for source in manifest], + 'mixed_source_manifest': manifest, + 'mixed_source_evidence': handoff.get('evidence_envelopes') or [], + 'generated_tabular_outputs': generated_tabular_outputs, + 'agent_citations': tabular_agent_citations, + } + +def _resolve_cross_format_comparison_manifest( + comparison_config, + user_id, + conversation_id, + cancel_requested=None, + request_correlation_id=None, +): + """Resolve the Source and ordered Targets once for the mixed Compare decision.""" + requested_ids, role_by_document_id = _get_document_action_source_ids(comparison_config) + manifest = resolve_authorized_source_manifest( + requested_ids, + user_id=user_id, + selection_mode=SELECTION_MODE_SELECTED, + conversation_id=conversation_id, + active_group_ids=comparison_config.get('active_group_ids'), + active_public_workspace_ids=comparison_config.get('active_public_workspace_id'), + doc_scope=comparison_config.get('doc_scope', 'all'), + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + for source in manifest: + source['comparison_role'] = role_by_document_id.get(source.get('document_id'), 'target') + return manifest + + +def _raise_legacy_cross_format_compare_limitation(comparison_config, user_id, conversation_id=''): + """Fail closed when rollback would otherwise route a table through chunk analysis.""" + manifest = _resolve_cross_format_comparison_manifest(comparison_config, user_id, conversation_id) + partitions = partition_source_manifest(manifest) + if partitions['narrative_sources'] and partitions['tabular_sources']: + raise ValueError( + 'Mixed narrative and tabular Compare is temporarily unavailable while cross-format Compare is disabled.' + ) + + +def _execute_cross_format_comparison_workflow( + workflow, + comparison_config, + settings, + invoke_prompt, + conversation_id='', + activity_callback=None, + thought_tracker=None, + live_thought_callback=None, + cancel_requested=None, + request_correlation_id=None, + token_usage_callback=None, +): + """Prepare native envelopes for a mixed Source/Target Compare, then reuse pairwise reduction.""" + started_at = time.perf_counter() + request_correlation_id = normalize_mixed_source_correlation_id( + request_correlation_id + ) + user_id = str(workflow.get('user_id') or '').strip() + raise_if_mixed_source_cancelled( + cancel_requested, + 'manifest', + request_correlation_id=request_correlation_id, + ) + manifest = _resolve_cross_format_comparison_manifest( + comparison_config, + user_id, + conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + partitions = partition_source_manifest(manifest) + if not (partitions['narrative_sources'] and partitions['tabular_sources']): + return None + target_sources = [source for source in manifest if source.get('comparison_role') == 'right'] + if len(target_sources) > 1 and not is_cross_format_compare_one_to_many_enabled(settings): + raise ValueError('Cross-format Compare currently supports one Target while one-to-many rollout is disabled.') + if callable(activity_callback): + activity_callback({'type': 'mixed_source_progress', 'phase': 'resolving_sources', 'label': 'Resolving Source and Targets'}) + + evidence_by_id = {} + generated_tabular_outputs = [] + tabular_agent_citations = [] + for source in partitions['narrative_sources']: + if callable(activity_callback): + activity_callback({'type': 'mixed_source_progress', 'phase': 'analyzing_narrative', 'label': 'Analyzing narrative source'}) + try: + narrative_result = run_document_analysis( + user_id=user_id, + analysis_prompt=workflow.get('task_prompt', ''), + document_ids=[source.get('document_id')], + invoke_prompt=invoke_prompt, + doc_scope=comparison_config.get('doc_scope'), + active_group_ids=comparison_config.get('active_group_ids'), + active_public_workspace_id=comparison_config.get('active_public_workspace_id'), + conversation_id=conversation_id, + window_unit=comparison_config.get('window_unit'), + window_size=comparison_config.get('window_size'), + window_percent=comparison_config.get('window_percent'), + max_retries_per_window=comparison_config.get('max_retries_per_window'), + activity_callback=activity_callback, + max_documents=1, + include_coverage_summary=False, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + document_coverage = list((narrative_result.get('coverage') or {}).get('documents') or [{}])[0] + status = EVIDENCE_STATUS_COMPLETED if not document_coverage.get('failed_windows') else 'partial' + evidence_by_id[source.get('document_id')] = build_evidence_envelope( + document_id=source.get('document_id'), source_kind='narrative', + engine=EVIDENCE_ENGINE_DOCUMENT_ANALYSIS, status=status, + summary=str(narrative_result.get('analysis_reply') or narrative_result.get('reply') or ''), + coverage={'terminal': True, 'source_version': source.get('source_version'), **document_coverage}, + ) + except MixedSourceCancellationError: + raise + except Exception: + evidence_by_id[source.get('document_id')] = build_evidence_envelope( + document_id=source.get('document_id'), source_kind='narrative', + engine=EVIDENCE_ENGINE_DOCUMENT_ANALYSIS, status=EVIDENCE_STATUS_FAILED, + summary='Narrative evidence could not be completed for this source.', + coverage={'terminal': True, 'source_version': source.get('source_version')}, + error='Narrative analysis could not be completed.', + ) + + for source in partitions['tabular_sources']: + if callable(activity_callback): + activity_callback({'type': 'mixed_source_progress', 'phase': 'analyzing_tabular', 'label': 'Analyzing tabular source'}) + tabular_config = dict(comparison_config) + tabular_payload = _maybe_execute_tabular_document_action( + DOCUMENT_ACTION_TYPE_ANALYZE, workflow, { + **tabular_config, + 'type': DOCUMENT_ACTION_TYPE_ANALYZE, + 'document_ids': [source.get('document_id')], + }, + settings, conversation_id=conversation_id, invoke_prompt=invoke_prompt, + thought_tracker=thought_tracker, live_thought_callback=live_thought_callback, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + token_usage_callback=token_usage_callback, + ) + tabular_result = (tabular_payload or {}).get('result') or {} + if tabular_result.get('analysis_reply'): + evidence_by_id[source.get('document_id')] = build_evidence_envelope( + document_id=source.get('document_id'), source_kind='tabular', + engine=EVIDENCE_ENGINE_TABULAR_TOOLS, status=EVIDENCE_STATUS_COMPLETED, + summary=str(tabular_result.get('analysis_reply') or ''), + citations=list((tabular_payload or {}).get('agent_citations') or []), + generated_artifacts=list((tabular_payload or {}).get('generated_tabular_outputs') or []), + coverage={'terminal': True, 'source_version': source.get('source_version'), 'tool_call_count': 1}, + ) + generated_tabular_outputs.extend((tabular_payload or {}).get('generated_tabular_outputs') or []) + tabular_agent_citations.extend((tabular_payload or {}).get('agent_citations') or []) + else: + evidence_by_id[source.get('document_id')] = build_evidence_envelope( + document_id=source.get('document_id'), source_kind='tabular', + engine=EVIDENCE_ENGINE_TABULAR_TOOLS, status=EVIDENCE_STATUS_FAILED, + summary='Tabular evidence could not be completed for this source.', + coverage={'terminal': True, 'source_version': source.get('source_version')}, + error='Tabular analysis could not be completed.', + ) + + handoff = build_mixed_source_evidence_handoff( + manifest, + list(evidence_by_id.values()), + SELECTION_MODE_SELECTED, + mode='compare', + telemetry_settings=settings, + request_correlation_id=request_correlation_id, + ) + mode_outcome = evaluate_mixed_source_mode_outcome( + 'compare', + { + 'entries': (handoff.get('mixed_source_coverage') or {}).get('terminal_ledger') or [], + 'partial_coverage': bool((handoff.get('mixed_source_coverage') or {}).get('partial_coverage')), + }, + ) + if not mode_outcome['should_reduce']: + if mode_outcome.get('reason') == 'source_preparation_failed': + raise RuntimeError('Cross-format Compare Source could not be prepared.') + raise RuntimeError('Cross-format Compare could not prepare any Target evidence.') + envelopes = {envelope.get('document_id'): envelope for envelope in handoff.get('evidence_envelopes') or []} + def source_payload(source): + envelope = dict(envelopes.get(source.get('document_id')) or {}) + envelope['document_name'] = source.get('display_name') or ('Source' if source.get('comparison_role') == 'left' else 'Target') + role_label = 'Source' if source.get('comparison_role') == 'left' else 'Target' + evidence_type = 'computed tabular facts' if envelope.get('source_kind') == 'tabular' else 'narrative document analysis' + envelope['summary'] = ( + f'Role: {role_label}. Evidence type: {evidence_type}.\n' + f"{str(envelope.get('summary') or '').strip()}" + ) + return envelope + left_source = next((source for source in manifest if source.get('comparison_role') == 'left'), {}) + comparison_result = run_evidence_document_comparison( + workflow.get('task_prompt', ''), source_payload(left_source), + [source_payload(source) for source in target_sources], invoke_prompt, activity_callback=activity_callback, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + pairwise_coverage = dict(comparison_result.get('coverage') or {}) + mixed_coverage = _build_mixed_source_analysis_coverage(handoff) + mixed_coverage['failed_targets'] = list(pairwise_coverage.get('failed_targets') or []) + mixed_coverage['partial_coverage'] = bool( + mixed_coverage.get('partial_coverage') + or pairwise_coverage.get('partial_coverage') + ) + comparison_result['coverage'] = mixed_coverage + comparison_result['mixed_source_manifest'] = manifest + comparison_result['mixed_source_evidence'] = handoff.get('evidence_envelopes') or [] + comparison_result['generated_tabular_outputs'] = generated_tabular_outputs + comparison_result['agent_citations'] = tabular_agent_citations + emit_mixed_source_telemetry( + settings, + 'reduction', + 'compare', + request_correlation_id=request_correlation_id, + metrics={ + 'engine_call_count': len(evidence_by_id), + 'model_request_count': len(comparison_result.get('comparison_items') or []), + 'latency_ms': round((time.perf_counter() - started_at) * 1000, 3), + }, + dimensions={ + 'selection_mode': SELECTION_MODE_SELECTED, + 'outcome_status': mode_outcome['status'], + }, + ) + return comparison_result + + def _build_workflow_generation_prompt(task_prompt): return append_proactive_chart_guidance(task_prompt) @@ -1807,6 +2640,9 @@ def _build_workflow_chat_messages(prompt_text, url_access_context=None, apply_ge source_review_content = _get_workflow_url_access_system_content(url_access_context) if source_review_content: messages.append({'role': 'system', 'content': source_review_content}) + generated_file_output_guidance = build_generated_file_output_guidance(prompt_text) + if generated_file_output_guidance: + messages.append({'role': 'system', 'content': generated_file_output_guidance}) messages.append({'role': 'user', 'content': user_content}) return messages @@ -1814,9 +2650,14 @@ def _build_workflow_chat_messages(prompt_text, url_access_context=None, apply_ge def _build_workflow_agent_messages(prompt_text, url_access_context=None, apply_generation_guidance=False): user_content = _build_workflow_generation_prompt(prompt_text) if apply_generation_guidance else str(prompt_text or '').strip() source_review_content = _get_workflow_url_access_system_content(url_access_context) - if source_review_content: - user_content = f'{source_review_content}\n\n[Workflow Task]\n{user_content}' - return [ChatMessageContent(role='user', content=user_content)] + generated_file_output_guidance = build_generated_file_output_guidance(prompt_text) + content_sections = [ + content + for content in (generated_file_output_guidance, source_review_content) + if content + ] + content_sections.append(f'[Workflow Task]\n{user_content}') + return [ChatMessageContent(role='user', content='\n\n'.join(content_sections))] def _format_workflow_url_access_error(validation_result): @@ -1999,16 +2840,51 @@ def _maybe_execute_tabular_document_action( invoke_prompt=None, thought_tracker=None, live_thought_callback=None, + cancel_requested=None, + request_correlation_id=None, + token_usage_callback=None, ): + raise_if_mixed_source_cancelled( + cancel_requested, + 'tabular', + request_correlation_id=request_correlation_id, + ) if action_type not in {DOCUMENT_ACTION_TYPE_ANALYZE, DOCUMENT_ACTION_TYPE_COMPARISON}: return None - if not callable(invoke_prompt) or not is_tabular_processing_enabled(settings): - return None user_id = str(workflow.get('user_id') or '').strip() if not user_id: return None + if is_mixed_source_manifest_enabled(settings): + requested_source_ids, _ = _get_document_action_source_ids(action_config) + if requested_source_ids: + try: + resolve_authorized_source_manifest( + requested_source_ids, + user_id=user_id, + selection_mode='selected', + conversation_id=conversation_id, + active_group_ids=action_config.get('active_group_ids'), + active_public_workspace_ids=action_config.get('active_public_workspace_id'), + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + raise + except Exception: + log_event( + '[MixedSourceManifest] Workflow shadow resolution failed.', + extra={ + 'requested_source_count': len(requested_source_ids), + 'selection_mode': 'selected', + }, + level=logging.WARNING, + ) + + if not callable(invoke_prompt) or not is_tabular_processing_enabled(settings): + return None + tabular_documents = _resolve_tabular_document_action_documents( action_config, user_id, @@ -2045,6 +2921,11 @@ def _maybe_execute_tabular_document_action( try: for tabular_document in tabular_documents: + raise_if_mixed_source_cancelled( + cancel_requested, + 'tabular', + request_correlation_id=request_correlation_id, + ) document_baseline_invocation_count = 0 if conversation_id: document_baseline_invocation_count = len( @@ -2075,8 +2956,14 @@ def _maybe_execute_tabular_document_action( }], thought_tracker=thought_tracker, live_thought_callback=live_thought_callback, + token_usage_callback=token_usage_callback, ) ) + raise_if_mixed_source_cancelled( + cancel_requested, + 'tabular', + request_correlation_id=request_correlation_id, + ) if not str(tabular_analysis or '').strip(): raise ValueError( f"Tabular analysis returned no computed results for {tabular_document.get('document_name') or tabular_document.get('file_name') or tabular_document.get('document_id')}." @@ -2114,10 +3001,20 @@ def _maybe_execute_tabular_document_action( conversation_id=conversation_id, thought_callback=tabular_post_processing_thought_callback, user_id=user_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + token_usage_callback=token_usage_callback, ) ) + raise_if_mixed_source_cancelled( + cancel_requested, + 'export', + request_correlation_id=request_correlation_id, + ) if generated_tabular_output: generated_tabular_outputs.append(generated_tabular_output) + except MixedSourceCancellationError: + raise except Exception as exc: log_event( f'[WorkflowDocumentAction] Tabular document-action helper skipped: {exc}', @@ -2141,6 +3038,11 @@ def _maybe_execute_tabular_document_action( tabular_agent_citations = _build_agent_citations_from_plugin_invocations(tabular_invocations) if action_type == DOCUMENT_ACTION_TYPE_ANALYZE: + raise_if_mixed_source_cancelled( + cancel_requested, + 'reduction', + request_correlation_id=request_correlation_id, + ) analysis_result = { 'reply': '', 'analysis_reply': str(invoke_prompt( @@ -2159,6 +3061,11 @@ def _maybe_execute_tabular_document_action( 'window_size': None, 'window_percent': None, } + raise_if_mixed_source_cancelled( + cancel_requested, + 'reduction', + request_correlation_id=request_correlation_id, + ) if not analysis_result['analysis_reply']: raise RuntimeError('Tabular analysis synthesis returned an empty response.') analysis_result['reply'] = analysis_result['analysis_reply'] @@ -2171,6 +3078,11 @@ def _maybe_execute_tabular_document_action( left_document = tabular_documents[0] if tabular_documents else {} right_documents = tabular_documents[1:] + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison_reduction', + request_correlation_id=request_correlation_id, + ) comparison_result = { 'reply': '', 'analysis_reply': str(invoke_prompt( @@ -2201,6 +3113,11 @@ def _maybe_execute_tabular_document_action( ], 'comparison_items': [], } + raise_if_mixed_source_cancelled( + cancel_requested, + 'comparison_reduction', + request_correlation_id=request_correlation_id, + ) if not comparison_result['analysis_reply']: raise RuntimeError('Tabular comparison synthesis returned an empty response.') comparison_result['reply'] = comparison_result['analysis_reply'] @@ -3434,33 +4351,178 @@ def _add_workflow_activity_thought( ) -def _create_assistant_message(conversation, workflow, result, trigger_source, run_id, user_message_doc, assistant_message_id=None): - assistant_message_id = assistant_message_id or str(uuid.uuid4()) - timestamp = _utc_now_iso() - user_thread_info = (user_message_doc.get('metadata') or {}).get('thread_info') or {} - document_action = _get_document_action_config(workflow) - workspace_type = _get_workflow_scope(workflow) - group_id = _get_workflow_group_id(workflow) - raw_agent_citations = list(result.get('agent_citations') or []) - web_search_citations = list(result.get('web_search_citations') or []) - source_review_metadata = result.get('source_review') if isinstance(result.get('source_review'), dict) else {} - url_access_metadata = result.get('url_access') if isinstance(result.get('url_access'), dict) else {} - prepared_agent_citations = _persist_agent_citation_artifacts( - conversation_id=conversation.get('id'), - assistant_message_id=assistant_message_id, - agent_citations=raw_agent_citations, - created_timestamp=timestamp, - user_info={ - 'user_id': str(workflow.get('user_id') or '').strip(), - }, +def _maybe_create_workflow_generated_file_output( + workflow, + conversation_id, + user_question, + assistant_content, + function_results=None, + existing_outputs=None, +): + """Persist a requested workflow CSV, DOCX, or PDF artifact.""" + output_format = get_requested_generated_file_format(user_question) + if not output_format: + return None + if output_format == 'csv' and has_generated_tabular_csv_output(existing_outputs): + return None + if has_generated_file_output(existing_outputs, output_format): + return None + + export_payload = build_generated_file_export( + user_question, + assistant_content, + function_results=function_results, ) - assistant_doc = { - 'id': assistant_message_id, - 'conversation_id': conversation.get('id'), - 'role': 'assistant', - 'content': result.get('reply', ''), - 'timestamp': timestamp, - 'model_deployment_name': result.get('model_deployment_name'), + if not export_payload: + return None + + normalized_workflow = workflow if isinstance(workflow, dict) else {} + user_id = str(normalized_workflow.get('user_id') or '').strip() + normalized_conversation_id = str(conversation_id or '').strip() + if not user_id or not normalized_conversation_id: + return None + + generated_file_name = str(export_payload.get('file_name') or '').strip() + row_count = int(export_payload.get('row_count') or 0) + settings = get_settings() + structured_rows = export_payload.get('_structured_rows') or [] + row_batches = [] + if output_format == 'csv': + row_batches = build_tabular_generated_output_row_batches(structured_rows, settings=settings) + if not generated_file_name: + return None + + if output_format == 'csv' and should_queue_tabular_generated_output_background( + row_count, + len(row_batches), + settings, + ): + try: + background_run = queue_tabular_generated_output_run( + user_id=user_id, + conversation_id=normalized_conversation_id, + user_question=user_question, + source_candidate={ + 'filename': generated_file_name, + 'selected_sheet': '', + 'source_authorization': { + 'source': 'chat', + }, + }, + output_format=output_format, + row_batches=row_batches, + gpt_model='', + settings=settings, + passthrough_input_rows=True, + ) + return build_background_tabular_generated_output_metadata(background_run) + except Exception as exc: + log_event( + '[Workflow Generated File Export] Failed to queue large CSV export', + { + 'workflow_id': normalized_workflow.get('id'), + 'conversation_id': normalized_conversation_id, + 'row_count': row_count, + 'output_format': output_format, + 'error': str(exc), + }, + level=logging.ERROR, + exceptionTraceback=True, + ) + return None + + try: + upload_result = upload_generated_analysis_artifact_for_user( + current_user_id=user_id, + conversation_id=normalized_conversation_id, + file_name=generated_file_name, + file_content=export_payload.get('file_content'), + capability=export_payload.get('capability') or 'file_export', + output_format=output_format, + summary=export_payload.get('summary'), + ) + except Exception as exc: + log_event( + '[Workflow Generated File Export] Failed to save generated file artifact', + { + 'workflow_id': normalized_workflow.get('id'), + 'conversation_id': normalized_conversation_id, + 'row_count': row_count, + 'output_format': output_format, + 'error': str(exc), + }, + debug_only=True, + ) + return None + + artifact_metadata = build_generated_file_artifact_metadata( + export_payload, + upload_result, + normalized_conversation_id, + ) + if not artifact_metadata: + return None + + log_event( + '[Workflow Generated File Export] Saved generated file artifact', + { + 'workflow_id': normalized_workflow.get('id'), + 'conversation_id': normalized_conversation_id, + 'artifact_message_id': artifact_metadata.get('artifact_message_id'), + 'row_count': row_count, + 'output_format': output_format, + }, + debug_only=True, + ) + return artifact_metadata + + +def _maybe_create_workflow_assistant_table_generated_output(*args, **kwargs): + """Backward-compatible wrapper for the generic workflow file-output finalizer.""" + return _maybe_create_workflow_generated_file_output(*args, **kwargs) + + +def _create_assistant_message(conversation, workflow, result, trigger_source, run_id, user_message_doc, assistant_message_id=None): + assistant_message_id = assistant_message_id or str(uuid.uuid4()) + timestamp = _utc_now_iso() + user_thread_info = (user_message_doc.get('metadata') or {}).get('thread_info') or {} + document_action = _get_document_action_config(workflow) + workspace_type = _get_workflow_scope(workflow) + group_id = _get_workflow_group_id(workflow) + generated_analysis_artifacts = list(result.get('generated_analysis_artifacts') or []) + generated_tabular_outputs = list(result.get('generated_tabular_outputs') or []) + raw_agent_citations = list(result.get('agent_citations') or []) + generated_file_output = _maybe_create_workflow_generated_file_output( + workflow=workflow, + conversation_id=conversation.get('id'), + user_question=workflow.get('task_prompt', ''), + assistant_content=get_generated_file_export_content(result), + function_results=raw_agent_citations, + existing_outputs=generated_analysis_artifacts + generated_tabular_outputs, + ) + if generated_file_output: + generated_analysis_artifacts.append(generated_file_output) + if generated_file_output.get('output_format') == 'csv': + generated_tabular_outputs.append(generated_file_output) + web_search_citations = list(result.get('web_search_citations') or []) + source_review_metadata = result.get('source_review') if isinstance(result.get('source_review'), dict) else {} + url_access_metadata = result.get('url_access') if isinstance(result.get('url_access'), dict) else {} + prepared_agent_citations = _persist_agent_citation_artifacts( + conversation_id=conversation.get('id'), + assistant_message_id=assistant_message_id, + agent_citations=raw_agent_citations, + created_timestamp=timestamp, + user_info={ + 'user_id': str(workflow.get('user_id') or '').strip(), + }, + ) + assistant_doc = { + 'id': assistant_message_id, + 'conversation_id': conversation.get('id'), + 'role': 'assistant', + 'content': result.get('reply', ''), + 'timestamp': timestamp, + 'model_deployment_name': result.get('model_deployment_name'), 'augmented': bool(result.get('augmented') or result.get('hybrid_citations')), 'hybrid_citations': list(result.get('hybrid_citations') or []), 'agent_citations': prepared_agent_citations, @@ -3474,6 +4536,8 @@ def _create_assistant_message(conversation, workflow, result, trigger_source, ru 'workspace_type': workspace_type, 'group_id': group_id or None, 'token_usage': result.get('token_usage'), + 'generated_analysis_artifacts': generated_analysis_artifacts, + 'generated_tabular_outputs': generated_tabular_outputs, 'source_review': source_review_metadata, 'workflow': { 'workflow_id': workflow.get('id'), @@ -3486,6 +4550,7 @@ def _create_assistant_message(conversation, workflow, result, trigger_source, ru 'model_binding_summary': workflow.get('model_binding_summary') or {}, 'document_action': document_action, 'document_search': result.get('document_search') or {}, + 'mixed_source_coverage': result.get('mixed_source_coverage') or {}, 'analyze': workflow.get('analyze') or {}, 'analysis_coverage': result.get('analysis_coverage') or {}, }, @@ -3845,6 +4910,80 @@ def _get_recent_workflow_document_limit(action_type, settings): return _get_workflow_search_max_documents(settings) +def _resolve_analyze_all_document_ids(workflow, action_config, settings, max_documents): + """Enumerate a bounded catalog before the manifest performs object authorization.""" + user_id = str(workflow.get('user_id') or '').strip() + workflow_group_id = _get_workflow_group_id(workflow) + doc_scope = normalize_search_scope(action_config.get('doc_scope')) + active_group_ids = normalize_search_id_list(action_config.get('active_group_ids')) + if workflow_group_id: + assert_group_role( + user_id, + workflow_group_id, + allowed_roles=("Owner", "Admin", "DocumentManager", "User"), + ) + active_group_ids = [workflow_group_id] + doc_scope = 'group' + else: + active_group_ids = _resolve_recent_authorized_group_ids(user_id, active_group_ids) + + active_public_workspace_ids = normalize_search_id_list( + action_config.get('active_public_workspace_id') + ) + if not workflow_group_id: + active_public_workspace_ids = _resolve_recent_authorized_public_workspace_ids( + user_id, + active_public_workspace_ids, + ) + + catalog_requests = [] + if doc_scope in {'personal', 'all'} and not workflow_group_id: + catalog_requests.append((DOCUMENT_ACCESS_SCOPE_PERSONAL, {'user_id': user_id})) + if doc_scope in {'group', 'all'} and active_group_ids: + catalog_requests.append((DOCUMENT_ACCESS_SCOPE_GROUP, {'group_ids': active_group_ids})) + if doc_scope in {'public', 'all'} and active_public_workspace_ids and not workflow_group_id: + catalog_requests.append(( + DOCUMENT_ACCESS_SCOPE_PUBLIC, + {'public_workspace_ids': active_public_workspace_ids}, + )) + if not catalog_requests: + raise ValueError('Analyze All Documents has no authorized workspace scopes to enumerate.') + + document_ids = [] + for source_scope, scope_arguments in catalog_requests: + catalog_result = enumerate_bounded_document_access_index_ids( + source_scope, + max_documents, + settings=settings, + **scope_arguments, + ) + if not catalog_result.get('success'): + if catalog_result.get('status') == 'document_limit_exceeded': + raise ValueError( + f'Analyze All Documents exceeds the configured {max_documents}-document limit.' + ) + raise RuntimeError( + 'Analyze All Documents is temporarily unavailable until the authorized document catalog is ready.' + ) + for document_id in catalog_result.get('document_ids') or []: + if document_id not in document_ids: + document_ids.append(document_id) + if len(document_ids) > max_documents: + raise ValueError( + f'Analyze All Documents exceeds the configured {max_documents}-document limit.' + ) + + if not document_ids: + raise ValueError('Analyze All Documents found no authorized documents in the selected scopes.') + return { + 'document_ids': document_ids, + 'doc_scope': doc_scope, + 'active_group_ids': active_group_ids, + 'active_public_workspace_id': active_public_workspace_ids, + 'target_mode': DOCUMENT_ACTION_TARGET_MODE_ALL, + } + + def _collect_recent_workflow_documents(workflow, action_config, settings, max_documents): user_id = str(workflow.get('user_id') or '').strip() workflow_group_id = _get_workflow_group_id(workflow) @@ -3982,20 +5121,90 @@ def _format_workflow_search_results(results): def _build_workflow_search_prompt(task_prompt, search_context): task_prompt = str(task_prompt or '').strip() retrieved_content = str((search_context or {}).get('retrieved_content') or '').strip() - if not retrieved_content: + evidence_messages = [ + str((message or {}).get('content') or '').strip() + for message in list((search_context or {}).get('evidence_messages') or []) + if str((message or {}).get('content') or '').strip() + ] + if not retrieved_content and not evidence_messages: return task_prompt - return ( + prompt_sections = [ '[Workflow document search context]\n' - 'Use the retrieved document excerpts below as grounding for the workflow task. ' - 'When the excerpts are insufficient, say what is missing instead of guessing.\n\n' - f'{retrieved_content}\n\n' - '[Workflow task]\n' - f'{task_prompt}' - ).strip() + 'Use the bounded native-engine evidence below as grounding for the workflow task. ' + 'When the evidence is insufficient or source coverage is partial, say what is missing instead of guessing.' + ] + if retrieved_content: + prompt_sections.append( + f'[Narrative excerpts]\n{retrieved_content}' + ) + if evidence_messages: + prompt_sections.append( + '[Computed and coverage evidence]\n' + + '\n\n'.join(evidence_messages) + ) + prompt_sections.append(f'[Workflow task]\n{task_prompt}') + return '\n\n'.join(prompt_sections).strip() + + +@contextmanager +def _workflow_mixed_source_execution_context(user_id, conversation_id, manifest): + """Bind tabular tools to scope IDs from a freshly authorized manifest.""" + group_ids = [] + public_workspace_ids = [] + for source in list(manifest or []): + if source.get('authorization_status') != 'authorized': + continue + group_id = str(source.get('group_id') or '').strip() + public_workspace_id = str(source.get('public_workspace_id') or '').strip() + if group_id and group_id not in group_ids: + group_ids.append(group_id) + if public_workspace_id and public_workspace_id not in public_workspace_ids: + public_workspace_ids.append(public_workspace_id) + with _ensure_execution_context(user_id): + previous_conversation_id = getattr(g, 'conversation_id', None) if hasattr(g, 'conversation_id') else None + previous_authorized_chat_context = getattr(g, 'authorized_chat_context', None) if hasattr(g, 'authorized_chat_context') else None + g.conversation_id = conversation_id + g.authorized_chat_context = { + 'user_id': user_id, + 'conversation_id': conversation_id, + 'active_group_ids': group_ids, + 'active_group_id': group_ids[0] if len(group_ids) == 1 else None, + 'active_public_workspace_ids': public_workspace_ids, + 'active_public_workspace_id': ( + public_workspace_ids[0] + if len(public_workspace_ids) == 1 + else None + ), + 'fact_memory_scope_id': group_ids[0] if len(group_ids) == 1 else user_id, + 'fact_memory_scope_type': 'group' if len(group_ids) == 1 else 'user', + } + try: + yield + finally: + if previous_conversation_id is None and hasattr(g, 'conversation_id'): + delattr(g, 'conversation_id') + else: + g.conversation_id = previous_conversation_id + if previous_authorized_chat_context is None and hasattr(g, 'authorized_chat_context'): + delattr(g, 'authorized_chat_context') + else: + g.authorized_chat_context = previous_authorized_chat_context -def _prepare_workflow_search_context(workflow, action_config, settings, thought_tracker=None, run_id=None): + +def _prepare_workflow_search_context( + workflow, + action_config, + settings, + conversation_id='', + thought_tracker=None, + run_id=None, + request_correlation_id=None, +): + request_correlation_id = normalize_mixed_source_correlation_id( + request_correlation_id + ) if not _is_document_search_workflow(action_config): return {'workflow': workflow, 'citations': [], 'result_count': 0, 'document_count': 0, 'query': None} @@ -4008,20 +5217,186 @@ def _prepare_workflow_search_context(workflow, action_config, settings, thought_ if not query: return {'workflow': workflow, 'citations': [], 'result_count': 0, 'document_count': 0, 'query': None} - search_top_n = normalize_search_top_n(max(12, len(document_ids) * 3 if document_ids else 12)) - search_result = search_documents( - query=query, - user_id=str(workflow.get('user_id') or '').strip(), - top_n=search_top_n, - doc_scope=resolved_action.get('doc_scope') or 'all', - document_ids=document_ids, - active_group_ids=resolved_action.get('active_group_ids'), - active_public_workspace_id=resolved_action.get('active_public_workspace_id'), + user_id = str(workflow.get('user_id') or '').strip() + manifest_doc_scope = resolved_action.get('doc_scope') or 'all' + manifest_group_ids = list(resolved_action.get('active_group_ids') or []) + manifest_public_workspace_ids = list( + resolved_action.get('active_public_workspace_id') or [] ) + workflow_group_id = _get_workflow_group_id(workflow) + if workflow_group_id: + assert_group_role( + user_id, + workflow_group_id, + allowed_roles=("Owner", "Admin", "DocumentManager", "User"), + ) + manifest_doc_scope = 'group' + manifest_group_ids = [workflow_group_id] + manifest_public_workspace_ids = [] + + scoped_action = dict(resolved_action) + scoped_action['doc_scope'] = manifest_doc_scope + scoped_action['active_group_ids'] = manifest_group_ids + scoped_action['active_public_workspace_id'] = manifest_public_workspace_ids + + if not is_mixed_source_chat_search_enabled(settings): + search_top_n = normalize_search_top_n(max(50, len(document_ids) * 3 if document_ids else 50)) + search_result = search_documents( + query=query, + user_id=user_id, + top_n=search_top_n, + doc_scope=manifest_doc_scope, + document_ids=document_ids, + active_group_ids=manifest_group_ids, + active_public_workspace_id=manifest_public_workspace_ids, + ) + retrieved_content, citations = _format_workflow_search_results(search_result.get('results') or []) + prepared_workflow = _apply_runtime_document_action_config(workflow, scoped_action) + prepared_workflow['task_prompt'] = _build_workflow_search_prompt(workflow.get('task_prompt', ''), { + 'retrieved_content': retrieved_content, + }) + + if thought_tracker and run_id: + _add_workflow_activity_thought( + thought_tracker, + prepared_workflow, + run_id, + step_type='document', + content='Searched selected workflow documents', + detail=( + f"results={search_result.get('result_count', 0)} | " + f"documents={search_result.get('document_count', 0)}" + ), + activity_key=f'search:{run_id}:documents', + kind='document_search', + title='Document search', + status='completed', + ) + + return { + 'workflow': prepared_workflow, + 'citations': citations, + 'agent_citations': [], + 'generated_tabular_outputs': [], + 'coverage': {}, + 'result_count': search_result.get('result_count', 0), + 'document_count': search_result.get('document_count', 0), + 'query': search_result.get('query'), + } + + manifest = resolve_authorized_source_manifest( + document_ids, + user_id=user_id, + selection_mode='selected', + conversation_id=conversation_id, + active_group_ids=manifest_group_ids, + active_public_workspace_ids=manifest_public_workspace_ids, + doc_scope=manifest_doc_scope, + request_correlation_id=request_correlation_id, + ) + partitions = partition_source_manifest(manifest) + narrative_sources = list(partitions.get('narrative_sources') or []) + tabular_sources = list(partitions.get('tabular_sources') or []) + narrative_document_ids = [ + str(source.get('document_id') or '').strip() + for source in narrative_sources + if str(source.get('document_id') or '').strip() + ] + authorized_document_ids = narrative_document_ids + [ + str(source.get('document_id') or '').strip() + for source in tabular_sources + if str(source.get('document_id') or '').strip() + ] + + manifest_group_ids = [] + manifest_public_workspace_ids = [] + for source in narrative_sources + tabular_sources: + group_id = str(source.get('group_id') or '').strip() + public_workspace_id = str(source.get('public_workspace_id') or '').strip() + if group_id and group_id not in manifest_group_ids: + manifest_group_ids.append(group_id) + if public_workspace_id and public_workspace_id not in manifest_public_workspace_ids: + manifest_public_workspace_ids.append(public_workspace_id) + + search_top_n = normalize_search_top_n( + max(50, len(narrative_document_ids) * 3 if narrative_document_ids else 50) + ) + search_result = { + 'results': [], + 'result_count': 0, + 'document_count': 0, + 'query': query, + } + narrative_retrieval_failed = False + if narrative_document_ids: + try: + search_result = search_documents( + query=query, + user_id=user_id, + top_n=search_top_n, + doc_scope=manifest_doc_scope, + document_ids=narrative_document_ids, + active_group_ids=manifest_group_ids, + active_public_workspace_id=manifest_public_workspace_ids, + include_all_public_workspaces=True, + ) + except Exception: + if not tabular_sources: + raise + narrative_retrieval_failed = True retrieved_content, citations = _format_workflow_search_results(search_result.get('results') or []) - prepared_workflow = _apply_runtime_document_action_config(workflow, resolved_action) + evidence_envelopes = ( + build_failed_narrative_evidence_envelopes( + narrative_sources, + 'selected', + ) + if narrative_retrieval_failed + else build_narrative_evidence_envelopes( + narrative_sources, + search_result.get('results') or [], + 'selected', + ) + ) + + from functions_tabular_analysis import execute_mixed_source_tabular_evidence + + native_token_usage_aggregate = _create_token_usage_aggregate() + with _workflow_mixed_source_execution_context(user_id, conversation_id, manifest): + tabular_result = execute_mixed_source_tabular_evidence( + tabular_sources=tabular_sources, + selection_mode='selected', + has_narrative_sources=bool(narrative_sources), + user_question=query, + user_id=user_id, + conversation_id=conversation_id, + gpt_model=_resolve_tabular_document_action_model_name(workflow, settings), + settings=settings, + thought_tracker=thought_tracker, + request_correlation_id=request_correlation_id, + ) + _accumulate_token_usage_summary( + native_token_usage_aggregate, + tabular_result.get('token_usage'), + ) + evidence_envelopes.extend(tabular_result.get('evidence_envelopes') or []) + mixed_source_handoff = build_mixed_source_evidence_handoff( + manifest, + evidence_envelopes, + 'selected', + mode='search', + telemetry_settings=settings, + request_correlation_id=request_correlation_id, + ) + + prepared_action = dict(resolved_action) + prepared_action['document_ids'] = authorized_document_ids + prepared_action['doc_scope'] = manifest_doc_scope + prepared_action['active_group_ids'] = manifest_group_ids + prepared_action['active_public_workspace_id'] = manifest_public_workspace_ids + prepared_workflow = _apply_runtime_document_action_config(workflow, prepared_action) prepared_workflow['task_prompt'] = _build_workflow_search_prompt(workflow.get('task_prompt', ''), { - 'retrieved_content': retrieved_content, + 'retrieved_content': '', + 'evidence_messages': [mixed_source_handoff], }) if thought_tracker and run_id: @@ -4033,7 +5408,8 @@ def _prepare_workflow_search_context(workflow, action_config, settings, thought_ content='Searched selected workflow documents', detail=( f"results={search_result.get('result_count', 0)} | " - f"documents={search_result.get('document_count', 0)}" + f"narrative_documents={len(narrative_sources)} | " + f"tabular_documents={len(tabular_sources)}" ), activity_key=f'search:{run_id}:documents', kind='document_search', @@ -4044,12 +5420,63 @@ def _prepare_workflow_search_context(workflow, action_config, settings, thought_ return { 'workflow': prepared_workflow, 'citations': citations, + 'agent_citations': list(tabular_result.get('agent_citations') or []), + 'generated_tabular_outputs': list(tabular_result.get('generated_outputs') or []), + 'coverage': mixed_source_handoff.get('mixed_source_coverage') or {}, + 'token_usage': _finalize_token_usage(native_token_usage_aggregate), 'result_count': search_result.get('result_count', 0), - 'document_count': search_result.get('document_count', 0), + 'document_count': len(authorized_document_ids), 'query': search_result.get('query'), } +def _attach_workflow_search_context(execution_result, workflow_search_context): + """Merge mixed Search evidence into either a workflow model or agent result.""" + execution_result = execution_result if isinstance(execution_result, dict) else {} + search_context = workflow_search_context if isinstance(workflow_search_context, dict) else {} + if not search_context: + return execution_result + + existing_agent_citations = deduplicate_mixed_source_references( + list(execution_result.get('agent_citations') or []) + + list(search_context.get('agent_citations') or []), + reference_type='citation', + ) + existing_outputs = deduplicate_mixed_source_references( + list(execution_result.get('generated_tabular_outputs') or []) + + list(search_context.get('generated_tabular_outputs') or []), + reference_type='artifact', + ) + hybrid_citations = deduplicate_mixed_source_references( + list(search_context.get('citations') or []), + reference_type='citation', + ) + coverage = search_context.get('coverage') if isinstance(search_context.get('coverage'), dict) else {} + merged_token_usage = _merge_token_usage_summaries([ + execution_result, + {'token_usage': search_context.get('token_usage')}, + ]) + execution_result.update({ + 'hybrid_citations': hybrid_citations, + 'agent_citations': existing_agent_citations, + 'generated_tabular_outputs': existing_outputs, + 'mixed_source_coverage': coverage, + 'token_usage': merged_token_usage or execution_result.get('token_usage'), + 'augmented': bool( + hybrid_citations + or existing_agent_citations + or coverage.get('requested_source_count') + ), + 'document_search': { + 'query': search_context.get('query'), + 'result_count': search_context.get('result_count', 0), + 'document_count': search_context.get('document_count', 0), + 'partial_coverage': bool(coverage.get('partial_coverage')), + }, + }) + return execution_result + + def _apply_runtime_document_action_config(workflow, action_config): prepared_workflow = dict(workflow or {}) prepared_workflow['document_action'] = dict(action_config or {}) @@ -4743,7 +6170,14 @@ def _execute_document_analysis_workflow( external_activity_callback=None, action_config=None, url_access_context=None, + cancel_requested=None, + request_correlation_id=None, ): + raise_if_mixed_source_cancelled( + cancel_requested, + 'manifest', + request_correlation_id=request_correlation_id, + ) analysis_config = action_config if isinstance(action_config, dict) else _get_document_action_config(workflow) if analysis_config.get('type') != DOCUMENT_ACTION_TYPE_ANALYZE: raise ValueError('Document analysis is not enabled for this workflow.') @@ -4807,6 +6241,8 @@ def _execute_document_analysis_workflow( external_activity_callback=external_activity_callback, action_config=per_document_action, url_access_context=url_access_context, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, ), }) @@ -4893,45 +6329,77 @@ def invoke_prompt(prompt_text, stage='window_analysis', metadata=None): _accumulate_token_usage(token_usage_aggregate, result) return str(result) - tabular_action_payload = _maybe_execute_tabular_document_action( - DOCUMENT_ACTION_TYPE_ANALYZE, + def record_native_token_usage(token_usage): + _accumulate_token_usage_summary(token_usage_aggregate, token_usage) + + analysis_result = _execute_mixed_source_analyze_workflow( + workflow, analysis_config, settings, invoke_prompt, + conversation_id=conversation_id, activity_callback=activity_callback, + thought_tracker=thought_tracker, live_thought_callback=external_activity_callback, + max_documents=workflow_analysis_max_documents, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + token_usage_callback=record_native_token_usage, + ) + primary_generated_outputs = list( + analysis_result.get('generated_tabular_outputs') + or [] + ) + _reauthorize_mixed_source_workflow_result( workflow, analysis_config, + analysis_result, settings, - conversation_id=conversation_id, - invoke_prompt=invoke_prompt, - thought_tracker=thought_tracker, - live_thought_callback=external_activity_callback, + conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + additional_generated_outputs=primary_generated_outputs, ) - if tabular_action_payload: - analysis_result = tabular_action_payload.get('result') or {} - else: - analysis_result = run_document_analysis( - user_id=user_id, - analysis_prompt=workflow.get('task_prompt', ''), - document_ids=analysis_config.get('document_ids'), - invoke_prompt=invoke_prompt, - doc_scope=analysis_config.get('doc_scope'), - active_group_ids=analysis_config.get('active_group_ids'), - active_public_workspace_id=analysis_config.get('active_public_workspace_id'), - window_unit=analysis_config.get('window_unit'), - window_size=analysis_config.get('window_size'), - window_percent=analysis_config.get('window_percent'), - max_retries_per_window=analysis_config.get('max_retries_per_window'), - activity_callback=activity_callback, - max_documents=workflow_analysis_max_documents, + try: + document_analysis_artifact_payload = _maybe_create_document_analysis_generated_artifacts( + analysis_result, + workflow.get('task_prompt', ''), + conversation_id=conversation_id, + primary_generated_outputs=primary_generated_outputs, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + _rollback_mixed_source_generated_outputs( + user_id, + conversation_id, + primary_generated_outputs, + reason='cancellation', ) - document_analysis_artifact_payload = _maybe_create_document_analysis_generated_artifacts( + raise + _reauthorize_mixed_source_workflow_result( + workflow, + analysis_config, analysis_result, - workflow.get('task_prompt', ''), - conversation_id=conversation_id, - primary_generated_outputs=list((tabular_action_payload or {}).get('generated_tabular_outputs') or []), + settings, + conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + additional_generated_outputs=( + primary_generated_outputs + + list(document_analysis_artifact_payload.get('artifacts') or []) + ), ) agent_citations = _build_agent_citations_from_invocations(user_id, conversation_id) if not agent_citations: - agent_citations = list((tabular_action_payload or {}).get('agent_citations') or []) + agent_citations = list( + analysis_result.get('agent_citations') + or [] + ) alert_targets = _collect_agent_alert_targets(user_id, conversation_id) token_usage = _finalize_token_usage(token_usage_aggregate) + if analysis_result.get('mixed_source_manifest'): + _emit_mixed_source_token_telemetry( + settings, + 'analyze', + token_usage, + request_correlation_id=request_correlation_id, + ) return { 'reply': ( @@ -4947,7 +6415,10 @@ def invoke_prompt(prompt_text, stage='window_analysis', metadata=None): 'agent_name': getattr(loaded_agent, 'name', None) or requested_name, 'agent_display_name': getattr(loaded_agent, 'display_name', None) or selected_agent.get('display_name') or requested_name, 'agent_citations': agent_citations, - 'generated_tabular_outputs': list((tabular_action_payload or {}).get('generated_tabular_outputs') or []), + 'generated_tabular_outputs': list( + analysis_result.get('generated_tabular_outputs') + or [] + ), 'alert_targets': alert_targets, } finally: @@ -5008,41 +6479,70 @@ def invoke_model_prompt(prompt_text, stage='window_analysis', metadata=None): return '' return _extract_message_text(completion.choices[0].message.content) - tabular_action_payload = _maybe_execute_tabular_document_action( - DOCUMENT_ACTION_TYPE_ANALYZE, + def record_native_token_usage(token_usage): + _accumulate_token_usage_summary(token_usage_aggregate, token_usage) + + analysis_result = _execute_mixed_source_analyze_workflow( + workflow, analysis_config, settings, invoke_model_prompt, + conversation_id=conversation_id, activity_callback=activity_callback, + thought_tracker=thought_tracker, live_thought_callback=external_activity_callback, + max_documents=workflow_analysis_max_documents, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + token_usage_callback=record_native_token_usage, + ) + primary_generated_outputs = list( + analysis_result.get('generated_tabular_outputs') + or [] + ) + _reauthorize_mixed_source_workflow_result( workflow, analysis_config, + analysis_result, settings, - conversation_id=conversation_id, - invoke_prompt=invoke_model_prompt, - thought_tracker=thought_tracker, - live_thought_callback=external_activity_callback, + conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + additional_generated_outputs=primary_generated_outputs, ) - if tabular_action_payload: - analysis_result = tabular_action_payload.get('result') or {} - else: - analysis_result = run_document_analysis( - user_id=user_id, - analysis_prompt=workflow.get('task_prompt', ''), - document_ids=analysis_config.get('document_ids'), - invoke_prompt=invoke_model_prompt, - doc_scope=analysis_config.get('doc_scope'), - active_group_ids=analysis_config.get('active_group_ids'), - active_public_workspace_id=analysis_config.get('active_public_workspace_id'), - window_unit=analysis_config.get('window_unit'), - window_size=analysis_config.get('window_size'), - window_percent=analysis_config.get('window_percent'), - max_retries_per_window=analysis_config.get('max_retries_per_window'), - activity_callback=activity_callback, - max_documents=workflow_analysis_max_documents, + try: + document_analysis_artifact_payload = _maybe_create_document_analysis_generated_artifacts( + analysis_result, + workflow.get('task_prompt', ''), + conversation_id=conversation_id, + primary_generated_outputs=primary_generated_outputs, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, ) - document_analysis_artifact_payload = _maybe_create_document_analysis_generated_artifacts( + except MixedSourceCancellationError: + _rollback_mixed_source_generated_outputs( + user_id, + conversation_id, + primary_generated_outputs, + reason='cancellation', + ) + raise + _reauthorize_mixed_source_workflow_result( + workflow, + analysis_config, analysis_result, - workflow.get('task_prompt', ''), - conversation_id=conversation_id, - primary_generated_outputs=list((tabular_action_payload or {}).get('generated_tabular_outputs') or []), + settings, + conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + additional_generated_outputs=( + primary_generated_outputs + + list(document_analysis_artifact_payload.get('artifacts') or []) + ), ) token_usage = _finalize_token_usage(token_usage_aggregate) + if analysis_result.get('mixed_source_manifest'): + _emit_mixed_source_token_telemetry( + settings, + 'analyze', + token_usage, + request_correlation_id=request_correlation_id, + ) debug_print( '[WorkflowDocumentAnalysis] Completed workflow action | ' f"workflow_id={workflow.get('id')} | " @@ -5064,8 +6564,14 @@ def invoke_model_prompt(prompt_text, stage='window_analysis', metadata=None): 'model_deployment_name': deployment_name, 'token_usage': token_usage, 'provider': provider, - 'agent_citations': list((tabular_action_payload or {}).get('agent_citations') or []), - 'generated_tabular_outputs': list((tabular_action_payload or {}).get('generated_tabular_outputs') or []), + 'agent_citations': list( + analysis_result.get('agent_citations') + or [] + ), + 'generated_tabular_outputs': list( + analysis_result.get('generated_tabular_outputs') + or [] + ), } @@ -5078,7 +6584,14 @@ def _execute_document_comparison_workflow( external_activity_callback=None, action_config=None, url_access_context=None, + cancel_requested=None, + request_correlation_id=None, ): + raise_if_mixed_source_cancelled( + cancel_requested, + 'manifest', + request_correlation_id=request_correlation_id, + ) comparison_config = action_config if isinstance(action_config, dict) else _get_document_action_config(workflow) if comparison_config.get('type') != DOCUMENT_ACTION_TYPE_COMPARISON: raise ValueError('Document comparison is not enabled for this workflow.') @@ -5165,17 +6678,33 @@ def invoke_prompt(prompt_text, stage='window_analysis', metadata=None): _accumulate_token_usage(token_usage_aggregate, result) return str(result) - tabular_action_payload = _maybe_execute_tabular_document_action( - DOCUMENT_ACTION_TYPE_COMPARISON, - workflow, - comparison_config, - settings, - conversation_id=conversation_id, - invoke_prompt=invoke_prompt, - thought_tracker=thought_tracker, - live_thought_callback=external_activity_callback, + def record_native_token_usage(token_usage): + _accumulate_token_usage_summary(token_usage_aggregate, token_usage) + + mixed_comparison_enabled = is_cross_format_compare_enabled(settings) + mixed_comparison_result = ( + _execute_cross_format_comparison_workflow( + workflow, comparison_config, settings, invoke_prompt, + conversation_id=conversation_id, activity_callback=activity_callback, + thought_tracker=thought_tracker, live_thought_callback=external_activity_callback, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + token_usage_callback=record_native_token_usage, + ) if mixed_comparison_enabled else None ) - if tabular_action_payload: + if not mixed_comparison_enabled: + _raise_legacy_cross_format_compare_limitation(comparison_config, user_id, conversation_id) + tabular_action_payload = None if mixed_comparison_result else _maybe_execute_tabular_document_action( + DOCUMENT_ACTION_TYPE_COMPARISON, workflow, comparison_config, settings, + conversation_id=conversation_id, invoke_prompt=invoke_prompt, + thought_tracker=thought_tracker, live_thought_callback=external_activity_callback, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + token_usage_callback=record_native_token_usage, + ) + if mixed_comparison_result: + comparison_result = mixed_comparison_result + elif tabular_action_payload: comparison_result = tabular_action_payload.get('result') or {} else: comparison_result = run_document_comparison( @@ -5185,17 +6714,69 @@ def invoke_prompt(prompt_text, stage='window_analysis', metadata=None): invoke_prompt=invoke_prompt, activity_callback=activity_callback, conversation_id=conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, ) - comparison_artifact_payload = _maybe_create_comparison_generated_artifacts( + agent_citations = _build_agent_citations_from_invocations(user_id, conversation_id) + agent_citations = deduplicate_mixed_source_references( + list(agent_citations or []) + + list((tabular_action_payload or {}).get('agent_citations') or []) + + list(comparison_result.get('agent_citations') or []), + reference_type='citation', + ) + generated_tabular_outputs = deduplicate_mixed_source_references( + list((tabular_action_payload or {}).get('generated_tabular_outputs') or []) + + list(comparison_result.get('generated_tabular_outputs') or []), + reference_type='artifact', + ) + _reauthorize_mixed_source_workflow_result( + workflow, + comparison_config, comparison_result, - workflow.get('task_prompt', ''), - conversation_id=conversation_id, + settings, + conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + additional_generated_outputs=generated_tabular_outputs, + ) + try: + comparison_artifact_payload = _maybe_create_comparison_generated_artifacts( + comparison_result, + workflow.get('task_prompt', ''), + conversation_id=conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + _rollback_mixed_source_generated_outputs( + user_id, + conversation_id, + generated_tabular_outputs, + reason='cancellation', + ) + raise + _reauthorize_mixed_source_workflow_result( + workflow, + comparison_config, + comparison_result, + settings, + conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + additional_generated_outputs=( + generated_tabular_outputs + + list(comparison_artifact_payload.get('artifacts') or []) + ), ) - agent_citations = _build_agent_citations_from_invocations(user_id, conversation_id) - if not agent_citations: - agent_citations = list((tabular_action_payload or {}).get('agent_citations') or []) alert_targets = _collect_agent_alert_targets(user_id, conversation_id) token_usage = _finalize_token_usage(token_usage_aggregate) + if comparison_result.get('mixed_source_manifest'): + _emit_mixed_source_token_telemetry( + settings, + 'compare', + token_usage, + request_correlation_id=request_correlation_id, + ) return { 'reply': ( @@ -5211,7 +6792,7 @@ def invoke_prompt(prompt_text, stage='window_analysis', metadata=None): 'agent_name': getattr(loaded_agent, 'name', None) or requested_name, 'agent_display_name': getattr(loaded_agent, 'display_name', None) or selected_agent.get('display_name') or requested_name, 'agent_citations': agent_citations, - 'generated_tabular_outputs': list((tabular_action_payload or {}).get('generated_tabular_outputs') or []), + 'generated_tabular_outputs': generated_tabular_outputs, 'alert_targets': alert_targets, } finally: @@ -5272,17 +6853,33 @@ def invoke_model_prompt(prompt_text, stage='window_analysis', metadata=None): return '' return _extract_message_text(completion.choices[0].message.content) - tabular_action_payload = _maybe_execute_tabular_document_action( - DOCUMENT_ACTION_TYPE_COMPARISON, - workflow, - comparison_config, - settings, - conversation_id=conversation_id, - invoke_prompt=invoke_model_prompt, - thought_tracker=thought_tracker, - live_thought_callback=external_activity_callback, + def record_native_token_usage(token_usage): + _accumulate_token_usage_summary(token_usage_aggregate, token_usage) + + mixed_comparison_enabled = is_cross_format_compare_enabled(settings) + mixed_comparison_result = ( + _execute_cross_format_comparison_workflow( + workflow, comparison_config, settings, invoke_model_prompt, + conversation_id=conversation_id, activity_callback=activity_callback, + thought_tracker=thought_tracker, live_thought_callback=external_activity_callback, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + token_usage_callback=record_native_token_usage, + ) if mixed_comparison_enabled else None + ) + if not mixed_comparison_enabled: + _raise_legacy_cross_format_compare_limitation(comparison_config, user_id, conversation_id) + tabular_action_payload = None if mixed_comparison_result else _maybe_execute_tabular_document_action( + DOCUMENT_ACTION_TYPE_COMPARISON, workflow, comparison_config, settings, + conversation_id=conversation_id, invoke_prompt=invoke_model_prompt, + thought_tracker=thought_tracker, live_thought_callback=external_activity_callback, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + token_usage_callback=record_native_token_usage, ) - if tabular_action_payload: + if mixed_comparison_result: + comparison_result = mixed_comparison_result + elif tabular_action_payload: comparison_result = tabular_action_payload.get('result') or {} else: comparison_result = run_document_comparison( @@ -5292,13 +6889,66 @@ def invoke_model_prompt(prompt_text, stage='window_analysis', metadata=None): invoke_prompt=invoke_model_prompt, activity_callback=activity_callback, conversation_id=conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, ) - comparison_artifact_payload = _maybe_create_comparison_generated_artifacts( + agent_citations = deduplicate_mixed_source_references( + list((tabular_action_payload or {}).get('agent_citations') or []) + + list(comparison_result.get('agent_citations') or []), + reference_type='citation', + ) + generated_tabular_outputs = deduplicate_mixed_source_references( + list((tabular_action_payload or {}).get('generated_tabular_outputs') or []) + + list(comparison_result.get('generated_tabular_outputs') or []), + reference_type='artifact', + ) + _reauthorize_mixed_source_workflow_result( + workflow, + comparison_config, comparison_result, - workflow.get('task_prompt', ''), - conversation_id=conversation_id, + settings, + conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + additional_generated_outputs=generated_tabular_outputs, + ) + try: + comparison_artifact_payload = _maybe_create_comparison_generated_artifacts( + comparison_result, + workflow.get('task_prompt', ''), + conversation_id=conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + _rollback_mixed_source_generated_outputs( + user_id, + conversation_id, + generated_tabular_outputs, + reason='cancellation', + ) + raise + _reauthorize_mixed_source_workflow_result( + workflow, + comparison_config, + comparison_result, + settings, + conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + additional_generated_outputs=( + generated_tabular_outputs + + list(comparison_artifact_payload.get('artifacts') or []) + ), ) token_usage = _finalize_token_usage(token_usage_aggregate) + if comparison_result.get('mixed_source_manifest'): + _emit_mixed_source_token_telemetry( + settings, + 'compare', + token_usage, + request_correlation_id=request_correlation_id, + ) debug_print( '[WorkflowDocumentComparison] Completed workflow action | ' f"workflow_id={workflow.get('id')} | " @@ -5320,8 +6970,8 @@ def invoke_model_prompt(prompt_text, stage='window_analysis', metadata=None): 'model_deployment_name': deployment_name, 'token_usage': token_usage, 'provider': provider, - 'agent_citations': list((tabular_action_payload or {}).get('agent_citations') or []), - 'generated_tabular_outputs': list((tabular_action_payload or {}).get('generated_tabular_outputs') or []), + 'agent_citations': agent_citations, + 'generated_tabular_outputs': generated_tabular_outputs, } @@ -5333,7 +6983,17 @@ def _execute_document_action_workflow( thought_tracker=None, external_activity_callback=None, url_access_context=None, + cancel_requested=None, + request_correlation_id=None, ): + request_correlation_id = normalize_mixed_source_correlation_id( + request_correlation_id + ) + raise_if_mixed_source_cancelled( + cancel_requested, + 'manifest', + request_correlation_id=request_correlation_id, + ) action_config = _get_document_action_config(workflow) action_config = _resolve_recent_document_action_targets(workflow, action_config, settings) workflow = _apply_runtime_document_action_config(workflow, action_config) @@ -5358,6 +7018,8 @@ def _execute_document_action_workflow( external_activity_callback=external_activity_callback, action_config=action_config, url_access_context=url_access_context, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, ) elif action_type == DOCUMENT_ACTION_TYPE_COMPARISON: result = _execute_document_comparison_workflow( @@ -5369,6 +7031,8 @@ def _execute_document_action_workflow( external_activity_callback=external_activity_callback, action_config=action_config, url_access_context=url_access_context, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, ) else: raise ValueError('No document action is enabled for this workflow.') @@ -5383,6 +7047,11 @@ def _execute_document_action_workflow( ) raise + raise_if_mixed_source_cancelled( + cancel_requested, + 'finalization', + request_correlation_id=request_correlation_id, + ) debug_print( '[WorkflowDocumentAction] Action completed | ' f"workflow_id={workflow.get('id')} | " @@ -5672,6 +7341,7 @@ def run_personal_workflow(workflow, trigger_source='manual', user_roles=None, ac execution_workflow, document_action, settings, + conversation_id=conversation.get('id'), thought_tracker=thought_tracker, run_id=run_id, ) @@ -5710,16 +7380,10 @@ def run_personal_workflow(workflow, trigger_source='manual', user_roles=None, ac thought_tracker=thought_tracker, url_access_context=url_access_context, ) - if workflow_search_context: - execution_result.update({ - 'hybrid_citations': workflow_search_context.get('citations') or [], - 'augmented': bool(workflow_search_context.get('citations')), - 'document_search': { - 'query': workflow_search_context.get('query'), - 'result_count': workflow_search_context.get('result_count', 0), - 'document_count': workflow_search_context.get('document_count', 0), - }, - }) + execution_result = _attach_workflow_search_context( + execution_result, + workflow_search_context, + ) else: execution_result = _execute_model_workflow( execution_workflow, @@ -5728,16 +7392,10 @@ def run_personal_workflow(workflow, trigger_source='manual', user_roles=None, ac thought_tracker=thought_tracker, url_access_context=url_access_context, ) - if workflow_search_context: - execution_result.update({ - 'hybrid_citations': workflow_search_context.get('citations') or [], - 'augmented': bool(workflow_search_context.get('citations')), - 'document_search': { - 'query': workflow_search_context.get('query'), - 'result_count': workflow_search_context.get('result_count', 0), - 'document_count': workflow_search_context.get('document_count', 0), - }, - }) + execution_result = _attach_workflow_search_context( + execution_result, + workflow_search_context, + ) execution_result = _attach_workflow_url_access_result(execution_result, url_access_context) assistant_doc = _create_assistant_message( diff --git a/application/single_app/route_backend_chats.py b/application/single_app/route_backend_chats.py index e73a05c03..2b1d9f737 100644 --- a/application/single_app/route_backend_chats.py +++ b/application/single_app/route_backend_chats.py @@ -12,7 +12,11 @@ from semantic_kernel_plugins.plugin_invocation_thoughts import ( register_plugin_invocation_thought_callback, ) -from semantic_kernel_plugins.plugin_invocation_logger import get_plugin_logger, sanitize_plugin_invocation_value +from semantic_kernel_plugins.plugin_invocation_logger import ( + PluginInvocationResult, + get_plugin_logger, + sanitize_plugin_invocation_value, +) from semantic_kernel_plugins.chart_plugin import ChartPlugin from foundry_agent_runtime import FoundryAgentInvocationError, FoundryAgentUserAuthenticationRequired, execute_foundry_agent, resolve_authority from model_endpoint_clients import ( @@ -33,6 +37,26 @@ build_model_endpoint_sync_chat_client, build_semantic_kernel_chat_service_for_model, ) +from functions_mixed_source_orchestration import ( + MixedSourceCancellationError, + MixedSourceFinalizationError, + build_failed_narrative_evidence_envelopes, + build_mixed_source_evidence_handoff, + build_narrative_evidence_envelopes, + build_tabular_file_contexts_from_manifest, + compare_reauthorized_source_manifests, + emit_mixed_source_telemetry, + execute_tabular_evidence_sources, + normalize_mixed_source_correlation_id, + normalize_document_context_request, + partition_source_manifest, + resolve_authorized_source_manifest, + raise_if_mixed_source_cancelled, + should_run_tabular_evidence, +) +from functions_tabular_analysis import ( + get_new_plugin_invocations as _shared_get_new_plugin_invocations, +) import builtins import asyncio, types import ast @@ -55,6 +79,7 @@ from flask import Response, copy_current_request_context, g, has_request_context, stream_with_context from functions_authentication import * from functions_search import * +from functions_search_service import search_relevant_tabular_candidates from functions_service_health import ( SEMANTIC_SEARCH_QUOTA_WARNING_TYPE, SemanticSearchQuotaExceededError, @@ -94,8 +119,18 @@ from functions_chat import * from functions_content import generate_embedding, generate_embeddings_batch from functions_assistant_table_exports import ( - TABLE_EXPORT_REQUEST_MARKERS, - build_assistant_table_csv_export, + assistant_table_export_requested, + build_safe_csv_headers, + has_generated_tabular_csv_output, + neutralize_csv_spreadsheet_formula, +) +from functions_generated_file_exports import ( + build_generated_file_artifact_metadata, + build_generated_file_export, + build_generated_file_output_guidance, + get_generated_file_export_content, + get_requested_generated_file_format, + has_generated_file_output, ) from functions_chart_operations import ( CORE_CHART_PLUGIN_NAME, @@ -149,14 +184,20 @@ normalize_document_action_config, ) from functions_thoughts import ThoughtTracker +from functions_tabular_csv_query import validate_tabular_csv_query_expression from functions_workflow_runner import _execute_document_action_workflow from functions_simplechat_operations import ( + delete_generated_chat_artifact_for_current_user, derive_conversation_title_from_message, upload_chat_image_bytes_for_user, upload_generated_analysis_artifact_for_current_user, ) from functions_tabular_generated_exports import ( + _normalize_generated_batch_entries, + _prepare_tabular_source_rows, build_background_tabular_generated_output_metadata, + build_tabular_generated_output_row_batches, + cancel_tabular_generated_output_run, get_tabular_generated_output_run_status, queue_tabular_generated_output_run, resume_tabular_generated_output_run, @@ -170,8 +211,9 @@ DOCUMENT_ACTION_TYPE_ANALYZE: ASSIGNED_KNOWLEDGE_USER_ACTION_ANALYZE, DOCUMENT_ACTION_TYPE_COMPARISON: ASSIGNED_KNOWLEDGE_USER_ACTION_COMPARE, } -ASSIGNED_KNOWLEDGE_CONTEXT_TOP_N = 12 +ASSIGNED_KNOWLEDGE_CONTEXT_TOP_N = 50 ASSIGNED_KNOWLEDGE_CONTEXT_EXCERPT_MAX_CHARS = 1800 +MIXED_SOURCE_CHAT_RELEVANCE_SOURCE_LIMIT = 48 FOUNDRY_SELECTED_AGENT_TYPES = {'aifoundry', 'new_foundry', 'foundry_workflow'} FOUNDRY_AGENT_PLUGIN_NAMES = { 'aifoundry': 'azure_ai_foundry', @@ -272,6 +314,37 @@ def _safe_metadata_int(value): return 0 +def _merge_chat_token_usage(*token_summaries): + """Merge observed token summaries without estimating missing provider usage.""" + merged = { + 'prompt_tokens': 0, + 'completion_tokens': 0, + 'total_tokens': 0, + 'request_count': 0, + } + has_usage = False + for token_summary in token_summaries: + if not isinstance(token_summary, dict): + continue + summary_has_usage = False + for key in ('prompt_tokens', 'completion_tokens', 'total_tokens'): + try: + value = max(0, int(token_summary.get(key) or 0)) + except (TypeError, ValueError): + value = 0 + merged[key] += value + summary_has_usage = summary_has_usage or bool(value) + try: + request_count = max(0, int(token_summary.get('request_count') or 0)) + except (TypeError, ValueError): + request_count = 0 + if not request_count and summary_has_usage: + request_count = 1 + merged['request_count'] += request_count + has_usage = has_usage or summary_has_usage or bool(request_count) + return merged if has_usage else None + + def _normalize_capability_action(document_action_type): normalized_action_type = str(document_action_type or DOCUMENT_ACTION_TYPE_NONE).strip().lower() if normalized_action_type == DOCUMENT_ACTION_TYPE_ANALYZE: @@ -281,6 +354,505 @@ def _normalize_capability_action(document_action_type): return ASSIGNED_KNOWLEDGE_USER_ACTION_SEARCH +def _maybe_resolve_chat_source_manifest( + settings, + user_id, + conversation_id, + selected_document_ids, + scope_context, +): + if not is_mixed_source_manifest_enabled(settings): + return [] + + requested_source_ids = _normalize_conversation_task_document_ids( + selected_document_ids + ) + if not requested_source_ids: + return [] + + scope_context = scope_context if isinstance(scope_context, dict) else {} + try: + return resolve_authorized_source_manifest( + requested_source_ids, + user_id=user_id, + selection_mode='selected', + conversation_id=conversation_id, + active_group_ids=scope_context.get('active_group_ids'), + active_public_workspace_ids=scope_context.get('active_public_workspace_ids'), + ) + except Exception: + log_event( + '[MixedSourceManifest] Chat shadow resolution failed.', + extra={ + 'requested_source_count': len(requested_source_ids), + 'selection_mode': 'selected', + }, + level=logging.WARNING, + ) + return [] + + +def _normalize_chat_document_context_contract( + settings, + data, + selected_document_ids, + hybrid_search_enabled, +): + """Normalize Phase 2 context intent while preserving flag-off compatibility.""" + normalized_document_ids = _normalize_conversation_task_document_ids( + selected_document_ids + ) + if not is_mixed_source_chat_search_enabled(settings): + return { + 'selection_mode': 'selected' if normalized_document_ids else 'relevance', + 'selected_document_ids': normalized_document_ids, + 'document_context_requested': bool(hybrid_search_enabled), + 'hybrid_search': bool(hybrid_search_enabled), + 'explicit_selection': False, + } + + return normalize_document_context_request( + selection_mode=data.get('selection_mode'), + selected_document_ids=normalized_document_ids, + document_context_requested=data.get('document_context_requested'), + hybrid_search=hybrid_search_enabled, + ) + + +def _resolve_chat_mixed_source_manifest( + settings, + user_id, + conversation_id, + document_ids, + selection_mode, + active_group_ids=None, + active_public_workspace_ids=None, + cancel_requested=None, + request_correlation_id=None, +): + """Resolve a fresh Phase 1 manifest for Phase 2 Chat evidence preparation.""" + if not is_mixed_source_chat_search_enabled(settings): + return [] + normalized_document_ids = _normalize_conversation_task_document_ids(document_ids) + if not normalized_document_ids: + return [] + return resolve_authorized_source_manifest( + normalized_document_ids, + user_id=user_id, + selection_mode=selection_mode, + conversation_id=conversation_id, + active_group_ids=active_group_ids, + active_public_workspace_ids=active_public_workspace_ids, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + + +def _get_manifest_partition_document_ids(manifest_partitions, partition_name): + partitions = manifest_partitions if isinstance(manifest_partitions, dict) else {} + return [ + str(source.get('document_id') or '').strip() + for source in partitions.get(partition_name) or [] + if str(source.get('document_id') or '').strip() + ] + + +def _resolve_chat_mixed_source_partition( + settings, + user_id, + conversation_id, + document_ids, + selection_mode, + active_group_ids=None, + active_public_workspace_ids=None, + cancel_requested=None, + request_correlation_id=None, +): + """Resolve and partition one current authorization-safe source manifest.""" + manifest = _resolve_chat_mixed_source_manifest( + settings, + user_id, + conversation_id, + document_ids, + selection_mode, + active_group_ids=active_group_ids, + active_public_workspace_ids=active_public_workspace_ids, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + partitions = partition_source_manifest(manifest) + return { + 'manifest': manifest, + 'partitions': partitions, + 'narrative_document_ids': _get_manifest_partition_document_ids( + partitions, + 'narrative_sources', + ), + 'tabular_sources': list(partitions.get('tabular_sources') or []), + } + + +def _build_mixed_source_continuity_refs(manifest, evidence_envelopes, selection_origin): + """Persist only compact source identity and terminal state for later reauthorization.""" + evidence_by_document_id = { + str(envelope.get('document_id') or '').strip(): envelope + for envelope in list(evidence_envelopes or []) + if isinstance(envelope, dict) and str(envelope.get('document_id') or '').strip() + } + continuity_refs = [] + for requested_order, source in enumerate(list(manifest or [])[:100]): + if not isinstance(source, dict) or source.get('authorization_status') != 'authorized': + continue + document_id = str(source.get('document_id') or '').strip() + scope = str(source.get('scope') or '').strip().lower() + scope_id = str(source.get('scope_id') or '').strip() + if not document_id or not scope or not scope_id: + continue + envelope = evidence_by_document_id.get(document_id) or {} + coverage = envelope.get('coverage') if isinstance(envelope.get('coverage'), dict) else {} + continuity_ref = { + 'document_id': document_id, + 'scope': scope, + 'scope_id': scope_id, + 'source_role': str(source.get('source_role') or 'selected'), + 'requested_order': requested_order, + 'source_kind': source.get('source_kind'), + 'engine': envelope.get('engine'), + 'source_version': source.get('source_version'), + 'status': envelope.get('status') or 'unavailable', + 'coverage': { + 'partial_coverage': bool(coverage.get('partial_coverage')), + 'evidence_envelope_truncated': bool(coverage.get('evidence_envelope_truncated')), + 'failed': str(envelope.get('status') or '').strip().lower() in {'failed', 'unavailable'}, + }, + 'selection_origin': selection_origin, + 'action_mode': 'chat', + 'citation_count': len(envelope.get('citations') or []), + 'artifact_count': len(envelope.get('generated_artifacts') or []), + } + if scope == 'group': + continuity_ref['group_id'] = scope_id + elif scope == 'public': + continuity_ref['public_workspace_id'] = scope_id + else: + continuity_ref['user_id'] = scope_id + continuity_refs.append(continuity_ref) + return continuity_refs + + +def _reauthorize_document_action_finalization( + normalized_action, + execution_result, + user_id, + conversation_id, + cancel_requested=None, + request_correlation_id=None, + settings=None, +): + """Reauthorize every action source and exact version before publishing output.""" + raise_if_mixed_source_cancelled( + cancel_requested, + 'finalization', + request_correlation_id=request_correlation_id, + ) + normalized_action = normalized_action if isinstance(normalized_action, dict) else {} + analysis_result = ( + execution_result.get('analysis_result') + if isinstance(execution_result, dict) + and isinstance(execution_result.get('analysis_result'), dict) + else {} + ) + execution_manifest = analysis_result.get('mixed_source_manifest') + requested_ids = [ + str(document_id or '').strip() + for document_id in list(normalized_action.get('document_ids') or []) + if str(document_id or '').strip() + ] + if not requested_ids and isinstance(execution_manifest, list): + requested_ids = [ + str(source.get('document_id') or '').strip() + for source in execution_manifest + if isinstance(source, dict) and str(source.get('document_id') or '').strip() + ] + if not requested_ids: + return + + finalization_selection_mode = str( + normalized_action.get('target_mode') or 'selected' + ).strip().lower() + if finalization_selection_mode not in {'selected', 'all', 'history', 'relevance'}: + finalization_selection_mode = 'selected' + fresh_manifest = resolve_authorized_source_manifest( + requested_ids, + user_id=user_id, + selection_mode=finalization_selection_mode, + conversation_id=conversation_id, + active_group_ids=normalized_action.get('active_group_ids'), + active_public_workspace_ids=normalized_action.get('active_public_workspace_id'), + doc_scope=normalized_action.get('doc_scope', 'all'), + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + if isinstance(execution_manifest, list): + _validate_reauthorized_manifest_finalization( + execution_manifest, + fresh_manifest, + settings=settings, + mode=( + 'compare' + if normalized_action.get('type') == DOCUMENT_ACTION_TYPE_COMPARISON + else 'analyze' + ), + request_correlation_id=request_correlation_id, + ) + elif len(fresh_manifest) != len(requested_ids) or any( + source.get('authorization_status') != 'authorized' + for source in fresh_manifest + ): + raise PermissionError('One or more selected sources are no longer available.') + + +def _validate_reauthorized_manifest_finalization( + execution_manifest, + fresh_manifest, + settings=None, + mode='chat', + request_correlation_id=None, +): + """Require every previously authorized source to retain identity and version.""" + finalization_result = compare_reauthorized_source_manifests( + execution_manifest, + fresh_manifest, + ) + authorization_failure_count = finalization_result['authorization_failure_count'] + source_version_changed_count = finalization_result['source_version_changed_count'] + + if authorization_failure_count: + emit_mixed_source_telemetry( + settings, + 'authorization_failure', + mode, + request_correlation_id=request_correlation_id, + metrics={ + 'authorization_failure_count': authorization_failure_count, + }, + dimensions={'outcome_status': 'failed'}, + ) + raise MixedSourceFinalizationError('authorization_lost') + if source_version_changed_count: + log_event( + '[MixedSourceLifecycle] Finalization source version changed.', + extra={ + 'request_correlation_id': request_correlation_id, + 'source_version_changed_count': source_version_changed_count, + }, + level=logging.WARNING, + ) + raise MixedSourceFinalizationError('source_version_changed') + + +def _build_reauthorized_continuity_decision(prior_refs, manifest, explicit_selection): + """Describe a fresh manifest decision without treating persisted refs as authority.""" + if explicit_selection: + return { + 'selection_origin': 'selected', + 'prior_source_count': 0, + 'reauthorized_source_count': 0, + 'unavailable_source_count': 0, + 'source_version_changed_count': 0, + 'requires_native_execution': False, + } + + prior_by_document_id = { + str(ref.get('document_id') or '').strip(): ref + for ref in list(prior_refs or []) + if isinstance(ref, dict) and str(ref.get('document_id') or '').strip() + } + unavailable_source_count = 0 + source_version_changed_count = 0 + incomplete_prior_source_count = 0 + reauthorized_source_count = 0 + for source in list(manifest or []): + document_id = str((source or {}).get('document_id') or '').strip() + prior_ref = prior_by_document_id.get(document_id) + if not prior_ref: + continue + if source.get('authorization_status') != 'authorized': + unavailable_source_count += 1 + continue + reauthorized_source_count += 1 + prior_status = str(prior_ref.get('status') or '').strip().lower() + prior_coverage = ( + prior_ref.get('coverage') + if isinstance(prior_ref.get('coverage'), dict) + else {} + ) + if ( + prior_status not in {'completed'} + or prior_coverage.get('partial_coverage') + or prior_coverage.get('failed') + or prior_coverage.get('evidence_envelope_truncated') + ): + incomplete_prior_source_count += 1 + prior_version = prior_ref.get('source_version') + current_version = source.get('source_version') + if prior_version is not None and current_version is not None and str(prior_version) != str(current_version): + source_version_changed_count += 1 + + return { + 'selection_origin': 'history', + 'prior_source_count': len(prior_by_document_id), + 'reauthorized_source_count': reauthorized_source_count, + 'unavailable_source_count': unavailable_source_count, + 'source_version_changed_count': source_version_changed_count, + 'incomplete_prior_source_count': incomplete_prior_source_count, + 'requires_native_execution': bool( + unavailable_source_count + or source_version_changed_count + or incomplete_prior_source_count + ), + } + + +def _resolve_reauthorized_continuity_decision( + settings, + user_id, + conversation_id, + prior_refs, + request_correlation_id=None, +): + """Resolve persisted continuity hints through a fresh authorized manifest.""" + if not is_mixed_source_conversation_continuity_enabled(settings): + return None + requested_document_ids = [ + str(ref.get('document_id') or '').strip() + for ref in list(prior_refs or []) + if isinstance(ref, dict) and str(ref.get('document_id') or '').strip() + ] + if not requested_document_ids: + return None + + search_parameters = build_prior_grounded_document_search_parameters(prior_refs) + search_parameters = revalidate_prior_grounded_document_search_parameters( + user_id, + search_parameters, + ) + manifest = _resolve_chat_mixed_source_manifest( + settings, + user_id, + conversation_id, + requested_document_ids, + 'history', + active_group_ids=search_parameters.get('active_group_ids'), + active_public_workspace_ids=search_parameters.get('active_public_workspace_ids'), + request_correlation_id=request_correlation_id, + ) + return _build_reauthorized_continuity_decision( + prior_refs, + manifest, + explicit_selection=False, + ) + + +def _can_reuse_prior_grounded_history(history_assessment, continuity_decision): + """Return whether existing history is complete enough to avoid fresh native evidence.""" + if ( + isinstance(continuity_decision, dict) + and continuity_decision.get('requires_native_execution') + ): + return False + return bool( + isinstance(history_assessment, dict) + and history_assessment.get('can_answer_from_history') + ) + + +def _resolve_chat_mixed_source_relevance_context( + *, + settings, + user_id, + conversation_id, + query, + search_results, + document_scope, + candidate_document_ids=None, + tags_filter=None, + active_group_ids=None, + active_public_workspace_ids=None, + cancel_requested=None, + request_correlation_id=None, +): + """Add bounded schema candidates and reauthorize all relevance-derived sources.""" + if not is_mixed_source_chat_search_enabled(settings): + return { + 'manifest': [], + 'partitions': {}, + 'narrative_document_ids': [], + 'tabular_sources': [], + 'search_results': list(search_results or []), + 'tabular_candidate_count': 0, + } + + relevance_candidates_enabled = is_mixed_source_relevance_candidates_enabled(settings) + candidate_result = { + 'document_ids': [], + 'candidate_count': 0, + } + if relevance_candidates_enabled: + candidate_result = search_relevant_tabular_candidates( + query=query, + user_id=user_id, + doc_scope=document_scope, + document_ids=candidate_document_ids, + tags_filter=tags_filter, + active_group_ids=active_group_ids, + active_public_workspace_id=active_public_workspace_ids, + ) + relevance_document_ids = [] + seen_document_ids = set() + for result in list(search_results or []): + document_id = str((result or {}).get('document_id') or '').strip() + if document_id and document_id not in seen_document_ids: + seen_document_ids.add(document_id) + relevance_document_ids.append(document_id) + if len(relevance_document_ids) >= ( + MIXED_SOURCE_CHAT_RELEVANCE_SOURCE_LIMIT - 6 + ): + break + for document_id in candidate_result.get('document_ids') or []: + normalized_document_id = str(document_id or '').strip() + if normalized_document_id and normalized_document_id not in seen_document_ids: + seen_document_ids.add(normalized_document_id) + relevance_document_ids.append(normalized_document_id) + if len(relevance_document_ids) >= MIXED_SOURCE_CHAT_RELEVANCE_SOURCE_LIMIT: + break + + resolved_context = _resolve_chat_mixed_source_partition( + settings, + user_id, + conversation_id, + relevance_document_ids, + 'relevance', + active_group_ids=active_group_ids, + active_public_workspace_ids=active_public_workspace_ids, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + narrative_document_id_set = set( + resolved_context.get('narrative_document_ids') or [] + ) + resolved_context['search_results'] = [ + result + for result in list(search_results or []) + if str((result or {}).get('document_id') or '').strip() + in narrative_document_id_set + ] + resolved_context['tabular_candidate_count'] = int( + candidate_result.get('candidate_count') or 0 + ) + return resolved_context + + def _source_review_metadata_used(source_review_result): if not isinstance(source_review_result, dict): return False @@ -1106,6 +1678,124 @@ def _strip_agent_citation_artifact_refs(agent_citations): return compact_citations +def _rollback_agent_citation_artifacts(conversation_id, compact_citations): + """Remove current-message citation artifacts and chunks after aborted publication.""" + normalized_conversation_id = str(conversation_id or '').strip() + artifact_ids = [] + for citation in list(compact_citations or []): + if not isinstance(citation, dict): + continue + artifact_id = str(citation.get('artifact_id') or '').strip() + if artifact_id and artifact_id not in artifact_ids: + artifact_ids.append(artifact_id) + + deleted_artifact_count = 0 + rollback_failure_count = 0 + for artifact_id in artifact_ids: + try: + chunk_documents = list(cosmos_messages_container.query_items( + query='SELECT c.id FROM c WHERE c.parent_message_id = @parent_message_id', + parameters=[{'name': '@parent_message_id', 'value': artifact_id}], + partition_key=normalized_conversation_id, + )) + for chunk_document in chunk_documents: + chunk_id = str((chunk_document or {}).get('id') or '').strip() + if chunk_id: + cosmos_messages_container.delete_item( + item=chunk_id, + partition_key=normalized_conversation_id, + ) + cosmos_messages_container.delete_item( + item=artifact_id, + partition_key=normalized_conversation_id, + ) + deleted_artifact_count += 1 + except Exception: + rollback_failure_count += 1 + + if artifact_ids: + log_event( + '[MixedSourceLifecycle] Citation artifact rollback completed.', + extra={ + 'citation_artifact_count': len(artifact_ids), + 'deleted_artifact_count': deleted_artifact_count, + 'rollback_failure_count': rollback_failure_count, + }, + level=logging.INFO if not rollback_failure_count else logging.WARNING, + ) + return { + 'deleted_artifact_count': deleted_artifact_count, + 'rollback_failure_count': rollback_failure_count, + } + + +def _rollback_mixed_source_chat_publication( + user_id, + conversation_id, + generated_outputs=None, + compact_citations=None, +): + """Cancel queued exports and remove artifacts created before mixed publication stopped.""" + normalized_user_id = str(user_id or '').strip() + normalized_conversation_id = str(conversation_id or '').strip() + export_run_ids = [] + artifact_message_ids = [] + for output in list(generated_outputs or []): + if not isinstance(output, dict): + continue + export_run_id = str(output.get('export_run_id') or '').strip() + artifact_message_id = str(output.get('artifact_message_id') or '').strip() + if export_run_id and export_run_id not in export_run_ids: + export_run_ids.append(export_run_id) + if artifact_message_id and artifact_message_id not in artifact_message_ids: + artifact_message_ids.append(artifact_message_id) + + canceled_export_count = 0 + deleted_generated_artifact_count = 0 + rollback_failure_count = 0 + for export_run_id in export_run_ids: + try: + cancel_result = cancel_tabular_generated_output_run( + normalized_user_id, + export_run_id, + ) + if isinstance(cancel_result, dict) and cancel_result.get('canceled'): + canceled_export_count += 1 + except Exception: + rollback_failure_count += 1 + for artifact_message_id in artifact_message_ids: + try: + if delete_generated_chat_artifact_for_current_user( + normalized_conversation_id, + artifact_message_id, + ): + deleted_generated_artifact_count += 1 + except Exception: + rollback_failure_count += 1 + + citation_rollback = _rollback_agent_citation_artifacts( + normalized_conversation_id, + compact_citations, + ) + rollback_failure_count += citation_rollback.get('rollback_failure_count', 0) + log_event( + '[MixedSourceLifecycle] Chat publication rollback completed.', + extra={ + 'canceled_export_count': canceled_export_count, + 'deleted_generated_artifact_count': deleted_generated_artifact_count, + 'deleted_citation_artifact_count': citation_rollback.get('deleted_artifact_count', 0), + 'rollback_failure_count': rollback_failure_count, + }, + level=logging.INFO if not rollback_failure_count else logging.WARNING, + ) + return { + 'canceled_export_count': canceled_export_count, + 'deleted_generated_artifact_count': deleted_generated_artifact_count, + 'deleted_citation_artifact_count': citation_rollback.get('deleted_artifact_count', 0), + 'rollback_failure_count': rollback_failure_count, + } + + FACT_MEMORY_TYPE_FACT = 'fact' FACT_MEMORY_TYPE_INSTRUCTION = 'instruction' FACT_MEMORY_TYPE_LEGACY_DESCRIBER = 'describer' @@ -1223,7 +1913,13 @@ def _normalize_generated_analysis_artifact_metadata(raw_artifact, default_capabi artifact_message_id = str(raw_artifact.get('artifact_message_id') or '').strip() document_id = str(raw_artifact.get('document_id') or '').strip() export_run_id = str(raw_artifact.get('export_run_id') or raw_artifact.get('run_id') or '').strip() - if not artifact_message_id and not document_id and not export_run_id: + terminal_status = str(raw_artifact.get('status') or '').strip().lower() + suppress_assistant_table_export = bool(raw_artifact.get('suppress_assistant_table_export')) + is_terminal_export_status = ( + terminal_status in {'failed', 'canceled'} + and suppress_assistant_table_export + ) + if not artifact_message_id and not document_id and not export_run_id and not is_terminal_export_status: return None normalized_artifact = dict(raw_artifact) @@ -1238,6 +1934,9 @@ def _normalize_generated_analysis_artifact_metadata(raw_artifact, default_capabi if export_run_id: normalized_artifact['export_run_id'] = export_run_id normalized_artifact['background_export'] = bool(raw_artifact.get('background_export', True)) + elif is_terminal_export_status: + normalized_artifact['background_export'] = True + normalized_artifact['suppress_assistant_table_export'] = True normalized_output_format = str(raw_artifact.get('output_format') or '').strip().lower() if normalized_output_format: @@ -1352,84 +2051,173 @@ def _maybe_create_deep_research_ledger_artifact(settings, conversation_id, ledge def _has_generated_tabular_csv_output(generated_outputs): - for generated_output in generated_outputs or []: - if not isinstance(generated_output, dict): - continue - - capability = str(generated_output.get('capability') or '').strip().lower() - output_format = str(generated_output.get('output_format') or '').strip().lower() - file_name = str(generated_output.get('file_name') or '').strip().lower() - if output_format == 'csv' or file_name.endswith('.csv'): - if not capability or capability == 'tabular': - return True - - return False + return has_generated_tabular_csv_output(generated_outputs) -def maybe_create_assistant_table_generated_output( +def maybe_create_generated_file_output( user_question, assistant_content, conversation_id, + function_results=None, existing_outputs=None, + cancel_requested=None, + request_correlation_id=None, ): - """Save a CSV artifact when a table-request answer contains a parseable table.""" - if _has_generated_tabular_csv_output(existing_outputs): + """Save a requested CSV, DOCX, or PDF artifact from response and action evidence.""" + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) + output_format = get_requested_generated_file_format(user_question) + if not output_format: + return None + if output_format == 'csv' and _has_generated_tabular_csv_output(existing_outputs): + return None + if has_generated_file_output(existing_outputs, output_format): return None - export_payload = build_assistant_table_csv_export(user_question, assistant_content) + export_payload = build_generated_file_export( + user_question, + assistant_content, + function_results=function_results, + ) if not export_payload: return None - generated_file_name = export_payload.get('file_name') + generated_file_name = str(export_payload.get('file_name') or '').strip() + if not generated_file_name: + return None row_count = _safe_int(export_payload.get('row_count')) + settings = get_settings() + structured_rows = export_payload.get('_structured_rows') or [] + row_batches = [] + if output_format == 'csv': + row_batches = _build_tabular_generated_output_row_batches( + structured_rows, + settings=settings, + ) + if output_format == 'csv' and should_queue_tabular_generated_output_background( + row_count, + len(row_batches), + settings, + ): + try: + background_run = queue_tabular_generated_output_run( + user_id=get_current_user_id(), + conversation_id=conversation_id, + user_question=user_question, + source_candidate={ + 'filename': generated_file_name, + 'selected_sheet': '', + 'source_authorization': { + 'source': 'chat', + }, + }, + output_format=output_format, + row_batches=row_batches, + gpt_model='', + settings=settings, + passthrough_input_rows=True, + ) + background_metadata = build_background_tabular_generated_output_metadata(background_run) + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'export', + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + cancel_tabular_generated_output_run( + get_current_user_id(), + background_metadata.get('export_run_id'), + ) + raise + return background_metadata + except MixedSourceCancellationError: + raise + except Exception as exc: + log_event( + '[Generated File Export] Failed to queue large CSV export', + { + 'conversation_id': conversation_id, + 'generated_file_name': generated_file_name, + 'row_count': row_count, + 'output_format': output_format, + 'error': str(exc), + }, + level=logging.ERROR, + exceptionTraceback=True, + ) + return None + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) upload_result = upload_generated_analysis_artifact_for_current_user( conversation_id=conversation_id, file_name=generated_file_name, file_content=export_payload.get('file_content'), - capability='tabular', - output_format='csv', + capability=export_payload.get('capability') or 'file_export', + output_format=output_format, summary=export_payload.get('summary'), ) + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + delete_generated_chat_artifact_for_current_user( + conversation_id, + (upload_result.get('message') or {}).get('id'), + ) + raise + except MixedSourceCancellationError: + raise except Exception as exc: log_event( - '[Assistant Table Export] Failed to save assistant table CSV artifact', + '[Generated File Export] Failed to save generated file artifact', { 'conversation_id': conversation_id, 'generated_file_name': generated_file_name, 'row_count': row_count, + 'output_format': output_format, 'error': str(exc), }, debug_only=True, ) return None - artifact_message_id = upload_result.get('message', {}).get('id') - if not artifact_message_id: + artifact_metadata = build_generated_file_artifact_metadata( + export_payload, + upload_result, + conversation_id, + ) + if not artifact_metadata: return None - uploaded_file_name = upload_result.get('message', {}).get('file_name') or generated_file_name log_event( - '[Assistant Table Export] Saved assistant table CSV artifact', + '[Generated File Export] Saved generated file artifact', { 'conversation_id': conversation_id, - 'artifact_message_id': artifact_message_id, - 'generated_file_name': uploaded_file_name, + 'artifact_message_id': artifact_metadata.get('artifact_message_id'), + 'generated_file_name': artifact_metadata.get('file_name'), 'row_count': row_count, + 'output_format': output_format, }, debug_only=True, ) - return { - 'capability': 'tabular', - 'artifact_message_id': artifact_message_id, - 'conversation_id': conversation_id, - 'storage_scope': 'chat', - 'file_name': uploaded_file_name, - 'output_format': 'csv', - 'row_count': row_count, - 'preview_rows': export_payload.get('preview_rows') or [], - 'summary': export_payload.get('summary'), - } + return artifact_metadata + + +def maybe_create_assistant_table_generated_output(*args, **kwargs): + """Backward-compatible wrapper for the generic generated-file finalizer.""" + return maybe_create_generated_file_output(*args, **kwargs) def _safe_int(value, default=0): @@ -2588,6 +3376,8 @@ def persist_agent_citation_artifacts( agent_citations, created_timestamp, user_info=None, + cancel_requested=None, + request_correlation_id=None, ): """Persist raw agent citation payloads outside the primary assistant message doc.""" if not agent_citations: @@ -2601,10 +3391,43 @@ def persist_agent_citation_artifacts( user_info=user_info, ) + persisted_artifact_ids = [] try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) for artifact_doc in artifact_docs: + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) cosmos_messages_container.upsert_item(artifact_doc) + artifact_id = str(artifact_doc.get('id') or '').strip() + if artifact_id: + persisted_artifact_ids.append(artifact_id) + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) return compact_citations + except MixedSourceCancellationError: + for artifact_id in reversed(persisted_artifact_ids): + try: + cosmos_messages_container.delete_item( + item=artifact_id, + partition_key=conversation_id, + ) + except Exception: + log_event( + '[Agent Citations] Citation rollback after cancellation failed.', + extra={'citation_artifact_rollback_failure_count': 1}, + level=logging.WARNING, + ) + raise except Exception as exc: log_event( f"[Agent Citations] Failed to persist assistant artifacts: {exc}", @@ -3492,7 +4315,7 @@ def _resolve_tabular_related_document_evidence(document_match, user_question, us search_payload = search_documents( query=search_query, user_id=user_id, - top_n=2, + top_n=10, doc_scope=doc_scope, document_ids=[document_id], active_group_ids=active_group_ids, @@ -3683,10 +4506,15 @@ def augment_tabular_invocations_with_related_document_evidence(invocations, user updated_result_payload['data'] = updated_rows updated_result_payload['referenced_document_row_count'] = augmented_rows_for_invocation updated_result_payload['referenced_document_match_count'] = augmented_documents_for_invocation - if isinstance(getattr(invocation, 'result', None), dict): + existing_result = getattr(invocation, 'result', None) + internal_metadata = getattr(existing_result, 'internal_metadata', {}) or {} + if isinstance(existing_result, dict): invocation.result = updated_result_payload else: - invocation.result = json.dumps(updated_result_payload, indent=2, default=str, ensure_ascii=False) + invocation.result = PluginInvocationResult( + json.dumps(updated_result_payload, indent=2, default=str, ensure_ascii=False), + internal_metadata=internal_metadata, + ) return { 'augmented_row_count': augmented_row_count, @@ -3815,11 +4643,9 @@ def get_tabular_generated_output_format(user_question): 'return json', 'valid json', ) - csv_markers = TABLE_EXPORT_REQUEST_MARKERS - if any(marker in normalized_question for marker in json_markers): return 'json' - if any(marker in normalized_question for marker in csv_markers): + if assistant_table_export_requested(user_question): return 'csv' return None @@ -3844,7 +4670,7 @@ def question_requests_tabular_generated_output(user_question): 'each object', 'each row', ) - if requested_format == 'csv' and any(marker in normalized_question for marker in TABLE_EXPORT_REQUEST_MARKERS): + if requested_format == 'csv' and assistant_table_export_requested(user_question): return True return any(marker in normalized_question for marker in exhaustive_markers) @@ -3913,13 +4739,13 @@ def _serialize_tabular_generated_output_value(value): if value is None: return '' if isinstance(value, (dict, list)): - return json.dumps(value, default=str, ensure_ascii=False) + return neutralize_csv_spreadsheet_formula(json.dumps(value, default=str, ensure_ascii=False)) if hasattr(value, 'isoformat') and not isinstance(value, str): try: - return value.isoformat() + return neutralize_csv_spreadsheet_formula(value.isoformat()) except TypeError: pass - return str(value) + return neutralize_csv_spreadsheet_formula(value) def _build_tabular_generated_output_csv(entries): @@ -3938,14 +4764,15 @@ def _build_tabular_generated_output_csv(entries): if not ordered_columns: ordered_columns = ['value'] + safe_ordered_columns = build_safe_csv_headers(ordered_columns) output_buffer = io.StringIO() - writer = csv.DictWriter(output_buffer, fieldnames=ordered_columns) + writer = csv.DictWriter(output_buffer, fieldnames=safe_ordered_columns) writer.writeheader() for entry in entries or []: serialized_row = {} if isinstance(entry, dict): - for field_name in ordered_columns: - serialized_row[field_name] = _serialize_tabular_generated_output_value(entry.get(field_name)) + for field_name, safe_field_name in zip(ordered_columns, safe_ordered_columns): + serialized_row[safe_field_name] = _serialize_tabular_generated_output_value(entry.get(field_name)) writer.writerow(serialized_row) return output_buffer.getvalue() @@ -4253,30 +5080,7 @@ def _build_tabular_generated_output_file_name(source_file_name, output_format): def _build_tabular_generated_output_row_batches(rows, settings=None): - budget = _get_tabular_generated_output_batch_budget(settings) - max_batch_rows = budget['max_rows'] - max_batch_chars = budget['max_chars'] - batches = [] - current_batch = [] - current_batch_chars = 0 - - for row in rows or []: - row_text = _dump_tabular_generated_output_json(row) - if current_batch and ( - len(current_batch) >= max_batch_rows - or current_batch_chars + len(row_text) > max_batch_chars - ): - batches.append(current_batch) - current_batch = [] - current_batch_chars = 0 - - current_batch.append(row) - current_batch_chars += len(row_text) - - if current_batch: - batches.append(current_batch) - - return batches + return build_tabular_generated_output_row_batches(rows, settings=settings) def _build_tabular_generated_output_candidate_diagnostic(invocation): @@ -4319,6 +5123,9 @@ def _build_tabular_generated_output_candidate_diagnostic(invocation): 'full_result_available': full_result_available, 'function_rank': function_rank, 'max_rows': result_payload.get('max_rows') if isinstance(result_payload, dict) else None, + 'start_row': _safe_int(result_payload.get('start_row')) if isinstance(result_payload, dict) else 0, + 'has_more': bool(result_payload.get('has_more')) if isinstance(result_payload, dict) else False, + 'next_start_row': result_payload.get('next_start_row') if isinstance(result_payload, dict) else None, 'filter_applied': result_payload.get('filter_applied') if isinstance(result_payload, dict) else None, 'normalized_match': result_payload.get('normalized_match') if isinstance(result_payload, dict) else None, 'skip_reason': skip_reason, @@ -4333,50 +5140,327 @@ def _build_tabular_generated_output_candidate_diagnostics(invocations): ] +def _build_tabular_generated_output_source_signature(invocation, result_payload): + invocation_parameters = getattr(invocation, 'parameters', {}) or {} + invocation_internal_metadata = getattr( + getattr(invocation, 'result', None), + 'internal_metadata', + {}, + ) or {} + source_descriptor = invocation_internal_metadata.get('tabular_generated_export_source') or {} + source_authorization = invocation_internal_metadata.get('tabular_source_authorization') or {} + source_identity = { + 'source': source_descriptor.get('source') or source_authorization.get('source'), + 'scope_id': source_descriptor.get('scope_id') or source_authorization.get('scope_id'), + 'container': source_descriptor.get('container') or source_authorization.get('container'), + 'blob_path': source_descriptor.get('blob_path') or source_authorization.get('blob_path'), + 'blob_etag': source_descriptor.get('blob_etag') or source_authorization.get('blob_etag'), + } + signature_parameters = { + str(parameter_name): parameter_value + for parameter_name, parameter_value in invocation_parameters.items() + if str(parameter_name) not in { + 'user_id', + 'conversation_id', + 'start_row', + 'max_rows', + } + } + signature_payload = { + 'plugin_name': str(getattr(invocation, 'plugin_name', '') or '').strip(), + 'function_name': str(getattr(invocation, 'function_name', '') or '').strip(), + 'filename': result_payload.get('filename'), + 'selected_sheet': result_payload.get('selected_sheet'), + 'parameters': signature_parameters, + 'source_identity': source_identity, + } + return json.dumps(signature_payload, sort_keys=True, default=str, separators=(',', ':')) + + +def _coalesce_tabular_generated_output_pages(pages, total_matches): + normalized_total_matches = _safe_int(total_matches) + ordered_pages = sorted(pages or [], key=lambda page: page.get('start_row', 0)) + merged_rows = [] + next_expected_row = 0 + + for page in ordered_pages: + start_row = _safe_int(page.get('start_row')) + page_rows = page.get('rows') if isinstance(page.get('rows'), list) else [] + returned_rows = _safe_int(page.get('returned_rows')) + if returned_rows != len(page_rows): + return { + 'rows': merged_rows, + 'row_count': len(merged_rows), + 'full_result_available': False, + 'validation_error': ( + f'Page at row {start_row} declared {returned_rows} row(s) but contained {len(page_rows)}' + ), + } + if start_row < next_expected_row: + return { + 'rows': merged_rows, + 'row_count': len(merged_rows), + 'full_result_available': False, + 'validation_error': ( + f'Page overlap at row {start_row}; next expected row was {next_expected_row}' + ), + } + if start_row > next_expected_row: + return { + 'rows': merged_rows, + 'row_count': len(merged_rows), + 'full_result_available': False, + 'validation_error': ( + f'Page gap from row {next_expected_row} through {start_row - 1}' + ), + } + + merged_rows.extend(page_rows) + next_expected_row += len(page_rows) + + full_result_available = normalized_total_matches > 0 and next_expected_row == normalized_total_matches + validation_error = None + if next_expected_row < normalized_total_matches: + validation_error = ( + f'Page gap from row {next_expected_row} through {normalized_total_matches - 1}' + ) + elif normalized_total_matches and next_expected_row > normalized_total_matches: + validation_error = ( + f'Page coverage returned {next_expected_row} row(s) for {normalized_total_matches} total match(es)' + ) + + return { + 'rows': merged_rows, + 'row_count': len(merged_rows), + 'full_result_available': full_result_available, + 'validation_error': validation_error, + } + + def _build_tabular_generated_output_source_candidate(invocations): - best_candidate = None - best_score = None + candidate_groups = {} + logical_source_identities = {} for invocation in invocations or []: diagnostic = _build_tabular_generated_output_candidate_diagnostic(invocation) if diagnostic.get('skip_reason'): continue + result_payload = get_tabular_invocation_result_payload(invocation) + invocation_internal_metadata = getattr( + getattr(invocation, 'result', None), + 'internal_metadata', + {}, + ) or {} + source_signature = _build_tabular_generated_output_source_signature(invocation, result_payload) + logical_signature = json.dumps({ + 'plugin_name': str(getattr(invocation, 'plugin_name', '') or '').strip(), + 'function_name': str(getattr(invocation, 'function_name', '') or '').strip(), + 'filename': result_payload.get('filename'), + 'selected_sheet': result_payload.get('selected_sheet'), + 'parameters': { + str(parameter_name): parameter_value + for parameter_name, parameter_value in (getattr(invocation, 'parameters', {}) or {}).items() + if str(parameter_name) not in { + 'user_id', + 'conversation_id', + 'start_row', + 'max_rows', + } + }, + }, sort_keys=True, default=str, separators=(',', ':')) + logical_source_identities.setdefault(logical_signature, set()).add(source_signature) + candidate_group = candidate_groups.setdefault(source_signature, { + 'function_name': diagnostic.get('function_name'), + 'filename': result_payload.get('filename'), + 'selected_sheet': result_payload.get('selected_sheet'), + 'source_parameters': dict(getattr(invocation, 'parameters', {}) or {}), + 'source_descriptor': invocation_internal_metadata.get('tabular_generated_export_source'), + 'source_authorization': invocation_internal_metadata.get('tabular_source_authorization'), + 'function_rank': diagnostic.get('function_rank') or 0, + 'total_matches': diagnostic.get('total_matches') or 0, + 'pages': [], + 'diagnostics': [], + 'logical_signature': logical_signature, + }) + candidate_group['diagnostics'].append(diagnostic) + candidate_group['pages'].append({ + 'start_row': diagnostic.get('start_row') or 0, + 'returned_rows': diagnostic.get('returned_rows') or diagnostic.get('data_row_count') or 0, + 'rows': result_payload.get('data'), + }) + if diagnostic.get('total_matches') != candidate_group['total_matches']: + candidate_group['total_count_mismatch'] = True + + best_candidate = None + best_score = None + for candidate_group in candidate_groups.values(): + coalesced_result = _coalesce_tabular_generated_output_pages( + candidate_group.get('pages'), + candidate_group.get('total_matches'), + ) + if candidate_group.get('total_count_mismatch'): + coalesced_result['full_result_available'] = False + coalesced_result['validation_error'] = 'Compatible pages reported inconsistent total row counts' + if len(logical_source_identities.get(candidate_group.get('logical_signature'), set())) > 1: + coalesced_result['full_result_available'] = False + coalesced_result['validation_error'] = ( + 'Compatible pages resolved to different source blobs or versions' + ) + score = ( - 1 if diagnostic.get('full_result_available') else 0, - diagnostic.get('returned_rows') or diagnostic.get('data_row_count') or 0, - diagnostic.get('function_rank') or 0, + 1 if coalesced_result.get('full_result_available') else 0, + coalesced_result.get('row_count') or 0, + candidate_group.get('function_rank') or 0, ) if best_score is not None and score <= best_score: continue - result_payload = get_tabular_invocation_result_payload(invocation) best_candidate = { - 'function_name': diagnostic.get('function_name'), - 'filename': result_payload.get('filename'), - 'selected_sheet': result_payload.get('selected_sheet'), - 'rows': result_payload.get('data'), - 'row_count': diagnostic.get('returned_rows') or diagnostic.get('data_row_count'), - 'total_matches': diagnostic.get('total_matches'), - 'full_result_available': diagnostic.get('full_result_available'), - 'diagnostics': diagnostic, + 'function_name': candidate_group.get('function_name'), + 'filename': candidate_group.get('filename'), + 'selected_sheet': candidate_group.get('selected_sheet'), + 'source_parameters': candidate_group.get('source_parameters'), + 'source_descriptor': candidate_group.get('source_descriptor'), + 'source_authorization': candidate_group.get('source_authorization'), + 'rows': coalesced_result.get('rows'), + 'row_count': coalesced_result.get('row_count'), + 'total_matches': candidate_group.get('total_matches'), + 'full_result_available': coalesced_result.get('full_result_available'), + 'validation_error': coalesced_result.get('validation_error'), + 'page_count': len(candidate_group.get('pages') or []), + 'diagnostics': candidate_group.get('diagnostics'), } best_score = score return best_candidate -def _build_tabular_generated_output_batch_prompt(user_question, batch_rows, batch_index, total_batches, source_candidate): +def _build_tabular_generated_output_query_descriptor( + source_candidate, + user_id, + conversation_id, + settings, +): + if not isinstance(source_candidate, dict): + return None + if source_candidate.get('function_name') != 'query_tabular_data': + return None + + source_parameters = source_candidate.get('source_parameters') or {} + query_expression = str(source_parameters.get('query_expression') or '').strip() + source_descriptor = source_candidate.get('source_descriptor') or {} + if ( + not query_expression + or source_descriptor.get('kind') != 'query_tabular_data' + or str(source_descriptor.get('filename') or '') != str(source_candidate.get('filename') or '') + or str(source_descriptor.get('query_expression') or '') != query_expression + ): + return None + validate_tabular_csv_query_expression(query_expression) + + batch_budget = _get_tabular_generated_output_batch_budget(settings) + descriptor = dict(source_descriptor) + descriptor['expected_row_count'] = _safe_int(source_candidate.get('total_matches')) + descriptor['batch_max_rows'] = batch_budget['max_rows'] + descriptor['batch_max_chars'] = batch_budget['max_chars'] + return descriptor + + +def _build_tabular_generated_output_source_authorization(source_candidate): + exact_source_authorization = (source_candidate or {}).get('source_authorization') or {} + if ( + exact_source_authorization.get('source') + and exact_source_authorization.get('container') + and exact_source_authorization.get('blob_path') + ): + return dict(exact_source_authorization) + + source_parameters = (source_candidate or {}).get('source_parameters') or {} + source = str(source_parameters.get('source') or 'chat').strip().lower() + if source not in {'chat', 'workspace', 'group', 'public'}: + source = 'chat' + + authorized_context = dict(getattr(g, 'authorized_chat_context', {}) or {}) + scope_id = None + if source == 'group': + scope_id = str( + source_parameters.get('group_id') + or authorized_context.get('active_group_id') + or '' + ).strip() or None + elif source == 'public': + scope_id = str( + source_parameters.get('public_workspace_id') + or authorized_context.get('active_public_workspace_id') + or '' + ).strip() or None + return { + 'source': source, + 'scope_id': scope_id, + } + + +def _build_failed_tabular_generated_output_metadata(source_candidate, output_format, reason): + normalized_output_format = str(output_format or 'json').strip().lower() or 'json' + row_count = _safe_int((source_candidate or {}).get('total_matches')) + failure_reason = str(reason or 'The exhaustive export could not be prepared.').strip() + return { + 'capability': 'tabular', + 'background_export': True, + 'status': 'failed', + 'status_label': 'Failed', + 'status_tone': 'danger', + 'status_detail': failure_reason, + 'retryable_failure': False, + 'can_resume': False, + 'can_cancel': False, + 'suppress_assistant_table_export': True, + 'file_name': _build_tabular_generated_output_file_name( + (source_candidate or {}).get('filename'), + normalized_output_format, + ), + 'output_format': normalized_output_format, + 'row_count': row_count, + 'processed_rows': 0, + 'source_file_name': (source_candidate or {}).get('filename'), + 'selected_sheet': (source_candidate or {}).get('selected_sheet'), + 'summary': failure_reason, + } + + +def _build_tabular_generated_output_batch_prompt( + user_question, + batch_rows, + batch_index, + total_batches, + source_candidate, + output_schema=None, +): source_file_name = str(source_candidate.get('filename') or 'unknown file').strip() or 'unknown file' selected_sheet = str(source_candidate.get('selected_sheet') or '').strip() batch_rows_json = _dump_tabular_generated_output_json(batch_rows) selected_sheet_line = f"Worksheet: {selected_sheet}\n" if selected_sheet else '' + model_output_schema = [ + field_name + for field_name in (output_schema or []) + if field_name not in {'source_row_number', 'source_row_identity'} + ] + output_schema_line = ( + f'Use exactly these output fields for every object, in this order: ' + f'{json.dumps(model_output_schema, ensure_ascii=False)}.\n' + if model_output_schema + else '' + ) return ( 'Transform the tabular input rows below into structured output for the user.\n\n' f'User instructions:\n{user_question}\n\n' 'Return ONLY a valid JSON array.\n' f'Return exactly {len(batch_rows)} JSON object(s), one per input row, in the same order.\n' + f'{output_schema_line}' + 'Copy __simplechat_source_row_token exactly from each input row into its matching output object. ' + 'Do not include the other fields beginning with __simplechat_source_ in generated objects.\n' 'Do not drop, merge, summarize, or cap rows.\n' 'Input rows may include normalized helper fields such as comment_id, body_text, source_file, attachment_present, attachment_names, and attachment_text. Use those normalized fields when they are present.\n' 'Input rows may include a referenced_documents array containing row-linked evidence from explicitly referenced non-tabular documents. Use that evidence as part of the source row context when it is relevant to the requested output.\n' @@ -4394,6 +5478,25 @@ def _build_tabular_generated_output_system_message(output_metadata): output_format = str(output_metadata.get('output_format') or 'json').upper() file_name = str(output_metadata.get('file_name') or 'generated output').strip() or 'generated output' row_count = _safe_int(output_metadata.get('row_count')) + output_status = str(output_metadata.get('status') or '').strip().lower() + + if output_status == 'failed': + failure_detail = str( + output_metadata.get('status_detail') + or output_metadata.get('summary') + or 'The exhaustive export failed validation.' + ).strip() + return ( + f'The requested exhaustive {output_format} export for {row_count} row(s) failed. ' + f'{failure_detail} Do not claim that a full or partial export is attached, and do not recreate ' + 'the assistant summary table as a CSV. Briefly report the failure and preserve any other requested analysis.' + ) + + if output_status == 'canceled': + return ( + f'The requested exhaustive {output_format} export for {row_count} row(s) was canceled. ' + 'Do not claim that the full export is attached, and do not recreate a partial assistant-table CSV.' + ) if output_metadata.get('background_export'): run_id = str(output_metadata.get('export_run_id') or output_metadata.get('run_id') or '').strip() @@ -4452,16 +5555,18 @@ async def _generate_tabular_structured_output_entries( user_id=None, conversation_id=None, model_context=None, + token_usage_callback=None, ): from semantic_kernel.contents.chat_history import ChatHistory as SKChatHistory - rows = [ + normalized_rows = [ _build_tabular_generated_output_input_row( row, source_file_name=source_candidate.get('filename'), ) for row in (source_candidate.get('rows') or []) ] + rows = _prepare_tabular_source_rows(normalized_rows) if not rows: return None @@ -4553,6 +5658,7 @@ async def _generate_tabular_structured_output_entries( ) merged_entries = [] + output_schema = None for batch_index, batch_rows in enumerate(row_batches): batch_number = batch_index + 1 log_event( @@ -4587,6 +5693,7 @@ async def _generate_tabular_structured_output_entries( batch_index, total_batches, source_candidate, + output_schema=output_schema, ) parsed_entries = None @@ -4621,10 +5728,27 @@ async def _generate_tabular_structured_output_entries( execution_settings = AzureChatPromptExecutionSettings(service_id='tabular-generated-output') result = await chat_service.get_chat_message_contents(chat_history, execution_settings) + if result: + _publish_tabular_response_token_usage( + result[0], + token_usage_callback, + ) raw_response_content = result[0].content if result and result[0].content else '' if result and result[0].content: parsed_entries = _parse_tabular_generated_json_entries(raw_response_content) parsed_entry_count = len(parsed_entries) if parsed_entries is not None else 0 + validation_error = None + if parsed_entries is not None and parsed_entry_count == len(batch_rows): + try: + parsed_entries, candidate_output_schema = _normalize_generated_batch_entries( + batch_rows, + parsed_entries, + expected_output_schema=output_schema, + ) + output_schema = candidate_output_schema + except ValueError as exc: + validation_error = str(exc) + parsed_entries = None if parsed_entries is None or parsed_entry_count != len(batch_rows): log_event( '[Tabular Generated Output] Structured export batch attempt mismatch', @@ -4636,6 +5760,7 @@ async def _generate_tabular_structured_output_entries( 'attempt_number': attempt_number, 'expected_row_count': len(batch_rows), 'parsed_row_count': parsed_entry_count, + 'validation_error': validation_error, 'response_char_count': len(raw_response_content), 'response_preview': _truncate_tabular_generated_output_response_preview(raw_response_content), }, @@ -4687,8 +5812,16 @@ async def maybe_create_tabular_generated_output( thought_callback=None, user_id=None, model_context=None, + cancel_requested=None, + request_correlation_id=None, + token_usage_callback=None, ): """Build, upload, and describe a generated tabular JSON/CSV export when requested.""" + raise_if_mixed_source_cancelled( + cancel_requested, + 'export', + request_correlation_id=request_correlation_id, + ) if not question_requests_tabular_generated_output(user_question): return None @@ -4718,6 +5851,208 @@ async def maybe_create_tabular_generated_output( debug_only=True, ) return None + + source_candidate['source_authorization'] = _build_tabular_generated_output_source_authorization( + source_candidate + ) + source_descriptor = None + source_descriptor_error = None + if ( + user_id + and conversation_id + and question_requests_tabular_structured_object_output(user_question) + ): + try: + source_descriptor = _build_tabular_generated_output_query_descriptor( + source_candidate, + user_id, + conversation_id, + settings, + ) + except Exception as exc: + source_descriptor_error = str(exc) + log_event( + '[Tabular Generated Output] Could not build durable source query descriptor', + { + 'conversation_id': conversation_id, + 'source_file_name': source_candidate.get('filename'), + 'function_name': source_candidate.get('function_name'), + 'error': str(exc), + }, + level=logging.WARNING, + ) + if source_descriptor: + source_candidate['source_authorization'] = { + field_name: source_descriptor.get(field_name) + for field_name in ('source', 'scope_id', 'container', 'blob_path') + } + + expected_row_count = _safe_int(source_candidate.get('total_matches')) + source_batch_rows = _safe_int((source_descriptor or {}).get('batch_max_rows')) or 1 + estimated_batch_count = max( + 1, + (expected_row_count + source_batch_rows - 1) // source_batch_rows, + ) + materialized_rows = None + materialized_batches = None + if ( + source_candidate.get('full_result_available') + and _safe_int(source_candidate.get('page_count')) > 1 + ): + materialized_rows = [ + _build_tabular_generated_output_input_row( + row, + source_file_name=source_candidate.get('filename'), + ) + for row in (source_candidate.get('rows') or []) + ] + materialized_batches = _build_tabular_generated_output_row_batches( + materialized_rows, + settings=settings, + ) + threshold_batch_count = len(materialized_batches) if materialized_batches else estimated_batch_count + exceeds_background_threshold = should_queue_tabular_generated_output_background( + expected_row_count, + threshold_batch_count, + settings, + ) + should_queue_materialized_pages = bool( + user_id + and conversation_id + and question_requests_tabular_structured_object_output(user_question) + and source_candidate.get('full_result_available') + and _safe_int(source_candidate.get('page_count')) > 1 + and not exceeds_background_threshold + ) + if should_queue_materialized_pages: + raise_if_mixed_source_cancelled( + cancel_requested, + 'export', + request_correlation_id=request_correlation_id, + ) + background_run = queue_tabular_generated_output_run( + user_id=user_id, + conversation_id=conversation_id, + user_question=user_question, + source_candidate=source_candidate, + output_format=output_format, + row_batches=materialized_batches, + gpt_model=gpt_model, + settings=settings, + model_context=model_context, + ) + background_metadata = build_background_tabular_generated_output_metadata(background_run) + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'export', + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + cancel_tabular_generated_output_run( + user_id, + background_metadata.get('export_run_id'), + ) + raise + await emit_tabular_post_processing_thought( + thought_callback, + f"Queued exhaustive {str(output_format or 'json').upper()} export from validated tabular pages", + detail=( + f"run_id={background_metadata.get('export_run_id')}; " + f"rows={expected_row_count}; batches={len(materialized_batches)}; checkpointed=true" + ), + activity=build_tabular_post_processing_activity_payload( + 'tabular.generated_output', + f"Exhaustive {str(output_format or 'json').upper()} export queued", + 'running', + phase='queued', + output_format=output_format, + file_name=source_candidate.get('filename'), + batch_index=0, + batch_count=len(materialized_batches), + ), + ) + return background_metadata + + should_queue_source_backed_run = bool( + source_descriptor + and expected_row_count > 0 + and ( + not source_candidate.get('full_result_available') + or exceeds_background_threshold + ) + ) + if should_queue_source_backed_run: + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'export', + request_correlation_id=request_correlation_id, + ) + background_run = queue_tabular_generated_output_run( + user_id=user_id, + conversation_id=conversation_id, + user_question=user_question, + source_candidate=source_candidate, + output_format=output_format, + row_batches=None, + gpt_model=gpt_model, + settings=settings, + model_context=model_context, + source_descriptor=source_descriptor, + ) + except MixedSourceCancellationError: + raise + except Exception as exc: + log_event( + '[Tabular Generated Output] Durable source-backed export queueing failed', + { + 'conversation_id': conversation_id, + 'source_file_name': source_candidate.get('filename'), + 'row_count': expected_row_count, + 'error': str(exc), + }, + level=logging.ERROR, + exceptionTraceback=True, + ) + return _build_failed_tabular_generated_output_metadata( + source_candidate, + output_format, + 'The exhaustive export could not be queued. No partial CSV was created.', + ) + background_metadata = build_background_tabular_generated_output_metadata(background_run) + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'export', + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + cancel_tabular_generated_output_run( + user_id, + background_metadata.get('export_run_id'), + ) + raise + await emit_tabular_post_processing_thought( + thought_callback, + f"Queued exhaustive {str(output_format or 'json').upper()} export from the authorized source query", + detail=( + f"run_id={background_metadata.get('export_run_id')}; " + f"rows={expected_row_count}; batches~={estimated_batch_count}; checkpointed=true" + ), + activity=build_tabular_post_processing_activity_payload( + 'tabular.generated_output', + f"Exhaustive {str(output_format or 'json').upper()} export queued", + 'running', + phase='queued', + output_format=output_format, + file_name=source_candidate.get('filename'), + batch_index=0, + batch_count=estimated_batch_count, + ), + ) + return background_metadata + if not source_candidate.get('full_result_available'): log_event( '[Tabular Generated Output] Selected source candidate is incomplete; skipping export', @@ -4725,10 +6060,22 @@ async def maybe_create_tabular_generated_output( 'conversation_id': conversation_id, 'output_format': output_format, 'selected_candidate': source_candidate.get('diagnostics'), + 'validation_error': source_candidate.get('validation_error'), + 'source_replay_available': bool(source_descriptor), }, debug_only=True, ) - return None + failure_detail = source_descriptor_error or source_candidate.get('validation_error') + return _build_failed_tabular_generated_output_metadata( + source_candidate, + output_format, + ( + f'The exhaustive export source could not be validated: {failure_detail}. ' + 'No partial CSV was created.' + if failure_detail + else 'The exhaustive export source was incomplete. No partial CSV was created.' + ), + ) log_event( '[Tabular Generated Output] Selected source candidate', { @@ -4743,6 +6090,11 @@ async def maybe_create_tabular_generated_output( return None if question_requests_tabular_structured_object_output(user_question): + raise_if_mixed_source_cancelled( + cancel_requested, + 'export', + request_correlation_id=request_correlation_id, + ) output_entries = await _generate_tabular_structured_output_entries( user_question, source_candidate, @@ -4754,8 +6106,17 @@ async def maybe_create_tabular_generated_output( conversation_id=conversation_id, model_context=model_context, ) + raise_if_mixed_source_cancelled( + cancel_requested, + 'export', + request_correlation_id=request_correlation_id, + ) if output_entries is None: - return None + return _build_failed_tabular_generated_output_metadata( + source_candidate, + output_format, + 'The exhaustive export failed output validation. No partial CSV was created.', + ) if isinstance(output_entries, dict) and output_entries.get('background_export'): return output_entries else: @@ -4806,6 +6167,18 @@ async def maybe_create_tabular_generated_output( 'in this chat as a downloadable export.' ), ) + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError: + delete_generated_chat_artifact_for_current_user( + conversation_id, + (upload_result.get('message') or {}).get('id'), + ) + raise preview_rows = output_entries[:TABULAR_GENERATED_OUTPUT_PREVIEW_ROWS] uploaded_file_name = upload_result.get('message', {}).get('file_name') or generated_file_name @@ -4835,6 +6208,7 @@ async def maybe_create_tabular_generated_output( ) return { 'capability': 'tabular', + 'suppress_assistant_table_export': True, 'artifact_message_id': upload_result.get('message', {}).get('id'), 'conversation_id': conversation_id, 'storage_scope': 'chat', @@ -5704,17 +7078,8 @@ def get_session(self, user_id, conversation_id, active_only=False): def get_new_plugin_invocations(invocations, baseline_count): - """Return only the plugin invocations created after the baseline count.""" - if not invocations: - return [] - - if baseline_count <= 0: - return list(invocations) - - if baseline_count >= len(invocations): - return [] - - return list(invocations[baseline_count:]) + """Compatibility shim for existing chat-route imports and call sites.""" + return _shared_get_new_plugin_invocations(invocations, baseline_count) def split_tabular_plugin_invocations(invocations): @@ -7542,6 +8907,42 @@ def derive_tabular_follow_up_calls_from_invocations(user_question, invocations): return follow_up_calls[:2] +def _extract_tabular_response_token_usage(response): + """Return observed token usage from a Semantic Kernel response metadata payload.""" + metadata = getattr(response, 'metadata', None) + metadata = metadata if isinstance(metadata, dict) else {} + usage = metadata.get('usage') or metadata.get('token_usage') + usage = usage if isinstance(usage, dict) else {} + + def as_nonnegative_int(value): + try: + return max(0, int(value or 0)) + except (TypeError, ValueError): + return 0 + + prompt_tokens = as_nonnegative_int(usage.get('prompt_tokens')) + completion_tokens = as_nonnegative_int(usage.get('completion_tokens')) + total_tokens = as_nonnegative_int(usage.get('total_tokens')) + if not total_tokens and (prompt_tokens or completion_tokens): + total_tokens = prompt_tokens + completion_tokens + if not any((prompt_tokens, completion_tokens, total_tokens)): + return None + return { + 'prompt_tokens': prompt_tokens, + 'completion_tokens': completion_tokens, + 'total_tokens': total_tokens, + 'request_count': 1, + } + + +def _publish_tabular_response_token_usage(response, token_usage_callback=None): + """Publish observed native model usage without changing existing return contracts.""" + token_usage = _extract_tabular_response_token_usage(response) + if token_usage and callable(token_usage_callback): + token_usage_callback(token_usage) + return token_usage + + async def maybe_recover_tabular_analysis_with_llm_reviewer(chat_service, kernel, tabular_plugin, plugin_logger, user_question, schema_context, @@ -7558,7 +8959,8 @@ async def maybe_recover_tabular_analysis_with_llm_reviewer(chat_service, kernel, discovery_feedback_messages=None, fallback_source_hint='workspace', fallback_group_id=None, - fallback_public_workspace_id=None): + fallback_public_workspace_id=None, + token_usage_callback=None): """Use an LLM reviewer to choose analytical tool calls when the main SK loop stalls.""" reviewer_allowed_function_names = [ function_name for function_name in (allowed_function_names or []) @@ -7644,6 +9046,10 @@ async def maybe_recover_tabular_analysis_with_llm_reviewer(chat_service, kernel, reviewer_text = '' if reviewer_result and reviewer_result[0].content: + _publish_tabular_response_token_usage( + reviewer_result[0], + token_usage_callback, + ) reviewer_text = reviewer_result[0].content.strip() reviewer_calls = parse_tabular_reviewer_plan(reviewer_text) @@ -8862,7 +10268,8 @@ async def run_tabular_sk_analysis(user_question, tabular_filenames, user_id, execution_mode='analysis', tabular_file_contexts=None, thought_callback=None, - model_context=None): + model_context=None, + token_usage_callback=None): """Run lightweight SK with tabular analysis and attachment follow-up support. Creates a temporary Kernel with TabularProcessingPlugin plus document-search @@ -8945,6 +10352,17 @@ async def run_tabular_sk_analysis(user_question, tabular_filenames, user_id, file_source_hint = file_context.get('source_hint', source_hint) file_group_id = file_context.get('group_id') file_public_workspace_id = file_context.get('public_workspace_id') + storage_locator = file_context.get('storage_locator') + if isinstance(storage_locator, dict): + locator_container = str(storage_locator.get('container') or '').strip() + locator_blob_path = str(storage_locator.get('blob_path') or '').strip() + if locator_container and locator_blob_path: + tabular_plugin.remember_resolved_blob_location( + file_source_hint, + fname, + locator_container, + locator_blob_path, + ) schema_source_context = {'source': file_source_hint} if file_group_id: schema_source_context['group_id'] = file_group_id @@ -9542,6 +10960,11 @@ def build_system_prompt(force_tool_use=False, tool_error_messages=None, result = await chat_service.get_chat_message_contents( chat_history, execution_settings, kernel=kernel ) + if result: + _publish_tabular_response_token_usage( + result[0], + token_usage_callback, + ) except Exception as exc: synthesis_exception = exc log_event( @@ -9886,6 +11309,7 @@ def build_system_prompt(force_tool_use=False, tool_error_messages=None, fallback_source_hint=source_hint, fallback_group_id=group_id, fallback_public_workspace_id=public_workspace_id, + token_usage_callback=token_usage_callback, ) if reviewer_recovery and reviewer_recovery.get('fallback'): return reviewer_recovery['fallback'] @@ -9897,12 +11321,8 @@ def build_system_prompt(force_tool_use=False, tool_error_messages=None, log_event(f"[Tabular SK Analysis] Error: {e}", level=logging.WARNING, exceptionTraceback=True) return None -def collect_tabular_sk_citations(user_id, conversation_id): - """Collect plugin invocations from the tabular SK analysis and convert to citation format.""" - from semantic_kernel_plugins.plugin_invocation_logger import get_plugin_logger - - plugin_logger = get_plugin_logger() - plugin_invocations = plugin_logger.get_invocations_for_conversation(user_id, conversation_id) +def _build_tabular_sk_citations_from_invocations(plugin_invocations): + """Convert selected tabular invocations to the existing tool-citation shape.""" plugin_invocations = filter_tabular_citation_invocations(plugin_invocations) if not plugin_invocations: @@ -9947,10 +11367,281 @@ def collect_tabular_sk_citations(user_id, conversation_id): } citations.append(citation) - log_event(f"[Tabular SK Citations] Collected {len(citations)} tool execution citations", level=logging.INFO) return citations +def collect_tabular_sk_citations(user_id, conversation_id): + """Collect plugin invocations from the tabular SK analysis and convert to citation format.""" + plugin_logger = get_plugin_logger() + plugin_invocations = plugin_logger.get_invocations_for_conversation(user_id, conversation_id) + citations = _build_tabular_sk_citations_from_invocations(plugin_invocations) + log_event( + f"[Tabular SK Citations] Collected {len(citations)} tool execution citations", + level=logging.INFO, + ) + return citations + + +def _execute_mixed_source_tabular_evidence( + *, + tabular_sources, + selection_mode, + has_narrative_sources, + user_question, + user_id, + conversation_id, + gpt_model, + settings, + thought_tracker=None, + live_thought_callback=None, + model_context=None, + cancel_requested=None, + request_correlation_id=None, +): + """Run the existing tabular engine once per manifest source with terminal coverage.""" + source_contexts = build_tabular_file_contexts_from_manifest(tabular_sources) + if has_request_context(): + authorized_context = dict(getattr(g, 'authorized_chat_context', {}) or {}) + authorized_blob_locations = [] + for file_context in source_contexts: + storage_locator = file_context.get('storage_locator') + if not isinstance(storage_locator, dict): + continue + locator_container = str(storage_locator.get('container') or '').strip() + locator_blob_path = str(storage_locator.get('blob_path') or '').strip() + if locator_container and locator_blob_path: + authorized_blob_locations.append([locator_container, locator_blob_path]) + authorized_context['authorized_blob_locations'] = authorized_blob_locations + g.authorized_chat_context = authorized_context + context_by_document_id = { + context['document_id']: context + for context in source_contexts + } + system_messages = [] + agent_citations = [] + generated_outputs = [] + all_invocations = [] + token_usage = { + 'prompt_tokens': 0, + 'completion_tokens': 0, + 'total_tokens': 0, + 'request_count': 0, + } + + def record_token_usage(usage): + if not isinstance(usage, dict): + return + for key in ('prompt_tokens', 'completion_tokens', 'total_tokens', 'request_count'): + try: + token_usage[key] += max(0, int(usage.get(key) or 0)) + except (TypeError, ValueError): + continue + if not is_tabular_processing_enabled(settings): + def fail_unavailable_tabular_source(source): + del source + raise RuntimeError('Tabular processing is unavailable') + + return { + 'evidence_envelopes': execute_tabular_evidence_sources( + tabular_sources, + fail_unavailable_tabular_source, + selection_mode, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ), + 'system_messages': [], + 'agent_citations': [], + 'generated_outputs': [], + 'invocations': [], + 'executed': False, + 'token_usage': None, + } + + execute_tabular = should_run_tabular_evidence( + user_question, + has_narrative_sources=has_narrative_sources, + ) + plugin_logger = get_plugin_logger() + + def publish_post_processing_thought(thought_payload): + payload = thought_payload if isinstance(thought_payload, dict) else {} + if thought_tracker is not None: + thought_tracker.add_thought( + payload.get('step_type', 'tabular_analysis'), + payload.get('content', ''), + detail=payload.get('detail'), + activity=payload.get('activity'), + ) + if callable(live_thought_callback): + live_payload = dict(payload) + if thought_tracker is not None: + live_payload['message_id'] = getattr(thought_tracker, 'message_id', None) + live_payload['step_index'] = thought_tracker.current_index - 1 + live_thought_callback(live_payload) + + def execute_source(source): + document_id = str(source.get('document_id') or '').strip() + file_context = context_by_document_id.get(document_id) + if not file_context: + raise ValueError('Authorized tabular source context is unavailable') + + baseline_invocation_count = len( + plugin_logger.get_invocations_for_conversation( + user_id, + conversation_id, + limit=1000, + ) + ) + execution_mode = get_tabular_execution_mode(user_question) + tabular_analysis, streamed_tool_thoughts = asyncio.run( + run_tabular_analysis_with_thought_tracking( + user_question=user_question, + tabular_filenames={file_context['file_name']}, + tabular_file_contexts=[file_context], + user_id=user_id, + conversation_id=conversation_id, + gpt_model=gpt_model, + settings=settings, + source_hint=file_context.get('source_hint', 'workspace'), + group_id=file_context.get('group_id'), + public_workspace_id=file_context.get('public_workspace_id'), + execution_mode=execution_mode, + thought_tracker=thought_tracker, + live_thought_callback=live_thought_callback, + model_context=model_context, + token_usage_callback=record_token_usage, + ) + ) + if not str(tabular_analysis or '').strip(): + raise ValueError('Tabular execution returned no computed results') + + invocations_after = plugin_logger.get_invocations_for_conversation( + user_id, + conversation_id, + limit=1000, + ) + source_invocations = get_new_plugin_invocations( + invocations_after, + baseline_invocation_count, + ) + if execution_mode == 'schema_summary': + successful_source_invocations = [ + invocation + for invocation in source_invocations + if getattr(invocation, 'function_name', '') == 'describe_tabular_file' + and not get_tabular_invocation_error_message(invocation) + ] + else: + successful_source_invocations, _ = split_tabular_analysis_invocations( + source_invocations + ) + if not successful_source_invocations: + raise ValueError('Tabular execution returned no successful native tool calls') + all_invocations.extend(source_invocations) + related_document_summary = '' + related_document_stats = augment_tabular_invocations_with_related_document_evidence( + source_invocations, + user_question, + user_id, + conversation_id=conversation_id, + ) + if related_document_stats.get('augmented_row_count'): + related_document_summary = build_tabular_related_document_evidence_summary( + source_invocations, + ) + + if thought_tracker is not None and not streamed_tool_thoughts: + for thought_content, thought_detail in get_tabular_tool_thought_payloads( + source_invocations + ): + thought_tracker.add_thought( + 'tabular_analysis', + thought_content, + thought_detail, + ) + if thought_tracker is not None: + for thought_content, thought_detail in get_tabular_status_thought_payloads( + source_invocations, + analysis_succeeded=True, + ): + thought_tracker.add_thought( + 'tabular_analysis', + thought_content, + thought_detail, + ) + + generated_output = asyncio.run(maybe_create_tabular_generated_output( + user_question=user_question, + invocations=source_invocations, + gpt_model=gpt_model, + settings=settings, + conversation_id=conversation_id, + thought_callback=publish_post_processing_thought, + user_id=user_id, + model_context=model_context, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + token_usage_callback=record_token_usage, + )) + if generated_output: + generated_outputs.append(generated_output) + + source_citations = _build_tabular_sk_citations_from_invocations( + source_invocations + ) + agent_citations.extend(source_citations) + chart_citations = build_tabular_inline_chart_citations( + user_question, + source_invocations, + ) + agent_citations.extend(chart_citations) + system_messages.append({ + 'role': 'system', + 'content': build_tabular_computed_results_system_message( + 'an authorized selected tabular source', + str(tabular_analysis).strip(), + related_document_evidence_summary=related_document_summary, + ), + }) + if generated_output: + system_messages.append({ + 'role': 'system', + 'content': _build_tabular_generated_output_system_message(generated_output), + }) + + return { + 'summary': str(tabular_analysis).strip(), + 'evidence': [ + get_tabular_invocation_compact_payload(invocation, max_rows=5) + for invocation in source_invocations + ], + 'citations': source_citations, + 'generated_artifacts': [generated_output] if generated_output else [], + 'coverage': { + 'tool_call_count': len(source_invocations), + 'execution_mode': execution_mode, + }, + } + + evidence_envelopes = execute_tabular_evidence_sources( + tabular_sources, + execute_source, + selection_mode, + execute=execute_tabular, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + return { + 'evidence_envelopes': evidence_envelopes, + 'system_messages': system_messages, + 'agent_citations': agent_citations, + 'generated_outputs': generated_outputs, + 'invocations': all_invocations, + 'executed': execute_tabular, + 'token_usage': token_usage if token_usage['request_count'] else None, + } + + def is_tabular_filename(filename): """Return True when the filename has a supported tabular extension.""" if not filename or not isinstance(filename, str): @@ -10437,7 +12128,8 @@ async def run_tabular_analysis_with_multi_file_support(user_question, tabular_fi execution_mode='analysis', tabular_file_contexts=None, thought_callback=None, - model_context=None): + model_context=None, + token_usage_callback=None): """Run deterministic multi-file helpers first, then fall back to the SK planner.""" analysis_file_contexts = normalize_tabular_file_contexts_for_analysis( tabular_filenames=tabular_filenames, @@ -10484,6 +12176,7 @@ async def run_tabular_analysis_with_multi_file_support(user_question, tabular_fi execution_mode=execution_mode, thought_callback=thought_callback, model_context=model_context, + token_usage_callback=token_usage_callback, ) @@ -10495,7 +12188,8 @@ async def run_tabular_analysis_with_thought_tracking(user_question, tabular_file tabular_file_contexts=None, thought_tracker=None, live_thought_callback=None, - model_context=None): + model_context=None, + token_usage_callback=None): """Run tabular analysis while streaming/persisting live tool thoughts when available.""" plugin_logger = get_plugin_logger() callback_key = None @@ -10554,6 +12248,7 @@ def record_and_publish_tabular_progress_thought(thought_payload): execution_mode=execution_mode, thought_callback=tabular_progress_callback, model_context=model_context, + token_usage_callback=token_usage_callback, ) if callable(tabular_progress_callback): @@ -11660,8 +13355,17 @@ def _load_or_create_analyze_conversation(user_id, conversation_id=None): invalidate_conversation_cache_for_item(conversation_item, reason="conversation_created") return conversation_item - def execute_document_action_chat_request(data=None, publish_background_event=None, forced_action_type=None): + def execute_document_action_chat_request( + data=None, + publish_background_event=None, + forced_action_type=None, + cancel_requested=None, + request_correlation_id=None, + ): settings = get_settings() + request_correlation_id = normalize_mixed_source_correlation_id( + request_correlation_id + ) data = data if isinstance(data, dict) else (request.get_json() or {}) user_id = get_current_user_id() if not user_id: @@ -12033,7 +13737,22 @@ def execute_document_action_chat_request(data=None, publish_background_event=Non run_id=assistant_message_id, thought_tracker=thought_tracker, external_activity_callback=stream_activity_callback, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, ) + except MixedSourceCancellationError as exc: + if thought_tracker.enabled: + thought_tracker.add_thought( + 'cancellation', + 'Document action canceled before final output publication', + detail=f'phase={exc.phase}', + ) + return { + 'canceled': True, + 'conversation_id': conversation_id, + 'user_message_id': user_message_id, + 'request_correlation_id': request_correlation_id, + }, 409 except Exception as exc: debug_print( '[ChatDocumentAction] Execution failed | ' @@ -12053,31 +13772,125 @@ def execute_document_action_chat_request(data=None, publish_background_event=Non level=logging.ERROR, exceptionTraceback=True, ) - return {'error': str(exc), 'conversation_id': conversation_id, 'user_message_id': user_message_id}, 500 + return {'error': str(exc), 'conversation_id': conversation_id, 'user_message_id': user_message_id}, 500 + + try: + _reauthorize_document_action_finalization( + normalized_action, + execution_result, + user_id, + conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + settings=settings, + ) + except MixedSourceCancellationError as exc: + if thought_tracker.enabled: + thought_tracker.add_thought( + 'cancellation', + 'Document action canceled before final output publication', + detail=f'phase={exc.phase}', + ) + return { + 'canceled': True, + 'conversation_id': conversation_id, + 'user_message_id': user_message_id, + 'request_correlation_id': request_correlation_id, + }, 409 + except PermissionError as exc: + return { + 'error': str(exc), + 'conversation_id': conversation_id, + 'user_message_id': user_message_id, + }, 403 + except RuntimeError as exc: + return { + 'error': str(exc), + 'conversation_id': conversation_id, + 'user_message_id': user_message_id, + }, 409 assistant_timestamp = datetime.utcnow().isoformat() hybrid_citations_list = _build_document_action_hybrid_citations(execution_result) if assigned_knowledge_context_citations: hybrid_citations_list.extend(assigned_knowledge_context_citations) hybrid_citations_list.sort(key=_build_hybrid_citation_sort_key, reverse=True) - prepared_agent_citations = persist_agent_citation_artifacts( - conversation_id=conversation_id, - assistant_message_id=assistant_message_id, - agent_citations=execution_result.get('agent_citations') or [], - created_timestamp=assistant_timestamp, - user_info=response_message_context.get('user_info'), - ) + prepared_agent_citations = [] document_generated_analysis_artifacts = list(execution_result.get('generated_analysis_artifacts') or []) document_generated_tabular_outputs = list(execution_result.get('generated_tabular_outputs') or []) - assistant_table_generated_output = maybe_create_assistant_table_generated_output( - user_question=user_message, - assistant_content=execution_result.get('reply', ''), - conversation_id=conversation_id, - existing_outputs=document_generated_analysis_artifacts + document_generated_tabular_outputs, - ) - if assistant_table_generated_output: - document_generated_analysis_artifacts.append(assistant_table_generated_output) - document_generated_tabular_outputs.append(assistant_table_generated_output) + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) + prepared_agent_citations = persist_agent_citation_artifacts( + conversation_id=conversation_id, + assistant_message_id=assistant_message_id, + agent_citations=execution_result.get('agent_citations') or [], + created_timestamp=assistant_timestamp, + user_info=response_message_context.get('user_info'), + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + raise_if_mixed_source_cancelled( + cancel_requested, + 'artifact_publication', + request_correlation_id=request_correlation_id, + ) + generated_file_output = maybe_create_generated_file_output( + user_question=user_message, + assistant_content=get_generated_file_export_content(execution_result), + conversation_id=conversation_id, + function_results=execution_result.get('agent_citations') or [], + existing_outputs=document_generated_analysis_artifacts + document_generated_tabular_outputs, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + ) + if generated_file_output: + document_generated_analysis_artifacts.append(generated_file_output) + if generated_file_output.get('output_format') == 'csv': + document_generated_tabular_outputs.append(generated_file_output) + _reauthorize_document_action_finalization( + normalized_action, + execution_result, + user_id, + conversation_id, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, + settings=settings, + ) + except MixedSourceCancellationError as exc: + _rollback_mixed_source_chat_publication( + user_id, + conversation_id, + document_generated_analysis_artifacts + document_generated_tabular_outputs, + compact_citations=prepared_agent_citations, + ) + if thought_tracker.enabled: + thought_tracker.add_thought( + 'cancellation', + 'Document action canceled before final output publication', + detail=f'phase={exc.phase}', + ) + return { + 'canceled': True, + 'conversation_id': conversation_id, + 'user_message_id': user_message_id, + 'request_correlation_id': request_correlation_id, + }, 409 + except MixedSourceFinalizationError: + _rollback_mixed_source_chat_publication( + user_id, + conversation_id, + document_generated_analysis_artifacts + document_generated_tabular_outputs, + compact_citations=prepared_agent_citations, + ) + return { + 'error': 'Selected source state changed before final output could be published.', + 'conversation_id': conversation_id, + 'user_message_id': user_message_id, + }, 409 generated_analysis_metadata = _build_generated_analysis_metadata( generated_analysis_artifacts=document_generated_analysis_artifacts, generated_tabular_outputs=document_generated_tabular_outputs, @@ -12130,6 +13943,42 @@ def execute_document_action_chat_request(data=None, publish_background_event=Non }, }) cosmos_messages_container.upsert_item(assistant_doc) + try: + raise_if_mixed_source_cancelled( + cancel_requested, + 'finalization', + request_correlation_id=request_correlation_id, + ) + except MixedSourceCancellationError as exc: + try: + cosmos_messages_container.delete_item( + item=assistant_message_id, + partition_key=conversation_id, + ) + except Exception: + log_event( + '[MixedSourceLifecycle] Assistant rollback after cancellation failed.', + extra={'assistant_message_rollback_failure_count': 1}, + level=logging.WARNING, + ) + _rollback_mixed_source_chat_publication( + user_id, + conversation_id, + document_generated_analysis_artifacts + document_generated_tabular_outputs, + compact_citations=prepared_agent_citations, + ) + if thought_tracker.enabled: + thought_tracker.add_thought( + 'cancellation', + 'Document action canceled before final output publication', + detail=f'phase={exc.phase}', + ) + return { + 'canceled': True, + 'conversation_id': conversation_id, + 'user_message_id': user_message_id, + 'request_correlation_id': request_correlation_id, + }, 409 token_usage = execution_result.get('token_usage') if isinstance(execution_result.get('token_usage'), dict) else None if token_usage and token_usage.get('total_tokens'): @@ -12236,11 +14085,18 @@ def execute_document_action_chat_request(data=None, publish_background_event=Non 'metadata': assistant_doc.get('metadata', {}), }), 200 - def execute_analyze_chat_request(data=None, publish_background_event=None): + def execute_analyze_chat_request( + data=None, + publish_background_event=None, + cancel_requested=None, + request_correlation_id=None, + ): return execute_document_action_chat_request( data=data, publish_background_event=publish_background_event, forced_action_type=DOCUMENT_ACTION_TYPE_ANALYZE, + cancel_requested=cancel_requested, + request_correlation_id=request_correlation_id, ) @bp.route('/api/chat/document-action', methods=['POST']) @@ -12269,6 +14125,7 @@ def chat_document_action_stream_api(): data['conversation_id'] = conversation_id g.conversation_id = conversation_id stream_session = CHAT_STREAM_REGISTRY.start_session(user_id, conversation_id) + request_correlation_id = normalize_mixed_source_correlation_id() def generate_document_action_response(publish_background_event=None): try: @@ -12282,7 +14139,17 @@ def generate_document_action_response(publish_background_event=None): payload, status_code = execute_document_action_chat_request( data=data, publish_background_event=publish_background_event, + cancel_requested=stream_session.is_cancel_requested, + request_correlation_id=request_correlation_id, ) + if payload.get('canceled'): + yield _build_stream_cancel_event( + payload.get('conversation_id') or conversation_id, + user_message_id=payload.get('user_message_id'), + reason=stream_session.get_cancel_reason(), + message_persisted=False, + ) + return if stream_session and stream_session.is_cancel_requested(): yield _build_stream_cancel_event( payload.get('conversation_id') or conversation_id, @@ -12330,6 +14197,7 @@ def chat_analyze_stream_api(): data['conversation_id'] = conversation_id g.conversation_id = conversation_id stream_session = CHAT_STREAM_REGISTRY.start_session(user_id, conversation_id) + request_correlation_id = normalize_mixed_source_correlation_id() def generate_analyze_response(publish_background_event=None): try: @@ -12343,7 +14211,17 @@ def generate_analyze_response(publish_background_event=None): payload, status_code = execute_analyze_chat_request( data=data, publish_background_event=publish_background_event, + cancel_requested=stream_session.is_cancel_requested, + request_correlation_id=request_correlation_id, ) + if payload.get('canceled'): + yield _build_stream_cancel_event( + payload.get('conversation_id') or conversation_id, + user_message_id=payload.get('user_message_id'), + reason=stream_session.get_cancel_reason(), + message_persisted=False, + ) + return if stream_session and stream_session.is_cancel_requested(): yield _build_stream_cancel_event( payload.get('conversation_id') or conversation_id, @@ -12615,7 +14493,14 @@ def result_requires_message_reload(result: Any) -> bool: generated_tabular_outputs_list = [] generated_analysis_artifacts_list = [] system_messages_for_augmentation = [] # Collect system messages from search + generated_file_output_guidance = build_generated_file_output_guidance(user_message) + if generated_file_output_guidance: + system_messages_for_augmentation.append({ + 'role': 'system', + 'content': generated_file_output_guidance, + }) search_results = [] + mixed_source_narrative_retrieval_failed = False selected_agent = None # Initialize selected_agent early to prevent NameError # --- Configuration --- # History / Summarization Settings @@ -12643,6 +14528,38 @@ def result_requires_message_reload(result: Any) -> bool: deep_research_enabled = deep_research_enabled.lower() == 'true' if isinstance(image_gen_enabled, str): image_gen_enabled = image_gen_enabled.lower() == 'true' + try: + document_context_contract = _normalize_chat_document_context_contract( + settings, + data, + selected_document_ids, + hybrid_search_enabled, + ) + except ValueError as contract_error: + return jsonify({'error': str(contract_error)}), 400 + selected_document_ids = list( + document_context_contract.get('selected_document_ids') or [] + ) + requested_selected_document_ids = list(selected_document_ids) + selected_document_id = ( + selected_document_ids[0] + if len(selected_document_ids) == 1 + else None + ) + selection_mode = document_context_contract.get('selection_mode') + document_context_requested = bool( + document_context_contract.get('document_context_requested') + ) + request_has_explicit_document_selection = bool( + document_context_contract.get('explicit_selection') + ) + request_document_context_enabled = bool( + hybrid_search_enabled + or ( + is_mixed_source_chat_search_enabled(settings) + and document_context_requested + ) + ) user_workspace_context_requested = data.get('user_workspace_context_enabled') if isinstance(user_workspace_context_requested, str): user_workspace_context_requested = user_workspace_context_requested.lower() == 'true' @@ -12692,6 +14609,7 @@ def result_requires_message_reload(result: Any) -> bool: history_grounded_search_used = False history_only_answerability = None prior_grounded_document_refs = [] + continuity_decision = None effective_document_scope = document_scope effective_selected_document_ids = list(selected_document_ids or []) effective_selected_document_id = selected_document_id @@ -12710,7 +14628,7 @@ def result_requires_message_reload(result: Any) -> bool: assigned_knowledge_deep_research_urls = [] if assigned_knowledge_filters: assigned_knowledge_user_context_active = ( - user_workspace_context_requested + (user_workspace_context_requested or document_context_requested) and _assigned_knowledge_allows_user_workspace_context(assigned_knowledge_filters) and _assigned_knowledge_allows_document_action( assigned_knowledge_filters, @@ -12770,6 +14688,15 @@ def result_requires_message_reload(result: Any) -> bool: g.assigned_knowledge_context = assigned_knowledge_filters g.assigned_knowledge_user_context_active = assigned_knowledge_user_context_active + mixed_source_explicit_selection = bool( + is_mixed_source_chat_search_enabled(settings) + and request_has_explicit_document_selection + and ( + not assigned_knowledge_filters + or assigned_knowledge_user_context_active + ) + ) + explicit_external_retrieval_requested = _is_explicit_external_retrieval_requested( web_search_enabled=web_search_enabled, url_access_enabled=url_access_enabled, @@ -12940,11 +14867,14 @@ def result_requires_message_reload(result: Any) -> bool: _set_authorized_chat_request_context(user_id, conversation_id, scope_context) auto_linked_chat_upload_document_ids = [] - auto_merge_chat_upload_workspace_context = _should_auto_merge_chat_upload_workspace_context( - explicit_external_retrieval_requested, - hybrid_search_enabled, - assigned_knowledge_filters=assigned_knowledge_filters, - assigned_knowledge_user_context_active=assigned_knowledge_user_context_active, + auto_merge_chat_upload_workspace_context = ( + not mixed_source_explicit_selection + and _should_auto_merge_chat_upload_workspace_context( + explicit_external_retrieval_requested, + hybrid_search_enabled, + assigned_knowledge_filters=assigned_knowledge_filters, + assigned_knowledge_user_context_active=assigned_knowledge_user_context_active, + ) ) if auto_merge_chat_upload_workspace_context: chat_upload_context = _resolve_chat_upload_workspace_context( @@ -13003,6 +14933,81 @@ def result_requires_message_reload(result: Any) -> bool: selected_document_id = effective_selected_document_id document_scope = effective_document_scope + mixed_source_manifest = [] + mixed_source_partitions = {} + mixed_source_narrative_document_ids = [] + mixed_source_tabular_sources = [] + mixed_source_evidence_envelopes = [] + mixed_source_native_token_usage = None + mixed_source_request_correlation_id = normalize_mixed_source_correlation_id() + if mixed_source_explicit_selection: + try: + mixed_source_manifest = _resolve_chat_mixed_source_manifest( + settings, + user_id, + conversation_id, + effective_selected_document_ids, + 'selected', + active_group_ids=effective_active_group_ids, + active_public_workspace_ids=effective_active_public_workspace_ids, + request_correlation_id=mixed_source_request_correlation_id, + ) + except ValueError as manifest_error: + return jsonify({'error': str(manifest_error)}), 400 + mixed_source_partitions = partition_source_manifest( + mixed_source_manifest + ) + mixed_source_narrative_document_ids = _get_manifest_partition_document_ids( + mixed_source_partitions, + 'narrative_sources', + ) + mixed_source_tabular_sources = list( + mixed_source_partitions.get('tabular_sources') or [] + ) + authorized_selected_document_ids = [ + str(source.get('document_id') or '').strip() + for source in ( + list(mixed_source_partitions.get('narrative_sources') or []) + + mixed_source_tabular_sources + ) + if str(source.get('document_id') or '').strip() + ] + effective_selected_document_ids = authorized_selected_document_ids + effective_selected_document_id = ( + authorized_selected_document_ids[0] + if len(authorized_selected_document_ids) == 1 + else None + ) + log_event( + '[MixedSourceChatSearch] Activated explicit selected-source context.', + extra={ + 'selection_mode': 'selected', + 'requested_source_count': len(mixed_source_manifest), + 'authorized_source_count': len(authorized_selected_document_ids), + 'narrative_source_count': len(mixed_source_narrative_document_ids), + 'tabular_source_count': len(mixed_source_tabular_sources), + 'omitted_source_count': len( + mixed_source_partitions.get('unresolved_sources') or [] + ), + }, + level=logging.INFO, + ) + else: + _maybe_resolve_chat_source_manifest( + settings, + user_id, + conversation_id, + effective_selected_document_ids, + scope_context, + ) + request_document_context_enabled = bool( + hybrid_search_enabled + or ( + is_mixed_source_chat_search_enabled(settings) + and document_context_requested + ) + ) + # Clear plugin invocations at start of message processing to ensure # each message only shows citations for tools executed during that specific interaction plugin_logger = get_plugin_logger() @@ -13100,13 +15105,13 @@ def result_requires_message_reload(result: Any) -> bool: # Button states and selections user_metadata['button_states'] = { 'image_generation': image_gen_enabled, - 'document_search': hybrid_search_enabled, + 'document_search': request_document_context_enabled, 'web_search': bool(web_search_enabled), 'url_access': bool(url_access_enabled), 'deep_research': bool(deep_research_enabled) } user_metadata['capability_usage'] = _build_capability_usage_metadata( - workspace_search_enabled=hybrid_search_enabled, + workspace_search_enabled=request_document_context_enabled, document_action_type=DOCUMENT_ACTION_TYPE_NONE, document_scope=effective_document_scope, selected_document_ids=effective_selected_document_ids, @@ -13119,12 +15124,18 @@ def result_requires_message_reload(result: Any) -> bool: ) # Document search scope and selections - if hybrid_search_enabled: + if request_document_context_enabled: user_metadata['workspace_search'] = { 'search_enabled': True, + 'selection_mode': selection_mode, + 'document_context_requested': document_context_requested, + 'hybrid_search_preference': bool(hybrid_search_enabled), 'document_scope': effective_document_scope, 'selected_document_id': effective_selected_document_id, 'selected_document_ids': effective_selected_document_ids, + 'requested_document_ids': requested_selected_document_ids, + 'active_group_ids': effective_active_group_ids, + 'active_public_workspace_ids': effective_active_public_workspace_ids, 'tags': tags_filter, 'classification': classifications_to_send } @@ -13313,7 +15324,7 @@ def result_requires_message_reload(result: Any) -> bool: conversation_id=conversation_id, message_type='user_message', message_length=len(user_message) if user_message else 0, - has_document_search=hybrid_search_enabled, + has_document_search=request_document_context_enabled, has_image_generation=image_gen_enabled, document_scope=document_scope, chat_context=actual_chat_type, @@ -13458,9 +15469,20 @@ def result_requires_message_reload(result: Any) -> bool: except Exception as ex: debug_print(f"[Content Safety] Unexpected error: {ex}") - if not original_hybrid_search_enabled and not explicit_external_retrieval_requested: + if ( + not original_hybrid_search_enabled + and not explicit_external_retrieval_requested + and not mixed_source_explicit_selection + ): prior_grounded_document_refs = _normalize_prior_grounded_document_refs(conversation_item) if prior_grounded_document_refs: + continuity_decision = _resolve_reauthorized_continuity_decision( + settings, + user_id, + conversation_id, + prior_grounded_document_refs, + request_correlation_id=mixed_source_request_correlation_id, + ) thought_tracker.add_thought( 'history_context', 'Checking whether prior conversation context already answers the question', @@ -13499,7 +15521,10 @@ def result_requires_message_reload(result: Any) -> bool: f"[History Fallback] History-only sufficiency assessment failed: {assessment_error}" ) - if history_only_answerability and history_only_answerability.get('can_answer_from_history'): + if _can_reuse_prior_grounded_history( + history_only_answerability, + continuity_decision, + ): thought_tracker.add_thought( 'history_context', 'Prior conversation context appears sufficient without new document retrieval', @@ -13572,13 +15597,71 @@ def result_requires_message_reload(result: Any) -> bool: 'history_context', 'No prior grounded documents were available; using conversation history only' ) + if ( + is_mixed_source_chat_search_enabled(settings) + and history_grounded_search_used + ): + history_context = _resolve_chat_mixed_source_partition( + settings, + user_id, + conversation_id, + effective_selected_document_ids, + 'history', + active_group_ids=effective_active_group_ids, + active_public_workspace_ids=effective_active_public_workspace_ids, + request_correlation_id=mixed_source_request_correlation_id, + ) + mixed_source_manifest = history_context.get('manifest') or [] + mixed_source_partitions = history_context.get('partitions') or {} + mixed_source_narrative_document_ids = list( + history_context.get('narrative_document_ids') or [] + ) + mixed_source_tabular_sources = list( + history_context.get('tabular_sources') or [] + ) + if is_mixed_source_conversation_continuity_enabled(settings): + continuity_decision = _build_reauthorized_continuity_decision( + prior_grounded_document_refs, + mixed_source_manifest, + explicit_selection=False, + ) + effective_selected_document_ids = ( + mixed_source_narrative_document_ids + + _get_manifest_partition_document_ids( + mixed_source_partitions, + 'tabular_sources', + ) + ) + effective_selected_document_id = ( + effective_selected_document_ids[0] + if len(effective_selected_document_ids) == 1 + else None + ) # region 4 - Augmentation # --------------------------------------------------------------------- # 4) Augmentation (Search, etc.) - Run *before* final history prep # --------------------------------------------------------------------- # Hybrid Search - if hybrid_search_enabled or history_grounded_search_used: + mixed_source_document_context_active = bool( + (hybrid_search_enabled or history_grounded_search_used) + if not is_mixed_source_chat_search_enabled(settings) + else ( + document_context_requested + or hybrid_search_enabled + or history_grounded_search_used + ) + ) + combined_documents = [] + mixed_source_narrative_search_active = bool( + mixed_source_document_context_active + and ( + not is_mixed_source_chat_search_enabled(settings) + or not mixed_source_manifest + or mixed_source_narrative_document_ids + ) + ) + if mixed_source_narrative_search_active: # Optional: Summarize recent history *for search* (uses its own limit) if hybrid_search_enabled and enable_summarize_content_history_for_search: @@ -13693,8 +15776,14 @@ def result_requires_message_reload(result: Any) -> bool: ): search_args["active_public_workspace_id"] = effective_active_public_workspace_id - if effective_selected_document_ids: - search_args["document_ids"] = effective_selected_document_ids + search_document_ids = ( + mixed_source_narrative_document_ids + if is_mixed_source_chat_search_enabled(settings) + and mixed_source_manifest + else effective_selected_document_ids + ) + if search_document_ids: + search_args["document_ids"] = search_document_ids elif effective_selected_document_id: search_args["document_id"] = effective_selected_document_id if auto_linked_chat_upload_document_ids: @@ -13729,21 +15818,70 @@ def result_requires_message_reload(result: Any) -> bool: else: # Public scope now automatically searches all visible public workspaces search_results = hybrid_search(**search_args) # Assuming hybrid_search handles None document_id + + if ( + is_mixed_source_chat_search_enabled(settings) + and not mixed_source_explicit_selection + and not history_grounded_search_used + ): + relevance_context = _resolve_chat_mixed_source_relevance_context( + settings=settings, + user_id=user_id, + conversation_id=conversation_id, + query=search_query, + search_results=search_results, + document_scope=effective_document_scope, + candidate_document_ids=( + assigned_knowledge_filters.get('document_ids') + if assigned_knowledge_filters + and assigned_knowledge_filters.get('has_workspace_knowledge') + else None + ), + tags_filter=tags_filter, + active_group_ids=effective_active_group_ids, + active_public_workspace_ids=effective_active_public_workspace_ids, + request_correlation_id=mixed_source_request_correlation_id, + ) + mixed_source_manifest = relevance_context.get('manifest') or [] + mixed_source_partitions = relevance_context.get('partitions') or {} + mixed_source_narrative_document_ids = list( + relevance_context.get('narrative_document_ids') or [] + ) + mixed_source_tabular_sources = list( + relevance_context.get('tabular_sources') or [] + ) + search_results = list( + relevance_context.get('search_results') or [] + ) except SemanticSearchQuotaExceededError as e: debug_print(f"Semantic search quota exceeded during hybrid search: {e}") - return jsonify({ - 'error': e.user_message, - 'warning_type': SEMANTIC_SEARCH_QUOTA_WARNING_TYPE, - 'service_health_warning': True, - }), 503 + if ( + is_mixed_source_chat_search_enabled(settings) + and mixed_source_manifest + and mixed_source_tabular_sources + ): + mixed_source_narrative_retrieval_failed = True + search_results = [] + else: + return jsonify({ + 'error': e.user_message, + 'warning_type': SEMANTIC_SEARCH_QUOTA_WARNING_TYPE, + 'service_health_warning': True, + }), 503 except Exception as e: debug_print(f"Error during hybrid search: {e}") - # Only treat as error if the exception is from embedding failure - return jsonify({ - 'error': 'There was an issue with the embedding process. Please check with an admin on embedding configuration.' - }), 500 + if ( + is_mixed_source_chat_search_enabled(settings) + and mixed_source_manifest + and mixed_source_tabular_sources + ): + mixed_source_narrative_retrieval_failed = True + search_results = [] + else: + return jsonify({ + 'error': 'There was an issue with the embedding process. Please check with an admin on embedding configuration.' + }), 500 - combined_documents = [] if search_results: unique_doc_names = set(doc.get('file_name', 'Unknown') for doc in search_results) thought_tracker.add_thought('search', f"Found {len(search_results)} results from {len(unique_doc_names)} documents") @@ -13804,11 +15942,12 @@ def result_requires_message_reload(result: Any) -> bool: # Construct system prompt for search results system_prompt_search = build_search_augmentation_system_prompt(retrieved_content) # Add this to a temporary list, don't save to DB yet - system_messages_for_augmentation.append({ - 'role': 'system', - 'content': system_prompt_search, - 'documents': combined_documents # Keep track of docs used - }) + if not is_mixed_source_chat_search_enabled(settings): + system_messages_for_augmentation.append({ + 'role': 'system', + 'content': system_prompt_search, + 'documents': combined_documents # Keep track of docs used + }) # Loop through each source document/chunk used for this message for source_doc in combined_documents: @@ -14024,7 +16163,13 @@ def result_requires_message_reload(result: Any) -> bool: # Update message-level chat_type based on actual document usage for this message # This must happen after document search is completed so search_results is populated message_chat_type = None - if (hybrid_search_enabled or history_grounded_search_used) and search_results and len(search_results) > 0: + mixed_source_has_authorized_evidence_sources = bool( + mixed_source_narrative_document_ids + or mixed_source_tabular_sources + ) + if mixed_source_document_context_active and ( + search_results or mixed_source_has_authorized_evidence_sources + ): # Documents were actually used for this message if effective_document_scope == 'group': message_chat_type = 'group' @@ -14082,8 +16227,10 @@ def result_requires_message_reload(result: Any) -> bool: source_review_used = _source_review_metadata_used(source_review_result) user_metadata['capability_usage'] = _build_capability_usage_metadata( - workspace_search_enabled=hybrid_search_enabled or history_grounded_search_used, - workspace_search_used=bool(search_results), + workspace_search_enabled=mixed_source_document_context_active, + workspace_search_used=bool( + search_results or mixed_source_has_authorized_evidence_sources + ), workspace_search_result_count=len(search_results or []), document_action_type=DOCUMENT_ACTION_TYPE_NONE, document_scope=effective_document_scope, @@ -14292,7 +16439,11 @@ def result_requires_message_reload(result: Any) -> bool: workspace_tabular_file_contexts = [] workspace_tabular_files = set() - if (hybrid_search_enabled or history_grounded_search_used) and is_tabular_processing_enabled(settings): + if ( + not is_mixed_source_chat_search_enabled(settings) + and (hybrid_search_enabled or history_grounded_search_used) + and is_tabular_processing_enabled(settings) + ): workspace_tabular_file_contexts = collect_workspace_tabular_file_contexts( combined_documents=combined_documents, selected_document_ids=effective_selected_document_ids, @@ -14316,7 +16467,138 @@ def record_tabular_post_processing_thought(thought_payload): activity=thought_payload.get('activity'), ) - if (hybrid_search_enabled or history_grounded_search_used) and workspace_tabular_files and is_tabular_processing_enabled(settings): + if ( + is_mixed_source_chat_search_enabled(settings) + and mixed_source_document_context_active + and mixed_source_manifest + ): + effective_mixed_source_selection_mode = ( + 'selected' + if mixed_source_explicit_selection + else 'history' + if history_grounded_search_used + else 'relevance' + ) + mixed_source_evidence_envelopes.extend( + build_failed_narrative_evidence_envelopes( + mixed_source_partitions.get('narrative_sources') or [], + effective_mixed_source_selection_mode, + ) + if mixed_source_narrative_retrieval_failed + else build_narrative_evidence_envelopes( + mixed_source_partitions.get('narrative_sources') or [], + search_results, + effective_mixed_source_selection_mode, + ) + ) + mixed_source_tabular_result = _execute_mixed_source_tabular_evidence( + tabular_sources=mixed_source_tabular_sources, + selection_mode=effective_mixed_source_selection_mode, + has_narrative_sources=bool(mixed_source_narrative_document_ids), + user_question=user_message, + user_id=user_id, + conversation_id=conversation_id, + gpt_model=gpt_model, + settings=settings, + thought_tracker=thought_tracker, + model_context=tabular_model_context, + request_correlation_id=mixed_source_request_correlation_id, + ) + mixed_source_evidence_envelopes.extend( + mixed_source_tabular_result.get('evidence_envelopes') or [] + ) + mixed_source_native_token_usage = mixed_source_tabular_result.get('token_usage') + agent_citations_list.extend( + mixed_source_tabular_result.get('agent_citations') or [] + ) + generated_tabular_outputs_list.extend( + mixed_source_tabular_result.get('generated_outputs') or [] + ) + generated_analysis_artifacts_list.extend( + mixed_source_tabular_result.get('generated_outputs') or [] + ) + mixed_source_handoff = build_mixed_source_evidence_handoff( + mixed_source_manifest, + mixed_source_evidence_envelopes, + effective_mixed_source_selection_mode, + mode='chat', + telemetry_settings=settings, + request_correlation_id=mixed_source_request_correlation_id, + ) + system_messages_for_augmentation.append(mixed_source_handoff) + mixed_source_coverage = mixed_source_handoff.get( + 'mixed_source_coverage', + {}, + ) + user_metadata['mixed_source_coverage'] = mixed_source_coverage + if continuity_decision: + user_metadata['source_continuity'] = continuity_decision + user_message_doc['metadata'] = user_metadata + cosmos_messages_container.upsert_item(user_message_doc) + if continuity_decision: + emit_mixed_source_telemetry( + settings, + 'continuity', + 'chat', + request_correlation_id=mixed_source_request_correlation_id, + metrics={ + 'history_source_count': continuity_decision.get('prior_source_count', 0), + 'history_rerun_count': int(bool(continuity_decision.get('requires_native_execution'))), + 'history_reuse_count': int(not continuity_decision.get('requires_native_execution')), + }, + dimensions={ + 'selection_mode': effective_mixed_source_selection_mode, + 'continuity_decision': ( + 'rerun' + if continuity_decision.get('requires_native_execution') + else 'reuse' + ), + }, + ) + emit_mixed_source_telemetry( + settings, + 'background_export', + 'chat', + request_correlation_id=mixed_source_request_correlation_id, + metrics={ + 'background_export_count': sum( + bool(output.get('background_export')) + for output in generated_tabular_outputs_list + if isinstance(output, dict) + ), + 'artifact_count': len(generated_analysis_artifacts_list), + 'citation_count': len(agent_citations_list) + len(hybrid_citations_list), + }, + dimensions={ + 'selection_mode': effective_mixed_source_selection_mode, + }, + ) + log_event( + '[MixedSourceChatSearch] Prepared bounded mixed-source synthesis evidence.', + extra={ + 'selection_mode': effective_mixed_source_selection_mode, + 'narrative_result_count': len(search_results or []), + 'tabular_candidate_count': len(mixed_source_tabular_sources), + 'tabular_completed_count': sum( + 1 + for envelope in mixed_source_evidence_envelopes + if envelope.get('source_kind') == 'tabular' + and envelope.get('status') == 'completed' + ), + 'mixed_synthesis_count': 1, + 'partial_coverage': bool( + mixed_source_coverage.get('partial_coverage') + ), + }, + level=logging.INFO, + ) + + if ( + not is_mixed_source_chat_search_enabled(settings) + and (hybrid_search_enabled or history_grounded_search_used) + and workspace_tabular_files + and is_tabular_processing_enabled(settings) + ): tabular_source_hint = determine_tabular_source_hint( effective_document_scope, active_group_id=effective_active_group_id, @@ -14579,7 +16861,10 @@ def record_tabular_post_processing_thought(thought_payload): gpt_model=gpt_model, user_message_id=user_message_id, fallback_user_message=user_message, - include_assistant_citation_context=not explicit_external_retrieval_requested, + include_assistant_citation_context=( + not explicit_external_retrieval_requested + and not mixed_source_explicit_selection + ), ) summary_of_older = history_segments['summary_of_older'] chat_tabular_files = history_segments['chat_tabular_files'] @@ -14656,7 +16941,14 @@ def record_tabular_post_processing_thought(thought_payload): final_api_source_refs.extend(history_debug_info.get('history_message_source_refs', [])) # --- Mini SK analysis for tabular files uploaded directly to chat --- - if chat_tabular_files and is_tabular_processing_enabled(settings): + if ( + chat_tabular_files + and is_tabular_processing_enabled(settings) + and not ( + is_mixed_source_chat_search_enabled(settings) + and (mixed_source_explicit_selection or mixed_source_tabular_sources) + ) + ): chat_tabular_filenames_str = ", ".join(chat_tabular_files) chat_tabular_execution_mode = get_tabular_execution_mode(user_message) log_event( @@ -14819,6 +17111,10 @@ def record_tabular_post_processing_thought(thought_payload): original_hybrid_search_enabled, prior_grounded_document_refs, explicit_external_retrieval_requested, + current_document_context_requested=( + is_mixed_source_chat_search_enabled(settings) + and document_context_requested + ), ): history_grounding_message = build_history_grounding_system_message() insert_idx = 0 @@ -15257,10 +17553,12 @@ def invoke_foundry_agent(): 'user_id': user_id, 'message_id': user_message_id, 'chat_type': chat_type, - 'document_scope': document_scope, + 'document_scope': effective_document_scope, 'group_id': active_group_id if chat_type == 'group' else None, 'hybrid_search_enabled': hybrid_search_enabled, - 'selected_document_id': selected_document_id, + 'selection_mode': selection_mode, + 'document_context_requested': mixed_source_document_context_active, + 'selected_document_id': effective_selected_document_id, 'selected_document_ids': effective_selected_document_ids, 'active_group_ids': effective_active_group_ids, 'active_public_workspace_ids': effective_active_public_workspace_ids, @@ -15586,7 +17884,6 @@ def gpt_error(e): else: ai_message, final_model_used, chat_mode, kernel_fallback_notice = fallback_result token_usage_data = None - ai_message = _append_inline_chart_blocks_to_message(ai_message, agent_citations_list) # Emit responded thought for non-agent paths (agent paths emit their own inside callbacks) @@ -15633,6 +17930,25 @@ def gpt_error(e): exceptionTraceback=True ) + token_usage_data = _merge_chat_token_usage( + mixed_source_native_token_usage, + token_usage_data, + ) + if mixed_source_manifest: + emit_mixed_source_telemetry( + settings, + 'native_execution', + 'chat', + request_correlation_id=mixed_source_request_correlation_id, + metrics={ + 'prompt_tokens': (token_usage_data or {}).get('prompt_tokens', 0), + 'completion_tokens': (token_usage_data or {}).get('completion_tokens', 0), + 'total_tokens': (token_usage_data or {}).get('total_tokens', 0), + 'token_request_count': (token_usage_data or {}).get('request_count', 0), + 'request_count': (token_usage_data or {}).get('request_count', 0), + }, + ) + # region 7 - Save GPT Response # --------------------------------------------------------------------- # 7) Save GPT response (or error message) @@ -15678,6 +17994,24 @@ def gpt_error(e): # Assistant message should be part of the same thread as the user message # Only system/augmentation messages create new threads within a conversation + if mixed_source_manifest: + fresh_finalization_manifest = resolve_authorized_source_manifest( + [source.get('document_id') for source in mixed_source_manifest], + user_id=user_id, + selection_mode=effective_mixed_source_selection_mode, + conversation_id=conversation_id, + active_group_ids=effective_active_group_ids, + active_public_workspace_ids=effective_active_public_workspace_ids, + doc_scope=effective_document_scope, + request_correlation_id=mixed_source_request_correlation_id, + ) + _validate_reauthorized_manifest_finalization( + mixed_source_manifest, + fresh_finalization_manifest, + settings=settings, + mode='chat', + request_correlation_id=mixed_source_request_correlation_id, + ) assistant_timestamp = datetime.utcnow().isoformat() prepared_agent_citations = persist_agent_citation_artifacts( conversation_id=conversation_id, @@ -15686,23 +18020,27 @@ def gpt_error(e): created_timestamp=assistant_timestamp, user_info=user_info_for_assistant, ) - assistant_table_generated_output = maybe_create_assistant_table_generated_output( + generated_file_output = maybe_create_generated_file_output( user_question=user_message, assistant_content=ai_message, conversation_id=conversation_id, + function_results=agent_citations_list, existing_outputs=generated_analysis_artifacts_list + generated_tabular_outputs_list, ) - if assistant_table_generated_output: - generated_analysis_artifacts_list.append(assistant_table_generated_output) - generated_tabular_outputs_list.append(assistant_table_generated_output) + if generated_file_output: + generated_analysis_artifacts_list.append(generated_file_output) + if generated_file_output.get('output_format') == 'csv': + generated_tabular_outputs_list.append(generated_file_output) generated_analysis_metadata = _build_generated_analysis_metadata( generated_analysis_artifacts=generated_analysis_artifacts_list, generated_tabular_outputs=generated_tabular_outputs_list, ) source_review_used = _source_review_metadata_used(source_review_result) assistant_capability_usage = _build_capability_usage_metadata( - workspace_search_enabled=hybrid_search_enabled or history_grounded_search_used, - workspace_search_used=bool(search_results), + workspace_search_enabled=mixed_source_document_context_active, + workspace_search_used=bool( + search_results or mixed_source_has_authorized_evidence_sources + ), workspace_search_result_count=len(hybrid_citations_list or []), document_action_type=DOCUMENT_ACTION_TYPE_NONE, document_scope=effective_document_scope, @@ -15848,6 +18186,16 @@ def gpt_error(e): selected_agent_name = None if selected_agent: selected_agent_name = getattr(selected_agent, 'name', None) + source_continuity_refs = None + if ( + is_mixed_source_conversation_continuity_enabled(settings) + and mixed_source_manifest + ): + source_continuity_refs = _build_mixed_source_continuity_refs( + mixed_source_manifest, + mixed_source_evidence_envelopes, + effective_mixed_source_selection_mode, + ) # Collect metadata for this conversation interaction conversation_item = collect_conversation_metadata( @@ -15859,7 +18207,7 @@ def gpt_error(e): document_scope=effective_document_scope, selected_document_id=effective_selected_document_id, model_deployment=actual_model_used, - hybrid_search_enabled=hybrid_search_enabled or history_grounded_search_used, + hybrid_search_enabled=mixed_source_document_context_active, image_gen_enabled=image_gen_enabled, selected_documents=combined_documents if 'combined_documents' in locals() else None, selected_agent=selected_agent_name, @@ -15867,7 +18215,8 @@ def gpt_error(e): search_results=search_results if 'search_results' in locals() else None, conversation_item=conversation_item, active_public_workspace_id=effective_active_public_workspace_id, - active_public_workspace_ids=effective_active_public_workspace_ids + active_public_workspace_ids=effective_active_public_workspace_ids, + source_continuity_refs=source_continuity_refs, ) except Exception as e: debug_print(f"Error collecting conversation metadata: {e}") @@ -16272,7 +18621,14 @@ def stream_cancel_requested(): generated_tabular_outputs_list = [] generated_analysis_artifacts_list = [] system_messages_for_augmentation = [] + generated_file_output_guidance = build_generated_file_output_guidance(user_message) + if generated_file_output_guidance: + system_messages_for_augmentation.append({ + 'role': 'system', + 'content': generated_file_output_guidance, + }) search_results = [] + mixed_source_narrative_retrieval_failed = False selected_agent = None # Configuration @@ -16296,6 +18652,39 @@ def stream_cancel_requested(): source_review_enabled = source_review_enabled.lower() == 'true' if isinstance(deep_research_enabled, str): deep_research_enabled = deep_research_enabled.lower() == 'true' + try: + document_context_contract = _normalize_chat_document_context_contract( + settings, + data, + selected_document_ids, + hybrid_search_enabled, + ) + except ValueError as contract_error: + yield f"data: {json.dumps({'error': str(contract_error)})}\n\n" + return + selected_document_ids = list( + document_context_contract.get('selected_document_ids') or [] + ) + requested_selected_document_ids = list(selected_document_ids) + selected_document_id = ( + selected_document_ids[0] + if len(selected_document_ids) == 1 + else None + ) + selection_mode = document_context_contract.get('selection_mode') + document_context_requested = bool( + document_context_contract.get('document_context_requested') + ) + request_has_explicit_document_selection = bool( + document_context_contract.get('explicit_selection') + ) + request_document_context_enabled = bool( + hybrid_search_enabled + or ( + is_mixed_source_chat_search_enabled(settings) + and document_context_requested + ) + ) user_workspace_context_requested = data.get('user_workspace_context_enabled') if isinstance(user_workspace_context_requested, str): user_workspace_context_requested = user_workspace_context_requested.lower() == 'true' @@ -16342,6 +18731,7 @@ def stream_cancel_requested(): history_grounded_search_used = False history_only_answerability = None prior_grounded_document_refs = [] + continuity_decision = None effective_document_scope = document_scope effective_selected_document_ids = list(selected_document_ids or []) effective_selected_document_id = selected_document_id @@ -16360,7 +18750,7 @@ def stream_cancel_requested(): assigned_knowledge_deep_research_urls = [] if assigned_knowledge_filters: assigned_knowledge_user_context_active = ( - user_workspace_context_requested + (user_workspace_context_requested or document_context_requested) and _assigned_knowledge_allows_user_workspace_context(assigned_knowledge_filters) and _assigned_knowledge_allows_document_action( assigned_knowledge_filters, @@ -16425,6 +18815,16 @@ def stream_cancel_requested(): f"public_workspaces={len(effective_active_public_workspace_ids)} | " f"tags={len(tags_filter)}" ) + mixed_source_explicit_selection = bool( + is_mixed_source_chat_search_enabled(settings) + and request_has_explicit_document_selection + and ( + not assigned_knowledge_filters + or assigned_knowledge_user_context_active + ) + ) + mixed_source_document_context_active = request_document_context_enabled + mixed_source_has_authorized_evidence_sources = False explicit_external_retrieval_requested = _is_explicit_external_retrieval_requested( web_search_enabled=web_search_enabled, url_access_enabled=url_access_enabled, @@ -16442,8 +18842,10 @@ def stream_cancel_requested(): def build_streaming_capability_usage(): source_review_was_used = _source_review_metadata_used(source_review_result) return _build_capability_usage_metadata( - workspace_search_enabled=hybrid_search_enabled or history_grounded_search_used, - workspace_search_used=bool(search_results), + workspace_search_enabled=mixed_source_document_context_active, + workspace_search_used=bool( + search_results or mixed_source_has_authorized_evidence_sources + ), workspace_search_result_count=len(hybrid_citations_list or []), document_action_type=DOCUMENT_ACTION_TYPE_NONE, document_scope=effective_document_scope, @@ -16610,11 +19012,14 @@ def build_streaming_capability_usage(): return auto_linked_chat_upload_document_ids = [] - auto_merge_chat_upload_workspace_context = _should_auto_merge_chat_upload_workspace_context( - explicit_external_retrieval_requested, - hybrid_search_enabled, - assigned_knowledge_filters=assigned_knowledge_filters, - assigned_knowledge_user_context_active=assigned_knowledge_user_context_active, + auto_merge_chat_upload_workspace_context = ( + not mixed_source_explicit_selection + and _should_auto_merge_chat_upload_workspace_context( + explicit_external_retrieval_requested, + hybrid_search_enabled, + assigned_knowledge_filters=assigned_knowledge_filters, + assigned_knowledge_user_context_active=assigned_knowledge_user_context_active, + ) ) if auto_merge_chat_upload_workspace_context: chat_upload_context = _resolve_chat_upload_workspace_context( @@ -16675,6 +19080,79 @@ def build_streaming_capability_usage(): selected_document_id = effective_selected_document_id document_scope = effective_document_scope + mixed_source_manifest = [] + mixed_source_partitions = {} + mixed_source_narrative_document_ids = [] + mixed_source_tabular_sources = [] + mixed_source_evidence_envelopes = [] + mixed_source_native_token_usage = None + mixed_source_request_correlation_id = normalize_mixed_source_correlation_id() + if mixed_source_explicit_selection: + try: + explicit_context = _resolve_chat_mixed_source_partition( + settings, + user_id, + conversation_id, + effective_selected_document_ids, + 'selected', + active_group_ids=effective_active_group_ids, + active_public_workspace_ids=effective_active_public_workspace_ids, + cancel_requested=stream_cancel_requested, + request_correlation_id=mixed_source_request_correlation_id, + ) + except ValueError as manifest_error: + yield f"data: {json.dumps({'error': str(manifest_error)})}\n\n" + return + mixed_source_manifest = explicit_context.get('manifest') or [] + mixed_source_partitions = explicit_context.get('partitions') or {} + mixed_source_narrative_document_ids = list( + explicit_context.get('narrative_document_ids') or [] + ) + mixed_source_tabular_sources = list( + explicit_context.get('tabular_sources') or [] + ) + effective_selected_document_ids = ( + mixed_source_narrative_document_ids + + _get_manifest_partition_document_ids( + mixed_source_partitions, + 'tabular_sources', + ) + ) + effective_selected_document_id = ( + effective_selected_document_ids[0] + if len(effective_selected_document_ids) == 1 + else None + ) + log_event( + '[MixedSourceChatSearch] Activated streaming explicit selected-source context.', + extra={ + 'selection_mode': 'selected', + 'requested_source_count': len(mixed_source_manifest), + 'authorized_source_count': len(effective_selected_document_ids), + 'narrative_source_count': len(mixed_source_narrative_document_ids), + 'tabular_source_count': len(mixed_source_tabular_sources), + 'omitted_source_count': len( + mixed_source_partitions.get('unresolved_sources') or [] + ), + }, + level=logging.INFO, + ) + else: + _maybe_resolve_chat_source_manifest( + settings, + user_id, + conversation_id, + effective_selected_document_ids, + scope_context, + ) + request_document_context_enabled = bool( + hybrid_search_enabled + or ( + is_mixed_source_chat_search_enabled(settings) + and document_context_requested + ) + ) + # Determine chat type actual_chat_type = 'personal_single_user' if conversation_item.get('chat_type'): @@ -16706,13 +19184,13 @@ def build_streaming_capability_usage(): user_metadata['button_states'] = { 'image_generation': False, - 'document_search': hybrid_search_enabled, + 'document_search': request_document_context_enabled, 'web_search': bool(web_search_enabled), 'url_access': bool(url_access_enabled), 'deep_research': bool(deep_research_enabled) } user_metadata['capability_usage'] = _build_capability_usage_metadata( - workspace_search_enabled=hybrid_search_enabled, + workspace_search_enabled=request_document_context_enabled, document_action_type=DOCUMENT_ACTION_TYPE_NONE, document_scope=effective_document_scope, selected_document_ids=effective_selected_document_ids, @@ -16725,12 +19203,16 @@ def build_streaming_capability_usage(): ) # Document search scope and selections - if hybrid_search_enabled: + if request_document_context_enabled: user_metadata['workspace_search'] = { 'search_enabled': True, + 'selection_mode': selection_mode, + 'document_context_requested': document_context_requested, + 'hybrid_search_preference': bool(hybrid_search_enabled), 'document_scope': effective_document_scope, 'selected_document_id': effective_selected_document_id, 'selected_document_ids': effective_selected_document_ids, + 'requested_document_ids': requested_selected_document_ids, 'active_group_ids': effective_active_group_ids, 'active_public_workspace_ids': effective_active_public_workspace_ids, 'classification': classifications_to_send @@ -16881,7 +19363,7 @@ def build_streaming_capability_usage(): conversation_id=conversation_id, message_type='user_message', message_length=len(user_message) if user_message else 0, - has_document_search=hybrid_search_enabled, + has_document_search=request_document_context_enabled, has_image_generation=False, document_scope=effective_document_scope, chat_context=actual_chat_type, @@ -17095,9 +19577,20 @@ def record_and_publish_streaming_thought(thought_payload): except Exception as ex: debug_print(f"[Content Safety - Streaming] Unexpected error: {ex}") - if not original_hybrid_search_enabled and not explicit_external_retrieval_requested: + if ( + not original_hybrid_search_enabled + and not explicit_external_retrieval_requested + and not mixed_source_explicit_selection + ): prior_grounded_document_refs = _normalize_prior_grounded_document_refs(conversation_item) if prior_grounded_document_refs: + continuity_decision = _resolve_reauthorized_continuity_decision( + settings, + user_id, + conversation_id, + prior_grounded_document_refs, + request_correlation_id=mixed_source_request_correlation_id, + ) yield emit_thought( 'history_context', 'Checking whether prior conversation context already answers the question', @@ -17136,7 +19629,10 @@ def record_and_publish_streaming_thought(thought_payload): f"[Streaming][History Fallback] History-only sufficiency assessment failed: {assessment_error}" ) - if history_only_answerability and history_only_answerability.get('can_answer_from_history'): + if _can_reuse_prior_grounded_history( + history_only_answerability, + continuity_decision, + ): yield emit_thought( 'history_context', 'Prior conversation context appears sufficient without new document retrieval', @@ -17210,9 +19706,68 @@ def record_and_publish_streaming_thought(thought_payload): 'No prior grounded documents were available; using conversation history only' ) + if ( + is_mixed_source_chat_search_enabled(settings) + and history_grounded_search_used + ): + history_context = _resolve_chat_mixed_source_partition( + settings, + user_id, + conversation_id, + effective_selected_document_ids, + 'history', + active_group_ids=effective_active_group_ids, + active_public_workspace_ids=effective_active_public_workspace_ids, + cancel_requested=stream_cancel_requested, + request_correlation_id=mixed_source_request_correlation_id, + ) + mixed_source_manifest = history_context.get('manifest') or [] + mixed_source_partitions = history_context.get('partitions') or {} + mixed_source_narrative_document_ids = list( + history_context.get('narrative_document_ids') or [] + ) + mixed_source_tabular_sources = list( + history_context.get('tabular_sources') or [] + ) + if is_mixed_source_conversation_continuity_enabled(settings): + continuity_decision = _build_reauthorized_continuity_decision( + prior_grounded_document_refs, + mixed_source_manifest, + explicit_selection=False, + ) + effective_selected_document_ids = ( + mixed_source_narrative_document_ids + + _get_manifest_partition_document_ids( + mixed_source_partitions, + 'tabular_sources', + ) + ) + effective_selected_document_id = ( + effective_selected_document_ids[0] + if len(effective_selected_document_ids) == 1 + else None + ) + # Hybrid search (if enabled) combined_documents = [] - if hybrid_search_enabled or history_grounded_search_used: + mixed_source_document_context_active = bool( + (hybrid_search_enabled or history_grounded_search_used) + if not is_mixed_source_chat_search_enabled(settings) + else ( + document_context_requested + or hybrid_search_enabled + or history_grounded_search_used + ) + ) + mixed_source_narrative_search_active = bool( + mixed_source_document_context_active + and ( + not is_mixed_source_chat_search_enabled(settings) + or not mixed_source_manifest + or mixed_source_narrative_document_ids + ) + ) + if mixed_source_narrative_search_active: debug_print( "[Streaming] Starting hybrid search | " f"conversation_id={conversation_id} | doc_scope={effective_document_scope} | " @@ -17232,7 +19787,7 @@ def record_and_publish_streaming_thought(thought_payload): search_args = { "query": search_query, "user_id": user_id, - "top_n": 12, + "top_n": 50, "doc_scope": effective_document_scope, } @@ -17255,8 +19810,14 @@ def record_and_publish_streaming_thought(thought_payload): ): search_args['active_public_workspace_id'] = effective_active_public_workspace_id - if effective_selected_document_ids: - search_args['document_ids'] = effective_selected_document_ids + search_document_ids = ( + mixed_source_narrative_document_ids + if is_mixed_source_chat_search_enabled(settings) + and mixed_source_manifest + else effective_selected_document_ids + ) + if search_document_ids: + search_args['document_ids'] = search_document_ids elif effective_selected_document_id: search_args['document_id'] = effective_selected_document_id if auto_linked_chat_upload_document_ids: @@ -17286,20 +19847,75 @@ def record_and_publish_streaming_thought(thought_payload): search_results = assigned_search_results else: search_results = hybrid_search(**search_args) + + if ( + is_mixed_source_chat_search_enabled(settings) + and not mixed_source_explicit_selection + and not history_grounded_search_used + ): + relevance_context = _resolve_chat_mixed_source_relevance_context( + settings=settings, + user_id=user_id, + conversation_id=conversation_id, + query=search_query, + search_results=search_results, + document_scope=effective_document_scope, + candidate_document_ids=( + assigned_knowledge_filters.get('document_ids') + if assigned_knowledge_filters + and assigned_knowledge_filters.get('has_workspace_knowledge') + else None + ), + tags_filter=tags_filter, + active_group_ids=effective_active_group_ids, + active_public_workspace_ids=effective_active_public_workspace_ids, + cancel_requested=stream_cancel_requested, + request_correlation_id=mixed_source_request_correlation_id, + ) + mixed_source_manifest = relevance_context.get('manifest') or [] + mixed_source_partitions = relevance_context.get('partitions') or {} + mixed_source_narrative_document_ids = list( + relevance_context.get('narrative_document_ids') or [] + ) + mixed_source_tabular_sources = list( + relevance_context.get('tabular_sources') or [] + ) + search_results = list( + relevance_context.get('search_results') or [] + ) debug_print( f"[Streaming] Hybrid search completed | results={len(search_results) if search_results else 0}" ) except SemanticSearchQuotaExceededError as e: debug_print(f"Semantic search quota exceeded during streaming hybrid search: {e}") - yield emit_thought( - 'search', - 'Workspace search warning: Semantic Ranker quota has been exceeded.', - detail=e.user_message, - ) - yield f"data: {json.dumps({'error': e.user_message, 'warning_type': SEMANTIC_SEARCH_QUOTA_WARNING_TYPE, 'service_health_warning': True})}\n\n" - return + if ( + is_mixed_source_chat_search_enabled(settings) + and mixed_source_manifest + and mixed_source_tabular_sources + ): + mixed_source_narrative_retrieval_failed = True + search_results = [] + yield emit_thought( + 'search', + 'Narrative search was unavailable; continuing with available table evidence.', + ) + else: + yield emit_thought( + 'search', + 'Workspace search warning: Semantic Ranker quota has been exceeded.', + detail=e.user_message, + ) + yield f"data: {json.dumps({'error': e.user_message, 'warning_type': SEMANTIC_SEARCH_QUOTA_WARNING_TYPE, 'service_health_warning': True})}\n\n" + return except Exception as e: debug_print(f"Error during hybrid search: {e}") + if ( + is_mixed_source_chat_search_enabled(settings) + and mixed_source_manifest + and mixed_source_tabular_sources + ): + mixed_source_narrative_retrieval_failed = True + search_results = [] if search_results: unique_doc_names_stream = set(doc.get('file_name', 'Unknown') for doc in search_results) @@ -17491,11 +20107,12 @@ def record_and_publish_streaming_thought(thought_payload): retrieved_content = "\n\n".join(retrieved_texts) system_prompt_search = build_search_augmentation_system_prompt(retrieved_content) - system_messages_for_augmentation.append({ - 'role': 'system', - 'content': system_prompt_search, - 'documents': combined_documents - }) + if not is_mixed_source_chat_search_enabled(settings): + system_messages_for_augmentation.append({ + 'role': 'system', + 'content': system_prompt_search, + 'documents': combined_documents + }) hybrid_citations_list.sort(key=_build_hybrid_citation_sort_key, reverse=True) elif history_grounded_search_used: @@ -17524,7 +20141,11 @@ def record_and_publish_streaming_thought(thought_payload): workspace_tabular_file_contexts = [] workspace_tabular_files = set() - if (hybrid_search_enabled or history_grounded_search_used) and is_tabular_processing_enabled(settings): + if ( + not is_mixed_source_chat_search_enabled(settings) + and (hybrid_search_enabled or history_grounded_search_used) + and is_tabular_processing_enabled(settings) + ): workspace_tabular_file_contexts = collect_workspace_tabular_file_contexts( combined_documents=combined_documents, selected_document_ids=effective_selected_document_ids, @@ -17540,7 +20161,144 @@ def record_and_publish_streaming_thought(thought_payload): file_context['file_name'] for file_context in workspace_tabular_file_contexts } - if (hybrid_search_enabled or history_grounded_search_used) and workspace_tabular_files and is_tabular_processing_enabled(settings): + if ( + is_mixed_source_chat_search_enabled(settings) + and mixed_source_document_context_active + and mixed_source_manifest + ): + effective_mixed_source_selection_mode = ( + 'selected' + if mixed_source_explicit_selection + else 'history' + if history_grounded_search_used + else 'relevance' + ) + mixed_source_evidence_envelopes.extend( + build_failed_narrative_evidence_envelopes( + mixed_source_partitions.get('narrative_sources') or [], + effective_mixed_source_selection_mode, + ) + if mixed_source_narrative_retrieval_failed + else build_narrative_evidence_envelopes( + mixed_source_partitions.get('narrative_sources') or [], + search_results, + effective_mixed_source_selection_mode, + ) + ) + mixed_source_tabular_result = _execute_mixed_source_tabular_evidence( + tabular_sources=mixed_source_tabular_sources, + selection_mode=effective_mixed_source_selection_mode, + has_narrative_sources=bool(mixed_source_narrative_document_ids), + user_question=user_message, + user_id=user_id, + conversation_id=conversation_id, + gpt_model=gpt_model, + settings=settings, + thought_tracker=thought_tracker, + live_thought_callback=publish_live_plugin_thought, + model_context=tabular_model_context, + cancel_requested=stream_cancel_requested, + request_correlation_id=mixed_source_request_correlation_id, + ) + mixed_source_evidence_envelopes.extend( + mixed_source_tabular_result.get('evidence_envelopes') or [] + ) + mixed_source_native_token_usage = mixed_source_tabular_result.get('token_usage') + agent_citations_list.extend( + mixed_source_tabular_result.get('agent_citations') or [] + ) + generated_tabular_outputs_list.extend( + mixed_source_tabular_result.get('generated_outputs') or [] + ) + generated_analysis_artifacts_list.extend( + mixed_source_tabular_result.get('generated_outputs') or [] + ) + mixed_source_handoff = build_mixed_source_evidence_handoff( + mixed_source_manifest, + mixed_source_evidence_envelopes, + effective_mixed_source_selection_mode, + mode='chat', + telemetry_settings=settings, + request_correlation_id=mixed_source_request_correlation_id, + ) + system_messages_for_augmentation.append(mixed_source_handoff) + mixed_source_coverage = mixed_source_handoff.get( + 'mixed_source_coverage', + {}, + ) + mixed_source_has_authorized_evidence_sources = bool( + mixed_source_narrative_document_ids + or mixed_source_tabular_sources + ) + user_metadata['mixed_source_coverage'] = mixed_source_coverage + if continuity_decision: + user_metadata['source_continuity'] = continuity_decision + user_message_doc['metadata'] = user_metadata + cosmos_messages_container.upsert_item(user_message_doc) + if continuity_decision: + emit_mixed_source_telemetry( + settings, + 'continuity', + 'chat', + request_correlation_id=mixed_source_request_correlation_id, + metrics={ + 'history_source_count': continuity_decision.get('prior_source_count', 0), + 'history_rerun_count': int(bool(continuity_decision.get('requires_native_execution'))), + 'history_reuse_count': int(not continuity_decision.get('requires_native_execution')), + }, + dimensions={ + 'selection_mode': effective_mixed_source_selection_mode, + 'continuity_decision': ( + 'rerun' + if continuity_decision.get('requires_native_execution') + else 'reuse' + ), + }, + ) + emit_mixed_source_telemetry( + settings, + 'background_export', + 'chat', + request_correlation_id=mixed_source_request_correlation_id, + metrics={ + 'background_export_count': sum( + bool(output.get('background_export')) + for output in generated_tabular_outputs_list + if isinstance(output, dict) + ), + 'artifact_count': len(generated_analysis_artifacts_list), + 'citation_count': len(agent_citations_list) + len(hybrid_citations_list), + }, + dimensions={ + 'selection_mode': effective_mixed_source_selection_mode, + }, + ) + log_event( + '[MixedSourceChatSearch] Prepared streaming bounded mixed-source synthesis evidence.', + extra={ + 'selection_mode': effective_mixed_source_selection_mode, + 'narrative_result_count': len(search_results or []), + 'tabular_candidate_count': len(mixed_source_tabular_sources), + 'tabular_completed_count': sum( + 1 + for envelope in mixed_source_evidence_envelopes + if envelope.get('source_kind') == 'tabular' + and envelope.get('status') == 'completed' + ), + 'mixed_synthesis_count': 1, + 'partial_coverage': bool( + mixed_source_coverage.get('partial_coverage') + ), + }, + level=logging.INFO, + ) + + if ( + not is_mixed_source_chat_search_enabled(settings) + and (hybrid_search_enabled or history_grounded_search_used) + and workspace_tabular_files + and is_tabular_processing_enabled(settings) + ): tabular_source_hint = determine_tabular_source_hint( effective_document_scope, active_group_id=effective_active_group_id, @@ -17812,7 +20570,9 @@ def record_and_publish_streaming_thought(thought_payload): # Update message chat type message_chat_type = None - if (hybrid_search_enabled or history_grounded_search_used) and search_results and len(search_results) > 0: + if mixed_source_document_context_active and ( + search_results or mixed_source_has_authorized_evidence_sources + ): if effective_document_scope == 'group': message_chat_type = 'group' elif effective_document_scope == 'public': @@ -17824,8 +20584,10 @@ def record_and_publish_streaming_thought(thought_payload): source_review_used = _source_review_metadata_used(source_review_result) user_metadata['capability_usage'] = _build_capability_usage_metadata( - workspace_search_enabled=hybrid_search_enabled or history_grounded_search_used, - workspace_search_used=bool(search_results), + workspace_search_enabled=mixed_source_document_context_active, + workspace_search_used=bool( + search_results or mixed_source_has_authorized_evidence_sources + ), workspace_search_result_count=len(search_results or []), document_action_type=DOCUMENT_ACTION_TYPE_NONE, document_scope=effective_document_scope, @@ -17867,7 +20629,10 @@ def record_and_publish_streaming_thought(thought_payload): gpt_model=gpt_model, user_message_id=user_message_id, fallback_user_message=user_message, - include_assistant_citation_context=not explicit_external_retrieval_requested, + include_assistant_citation_context=( + not explicit_external_retrieval_requested + and not mixed_source_explicit_selection + ), ) summary_of_older = history_segments['summary_of_older'] chat_tabular_files = history_segments['chat_tabular_files'] @@ -17894,7 +20659,14 @@ def record_and_publish_streaming_thought(thought_payload): final_api_source_refs.extend(history_debug_info.get('history_message_source_refs', [])) # --- Mini SK analysis for tabular files uploaded directly to chat --- - if chat_tabular_files and is_tabular_processing_enabled(settings): + if ( + chat_tabular_files + and is_tabular_processing_enabled(settings) + and not ( + is_mixed_source_chat_search_enabled(settings) + and (mixed_source_explicit_selection or mixed_source_tabular_sources) + ) + ): chat_tabular_filenames_str = ", ".join(chat_tabular_files) chat_tabular_execution_mode = get_tabular_execution_mode(user_message) log_event( @@ -18066,6 +20838,10 @@ def record_and_publish_streaming_thought(thought_payload): original_hybrid_search_enabled, prior_grounded_document_refs, explicit_external_retrieval_requested, + current_document_context_requested=( + is_mixed_source_chat_search_enabled(settings) + and document_context_requested + ), ): history_grounding_message = build_history_grounding_system_message() insert_idx = 0 @@ -18214,6 +20990,24 @@ def finalize_cancelled_stream_response(): 'cancel_reason': cancel_reason, } + if mixed_source_manifest: + _rollback_mixed_source_chat_publication( + user_id, + conversation_id, + generated_analysis_artifacts_list + generated_tabular_outputs_list, + ) + return _build_stream_cancel_event( + conversation_id, + user_message_id=user_message_id, + reason=cancel_reason, + message_persisted=False, + extra_payload={ + 'augmented': bool(system_messages_for_augmentation), + 'metadata': cancel_metadata, + 'thoughts_enabled': thought_tracker.enabled, + }, + ) + if partial_content: assistant_timestamp = datetime.utcnow().isoformat() prepared_agent_citations = persist_agent_citation_artifacts( @@ -18408,6 +21202,8 @@ def finalize_cancelled_agent_stream_response(): 'document_scope': effective_document_scope, 'group_id': effective_active_group_id if chat_type == 'group' else None, 'hybrid_search_enabled': hybrid_search_enabled, + 'selection_mode': selection_mode, + 'document_context_requested': mixed_source_document_context_active, 'selected_document_id': effective_selected_document_id, 'selected_document_ids': effective_selected_document_ids, 'active_group_ids': effective_active_group_ids, @@ -18775,6 +21571,25 @@ def finalize_cancelled_agent_stream_response(): yield finalize_cancelled_stream_response() return + token_usage_data = _merge_chat_token_usage( + mixed_source_native_token_usage, + token_usage_data, + ) + if mixed_source_manifest: + emit_mixed_source_telemetry( + settings, + 'native_execution', + 'chat', + request_correlation_id=mixed_source_request_correlation_id, + metrics={ + 'prompt_tokens': (token_usage_data or {}).get('prompt_tokens', 0), + 'completion_tokens': (token_usage_data or {}).get('completion_tokens', 0), + 'total_tokens': (token_usage_data or {}).get('total_tokens', 0), + 'token_request_count': (token_usage_data or {}).get('request_count', 0), + 'request_count': (token_usage_data or {}).get('request_count', 0), + }, + ) + # Stream complete - save message and send final metadata accumulated_content_before_chart_append = accumulated_content accumulated_content = _append_inline_chart_blocks_to_message(accumulated_content, agent_citations_list) @@ -18787,6 +21602,30 @@ def finalize_cancelled_agent_stream_response(): user_info_for_assistant = response_message_context.get('user_info') user_thread_id = response_message_context.get('thread_id') user_previous_thread_id = response_message_context.get('previous_thread_id') + if mixed_source_manifest: + raise_if_mixed_source_cancelled( + stream_cancel_requested, + 'finalization', + request_correlation_id=mixed_source_request_correlation_id, + ) + fresh_finalization_manifest = resolve_authorized_source_manifest( + [source.get('document_id') for source in mixed_source_manifest], + user_id=user_id, + selection_mode=effective_mixed_source_selection_mode, + conversation_id=conversation_id, + active_group_ids=effective_active_group_ids, + active_public_workspace_ids=effective_active_public_workspace_ids, + doc_scope=effective_document_scope, + cancel_requested=stream_cancel_requested, + request_correlation_id=mixed_source_request_correlation_id, + ) + _validate_reauthorized_manifest_finalization( + mixed_source_manifest, + fresh_finalization_manifest, + settings=settings, + mode='chat', + request_correlation_id=mixed_source_request_correlation_id, + ) assistant_timestamp = datetime.utcnow().isoformat() prepared_agent_citations = persist_agent_citation_artifacts( conversation_id=conversation_id, @@ -18794,16 +21633,46 @@ def finalize_cancelled_agent_stream_response(): agent_citations=agent_citations_list, created_timestamp=assistant_timestamp, user_info=user_info_for_assistant, + cancel_requested=stream_cancel_requested, + request_correlation_id=mixed_source_request_correlation_id, ) - assistant_table_generated_output = maybe_create_assistant_table_generated_output( + raise_if_mixed_source_cancelled( + stream_cancel_requested, + 'artifact_publication', + request_correlation_id=mixed_source_request_correlation_id, + ) + generated_file_output = maybe_create_generated_file_output( user_question=user_message, assistant_content=accumulated_content, conversation_id=conversation_id, + function_results=agent_citations_list, existing_outputs=generated_analysis_artifacts_list + generated_tabular_outputs_list, + cancel_requested=stream_cancel_requested, + request_correlation_id=mixed_source_request_correlation_id, ) - if assistant_table_generated_output: - generated_analysis_artifacts_list.append(assistant_table_generated_output) - generated_tabular_outputs_list.append(assistant_table_generated_output) + if generated_file_output: + generated_analysis_artifacts_list.append(generated_file_output) + if generated_file_output.get('output_format') == 'csv': + generated_tabular_outputs_list.append(generated_file_output) + if mixed_source_manifest: + fresh_finalization_manifest = resolve_authorized_source_manifest( + [source.get('document_id') for source in mixed_source_manifest], + user_id=user_id, + selection_mode=effective_mixed_source_selection_mode, + conversation_id=conversation_id, + active_group_ids=effective_active_group_ids, + active_public_workspace_ids=effective_active_public_workspace_ids, + doc_scope=effective_document_scope, + cancel_requested=stream_cancel_requested, + request_correlation_id=mixed_source_request_correlation_id, + ) + _validate_reauthorized_manifest_finalization( + mixed_source_manifest, + fresh_finalization_manifest, + settings=settings, + mode='chat', + request_correlation_id=mixed_source_request_correlation_id, + ) generated_analysis_metadata = _build_generated_analysis_metadata( generated_analysis_artifacts=generated_analysis_artifacts_list, generated_tabular_outputs=generated_tabular_outputs_list, @@ -18854,6 +21723,11 @@ def finalize_cancelled_agent_stream_response(): } }) cosmos_messages_container.upsert_item(assistant_doc) + raise_if_mixed_source_cancelled( + stream_cancel_requested, + 'finalization', + request_correlation_id=mixed_source_request_correlation_id, + ) if use_agent_streaming and agent_name_used: agent_scope_for_usage = 'personal' agent_group_id_for_usage = None @@ -18931,6 +21805,16 @@ def finalize_cancelled_agent_stream_response(): debug_print(f"Warning: Could not update streaming user message metadata: {e}") try: + source_continuity_refs = None + if ( + is_mixed_source_conversation_continuity_enabled(settings) + and mixed_source_manifest + ): + source_continuity_refs = _build_mixed_source_continuity_refs( + mixed_source_manifest, + mixed_source_evidence_envelopes, + effective_mixed_source_selection_mode, + ) conversation_item = collect_conversation_metadata( user_message=user_message, conversation_id=conversation_id, @@ -18940,7 +21824,7 @@ def finalize_cancelled_agent_stream_response(): document_scope=effective_document_scope, selected_document_id=effective_selected_document_id, model_deployment=final_model_used if use_agent_streaming else gpt_model, - hybrid_search_enabled=hybrid_search_enabled or history_grounded_search_used, + hybrid_search_enabled=mixed_source_document_context_active, image_gen_enabled=False, selected_documents=combined_documents if combined_documents else None, selected_agent=agent_name_used if use_agent_streaming else None, @@ -18948,7 +21832,8 @@ def finalize_cancelled_agent_stream_response(): search_results=search_results if search_results else None, conversation_item=conversation_item, active_public_workspace_id=effective_active_public_workspace_id, - active_public_workspace_ids=effective_active_public_workspace_ids + active_public_workspace_ids=effective_active_public_workspace_ids, + source_continuity_refs=source_continuity_refs, ) except Exception as e: debug_print(f"Error collecting conversation metadata: {e}") @@ -19016,6 +21901,52 @@ def finalize_cancelled_agent_stream_response(): ) yield f"data: {json.dumps(final_data)}\n\n" + except MixedSourceCancellationError: + if mixed_source_manifest: + try: + cosmos_messages_container.delete_item( + item=assistant_message_id, + partition_key=conversation_id, + ) + except Exception: + pass + _rollback_mixed_source_chat_publication( + user_id, + conversation_id, + generated_analysis_artifacts_list + generated_tabular_outputs_list, + compact_citations=locals().get('prepared_agent_citations') or [], + ) + yield _build_stream_cancel_event( + conversation_id, + user_message_id=user_message_id, + reason=stream_session.get_cancel_reason() if stream_session else 'user_requested', + message_persisted=False, + extra_payload={ + 'augmented': bool(system_messages_for_augmentation), + 'thoughts_enabled': thought_tracker.enabled, + }, + ) + return + yield finalize_cancelled_stream_response() + return + except MixedSourceFinalizationError: + if mixed_source_manifest: + try: + cosmos_messages_container.delete_item( + item=assistant_message_id, + partition_key=conversation_id, + ) + except Exception: + pass + _rollback_mixed_source_chat_publication( + user_id, + conversation_id, + generated_analysis_artifacts_list + generated_tabular_outputs_list, + compact_citations=locals().get('prepared_agent_citations') or [], + ) + yield f"data: {json.dumps({'error': 'Selected source state changed before final output could be published.', 'conversation_id': conversation_id})}\n\n" + return + raise except Exception as e: error_msg = str(e) debug_print(f"Error during streaming: {error_msg}") @@ -19149,10 +22080,29 @@ def tabular_generated_output_run_resume_api(run_id): resume_result = resume_tabular_generated_output_run(user_id, run_id) if not resume_result: return jsonify({'error': 'Tabular generated-output run not found'}), 404 + if resume_result.get('authorization_failed'): + return jsonify(resume_result), 403 if not resume_result.get('success'): return jsonify(resume_result), 409 return jsonify(resume_result) + @bp.route('/api/tabular/generated-output/runs//cancel', methods=['POST']) + @swagger_route(security=get_auth_security()) + @login_required + @user_required + def tabular_generated_output_run_cancel_api(run_id): + """Cancel a generated-output run owned by the current user.""" + user_id = get_current_user_id() + if not user_id: + return jsonify({'error': 'User not authenticated'}), 401 + + cancel_result = cancel_tabular_generated_output_run(user_id, run_id) + if not cancel_result: + return jsonify({'error': 'Tabular generated-output run not found'}), 404 + if not cancel_result.get('success'): + return jsonify(cancel_result), 409 + return jsonify(cancel_result) + @bp.route('/api/chat/stream/reattach/', methods=['GET']) @swagger_route(security=get_auth_security()) @login_required @@ -19717,6 +22667,17 @@ def add_ref(raw_ref): 'scope_id': scope_id, 'file_name': raw_ref.get('file_name') or raw_ref.get('title'), 'classification': raw_ref.get('classification'), + 'source_role': raw_ref.get('source_role'), + 'requested_order': raw_ref.get('requested_order'), + 'source_kind': raw_ref.get('source_kind'), + 'engine': raw_ref.get('engine'), + 'source_version': raw_ref.get('source_version'), + 'status': raw_ref.get('status'), + 'coverage': dict(raw_ref.get('coverage') or {}), + 'selection_origin': raw_ref.get('selection_origin'), + 'action_mode': raw_ref.get('action_mode'), + 'citation_count': _safe_metadata_int(raw_ref.get('citation_count')), + 'artifact_count': _safe_metadata_int(raw_ref.get('artifact_count')), } if scope == 'group': @@ -19755,6 +22716,8 @@ def build_prior_grounded_document_search_parameters(grounded_refs): document_ids = [] group_ids = [] public_workspace_ids = [] + chat_conversation_ids = [] + document_scopes = {} scope_types = set() for ref in grounded_refs or []: @@ -19769,6 +22732,10 @@ def build_prior_grounded_document_search_parameters(grounded_refs): if not scope: continue scope_types.add(scope) + document_scopes[document_id] = { + 'scope': scope, + 'scope_id': str(ref.get('scope_id') or '').strip(), + } if scope == 'group': group_id = str(ref.get('group_id') or ref.get('scope_id') or '').strip() @@ -19778,6 +22745,10 @@ def build_prior_grounded_document_search_parameters(grounded_refs): public_workspace_id = str(ref.get('public_workspace_id') or ref.get('scope_id') or '').strip() if public_workspace_id and public_workspace_id not in public_workspace_ids: public_workspace_ids.append(public_workspace_id) + elif scope == 'chat': + chat_conversation_id = str(ref.get('scope_id') or '').strip() + if chat_conversation_id and chat_conversation_id not in chat_conversation_ids: + chat_conversation_ids.append(chat_conversation_id) if len(scope_types) == 1: doc_scope = next(iter(scope_types)) @@ -19791,6 +22762,8 @@ def build_prior_grounded_document_search_parameters(grounded_refs): 'active_group_id': group_ids[0] if group_ids else None, 'active_public_workspace_ids': public_workspace_ids, 'active_public_workspace_id': public_workspace_ids[0] if public_workspace_ids else None, + 'chat_conversation_ids': chat_conversation_ids, + 'document_scopes': document_scopes, 'scope_types': sorted(scope_types), } @@ -19806,6 +22779,13 @@ def revalidate_prior_grounded_document_search_parameters(user_id, search_paramet ) allowed_group_ids = scope_context['active_group_ids'] allowed_public_workspace_ids = scope_context['active_public_workspace_ids'] + allowed_chat_conversation_ids = [] + for conversation_id in normalized_parameters.get('chat_conversation_ids') or []: + try: + _authorize_personal_conversation_access(user_id, conversation_id) + allowed_chat_conversation_ids.append(conversation_id) + except (LookupError, PermissionError): + continue allowed_scope_types = [] if 'personal' in scope_types: @@ -19814,12 +22794,25 @@ def revalidate_prior_grounded_document_search_parameters(user_id, search_paramet allowed_scope_types.append('group') if allowed_public_workspace_ids: allowed_scope_types.append('public') + if allowed_chat_conversation_ids: + allowed_scope_types.append('chat') normalized_parameters['active_group_ids'] = allowed_group_ids normalized_parameters['active_group_id'] = scope_context['active_group_id'] normalized_parameters['active_public_workspace_ids'] = allowed_public_workspace_ids normalized_parameters['active_public_workspace_id'] = scope_context['active_public_workspace_id'] + normalized_parameters['chat_conversation_ids'] = allowed_chat_conversation_ids normalized_parameters['scope_types'] = allowed_scope_types + document_scopes = normalized_parameters.get('document_scopes') or {} + normalized_parameters['document_ids'] = [ + document_id + for document_id in normalized_parameters.get('document_ids') or [] + if ( + (document_scopes.get(document_id) or {}).get('scope') != 'chat' + or (document_scopes.get(document_id) or {}).get('scope_id') + in allowed_chat_conversation_ids + ) + ] if not allowed_scope_types: normalized_parameters['document_ids'] = [] @@ -19827,7 +22820,9 @@ def revalidate_prior_grounded_document_search_parameters(user_id, search_paramet return normalized_parameters normalized_parameters['doc_scope'] = ( - allowed_scope_types[0] if len(allowed_scope_types) == 1 else 'all' + allowed_scope_types[0] + if len(allowed_scope_types) == 1 and allowed_scope_types[0] != 'chat' + else 'all' ) return normalized_parameters @@ -19918,11 +22913,13 @@ def should_apply_history_grounding_message( original_hybrid_search_enabled, prior_grounded_document_refs, explicit_external_retrieval_requested=False, + current_document_context_requested=False, ): """Apply bounded grounding only when prior grounded docs exist for this conversation.""" return ( not bool(original_hybrid_search_enabled) and not bool(explicit_external_retrieval_requested) + and not bool(current_document_context_requested) and bool(prior_grounded_document_refs) ) diff --git a/application/single_app/route_backend_collaboration.py b/application/single_app/route_backend_collaboration.py index 1d00cec85..0733372e8 100644 --- a/application/single_app/route_backend_collaboration.py +++ b/application/single_app/route_backend_collaboration.py @@ -307,6 +307,8 @@ def _build_collaboration_stream_request_payload(data, source_conversation_id, me 'message': message_content, 'conversation_id': source_conversation_id, 'hybrid_search': bool(data.get('hybrid_search')), + 'selection_mode': data.get('selection_mode'), + 'document_context_requested': data.get('document_context_requested'), 'web_search_enabled': bool(data.get('web_search_enabled')), 'selected_document_id': data.get('selected_document_id'), 'selected_document_ids': data.get('selected_document_ids') or [], diff --git a/application/single_app/route_backend_conversations.py b/application/single_app/route_backend_conversations.py index d53fbb432..413de7bd1 100644 --- a/application/single_app/route_backend_conversations.py +++ b/application/single_app/route_backend_conversations.py @@ -171,6 +171,67 @@ def _normalize_workspace_document_delete_ids(raw_document_ids): return normalized_document_ids +def _build_replayed_document_context(original_metadata): + """Rebuild document-context intent from stored metadata for retry and edit.""" + metadata = original_metadata if isinstance(original_metadata, dict) else {} + workspace_search = metadata.get('workspace_search') + if not isinstance(workspace_search, dict): + workspace_search = metadata.get('document_search') + workspace_search = workspace_search if isinstance(workspace_search, dict) else {} + + selected_document_ids = ( + workspace_search.get('requested_document_ids') + or workspace_search.get('selected_document_ids') + or [] + ) + if not isinstance(selected_document_ids, list): + selected_document_ids = [selected_document_ids] + selected_document_ids = [ + str(document_id or '').strip() + for document_id in selected_document_ids + if str(document_id or '').strip() + ] + selected_document_id = str( + workspace_search.get('selected_document_id') + or workspace_search.get('document_id') + or '' + ).strip() + if selected_document_id and selected_document_id not in selected_document_ids: + selected_document_ids.insert(0, selected_document_id) + + selection_mode = str(workspace_search.get('selection_mode') or '').strip().lower() + if selection_mode not in {'selected', 'all', 'history', 'relevance'}: + selection_mode = 'selected' if selected_document_ids else 'relevance' + document_context_requested = workspace_search.get('document_context_requested') + if not isinstance(document_context_requested, bool): + document_context_requested = bool( + workspace_search.get('search_enabled') + or workspace_search.get('enabled') + or selected_document_ids + ) + + return { + 'hybrid_search': bool( + workspace_search.get('hybrid_search_preference') + if 'hybrid_search_preference' in workspace_search + else workspace_search.get('enabled') + or workspace_search.get('search_enabled') + ), + 'selection_mode': selection_mode, + 'document_context_requested': document_context_requested, + 'selected_document_id': selected_document_ids[0] if selected_document_ids else None, + 'selected_document_ids': selected_document_ids, + 'doc_scope': workspace_search.get('document_scope') or workspace_search.get('scope'), + 'top_n': workspace_search.get('top_n'), + 'classifications': ( + workspace_search.get('classification') + or workspace_search.get('classifications') + ), + 'active_group_ids': workspace_search.get('active_group_ids') or [], + 'active_public_workspace_ids': workspace_search.get('active_public_workspace_ids') or [], + } + + def _get_requested_workspace_document_delete_ids_for_conversation(payload, conversation_id): if not isinstance(payload, dict): return [] @@ -2641,16 +2702,13 @@ def retry_message(message_id): "message_retry_created", ) # Build chat request parameters from original message metadata + replayed_document_context = _build_replayed_document_context(original_metadata) chat_request = { 'message': user_content, 'conversation_id': conversation_id, 'model_deployment': selected_model or original_metadata.get('model_selection', {}).get('selected_model'), 'reasoning_effort': reasoning_effort or original_metadata.get('reasoning_effort'), - 'hybrid_search': original_metadata.get('document_search', {}).get('enabled', False), - 'selected_document_id': original_metadata.get('document_search', {}).get('document_id'), - 'doc_scope': original_metadata.get('document_search', {}).get('scope'), - 'top_n': original_metadata.get('document_search', {}).get('top_n'), - 'classifications': original_metadata.get('document_search', {}).get('classifications'), + **replayed_document_context, 'image_generation': original_metadata.get('image_generation', {}).get('enabled', False), 'active_group_id': original_metadata.get('chat_context', {}).get('group_id'), 'active_public_workspace_id': original_metadata.get('chat_context', {}).get('public_workspace_id'), @@ -2864,16 +2922,13 @@ def edit_message(message_id): ) # Build chat request parameters from original message metadata # Keep all original settings (model, reasoning, doc search, etc.) + replayed_document_context = _build_replayed_document_context(original_metadata) chat_request = { 'message': edited_content, # Use edited content 'conversation_id': conversation_id, 'model_deployment': original_metadata.get('model_selection', {}).get('selected_model'), 'reasoning_effort': original_metadata.get('reasoning_effort'), - 'hybrid_search': original_metadata.get('document_search', {}).get('enabled', False), - 'selected_document_id': original_metadata.get('document_search', {}).get('document_id'), - 'doc_scope': original_metadata.get('document_search', {}).get('scope'), - 'top_n': original_metadata.get('document_search', {}).get('top_n'), - 'classifications': original_metadata.get('document_search', {}).get('classifications'), + **replayed_document_context, 'image_generation': original_metadata.get('image_generation', {}).get('enabled', False), 'active_group_id': original_metadata.get('chat_context', {}).get('group_id'), 'active_public_workspace_id': original_metadata.get('chat_context', {}).get('public_workspace_id'), diff --git a/application/single_app/route_backend_data_management.py b/application/single_app/route_backend_data_management.py index e1e1de7b0..c1992ee4a 100644 --- a/application/single_app/route_backend_data_management.py +++ b/application/single_app/route_backend_data_management.py @@ -3,7 +3,7 @@ import logging -from flask import current_app, jsonify, request, session +from flask import Response, current_app, jsonify, request, session from functions_activity_logging import log_general_admin_action from functions_appinsights import log_event @@ -15,17 +15,23 @@ DATA_MANAGEMENT_OPERATION_RESTORE, DataManagementCosmosEditorError, DataManagementSettingsValidationError, + export_data_management_migration_manifest, generate_data_management_encryption_key, get_data_management_cosmos_editor_containers, get_data_management_cosmos_editor_document, get_data_management_backup_summary, get_data_management_job_detail, + get_data_management_job_progress, get_data_management_jobs, get_data_management_migration_catalog, get_data_management_settings, log_data_management_cosmos_editor_activity, + preview_data_management_migration_plan, queue_data_management_job, query_data_management_cosmos_editor_documents, + request_data_management_migration_cancellation, + resolve_data_management_migration_manifest_item, + retry_data_management_migration_job, sanitize_data_management_job_for_admin, sanitize_data_management_settings_for_admin, save_data_management_cosmos_editor_document, @@ -182,8 +188,12 @@ def test_admin_data_management_storage(): def test_admin_data_management_target_cosmos(): payload = request.get_json(silent=True) or {} settings_payload = payload.get("settings") if isinstance(payload.get("settings"), dict) else None + migration_plan = payload.get("migration_plan") if isinstance(payload.get("migration_plan"), dict) else None try: - result = test_target_cosmos_connection(settings=settings_payload) + result = test_target_cosmos_connection( + settings=settings_payload, + migration_plan=migration_plan, + ) except Exception as exc: log_event( "[DataManagement] Target Cosmos connection test failed.", @@ -406,6 +416,116 @@ def get_admin_data_management_job_detail(job_id): return jsonify({"success": False, "error": "Data Management job was not found."}), 404 return jsonify({"success": True, **detail}), 200 + @bp.route("/api/admin/data-management/jobs//progress", methods=["GET"]) + @swagger_route(security=get_auth_security()) + @login_required + @admin_required + def get_admin_data_management_job_progress(job_id): + progress = get_data_management_job_progress(job_id) + if not progress: + return jsonify({"success": False, "error": "Data Management job was not found."}), 404 + return jsonify({"success": True, "job": progress}), 200 + + @bp.route("/api/admin/data-management/jobs//migration-manifest", methods=["GET"]) + @swagger_route(security=get_auth_security()) + @login_required + @admin_required + def download_admin_data_management_migration_manifest(job_id): + statuses = { + status.strip().lower() + for status in str(request.args.get("statuses") or "").split(",") + if status.strip() + } + try: + export = export_data_management_migration_manifest(job_id, statuses=statuses) + except DataManagementSettingsValidationError as exc: + return jsonify({"success": False, "error": str(exc)}), 404 + response = Response( + export["content"], + status=200, + content_type="application/x-ndjson; charset=utf-8", + ) + suffix = "-failures" if statuses else "" + response.headers["Content-Disposition"] = ( + f'attachment; filename="migration-{job_id}{suffix}.jsonl"' + ) + response.headers["X-Migration-Manifest-Entries"] = str(export["entry_count"]) + return response + + @bp.route( + "/api/admin/data-management/jobs//migration-manifest/items/", + methods=["GET"], + ) + @swagger_route(security=get_auth_security()) + @login_required + @admin_required + def resolve_admin_data_management_migration_manifest_item(job_id, item_ref): + try: + item = resolve_data_management_migration_manifest_item(job_id, item_ref) + except DataManagementSettingsValidationError as exc: + return jsonify({"success": False, "error": str(exc)}), 404 + return jsonify({"success": True, "item": item}), 200 + + @bp.route("/api/admin/data-management/jobs//retry", methods=["POST"]) + @swagger_route(security=get_auth_security()) + @login_required + @admin_required + def retry_admin_data_management_migration_job(job_id): + try: + job = retry_data_management_migration_job(job_id) + submitted = submit_data_management_job(current_app._get_current_object(), job.get("id")) + except DataManagementSettingsValidationError as exc: + return jsonify({"success": False, "error": str(exc)}), 400 + except Exception as exc: + log_event( + "[DataManagement] Failed to retry migration job.", + {"job_id": job_id, "error": str(exc)}, + level=logging.ERROR, + exceptionTraceback=True, + ) + return jsonify({"success": False, "error": "Migration retry could not be queued."}), 400 + + _log_data_management_admin_action( + "data_management_migration_retry_queued", + "Queued a Data Management migration retry from durable checkpoints.", + {"job_id": job.get("id"), "submitted": submitted}, + ) + public_job = sanitize_data_management_job_for_admin(job) + public_job["submitted_to_executor"] = submitted + return jsonify({"success": True, "job": public_job}), 202 + + @bp.route("/api/admin/data-management/jobs//cancel", methods=["POST"]) + @swagger_route(security=get_auth_security()) + @login_required + @admin_required + def cancel_admin_data_management_migration_job(job_id): + payload = request.get_json(silent=True) or {} + admin_user_id, admin_email = _get_admin_context() + try: + job = request_data_management_migration_cancellation( + job_id, + requested_by=admin_user_id, + requested_by_email=admin_email, + reason=payload.get("reason"), + ) + except DataManagementSettingsValidationError as exc: + return jsonify({"success": False, "error": str(exc)}), 400 + except Exception as exc: + log_event( + "[DataManagement] Failed to request migration cancellation.", + {"job_id": job_id, "error": str(exc)}, + level=logging.ERROR, + exceptionTraceback=True, + ) + return jsonify({"success": False, "error": "Migration cancellation could not be requested."}), 400 + + _log_data_management_admin_action( + "data_management_migration_cancel_requested", + "Requested cancellation of a Data Management migration.", + {"job_id": job.get("id"), "status": job.get("status")}, + ) + return jsonify({"success": True, "job": sanitize_data_management_job_for_admin(job)}), 202 + @bp.route("/api/admin/data-management/backups", methods=["GET"]) @swagger_route(security=get_auth_security()) @login_required @@ -436,9 +556,26 @@ def summarize_admin_data_management_migration(): payload = request.get_json(silent=True) or {} try: summary = summarize_data_management_migration_plan(payload) + preview = None + if payload.get("include_inventory") is True: + preview = preview_data_management_migration_plan( + get_data_management_settings(), + payload, + ) except DataManagementSettingsValidationError as exc: return jsonify({"success": False, "error": str(exc)}), 400 - return jsonify({"success": True, "summary": summary}), 200 + except Exception as exc: + log_event( + "[DataManagement] Migration inventory preview failed.", + {"error": str(exc)}, + level=logging.ERROR, + exceptionTraceback=True, + ) + return jsonify({ + "success": False, + "error": "Migration inventory preview could not be completed.", + }), 400 + return jsonify({"success": True, "summary": summary, "preview": preview}), 200 @bp.route("/api/admin/data-management/jobs", methods=["POST"]) @swagger_route(security=get_auth_security()) @@ -458,12 +595,17 @@ def create_admin_data_management_job(): admin_user_id, admin_email = _get_admin_context() try: + job_options = ( + dict(payload.get("options")) + if isinstance(payload.get("options"), dict) + else {} + ) job = queue_data_management_job( operation, backup_type=backup_type, requested_by=admin_user_id, requested_by_email=admin_email, - options=payload.get("options") if isinstance(payload.get("options"), dict) else {}, + options=job_options, ) submitted = submit_data_management_job(current_app._get_current_object(), job.get("id")) except Exception as exc: diff --git a/application/single_app/route_backend_documents.py b/application/single_app/route_backend_documents.py index f7d0a6b77..18ac783d2 100644 --- a/application/single_app/route_backend_documents.py +++ b/application/single_app/route_backend_documents.py @@ -2052,6 +2052,10 @@ def api_unshare_document(document_id): return jsonify({'message': 'Document unshared successfully'}), 200 else: return jsonify({'error': 'Failed to unshare document'}), 500 + except DocumentSearchAclProjectionDeferredError as exc: + response = jsonify({'error': str(exc)}) + response.headers['Retry-After'] = '150' + return response, 503 except exceptions.CosmosResourceNotFoundError: return jsonify({'error': 'Document not found or access denied'}), 404 except Exception as e: diff --git a/application/single_app/semantic_kernel_plugins/plugin_invocation_logger.py b/application/single_app/semantic_kernel_plugins/plugin_invocation_logger.py index 92dedaf5f..ffdba3a59 100644 --- a/application/single_app/semantic_kernel_plugins/plugin_invocation_logger.py +++ b/application/single_app/semantic_kernel_plugins/plugin_invocation_logger.py @@ -64,6 +64,15 @@ AUTHORIZATION_VALUE_RE = re.compile(r"(?i)\b(Bearer|Basic)\s+[A-Za-z0-9._~+/=-]+") +class PluginInvocationResult(str): + """String tool result with server-only metadata retained in invocation history.""" + + def __new__(cls, value: str, internal_metadata: Optional[Dict[str, Any]] = None): + instance = super().__new__(cls, value) + instance.internal_metadata = dict(internal_metadata or {}) + return instance + + def _normalize_sensitive_key(key: Any) -> str: return re.sub(r"[^a-z0-9]", "", str(key or "").strip().lower()) diff --git a/application/single_app/semantic_kernel_plugins/tabular_processing_plugin.py b/application/single_app/semantic_kernel_plugins/tabular_processing_plugin.py index 99ce381e2..a041c12f1 100644 --- a/application/single_app/semantic_kernel_plugins/tabular_processing_plugin.py +++ b/application/single_app/semantic_kernel_plugins/tabular_processing_plugin.py @@ -13,15 +13,18 @@ import json import logging import re +import tempfile import warnings import pandas +from azure.core import MatchConditions from flask import g, has_request_context from typing import Annotated, Dict, List, Optional, Set from urllib.parse import urlsplit, urlunsplit from semantic_kernel.functions import kernel_function -from semantic_kernel_plugins.plugin_invocation_logger import plugin_function_logger +from semantic_kernel_plugins.plugin_invocation_logger import PluginInvocationResult, plugin_function_logger from functions_appinsights import log_event from functions_authentication import get_current_user_id +from functions_tabular_csv_query import iter_tabular_csv_query_rows from functions_group import find_group_by_id, get_user_role_in_group from functions_public_workspaces import get_user_visible_public_workspace_ids_from_settings from config import ( @@ -124,6 +127,7 @@ def __init__(self): self._df_cache = {} # Per-instance cache: (container, blob_name, sheet_name) -> DataFrame self._blob_data_cache = {} # Per-instance cache: (container, blob_name) -> raw bytes self._workbook_metadata_cache = {} # Per-instance cache: (container, blob_name) -> workbook metadata + self._blob_version_cache = {} # Per-instance cache: (container, blob_name) -> immutable blob version metadata self._default_sheet_overrides = {} # (container, blob_name) -> default sheet name self._resolved_blob_location_overrides = {} # (source, filename) -> (container, blob_name) @@ -230,6 +234,14 @@ def _get_authorized_chat_context(self) -> dict: for workspace_id in (authorized_context.get('active_public_workspace_ids') or []) if str(workspace_id or '').strip() ] + authorized_blob_locations = [ + [str(location[0]), str(location[1])] + for location in authorized_context.get('authorized_blob_locations') or [] + if isinstance(location, (list, tuple)) + and len(location) == 2 + and str(location[0] or '').strip() + and str(location[1] or '').strip() + ] return { 'user_id': authorized_user_id, @@ -240,6 +252,7 @@ def _get_authorized_chat_context(self) -> dict: 'active_public_workspace_id': ( str(authorized_context.get('active_public_workspace_id') or '').strip() or None ), + 'authorized_blob_locations': authorized_blob_locations, } def _resolve_authorized_scope_arguments( @@ -342,6 +355,14 @@ def _is_authorized_public_workspace_scope( def _is_authorized_blob_location(self, container_name: str, blob_path: str, authorized_context: dict) -> bool: """Ensure remembered blob locations still fall within the caller's authorized request scope.""" + exact_authorized_locations = { + (str(location[0]), str(location[1])) + for location in authorized_context.get('authorized_blob_locations') or [] + if isinstance(location, (list, tuple)) and len(location) == 2 + } + if (str(container_name or ''), str(blob_path or '')) in exact_authorized_locations: + return True + source = self._infer_source_from_container(container_name) blob_parts = [part for part in str(blob_path or '').split('/') if part] if not source or not blob_parts: @@ -392,11 +413,45 @@ def _download_tabular_blob_bytes(self, container_name: str, blob_name: str) -> b client = self._get_blob_service_client() blob_client = client.get_blob_client(container=container_name, blob=blob_name) - stream = blob_client.download_blob() + blob_version = self._get_tabular_blob_version( + container_name, + blob_name, + refresh=True, + ) + stream = blob_client.download_blob( + etag=blob_version['blob_etag'], + match_condition=MatchConditions.IfNotModified, + ) data = stream.readall() self._blob_data_cache[cache_key] = data return data + def _get_tabular_blob_version(self, container_name: str, blob_name: str, refresh: bool = False) -> dict: + """Return the exact blob version used by this plugin instance.""" + cache_key = (container_name, blob_name) + if not refresh and cache_key in self._blob_version_cache: + return dict(self._blob_version_cache[cache_key]) + + blob_client = self._get_blob_service_client().get_blob_client( + container=container_name, + blob=blob_name, + ) + blob_properties = blob_client.get_blob_properties() + blob_etag = getattr(blob_properties, 'etag', None) + blob_size = getattr(blob_properties, 'size', None) + if isinstance(blob_properties, dict): + blob_etag = blob_etag or blob_properties.get('etag') + blob_size = blob_size if blob_size is not None else blob_properties.get('size') + if not blob_etag: + raise ValueError('Tabular source version could not be determined') + + blob_version = { + 'blob_etag': str(blob_etag), + 'blob_size': int(blob_size or 0), + } + self._blob_version_cache[cache_key] = blob_version + return dict(blob_version) + def _get_excel_engine(self, blob_name: str) -> Optional[str]: """Return the pandas Excel engine for a workbook, or None for CSV files.""" name_lower = blob_name.lower() @@ -3384,6 +3439,169 @@ def _resolve_blob_location_with_fallback(self, user_id: str, conversation_id: st return attempts[0] raise ValueError(f"Could not resolve blob location for {filename}") + def build_generated_export_query_descriptor( + self, + user_id: str, + conversation_id: str, + filename: str, + query_expression: str, + source: str = 'chat', + return_columns: Optional[str] = None, + group_id: Optional[str] = None, + public_workspace_id: Optional[str] = None, + expected_row_count: int = 0, + ) -> dict: + """Resolve an authorized, version-pinned CSV query for durable export replay.""" + container_name, blob_path = self._resolve_blob_location_with_fallback( + user_id, + conversation_id, + filename, + source, + group_id=group_id, + public_workspace_id=public_workspace_id, + ) + return self._build_generated_export_query_descriptor_from_location( + container_name=container_name, + blob_path=blob_path, + filename=filename, + query_expression=query_expression, + return_columns=return_columns, + expected_row_count=expected_row_count, + ) + + def _build_generated_export_query_descriptor_from_location( + self, + container_name: str, + blob_path: str, + filename: str, + query_expression: str, + return_columns: Optional[str] = None, + expected_row_count: int = 0, + blob_version: Optional[dict] = None, + ) -> dict: + """Pin a durable query descriptor to an already-authorized blob location.""" + if not str(blob_path or '').lower().endswith('.csv'): + raise ValueError('Durable source-backed generated exports currently require a CSV source') + + blob_version = dict(blob_version or self._get_tabular_blob_version(container_name, blob_path)) + blob_etag = blob_version.get('blob_etag') + blob_size = blob_version.get('blob_size') + + resolved_source = self._infer_source_from_container(container_name) + blob_parts = [part for part in str(blob_path or '').split('/') if part] + scope_id = blob_parts[0] if resolved_source in {'group', 'public'} and blob_parts else None + return { + 'version': 1, + 'kind': 'query_tabular_data', + 'source': resolved_source, + 'scope_id': scope_id, + 'container': container_name, + 'blob_path': blob_path, + 'blob_etag': str(blob_etag), + 'blob_size': int(blob_size or 0), + 'filename': str(filename or '').strip(), + 'query_expression': str(query_expression or '').strip(), + 'return_columns': return_columns, + 'expected_row_count': max(0, int(expected_row_count or 0)), + } + + def _build_source_authorization_from_location( + self, + container_name: str, + blob_path: str, + blob_version: Optional[dict] = None, + ) -> dict: + """Build exact server-only scope metadata for later worker revalidation.""" + resolved_source = self._infer_source_from_container(container_name) + blob_parts = [part for part in str(blob_path or '').split('/') if part] + scope_id = blob_parts[0] if resolved_source in {'group', 'public'} and blob_parts else None + return { + 'source': resolved_source, + 'scope_id': scope_id, + 'container': container_name, + 'blob_path': blob_path, + 'blob_etag': (blob_version or self._get_tabular_blob_version(container_name, blob_path)).get( + 'blob_etag' + ), + } + + def _query_csv_data_in_bounded_chunks( + self, + container_name: str, + blob_path: str, + filename: str, + query_expression: str, + return_columns: Optional[str], + start_row, + max_rows, + ) -> PluginInvocationResult: + """Execute foreground CSV pagination through the durable replay query engine.""" + start, limit = self._parse_row_page_arguments(start_row, max_rows) + replay_stats = {'used_reviewer_style_fallback': False} + matched_row_count = 0 + page_rows = [] + blob_client = self._get_blob_service_client().get_blob_client( + container=container_name, + blob=blob_path, + ) + blob_version = self._get_tabular_blob_version( + container_name, + blob_path, + refresh=True, + ) + with tempfile.SpooledTemporaryFile(max_size=1024 * 1024, mode='w+b') as source_stream: + blob_client.download_blob( + etag=blob_version['blob_etag'], + match_condition=MatchConditions.IfNotModified, + ).readinto(source_stream) + source_stream.seek(0) + for _, source_row in iter_tabular_csv_query_rows( + csv_stream=source_stream, + query_expression=query_expression, + return_columns=return_columns, + source_chunk_rows=1000, + tabular_plugin=self, + replay_stats=replay_stats, + ): + if start <= matched_row_count < start + limit: + page_rows.append(source_row) + matched_row_count += 1 + + response_payload = { + 'filename': filename, + 'selected_sheet': None, + 'query_expression': query_expression, + 'query_expression_fallback': replay_stats['used_reviewer_style_fallback'], + 'total_matches': matched_row_count, + } + response_payload.update(self._build_tabular_row_page_payload( + page_rows, + matched_row_count, + start_row=start, + max_rows=limit, + return_columns=self._parse_optional_column_list_argument(return_columns), + )) + source_descriptor = self._build_generated_export_query_descriptor_from_location( + container_name=container_name, + blob_path=blob_path, + filename=filename, + query_expression=query_expression, + return_columns=return_columns, + expected_row_count=matched_row_count, + blob_version=blob_version, + ) + return PluginInvocationResult( + json.dumps(response_payload, indent=2, default=str), + internal_metadata={ + 'tabular_generated_export_source': source_descriptor, + 'tabular_source_authorization': self._build_source_authorization_from_location( + container_name, + blob_path, + blob_version=blob_version, + ), + }, + ) + @kernel_function( description=( "List all tabular data files available for a user. Checks workspace documents " @@ -4462,6 +4680,16 @@ def _sync_work(): user_id, conversation_id, filename, source, group_id=group_id, public_workspace_id=public_workspace_id ) + if str(blob_path or '').lower().endswith('.csv'): + return self._query_csv_data_in_bounded_chunks( + container_name=container, + blob_path=blob_path, + filename=filename, + query_expression=query_expression, + return_columns=return_columns, + start_row=start_row, + max_rows=max_rows, + ) # When no explicit sheet_name is given, try cross-sheet query first normalized_sheet = (sheet_name or '').strip() normalized_sheet_idx = None if sheet_index is None else str(sheet_index).strip() @@ -4510,7 +4738,31 @@ def _sync_work(): max_rows=limit, return_columns=parsed_return_columns, )) - return json.dumps(response_payload, indent=2, default=str) + response_json = json.dumps(response_payload, indent=2, default=str) + source_blob_version = self._get_tabular_blob_version(container, blob_path) + internal_metadata = { + 'tabular_source_authorization': self._build_source_authorization_from_location( + container, + blob_path, + blob_version=source_blob_version, + ), + } + if str(blob_path or '').lower().endswith('.csv'): + internal_metadata['tabular_generated_export_source'] = ( + self._build_generated_export_query_descriptor_from_location( + container_name=container, + blob_path=blob_path, + filename=filename, + query_expression=query_expression, + return_columns=return_columns, + expected_row_count=len(result_df), + blob_version=source_blob_version, + ) + ) + return PluginInvocationResult( + response_json, + internal_metadata=internal_metadata, + ) except Exception as e: log_event(f"[TabularProcessingPlugin] Error querying data: {e}", level=logging.WARNING) return json.dumps({"error": f"Query error: {str(e)}. Ensure column names and values are correct."}) diff --git a/application/single_app/static/js/admin/admin_data_management.js b/application/single_app/static/js/admin/admin_data_management.js index 6476ec031..5c79699f9 100644 --- a/application/single_app/static/js/admin/admin_data_management.js +++ b/application/single_app/static/js/admin/admin_data_management.js @@ -7,6 +7,7 @@ const backupStorageAuthManagedIdentity = "managed_identity"; const backupStorageAuthConnectionString = "connection_string"; const targetCosmosDatabaseName = "SimpleChat"; const cosmosEditorConfirmationPhrase = "I understand this can damage system data"; +const migrationMirrorConfirmationPhrase = "MIRROR WITH DELETIONS"; const elements = {}; let dataManagementModified = false; let storedBackupConnectionStringAvailable = false; @@ -22,8 +23,9 @@ let cosmosEditorContinuationToken = null; let cosmosEditorResultCount = 0; let cosmosEditorSelectedDocument = null; let cosmosEditorPendingDocument = null; +let pendingMigrationCancellationJobId = null; -const jobDetailRefreshIntervalMs = 4000; +const jobDetailRefreshIntervalMs = 2000; const activeJobStatuses = new Set(["queued", "running"]); const migrationTargetTypes = ["users", "groups", "public_workspaces"]; const migrationSelections = { @@ -83,6 +85,8 @@ function bindElements() { "data_management_target_cosmos_database", "data-management-target-cosmos-key-field", "data_management_target_cosmos_key", + "data_management_target_cosmos_subscription_id", + "data_management_target_cosmos_resource_group", "data-management-test-target-cosmos-btn", "data-management-target-ai-search-section", "data_management_target_ai_search_auth", @@ -98,6 +102,15 @@ function bindElements() { "data_management_target_ec_connection_string", "data-management-test-target-ec-storage-btn", "data-management-refresh-migration-summary-btn", + "data-management-migration-mode-section", + "data_management_migration_mode_new_only", + "data_management_migration_mode_delta_upsert", + "data_management_migration_mode_mirror_with_deletions", + "data-management-migration-mode-description", + "data-management-migration-baseline-field", + "data_management_migration_baseline_job_id", + "data-management-migration-mirror-confirmation", + "data_management_migration_mirror_confirmation_phrase", "data_management_migration_users_mode", "data_management_migration_users_search", "data-management-search-users-btn", @@ -117,7 +130,16 @@ function bindElements() { "data-management-migration-public-workspaces-selected", "data_management_migration_public_workspaces_documents", "data_management_migration_include_ai_search", + "data-management-migration-search-write-freeze", + "data_management_migration_target_search_writes_frozen", "data_management_migration_include_source_blobs", + "data-management-test-migration-access-btn", + "data_management_migration_max_parallel_operations", + "data_management_migration_retry_count", + "data_management_migration_skip_recent_within_hours", + "data_management_migration_temporary_destination_ru_enabled", + "data-management-migration-temporary-ru-field", + "data_management_migration_temporary_destination_ru", "data-management-migration-summary", "data-management-migration-preview-btn", "data-management-execute-migration-btn", @@ -139,11 +161,15 @@ function bindElements() { "data-management-job-detail-subtitle", "data-management-job-detail-refresh-state", "data-management-job-detail-summary", + "data-management-job-detail-actions", "data-management-job-detail-progress", "data-management-job-items-tbody", "data-management-job-artifacts-tbody", "data-management-job-manifest-detail", "data-management-job-warnings", + "data-management-migration-cancel-modal", + "data-management-migration-cancel-message", + "data-management-confirm-migration-cancel-btn", "data-management-cosmos-editor-section", "data-management-cosmos-editor-open-danger-btn", "data-management-cosmos-editor-locked-message", @@ -205,6 +231,7 @@ function bindEvents() { elements.dataManagementGenerateKeyBtn?.addEventListener("click", generateEncryptionKey); elements.dataManagementTestStorageBtn?.addEventListener("click", testBackupStorage); elements.dataManagementTestTargetCosmosBtn?.addEventListener("click", testTargetCosmos); + elements.dataManagementTestMigrationAccessBtn?.addEventListener("click", testMigrationAccess); elements.dataManagementTestTargetSearchBtn?.addEventListener("click", testTargetSearch); elements.dataManagementTestTargetEcStorageBtn?.addEventListener("click", testTargetEnhancedCitationStorage); elements.dataManagementRunFullBackupBtn?.addEventListener("click", () => queueBackup("full")); @@ -219,6 +246,10 @@ function bindEvents() { elements.dataManagementViewAllBackupsBtn?.addEventListener("click", () => setBackupFilter("all")); elements.dataManagementRefreshJobsBtn?.addEventListener("click", loadDataManagementJobs); elements.dataManagementJobDetailModal?.addEventListener("hidden.bs.modal", () => stopJobDetailAutoRefresh({ clearJob: true })); + elements.dataManagementMigrationCancelModal?.addEventListener("hidden.bs.modal", () => { + pendingMigrationCancellationJobId = null; + }); + elements.dataManagementConfirmMigrationCancelBtn?.addEventListener("click", requestMigrationCancellation); elements.dataManagementKeyVaultLink?.addEventListener("click", openKeyVaultSettings); elements.dataManagementCosmosEditorOpenDangerBtn?.addEventListener("click", showCosmosEditorDangerModal); elements.datamanagementcosmoseditordangeraccept?.addEventListener("change", updateCosmosEditorDangerAcceptState); @@ -231,9 +262,25 @@ function bindEvents() { elements.dataManagementCosmosEditorSaveBtn?.addEventListener("click", openCosmosEditorSaveModal); elements.datamanagementcosmoseditorconfirmationphrase?.addEventListener("input", updateCosmosEditorConfirmSaveState); elements.dataManagementCosmosEditorConfirmSaveBtn?.addEventListener("click", saveCosmosEditorDocument); + elements.datamanagementmigrationtemporarydestinationruenabled?.addEventListener("change", updateMigrationCapacityVisibility); + [ + elements.datamanagementmigrationmodenewonly, + elements.datamanagementmigrationmodedeltaupsert, + elements.datamanagementmigrationmodemirrorwithdeletions, + ].forEach((modeElement) => modeElement?.addEventListener("change", () => { + updateMigrationModeVisibility(); + loadMigrationSummary(); + })); + elements.datamanagementmigrationbaselinejobid?.addEventListener("change", loadMigrationSummary); + elements.datamanagementmigrationmirrorconfirmationphrase?.addEventListener("input", () => { + elements.datamanagementmigrationmirrorconfirmationphrase?.classList.remove("is-invalid"); + }); bindDataManagementChangeTracking(); setStorageAuthVisibility(); setMigrationTargetVisibility(); + updateMigrationCapacityVisibility(); + updateMigrationModeVisibility(); + updateMigrationSearchWriteFreezeVisibility(); updateConnectionStringStatus(); updateDataManagementSaveButtonState(); } @@ -253,7 +300,10 @@ function bindMigrationPickerEvents() { } }); }); - elements.datamanagementmigrationincludeaisearch?.addEventListener("change", loadMigrationSummary); + elements.datamanagementmigrationincludeaisearch?.addEventListener("change", () => { + updateMigrationSearchWriteFreezeVisibility(); + loadMigrationSummary(); + }); elements.datamanagementmigrationincludesourceblobs?.addEventListener("change", loadMigrationSummary); } @@ -419,12 +469,19 @@ function populateSettings(settings) { setValue(elements.datamanagementtargetcosmosendpoint, settings.target_cosmos_endpoint || ""); setValue(elements.datamanagementtargetcosmosdatabase, targetCosmosDatabaseName); setValue(elements.datamanagementtargetcosmoskey, settings.target_cosmos_key || ""); + setValue(elements.datamanagementtargetcosmossubscriptionid, settings.target_cosmos_subscription_id || ""); + setValue(elements.datamanagementtargetcosmosresourcegroup, settings.target_cosmos_resource_group || ""); setValue(elements.datamanagementtargetaisearchauth, settings.target_ai_search_authentication_type || "managed_identity"); setValue(elements.datamanagementtargetaisearchendpoint, settings.target_ai_search_endpoint || ""); setValue(elements.datamanagementtargetaisearchkey, settings.target_ai_search_key || ""); setValue(elements.datamanagementtargetecstorageauth, settings.target_enhanced_citations_storage_authentication_type || "managed_identity"); setValue(elements.datamanagementtargetecblobendpoint, settings.target_enhanced_citations_storage_blob_endpoint || ""); setValue(elements.datamanagementtargetecconnectionstring, settings.target_enhanced_citations_storage_connection_string || ""); + setValue(elements.datamanagementmigrationmaxparalleloperations, settings.migration_max_parallel_operations || 8); + setValue(elements.datamanagementmigrationretrycount, settings.migration_retry_count || 5); + setValue(elements.datamanagementmigrationskiprecentwithinhours, settings.migration_skip_recent_within_hours || 0); + setChecked(elements.datamanagementmigrationtemporarydestinationruenabled, Boolean(settings.migration_temporary_destination_ru_enabled)); + setValue(elements.datamanagementmigrationtemporarydestinationru, settings.migration_temporary_destination_ru || 10000); if (elements.dataManagementKeyStorage) { elements.dataManagementKeyStorage.textContent = formatKeyStorage(settings.encryption_key_storage); @@ -436,6 +493,7 @@ function populateSettings(settings) { updateKeyStorageExperience(settings); setStorageAuthVisibility(); setMigrationTargetVisibility(); + updateMigrationCapacityVisibility(); migrationTargetTypes.forEach(updateMigrationPickerVisibility); updateConnectionStringStatus(); resetDataManagementModified(); @@ -529,6 +587,14 @@ function setMigrationTargetVisibility() { } } +function updateMigrationCapacityVisibility() { + const enabled = Boolean(elements.datamanagementmigrationtemporarydestinationruenabled?.checked); + setElementVisible(elements.dataManagementMigrationTemporaryRuField, enabled); + if (elements.datamanagementmigrationtemporarydestinationru) { + elements.datamanagementmigrationtemporarydestinationru.disabled = !enabled; + } +} + function setKeyStorageAlert(variant, iconClass, title, message, linkText) { const alertElement = elements.dataManagementKeyStorageAlert; const iconElement = elements.dataManagementKeyStorageAlertIcon; @@ -581,12 +647,19 @@ function collectSettings() { target_cosmos_endpoint: getValue(elements.datamanagementtargetcosmosendpoint), target_cosmos_database_name: targetCosmosDatabaseName, target_cosmos_key: getValue(elements.datamanagementtargetcosmoskey), + target_cosmos_subscription_id: getValue(elements.datamanagementtargetcosmossubscriptionid), + target_cosmos_resource_group: getValue(elements.datamanagementtargetcosmosresourcegroup), target_ai_search_authentication_type: getValue(elements.datamanagementtargetaisearchauth) || "managed_identity", target_ai_search_endpoint: getValue(elements.datamanagementtargetaisearchendpoint), target_ai_search_key: getValue(elements.datamanagementtargetaisearchkey), target_enhanced_citations_storage_authentication_type: getValue(elements.datamanagementtargetecstorageauth) || backupStorageAuthManagedIdentity, target_enhanced_citations_storage_blob_endpoint: getValue(elements.datamanagementtargetecstorageauth) === backupStorageAuthManagedIdentity ? getValue(elements.datamanagementtargetecblobendpoint) : "", target_enhanced_citations_storage_connection_string: getValue(elements.datamanagementtargetecstorageauth) === backupStorageAuthConnectionString ? getValue(elements.datamanagementtargetecconnectionstring) : "", + migration_max_parallel_operations: getNumberValue(elements.datamanagementmigrationmaxparalleloperations, 8), + migration_retry_count: getNumberValue(elements.datamanagementmigrationretrycount, 5), + migration_skip_recent_within_hours: getNumberValue(elements.datamanagementmigrationskiprecentwithinhours, 0), + migration_temporary_destination_ru_enabled: Boolean(elements.datamanagementmigrationtemporarydestinationruenabled?.checked), + migration_temporary_destination_ru: getNumberValue(elements.datamanagementmigrationtemporarydestinationru, 10000), }; } @@ -1165,9 +1238,12 @@ async function testTargetCosmos() { try { const data = await requestJson("/api/admin/data-management/target/cosmos/test", { method: "POST", - body: JSON.stringify({ settings: collectSettings() }), + body: JSON.stringify({ settings: collectSettings(), migration_plan: buildMigrationPlan() }), }); - setStatus(`Target Cosmos connection succeeded. Database: ${data.database_name || targetCosmosDatabaseName}.`, "success"); + const verifiedCount = Number(data.migration_access?.container_count || 0); + const accessText = verifiedCount ? ` Verified ${formatNumber(verifiedCount)} planned container(s).` : ""; + const capacityText = data.capacity?.target_ru ? ` Temporary capacity can reach ${formatNumber(data.capacity.target_ru)} RU/s.` : ""; + setStatus(`Target Cosmos connection succeeded. Database: ${data.database_name || targetCosmosDatabaseName}.${accessText}${capacityText}`, "success"); showToast("Target Cosmos connection succeeded.", "success"); } catch (error) { setStatus(error.message || "Target Cosmos connection test failed.", "danger"); @@ -1177,6 +1253,25 @@ async function testTargetCosmos() { } } +async function testMigrationAccess() { + setBusy(elements.dataManagementTestMigrationAccessBtn, true, "Validating..."); + try { + const data = await requestJson("/api/admin/data-management/target/cosmos/test", { + method: "POST", + body: JSON.stringify({ settings: collectSettings(), migration_plan: buildMigrationPlan() }), + }); + const verifiedCount = Number(data.migration_access?.container_count || 0); + const capacityText = data.capacity?.target_ru ? ` Destination capacity management is ready up to ${formatNumber(data.capacity.target_ru)} RU/s.` : ""; + setStatus(`Cosmos migration access validation succeeded. ${formatNumber(verifiedCount)} planned Cosmos container(s) can be read and written.${capacityText}`, "success"); + showToast("Cosmos migration access validation succeeded.", "success"); + } catch (error) { + setStatus(error.message || "Cosmos migration access validation failed.", "danger"); + showToast(error.message || "Cosmos migration access validation failed.", "danger"); + } finally { + setBusy(elements.dataManagementTestMigrationAccessBtn, false); + } +} + async function testTargetSearch() { setBusy(elements.dataManagementTestTargetSearchBtn, true, "Testing..."); try { @@ -1247,9 +1342,31 @@ async function queueOperation(operation, backupType = null, options = {}, trigge function queueMigration(dryRun) { const button = elements.dataManagementExecuteMigrationBtn; + const migrationPlan = buildMigrationPlan(); + if ( + migrationPlan.include_ai_search && + !migrationPlan.target_ai_search_writes_frozen + ) { + const message = "Confirm that external destination AI Search writers are frozen before migrating Search documents."; + elements.datamanagementmigrationtargetsearchwritesfrozen?.focus(); + setStatus(message, "danger"); + showToast(message, "danger"); + return Promise.resolve(); + } + if ( + migrationPlan.migration_mode === "mirror_with_deletions" && + migrationPlan.mirror_confirmation !== migrationMirrorConfirmationPhrase + ) { + const message = `Type ${migrationMirrorConfirmationPhrase} before running a mirror migration.`; + elements.datamanagementmigrationmirrorconfirmationphrase?.classList.add("is-invalid"); + elements.datamanagementmigrationmirrorconfirmationphrase?.focus(); + setStatus(message, "danger"); + showToast(message, "danger"); + return Promise.resolve(); + } return queueOperation("migration", null, { dry_run: Boolean(dryRun), - migration_plan: buildMigrationPlan(), + migration_plan: migrationPlan, }, button); } @@ -1403,13 +1520,46 @@ function createSmallMutedElement(text) { } function buildMigrationPlan() { + const migrationMode = getMigrationTransferMode(); return { users: buildMigrationSelection("users"), groups: buildMigrationSelection("groups"), public_workspaces: buildMigrationSelection("public_workspaces"), include_ai_search: Boolean(elements.datamanagementmigrationincludeaisearch?.checked), + target_ai_search_writes_frozen: Boolean(elements.datamanagementmigrationtargetsearchwritesfrozen?.checked), include_source_blobs: Boolean(elements.datamanagementmigrationincludesourceblobs?.checked), + migration_mode: migrationMode, + baseline_job_id: migrationMode === "new_only" ? "" : getValue(elements.datamanagementmigrationbaselinejobid), + mirror_confirmation: migrationMode === "mirror_with_deletions" ? getValue(elements.datamanagementmigrationmirrorconfirmationphrase) : "", + }; +} + +function updateMigrationSearchWriteFreezeVisibility() { + const includesAiSearch = Boolean(elements.datamanagementmigrationincludeaisearch?.checked); + setElementVisible(elements.dataManagementMigrationSearchWriteFreeze, includesAiSearch); + if (!includesAiSearch && elements.datamanagementmigrationtargetsearchwritesfrozen) { + elements.datamanagementmigrationtargetsearchwritesfrozen.checked = false; + } +} + +function getMigrationTransferMode() { + return document.querySelector("input[name='data_management_migration_mode']:checked")?.value || "new_only"; +} + +function updateMigrationModeVisibility() { + const migrationMode = getMigrationTransferMode(); + const usesBaseline = migrationMode !== "new_only"; + const isMirror = migrationMode === "mirror_with_deletions"; + elements.dataManagementMigrationBaselineField?.classList.toggle("d-none", !usesBaseline); + elements.dataManagementMigrationMirrorConfirmation?.classList.toggle("d-none", !isMirror); + const descriptions = { + new_only: "Copies source items that are absent from the destination. Existing destination data is never updated or deleted.", + delta_upsert: "Copies new items and updates changed migration-owned items since the prior successful watermark. Destination-only data is retained.", + mirror_with_deletions: "Runs delta/upsert, then deletes destination-only items that carry successful SimpleChat migration ownership. Unowned data is retained.", }; + if (elements.dataManagementMigrationModeDescription) { + elements.dataManagementMigrationModeDescription.textContent = descriptions[migrationMode] || descriptions.new_only; + } } function buildMigrationSelection(targetType) { @@ -1430,9 +1580,12 @@ async function loadMigrationSummary(triggerButton = null, showSuccess = false) { try { const data = await requestJson("/api/admin/data-management/migration/summary", { method: "POST", - body: JSON.stringify({ migration_plan: buildMigrationPlan() }), + body: JSON.stringify({ + migration_plan: buildMigrationPlan(), + include_inventory: Boolean(showSuccess), + }), }); - renderMigrationSummary(data.summary || {}); + renderMigrationSummary(data.summary || {}, data.preview || null); if (showSuccess) { setStatus("Migration preview refreshed.", "success"); showToast("Migration preview refreshed.", "success"); @@ -1448,7 +1601,7 @@ async function loadMigrationSummary(triggerButton = null, showSuccess = false) { } } -function renderMigrationSummary(summary) { +function renderMigrationSummary(summary, preview = null) { const summaryElement = elements.dataManagementMigrationSummary; if (!summaryElement) { return; @@ -1465,11 +1618,42 @@ function renderMigrationSummary(summary) { options.className = "d-flex flex-wrap gap-2"; options.appendChild(createBadge(summary.include_ai_search ? "AI Search included" : "AI Search skipped", summary.include_ai_search ? "bg-info text-dark" : "bg-secondary")); options.appendChild(createBadge(summary.include_source_blobs ? "Source blobs included" : "Source blobs skipped", summary.include_source_blobs ? "bg-info text-dark" : "bg-secondary")); + options.appendChild(createBadge(`Mode: ${formatActivityLabel(summary.migration_mode || "new_only")}`, summary.migration_mode === "mirror_with_deletions" ? "bg-danger" : "bg-primary")); optionsColumn.appendChild(options); container.appendChild(optionsColumn); + if (preview?.estimated_outcomes) { + container.appendChild(createMigrationPreviewOutcomes(preview)); + } summaryElement.replaceChildren(container); } +function createMigrationPreviewOutcomes(preview) { + const column = document.createElement("div"); + column.className = "col-12 border-top pt-3"; + const heading = document.createElement("div"); + heading.className = "fw-semibold mb-2"; + heading.textContent = "Estimated destination changes"; + const outcomes = document.createElement("div"); + outcomes.className = "d-flex flex-wrap gap-2"; + [ + ["Create", "create_count", "bg-success"], + ["Update", "update_count", "bg-primary"], + ["Unchanged", "unchanged_count", "bg-secondary"], + ["Delete", "delete_count", "bg-danger"], + ["Not applicable", "not_applicable_count", "bg-light text-dark border"], + ["Missing", "missing_count", "bg-warning text-dark"], + ["Conflicts", "conflict_count", "bg-warning text-dark"], + ].forEach(([label, fieldName, badgeClass]) => { + outcomes.appendChild(createBadge(`${label}: ${formatNumber(preview.estimated_outcomes[fieldName] || 0)}`, badgeClass)); + }); + const metadata = document.createElement("div"); + metadata.className = "small text-muted mt-2"; + const baselineText = preview.baseline_job_id ? ` Baseline ${preview.baseline_job_id}.` : ""; + metadata.textContent = `Captured ${formatDate(preview.captured_at)}.${baselineText} Actual completion counts will report any divergence from this snapshot.`; + column.append(heading, outcomes, metadata); + return column; +} + function createMigrationSummaryCard(targetType, targetSummary) { const column = document.createElement("div"); column.className = "col-md-4"; @@ -1707,7 +1891,7 @@ function createJobRow(job) { row.appendChild(createStatusCell(job.status || "unknown")); row.appendChild(createCell(formatProgress(job.progress))); row.appendChild(createCell(job.last_message || job.last_error || "")); - row.appendChild(createJobActionCell(job.id)); + row.appendChild(createJobActionCell(job.id, job)); return row; } @@ -1721,7 +1905,7 @@ function createMessageRow(message, variant, columnSpan) { return row; } -function createJobActionCell(jobId) { +function createJobActionCell(jobId, job = {}) { const cell = document.createElement("td"); const button = document.createElement("button"); button.type = "button"; @@ -1733,9 +1917,98 @@ function createJobActionCell(jobId) { button.appendChild(document.createTextNode("View Log")); button.addEventListener("click", () => loadDataManagementJobDetail(jobId)); cell.appendChild(button); + if (job.operation === "migration" && job.can_retry === true) { + const retryButton = document.createElement("button"); + retryButton.type = "button"; + retryButton.className = "btn btn-outline-warning btn-sm ms-1"; + retryButton.disabled = !jobId; + const retryIcon = document.createElement("i"); + retryIcon.className = "bi bi-arrow-repeat me-1"; + retryButton.append(retryIcon, document.createTextNode(job.status === "running" ? "Resume" : "Retry")); + retryButton.addEventListener("click", () => retryMigrationJob(jobId, retryButton)); + cell.appendChild(retryButton); + } + if (job.operation === "migration" && job.can_cancel === true) { + const cancelButton = document.createElement("button"); + cancelButton.type = "button"; + cancelButton.className = "btn btn-outline-danger btn-sm ms-1"; + cancelButton.disabled = !jobId; + const cancelIcon = document.createElement("i"); + cancelIcon.className = "bi bi-stop-circle me-1"; + cancelButton.append(cancelIcon, document.createTextNode("Cancel")); + cancelButton.addEventListener("click", () => openMigrationCancellationModal(job)); + cell.appendChild(cancelButton); + } return cell; } +function openMigrationCancellationModal(job) { + if (!job?.id || !elements.dataManagementMigrationCancelModal || !window.bootstrap?.Modal) { + return; + } + pendingMigrationCancellationJobId = job.id; + setText( + elements.dataManagementMigrationCancelMessage, + `Request cancellation for migration ${job.id}?` + ); + window.bootstrap.Modal.getOrCreateInstance(elements.dataManagementMigrationCancelModal).show(); +} + +async function requestMigrationCancellation() { + const jobId = pendingMigrationCancellationJobId; + if (!jobId) { + return; + } + setBusy(elements.dataManagementConfirmMigrationCancelBtn, true, "Requesting..."); + try { + const data = await requestJson( + `/api/admin/data-management/jobs/${encodeURIComponent(jobId)}/cancel`, + { method: "POST", body: JSON.stringify({}) } + ); + const canceledImmediately = data.job?.status === "canceled"; + setStatus( + canceledImmediately + ? "Queued migration canceled before execution started." + : "Migration cancellation requested. The worker will stop at its next durable checkpoint.", + "success" + ); + showToast( + canceledImmediately ? "Migration canceled." : "Migration cancellation requested.", + "success" + ); + window.bootstrap.Modal.getOrCreateInstance(elements.dataManagementMigrationCancelModal).hide(); + renderJobs([data.job]); + loadDataManagementJobs(); + if (currentJobDetailId === jobId) { + loadDataManagementJobDetail(jobId, { showModal: false, liveRefresh: true }); + } + } catch (error) { + setStatus(error.message || "Migration cancellation could not be requested.", "danger"); + showToast(error.message || "Migration cancellation could not be requested.", "danger"); + } finally { + setBusy(elements.dataManagementConfirmMigrationCancelBtn, false); + } +} + +async function retryMigrationJob(jobId, button) { + setBusy(button, true, "Retrying..."); + try { + const data = await requestJson(`/api/admin/data-management/jobs/${encodeURIComponent(jobId)}/retry`, { + method: "POST", + }); + setStatus("Migration retry queued from durable checkpoints.", "success"); + showToast("Migration retry queued.", "success"); + renderJobs([data.job]); + loadDataManagementJobs(); + loadDataManagementJobDetail(jobId, { showModal: false, liveRefresh: true }); + } catch (error) { + setStatus(error.message || "Migration retry could not be queued.", "danger"); + showToast(error.message || "Migration retry could not be queued.", "danger"); + } finally { + setBusy(button, false); + } +} + function createCell(text) { const cell = document.createElement("td"); cell.textContent = text ?? ""; @@ -1777,6 +2050,7 @@ function renderJobDetailModal(job, items) { setText(elements.dataManagementJobDetailTitle, formatOperation(job.operation || "job", job.backup_type || "")); setText(elements.dataManagementJobDetailSubtitle, job.id ? `Job ID: ${job.id}` : "Job ID not available"); renderJobDetailSummary(job); + renderJobDetailActions(job); renderJobDetailProgress(job); renderJobItems(items); const artifacts = getJobArtifacts(job, items); @@ -1785,12 +2059,82 @@ function renderJobDetailModal(job, items) { renderJobWarnings(job, artifacts); } +function renderJobDetailActions(job) { + const container = elements.dataManagementJobDetailActions; + if (!container) { + return; + } + container.replaceChildren(); + if (job.operation !== "migration") { + container.classList.add("d-none"); + return; + } + const fullManifestLink = createMigrationManifestDownloadLink( + job.id, + "Download manifest", + "", + "btn-outline-primary" + ); + const failuresLink = createMigrationManifestDownloadLink( + job.id, + "Download failures", + "failed,missing,collision", + "btn-outline-danger" + ); + container.append(fullManifestLink, failuresLink); + if (job.can_retry === true) { + const retryButton = document.createElement("button"); + retryButton.type = "button"; + retryButton.className = "btn btn-outline-warning btn-sm"; + const retryIcon = document.createElement("i"); + retryIcon.className = "bi bi-arrow-repeat me-1"; + retryButton.append( + retryIcon, + document.createTextNode(job.status === "running" ? "Resume" : "Retry") + ); + retryButton.addEventListener("click", () => retryMigrationJob(job.id, retryButton)); + container.appendChild(retryButton); + } + if (job.can_cancel === true) { + const cancelButton = document.createElement("button"); + cancelButton.type = "button"; + cancelButton.className = "btn btn-outline-danger btn-sm"; + const cancelIcon = document.createElement("i"); + cancelIcon.className = "bi bi-stop-circle me-1"; + cancelButton.append(cancelIcon, document.createTextNode("Cancel")); + cancelButton.addEventListener("click", () => openMigrationCancellationModal(job)); + container.appendChild(cancelButton); + } + container.classList.toggle("d-none", container.childElementCount === 0); +} + +function createMigrationManifestDownloadLink(jobId, label, statuses, buttonClass) { + const link = document.createElement("a"); + link.className = `btn ${buttonClass} btn-sm`; + const encodedJobId = encodeURIComponent(jobId || ""); + const query = statuses ? `?statuses=${encodeURIComponent(statuses)}` : ""; + link.href = `/api/admin/data-management/jobs/${encodedJobId}/migration-manifest${query}`; + link.download = ""; + const icon = document.createElement("i"); + icon.className = "bi bi-download me-1"; + link.append(icon, document.createTextNode(label)); + return link; +} + function renderJobDetailSummary(job) { const container = elements.dataManagementJobDetailSummary; if (!container) { return; } const result = job.result && typeof job.result === "object" ? job.result : {}; + const migrationState = job.migration_state && typeof job.migration_state === "object" ? job.migration_state : {}; + const migrationTotals = migrationState.totals && typeof migrationState.totals === "object" ? migrationState.totals : {}; + const migrationTiles = job.operation === "migration" ? [ + createSummaryTile("Migration ID", migrationState.migration_id || job.id || "N/A"), + createSummaryTile("Processed", formatNumber(migrationTotals.processed_count || 0)), + createSummaryTile("Transferred", formatBytes(migrationTotals.bytes || 0)), + createSummaryTile("Request units", formatNumber(migrationTotals.request_units || 0)), + ] : []; container.replaceChildren( createSummaryTile("Status", createStatusBadge(job.status || "unknown")), createSummaryTile("Progress", formatProgress(job.progress)), @@ -1800,7 +2144,8 @@ function renderJobDetailSummary(job) { createSummaryTile("Created", formatDate(job.created_at)), createSummaryTile("Started", formatDate(job.started_at)), createSummaryTile("Completed", formatDate(job.completed_at)), - createSummaryTile("Artifacts", formatNumber(result.artifact_count || getArtifactTotals(result.artifacts).artifactCount || 0)) + createSummaryTile("Artifacts", formatNumber(result.artifact_count || getArtifactTotals(result.artifacts).artifactCount || 0)), + ...migrationTiles ); } @@ -1830,6 +2175,11 @@ function renderJobDetailProgress(job) { const percent = getProgressPercent(progress); const completedSteps = Number(progress.completed_steps || 0); const totalSteps = Number(progress.total_steps || 0); + const isMigration = job.operation === "migration"; + const isMigrationActive = isMigration && activeJobStatuses.has(job.status); + const displayedStage = isMigrationActive + ? Math.min(totalSteps, completedSteps + 1) + : completedSteps; const card = document.createElement("div"); card.className = "card border-info-subtle"; const body = document.createElement("div"); @@ -1842,29 +2192,124 @@ function renderJobDetailProgress(job) { title.textContent = progress.current_step ? formatActivityLabel(progress.current_step) : "Job progress"; const stepText = document.createElement("small"); stepText.className = "text-muted"; - stepText.textContent = totalSteps > 0 ? `${formatNumber(completedSteps)} of ${formatNumber(totalSteps)} steps complete` : "Waiting for step details"; + stepText.textContent = isMigration + ? totalSteps > 0 ? `Stage ${formatNumber(displayedStage)} of ${formatNumber(totalSteps)}` : "Waiting for stage details" + : totalSteps > 0 ? `${formatNumber(completedSteps)} of ${formatNumber(totalSteps)} steps complete` : "Waiting for step details"; header.append(title, stepText); const progressWrapper = document.createElement("div"); progressWrapper.className = "progress"; progressWrapper.setAttribute("role", "progressbar"); - progressWrapper.setAttribute("aria-valuenow", String(percent)); progressWrapper.setAttribute("aria-valuemin", "0"); progressWrapper.setAttribute("aria-valuemax", "100"); const progressBar = document.createElement("div"); - progressBar.className = `progress-bar ${job.status === "failed" ? "bg-danger" : "bg-primary"}`; - progressBar.style.width = `${percent}%`; - progressBar.textContent = `${percent}%`; + if (isMigrationActive) { + progressWrapper.setAttribute("aria-valuetext", "Migration stage is active; measured throughput is shown below."); + progressBar.className = "progress-bar progress-bar-striped progress-bar-animated bg-primary"; + progressBar.style.width = "100%"; + progressBar.textContent = "Active stage"; + } else { + progressWrapper.setAttribute("aria-valuenow", String(percent)); + progressBar.className = `progress-bar ${job.status === "failed" ? "bg-danger" : "bg-primary"}`; + progressBar.style.width = `${percent}%`; + progressBar.textContent = `${percent}%`; + } progressWrapper.appendChild(progressBar); const message = document.createElement("p"); message.className = "text-muted small mt-2 mb-0"; message.textContent = job.last_message || job.last_error || "No job message has been recorded yet."; body.append(header, progressWrapper, message); + if (isMigration) { + const migrationMetrics = getMigrationLiveMetrics(job); + if (migrationMetrics.length) { + const metricGrid = document.createElement("div"); + metricGrid.className = "row g-2 small mt-3"; + migrationMetrics.forEach((metric) => { + metricGrid.appendChild(createDetailMetric(metric.label, metric.value)); + }); + body.appendChild(metricGrid); + } + } card.appendChild(body); container.replaceChildren(card); } +function getMigrationLiveMetrics(job) { + const migrationState = job?.migration_state; + if (!migrationState || typeof migrationState !== "object") { + return []; + } + const totals = migrationState.totals && typeof migrationState.totals === "object" ? migrationState.totals : {}; + const resources = migrationState.resources && typeof migrationState.resources === "object" ? Object.values(migrationState.resources) : []; + const activeResource = resources.find((resource) => resource?.status === "in_progress") || {}; + const activeMetrics = activeResource.progress && typeof activeResource.progress === "object" ? activeResource.progress : activeResource.result || {}; + const capacity = migrationState.capacity && typeof migrationState.capacity === "object" ? migrationState.capacity : {}; + const metrics = [ + { label: "Processed", value: formatNumber(totals.processed_count || 0) }, + { label: "Transferred", value: formatBytes(totals.bytes || 0) }, + { label: "Copied / skipped", value: `${formatNumber(totals.copied_count || 0)} / ${formatNumber(totals.skipped_count || 0)}` }, + { label: "Failures", value: formatNumber(totals.failed_count || 0) }, + { label: "Collisions", value: formatNumber(totals.collision_count || 0) }, + ]; + if (activeMetrics.observed_bytes !== undefined) { + metrics.push({ label: "Observed transferred", value: formatBytes(activeMetrics.observed_bytes || 0) }); + } + if (activeMetrics.items_per_second !== undefined) { + metrics.push({ label: "Items per second", value: formatNumber(activeMetrics.items_per_second) }); + } + const observedRate = activeMetrics.observed_bytes_per_second ?? activeMetrics.bytes_per_second; + if (observedRate !== undefined) { + metrics.push({ label: "Transfer rate", value: `${formatBytes(observedRate || 0)}/s` }); + } + if (activeMetrics.request_units_per_second !== undefined) { + metrics.push({ label: "RU per second", value: formatNumber(activeMetrics.request_units_per_second) }); + } + if (capacity.status) { + metrics.push({ label: "Destination capacity", value: formatActivityLabel(capacity.status) }); + } + const heartbeatAge = formatTimestampAge(job?.last_heartbeat_at); + const progressAge = formatTimestampAge(job?.last_progress_at || migrationState.last_progress_at); + if (heartbeatAge) { + metrics.push({ label: "Last heartbeat", value: heartbeatAge }); + } + if (progressAge) { + metrics.push({ label: "Last progress", value: progressAge }); + } + if (heartbeatAge) { + const hasRecentProgress = timestampAgeSeconds( + job?.last_progress_at || migrationState.last_progress_at + ) <= 10; + metrics.push({ + label: "Liveness", + value: hasRecentProgress ? "Running - progress active" : "Running - alive, no recent progress", + }); + } + return metrics; +} + +function timestampAgeSeconds(value) { + const timestamp = Date.parse(value || ""); + if (!Number.isFinite(timestamp)) { + return Number.POSITIVE_INFINITY; + } + return Math.max(0, Math.floor((Date.now() - timestamp) / 1000)); +} + +function formatTimestampAge(value) { + const ageSeconds = timestampAgeSeconds(value); + if (!Number.isFinite(ageSeconds)) { + return ""; + } + if (ageSeconds < 60) { + return `${ageSeconds}s ago`; + } + if (ageSeconds < 3600) { + return `${Math.floor(ageSeconds / 60)}m ago`; + } + return `${Math.floor(ageSeconds / 3600)}h ago`; +} + function renderJobItems(items) { const container = elements.dataManagementJobItemsTbody; if (!container) { @@ -1956,15 +2401,41 @@ function createArtifactCard(artifact) { const metrics = document.createElement("div"); metrics.className = "row g-2 small mb-2"; + const copiedCount = artifact.copied_count ?? artifact.item_count ?? 0; metrics.append( - createDetailMetric("Records", formatNumber(artifact.item_count || 0)), + createDetailMetric("Copied", formatNumber(copiedCount)), + createDetailMetric("Created", formatNumber(artifact.created_count || 0)), + createDetailMetric("Updated", formatNumber(artifact.updated_count || 0)), + createDetailMetric("Unchanged", formatNumber(artifact.unchanged_count || 0)), + createDetailMetric("Skipped", formatNumber(artifact.skipped_count || 0)), + createDetailMetric("Failed", formatNumber(artifact.failed_count || 0)), + createDetailMetric("Collisions", formatNumber(artifact.collision_count || 0)), createDetailMetric("Blobs", formatNumber(artifact.blob_count || 0)), createDetailMetric("Size", formatBytes(artifact.bytes || 0)) ); + if (artifact.type === "migration_reconciliation") { + metrics.append( + createDetailMetric("Readiness", formatActivityLabel(artifact.readiness || "unknown")), + createDetailMetric("Deleted", formatNumber(artifact.deleted_count || 0)), + createDetailMetric("Delete candidates", formatNumber(artifact.delete_candidate_count || 0)) + ); + } + if (artifact.items_per_second !== undefined || artifact.bytes_per_second !== undefined) { + metrics.append( + createDetailMetric("Items / second", formatNumber(artifact.items_per_second || 0)), + createDetailMetric("Transfer rate", `${formatBytes(artifact.bytes_per_second || 0)}/s`) + ); + } const location = createDetailBlock("Location", artifact.path || artifact.prefix || "Not recorded"); const notes = createArtifactNoteGroup(artifact); body.append(header, metrics, location, notes); + if (artifact.preview_actual_divergence && typeof artifact.preview_actual_divergence === "object") { + const divergenceTitle = document.createElement("div"); + divergenceTitle.className = "text-muted small mt-2"; + divergenceTitle.textContent = "Preview versus actual"; + body.append(divergenceTitle, createDetailChipGroup(artifact.preview_actual_divergence)); + } card.appendChild(body); return card; } @@ -2094,6 +2565,24 @@ function createArtifactNoteGroup(artifact) { { label: "Index", value: artifact.index_name }, { label: "Partial since", value: artifact.partial_since_epoch }, { label: "Filter", value: artifact.partial_filter }, + { label: "Request units", value: artifact.request_units }, + { label: "RU per second", value: artifact.request_units_per_second }, + { label: "Workers", value: artifact.parallel_operations }, + { label: "Retries", value: artifact.retry_count }, + { label: "Retry delays", value: artifact.retry_attempt_count }, + { label: "Prior failures", value: artifact.prior_failed_count }, + { label: "Checkpoints", value: artifact.checkpoint_count }, + { label: "Readiness", value: artifact.readiness, className: artifact.readiness === "ready" ? "bg-success" : "bg-warning text-dark" }, + { label: "Deletion", value: artifact.deletion_status }, + { label: "Remaining owned extras", value: artifact.remaining_destination_only_owned_count }, + { label: "Unowned extras", value: artifact.destination_only_unowned_count }, + { label: "Unresolved scope", value: artifact.unresolved_scope_count }, + { label: "Stale", value: artifact.stale_count }, + { label: "Deletion blockers", value: Array.isArray(artifact.deletion_blockers) ? artifact.deletion_blockers.join("; ") : artifact.deletion_blockers, className: "bg-warning text-dark" }, + { label: "Prior migration skips", value: artifact.destination_provenance_skip_count }, + { label: "Missing blobs", value: artifact.missing_count }, + { label: "No source blob", value: artifact.not_applicable_count }, + { label: "Collisions", value: artifact.collision_count }, { label: "Warning", value: artifact.warning, className: "bg-warning text-dark" }, ].filter((note) => isPresent(note.value)); @@ -2184,7 +2673,7 @@ function getArtifactTotals(artifacts) { } totals.artifactCount += 1; totals.bytes += Number(artifact.bytes || 0); - totals.recordCount += Number(artifact.item_count || 0); + totals.recordCount += Number(artifact.item_count ?? artifact.copied_count ?? 0); totals.blobCount += Number(artifact.blob_count || 0); }); return totals; @@ -2261,7 +2750,23 @@ async function refreshOpenJobDetail() { } jobDetailRefreshInFlight = true; try { - await loadDataManagementJobDetail(currentJobDetailId, { showModal: false, liveRefresh: true }); + const data = await requestJson( + `/api/admin/data-management/jobs/${encodeURIComponent(currentJobDetailId)}/progress`, + { method: "GET" } + ); + const job = data.job || {}; + renderJobDetailSummary(job); + renderJobDetailActions(job); + renderJobDetailProgress(job); + updateJobDetailAutoRefresh(job); + if (!activeJobStatuses.has(String(job.status || "unknown"))) { + await loadDataManagementJobDetail( + currentJobDetailId, + { showModal: false, liveRefresh: true } + ); + } + } catch (error) { + stopJobDetailAutoRefresh({ message: "Live refresh paused after a request failure." }); } finally { jobDetailRefreshInFlight = false; } diff --git a/application/single_app/static/js/chat/chat-messages.js b/application/single_app/static/js/chat/chat-messages.js index 021d38646..3dd73e4a4 100644 --- a/application/single_app/static/js/chat/chat-messages.js +++ b/application/single_app/static/js/chat/chat-messages.js @@ -3196,7 +3196,12 @@ function renderReplyQuoteHtml(fullMessageObject = null) { const normalizedDocumentId = String(output.document_id || '').trim(); const normalizedExportRunId = String(output.export_run_id || output.run_id || '').trim(); const isBackgroundExport = Boolean(output.background_export) && Boolean(normalizedExportRunId); - if (!normalizedArtifactMessageId && !normalizedDocumentId && !isBackgroundExport) { + const terminalStatus = String(output.status || '').trim().toLowerCase(); + const isTerminalExportStatus = Boolean( + output.suppress_assistant_table_export + && ['failed', 'canceled'].includes(terminalStatus) + ); + if (!normalizedArtifactMessageId && !normalizedDocumentId && !isBackgroundExport && !isTerminalExportStatus) { return null; } @@ -3207,7 +3212,8 @@ function renderReplyQuoteHtml(fullMessageObject = null) { document_id: normalizedDocumentId, export_run_id: normalizedExportRunId, run_id: normalizedExportRunId, - background_export: isBackgroundExport, + background_export: isBackgroundExport || isTerminalExportStatus, + suppress_assistant_table_export: Boolean(output.suppress_assistant_table_export), }; } @@ -3393,6 +3399,36 @@ function renderReplyQuoteHtml(fullMessageObject = null) { } } + function canCancelBackgroundGeneratedOutput(outputMetadata) { + return Boolean(outputMetadata?.background_export && outputMetadata?.can_cancel); + } + + function setBackgroundGeneratedOutputCancelButtonLabel(cancelButton, label) { + if (!(cancelButton instanceof HTMLElement)) { + return; + } + + const icon = document.createElement('i'); + icon.className = 'bi bi-x-circle me-1'; + icon.setAttribute('aria-hidden', 'true'); + const labelText = document.createElement('span'); + labelText.textContent = label; + cancelButton.replaceChildren(icon, labelText); + } + + function updateBackgroundGeneratedOutputCancelButton(cancelButton, outputMetadata) { + if (!(cancelButton instanceof HTMLElement)) { + return; + } + + const canCancel = canCancelBackgroundGeneratedOutput(outputMetadata); + cancelButton.classList.toggle('d-none', !canCancel); + cancelButton.disabled = !canCancel; + if (cancelButton.dataset.busy !== 'true') { + setBackgroundGeneratedOutputCancelButtonLabel(cancelButton, 'Cancel'); + } + } + function formatGeneratedTabularPreviewValue(value, maxLength = 120) { let formattedValue = ''; @@ -4227,6 +4263,7 @@ function renderReplyQuoteHtml(fullMessageObject = null) { updateBackgroundGeneratedOutputStatusCard(statusElements, outputMetadata); updateBackgroundGeneratedOutputContinueButton(statusElements.continueButton, outputMetadata); + updateBackgroundGeneratedOutputCancelButton(statusElements.cancelButton, outputMetadata); } catch (error) { if (statusElements.detailText) { statusElements.detailText.textContent = error.message || 'Could not refresh export progress.'; @@ -4298,6 +4335,53 @@ function renderReplyQuoteHtml(fullMessageObject = null) { } } + async function cancelBackgroundGeneratedOutputRun(outputMetadata, card, statusElements = {}, cancelButton = null) { + const runId = String(outputMetadata?.export_run_id || outputMetadata?.run_id || '').trim(); + if (!runId || !(card instanceof HTMLElement) || !document.body.contains(card)) { + return; + } + + if (cancelButton) { + cancelButton.dataset.busy = 'true'; + cancelButton.disabled = true; + setBackgroundGeneratedOutputCancelButtonLabel(cancelButton, 'Canceling...'); + } + + try { + const response = await fetch(`/api/tabular/generated-output/runs/${encodeURIComponent(runId)}/cancel`, { + method: 'POST', + headers: { + 'Accept': 'application/json', + }, + }); + const responseData = await response.json().catch(() => ({})); + if (!response.ok) { + throw new Error(responseData?.message || responseData?.error || `Server responded with status ${response.status}`); + } + + const runStatus = responseData?.run || {}; + Object.assign(outputMetadata, runStatus, { + export_run_id: runStatus.run_id || runId, + run_id: runStatus.run_id || runId, + background_export: true, + }); + updateBackgroundGeneratedOutputStatusCard(statusElements, outputMetadata); + updateBackgroundGeneratedOutputContinueButton(statusElements.continueButton, outputMetadata); + updateBackgroundGeneratedOutputCancelButton(cancelButton, outputMetadata); + showToast(responseData?.message || 'Background export canceled.', 'success'); + } catch (error) { + if (statusElements.detailText) { + statusElements.detailText.textContent = error.message || 'Could not cancel background export.'; + } + showToast(error.message || 'Could not cancel background export.', 'danger'); + } finally { + if (cancelButton) { + delete cancelButton.dataset.busy; + updateBackgroundGeneratedOutputCancelButton(cancelButton, outputMetadata); + } + } + } + function shouldPollBackgroundGeneratedOutput(outputMetadata) { if (!outputMetadata?.background_export) { return false; @@ -4451,6 +4535,11 @@ function renderReplyQuoteHtml(fullMessageObject = null) { actions.className = 'd-flex flex-wrap gap-2 mt-3'; if (outputMetadata?.background_export) { + const backgroundRunId = String(outputMetadata?.export_run_id || outputMetadata?.run_id || '').trim(); + if (!backgroundRunId) { + return card; + } + const continueButton = document.createElement('button'); continueButton.type = 'button'; continueButton.className = 'btn btn-sm btn-outline-primary generated-tabular-continue-btn d-none'; @@ -4464,18 +4553,20 @@ function renderReplyQuoteHtml(fullMessageObject = null) { updateBackgroundGeneratedOutputContinueButton(continueButton, outputMetadata); actions.appendChild(continueButton); - const refreshStatusButton = document.createElement('button'); - refreshStatusButton.type = 'button'; - refreshStatusButton.className = 'btn btn-sm btn-outline-secondary generated-tabular-refresh-status-btn'; - refreshStatusButton.textContent = 'Refresh Status'; - refreshStatusButton.addEventListener('click', async () => { - refreshStatusButton.disabled = true; - refreshStatusButton.textContent = 'Refreshing...'; - await refreshBackgroundGeneratedOutputStatus(outputMetadata, card, backgroundStatusElements || {}); - refreshStatusButton.disabled = false; - refreshStatusButton.textContent = 'Refresh Status'; + const cancelButton = document.createElement('button'); + cancelButton.type = 'button'; + cancelButton.className = 'btn btn-sm btn-outline-danger generated-tabular-cancel-btn d-none'; + cancelButton.setAttribute('aria-label', 'Cancel background export'); + setBackgroundGeneratedOutputCancelButtonLabel(cancelButton, 'Cancel'); + cancelButton.addEventListener('click', async () => { + await cancelBackgroundGeneratedOutputRun(outputMetadata, card, backgroundStatusElements || {}, cancelButton); }); - actions.appendChild(refreshStatusButton); + if (backgroundStatusElements) { + backgroundStatusElements.cancelButton = cancelButton; + } + updateBackgroundGeneratedOutputCancelButton(cancelButton, outputMetadata); + actions.appendChild(cancelButton); + card.appendChild(actions); scheduleBackgroundGeneratedOutputStatusPolling(outputMetadata, card, backgroundStatusElements || {}); return card; @@ -6011,6 +6102,8 @@ export function buildChatRequestPayload(finalMessageToSend, conversationId = cur .filter(value => value); selectedDocumentId = selectedDocumentIds.length > 0 ? selectedDocumentIds[0] : null; } + const selectionMode = selectedDocumentIds.length > 0 ? 'selected' : 'relevance'; + const documentContextRequested = hybridSearchEnabled || selectedDocumentIds.length > 0; let imageGenEnabled = false; const igbtn = document.getElementById('image-generate-btn'); @@ -6125,6 +6218,8 @@ export function buildChatRequestPayload(finalMessageToSend, conversationId = cur message: finalMessageToSend, conversation_id: conversationId, hybrid_search: hybridSearchEnabled, + selection_mode: selectionMode, + document_context_requested: documentContextRequested, user_workspace_context_enabled: userWorkspaceContextEnabled, web_search_enabled: webSearchEnabled, url_access_enabled: urlAccessEnabled, @@ -6186,6 +6281,13 @@ export function buildCollaborativeInvocationTarget(messageData = {}, explicitInv messageData.agent_info && (messageData.agent_info.id || messageData.agent_info.name || messageData.agent_info.display_name) ); + const workspaceContextInvocationRequested = Boolean( + messageData.hybrid_search + || ( + messageData.document_context_requested + && window.appSettings?.enable_mixed_source_chat_search + ) + ); const sourceMode = messageData.image_generation ? 'image_generation' : hasAgentTarget @@ -6196,7 +6298,7 @@ export function buildCollaborativeInvocationTarget(messageData = {}, explicitInv ? 'url_access' : messageData.web_search_enabled ? 'web_search' - : messageData.hybrid_search + : workspaceContextInvocationRequested ? 'workspace' : messageData.prompt_info ? 'prompt' diff --git a/application/single_app/static/js/plugin_modal_stepper.js b/application/single_app/static/js/plugin_modal_stepper.js index 39cea5cd0..0db221701 100644 --- a/application/single_app/static/js/plugin_modal_stepper.js +++ b/application/single_app/static/js/plugin_modal_stepper.js @@ -2193,7 +2193,7 @@ export class PluginModalStepper { populateDocumentSearchForm(additionalFields = {}) { document.getElementById('document-search-scope').value = additionalFields.default_doc_scope || 'all'; - document.getElementById('document-search-top-n').value = additionalFields.default_top_n || 12; + document.getElementById('document-search-top-n').value = additionalFields.default_top_n || 25; document.getElementById('document-search-window-unit').value = additionalFields.default_window_unit || 'pages'; document.getElementById('document-search-window-size').value = additionalFields.default_window_size || ''; document.getElementById('document-search-window-percent').value = additionalFields.default_window_percent || ''; @@ -5103,7 +5103,7 @@ export class PluginModalStepper { const config = this.getDocumentSearchAdditionalFields(); document.getElementById('summary-search-scope').textContent = this.formatDocumentScope(config.default_doc_scope); - document.getElementById('summary-search-top-n').textContent = String(config.default_top_n || 12); + document.getElementById('summary-search-top-n').textContent = String(config.default_top_n || 25); document.getElementById('summary-search-chunk-behavior').textContent = 'Returns all chunks by default'; document.getElementById('summary-search-windowing').textContent = this.formatDocumentSearchWindowing(config); document.getElementById('summary-search-window-target-length').textContent = config.default_window_target_length || '2 pages'; diff --git a/application/single_app/static/json/schemas/document_search_plugin.additional_settings.schema.json b/application/single_app/static/json/schemas/document_search_plugin.additional_settings.schema.json index d147ddd5f..a8c8a3384 100644 --- a/application/single_app/static/json/schemas/document_search_plugin.additional_settings.schema.json +++ b/application/single_app/static/json/schemas/document_search_plugin.additional_settings.schema.json @@ -17,7 +17,7 @@ "type": "integer", "minimum": 1, "maximum": 500, - "default": 12 + "default": 50 }, "default_window_unit": { "type": "string", diff --git a/application/single_app/templates/admin_settings.html b/application/single_app/templates/admin_settings.html index e9b758218..7165117b8 100644 --- a/application/single_app/templates/admin_settings.html +++ b/application/single_app/templates/admin_settings.html @@ -5071,6 +5071,16 @@
Target Cosmos Database
+
+
+ + +
+
+ + +
+
@@ -5141,6 +5151,32 @@
Migration Workflow
Refresh Summary
+
+ Synchronization mode +
+ + + + + + +
+

+ Copies source items that are absent from the destination. Existing destination data is never updated or deleted. +

+
+
+ + +
The baseline fixes the prior service watermarks used for this catch-up run.
+
+
+ +
@@ -5215,6 +5251,11 @@
Public Workspaces
+
+ + +
Azure AI Search does not support conditional document writes. SimpleChat target indexing pauses automatically; freeze all external writers before copying Search documents.
+
@@ -5223,6 +5264,44 @@
Public Workspaces
+
+
+
+
Performance And Resume
+

Migration uses durable per-resource checkpoints and the same migration ID after a retry.

+
+ +
+
+
+ + +
+
+ + +
+
+ + +
+
+
+ + +
+
+
+ + +
+
+
The migration records the prior database or dedicated-container capacity, raises only eligible targets up to 10,000 RU/s, and restores it after completion or failure. Capacity changes require destination ARM permission.
+
+
+

Choose migration targets to see the summary.

@@ -5521,6 +5600,7 @@
+
@@ -5568,6 +5648,25 @@
Warnings
+ +
diff --git a/application/single_app/templates/chats.html b/application/single_app/templates/chats.html index 460a49b16..38dd6d3eb 100644 --- a/application/single_app/templates/chats.html +++ b/application/single_app/templates/chats.html @@ -1487,6 +1487,7 @@