From d2656e95109c2d66365897cf181531f7f6d9b50e Mon Sep 17 00:00:00 2001 From: Aryan Mathur <126969847+neverthesameagain@users.noreply.github.com> Date: Tue, 3 Feb 2026 03:38:35 +0530 Subject: [PATCH] bg fixes and feature addition --- .claude/hooks/capture_session_event.py | 112 ++ .claude/hooks/claude_code_capture_utils.py | 93 ++ .claude/hooks/process_transcript.py | 720 ++++++++++++ .claude/settings.local.json | 55 + fastapi/__init__.py | 1 + fastapi/_compat/v2.py | 155 ++- fastapi/applications.py | 91 ++ fastapi/background.py | 104 +- fastapi/lifespan.py | 203 ++++ fastapi/openapi/_profiling.py | 205 ++++ fastapi/openapi/plugins.py | 376 ++++++ fastapi/openapi/utils.py | 32 +- fastapi/routing.py | 9 + tests/benchmarks/test_openapi_performance.py | 1019 +++++++++++++++++ tests/test_background_tasks_async_race.py | 370 ++++++ tests/test_background_tasks_bugs.py | 287 +++++ ...test_background_tasks_dependency_access.py | 344 ++++++ tests/test_background_tasks_fixes.py | 411 +++++++ tests/test_background_tasks_yield_scope.py | 355 ++++++ tests/test_lifespan_unified.py | 377 ++++++ tests/test_openapi_plugins.py | 488 ++++++++ 21 files changed, 5769 insertions(+), 38 deletions(-) create mode 100755 .claude/hooks/capture_session_event.py create mode 100755 .claude/hooks/claude_code_capture_utils.py create mode 100755 .claude/hooks/process_transcript.py create mode 100644 .claude/settings.local.json create mode 100644 fastapi/lifespan.py create mode 100644 fastapi/openapi/_profiling.py create mode 100644 fastapi/openapi/plugins.py create mode 100644 tests/benchmarks/test_openapi_performance.py create mode 100644 tests/test_background_tasks_async_race.py create mode 100644 tests/test_background_tasks_bugs.py create mode 100644 tests/test_background_tasks_dependency_access.py create mode 100644 tests/test_background_tasks_fixes.py create mode 100644 tests/test_background_tasks_yield_scope.py create mode 100644 tests/test_lifespan_unified.py create mode 100644 tests/test_openapi_plugins.py diff --git a/.claude/hooks/capture_session_event.py b/.claude/hooks/capture_session_event.py new file mode 100755 index 000000000..46f7dfeac --- /dev/null +++ b/.claude/hooks/capture_session_event.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +""" +Hook to capture session start/end events. +Cross-platform support for Windows, macOS, and Linux. +""" +import json +import sys +import os +import subprocess +from datetime import datetime, timezone +from claude_code_capture_utils import get_log_file_path, add_ab_metadata + +def get_git_metadata(repo_dir): + """Get current git commit and branch.""" + try: + # Get current commit hash + commit_result = subprocess.run( + ['git', 'rev-parse', 'HEAD'], + cwd=repo_dir, + capture_output=True, + text=True, + timeout=10 + ) + + # Get current branch + branch_result = subprocess.run( + ['git', 'branch', '--show-current'], + cwd=repo_dir, + capture_output=True, + text=True, + timeout=10 + ) + + git_metadata = { + "base_commit": commit_result.stdout.strip() if commit_result.returncode == 0 else None, + "branch": branch_result.stdout.strip() if branch_result.returncode == 0 else None, + "timestamp": datetime.now(timezone.utc).isoformat() + } + + if git_metadata["base_commit"]: + return git_metadata + else: + return None + + except Exception as e: + print(f"Warning: Could not capture git metadata: {e}", file=sys.stderr) + return None + +def main(): + try: + if len(sys.argv) < 2: + print("Usage: capture_session_event.py [start|end]", file=sys.stderr) + sys.exit(1) + + event_type = sys.argv[1].lower() + if event_type not in ["start", "end"]: + print("Event type must be 'start' or 'end'", file=sys.stderr) + sys.exit(1) + + input_data = json.load(sys.stdin) + + session_id = input_data.get("session_id", "unknown") + transcript_path = input_data.get("transcript_path", "") + cwd = input_data.get("cwd", "") + + if event_type == "start": + # Session start: capture git metadata + git_metadata = get_git_metadata(cwd) + + log_entry = { + "type": "session_start", + "timestamp": datetime.now(timezone.utc).isoformat(), + "session_id": session_id, + "transcript_path": transcript_path, + "cwd": cwd, + "git_metadata": git_metadata + } + + log_entry = add_ab_metadata(log_entry, cwd) + + if git_metadata: + print(f"[OK] Captured git metadata: {git_metadata['base_commit'][:8]} on {git_metadata['branch']}") + + # Write session_start event + log_file = get_log_file_path(session_id, cwd) + with open(log_file, "a", encoding="utf-8") as f: + f.write(json.dumps(log_entry) + "\n") + + elif event_type == "end": + # Session end: log the event + log_entry = { + "type": "session_end", + "timestamp": datetime.now(timezone.utc).isoformat(), + "session_id": session_id, + "transcript_path": transcript_path, + "cwd": cwd, + "reason": input_data.get("reason", "") + } + + log_entry = add_ab_metadata(log_entry, cwd) + + # Write session_end event + log_file = get_log_file_path(session_id, cwd) + with open(log_file, "a", encoding="utf-8") as f: + f.write(json.dumps(log_entry) + "\n") + + except Exception as e: + print(f"[ERROR] Session {event_type}: {e}", file=sys.stderr) + sys.exit(1) + +if __name__ == "__main__": + main() diff --git a/.claude/hooks/claude_code_capture_utils.py b/.claude/hooks/claude_code_capture_utils.py new file mode 100755 index 000000000..fe2d754da --- /dev/null +++ b/.claude/hooks/claude_code_capture_utils.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +""" +Utility functions for A/B testing hooks. +Cross-platform support for Windows, macOS, and Linux. +""" +import json +import os +from pathlib import Path + +def detect_model_lane(cwd): + """Detect if we're in model_a or model_b directory.""" + path_parts = Path(cwd).parts + if 'model_a' in path_parts: + return 'model_a' + elif 'model_b' in path_parts: + return 'model_b' + return None + +def get_experiment_root(cwd): + """Get the experiment root directory (parent of model_a/model_b).""" + current_path = Path(cwd) + + # Check if we're inside a model_a or model_b directory + # Look for the parent that contains both model_a and model_b + for parent in [current_path] + list(current_path.parents): + if (parent / 'model_a').exists() and (parent / 'model_b').exists(): + return str(parent) + # Also check one level up (in case we're inside the cloned repo) + parent_up = parent.parent + if (parent_up / 'model_a').exists() and (parent_up / 'model_b').exists(): + return str(parent_up) + return None + +def read_manifest(experiment_root): + """Read the manifest.json file to get task_id and model assignments.""" + try: + manifest_path = os.path.join(experiment_root, 'manifest.json') + if os.path.exists(manifest_path): + with open(manifest_path, 'r', encoding='utf-8') as f: + return json.load(f) + except Exception: + pass + return {} + +def get_ab_metadata(cwd): + """Get A/B testing metadata (task_id, model_lane, model_name) from current directory.""" + model_lane = detect_model_lane(cwd) + experiment_root = get_experiment_root(cwd) + + if not model_lane or not experiment_root: + return {} + + manifest = read_manifest(experiment_root) + + metadata = { + "task_id": manifest.get("task_id"), + "model_lane": model_lane, + "experiment_root": experiment_root + } + + # Get model name from assignments + assignments = manifest.get("assignments", {}) + if model_lane in assignments: + metadata["model_name"] = assignments[model_lane] + + return metadata + +def get_log_file_path(session_id, cwd): + """Get the correct log file path for A/B testing (routes to model-specific directory).""" + model_lane = detect_model_lane(cwd) + experiment_root = get_experiment_root(cwd) + + if model_lane and experiment_root: + # Route to model-specific logs directory + logs_dir = os.path.join(experiment_root, "logs", model_lane) + os.makedirs(logs_dir, exist_ok=True) + return os.path.join(logs_dir, f"session_{session_id}.jsonl") + else: + # Fallback to current behavior + project_dir = os.environ.get("CLAUDE_PROJECT_DIR", os.getcwd()) + logs_dir = os.path.join(project_dir, "logs") + os.makedirs(logs_dir, exist_ok=True) + return os.path.join(logs_dir, f"session_{session_id}.jsonl") + +def add_ab_metadata(event, cwd): + """Add A/B testing metadata to an event.""" + ab_metadata = get_ab_metadata(cwd) + if ab_metadata: + for key, value in ab_metadata.items(): + if value is not None: + event[key] = value + return event + diff --git a/.claude/hooks/process_transcript.py b/.claude/hooks/process_transcript.py new file mode 100755 index 000000000..955560dfb --- /dev/null +++ b/.claude/hooks/process_transcript.py @@ -0,0 +1,720 @@ +#!/usr/bin/env python3 +""" +Process raw transcript to extract all messages and generate summaries. +Handles both incremental (Stop event) and final (SessionEnd) processing. +Cross-platform support for Windows, macOS, and Linux. +""" +import json +import sys +import os +import shutil +import subprocess +from datetime import datetime, timezone +from pathlib import Path +from collections import defaultdict +from claude_code_capture_utils import get_log_file_path, add_ab_metadata, detect_model_lane, get_experiment_root + +def read_and_process_raw_transcript(transcript_path): + """ + Read raw transcript and extract all unique messages. + Returns deduplicated messages with last occurrence (final state). + Also extracts thinking blocks as separate entries. + """ + if not os.path.exists(transcript_path): + return [] + + # Track assistant messages by message.id (they have IDs, can have duplicates) + assistant_messages = {} + # Track thinking blocks separately (won't be counted in token usage) + thinking_blocks = {} + # Track user messages by uuid (they don't have message.id, use uuid) + user_messages = {} + + try: + with open(transcript_path, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + + try: + event = json.loads(line) + event_type = event.get('type') + message = event.get('message', {}) + + # Process assistant messages (have message.id) + if event_type == 'assistant': + msg_id = message.get('id') + if msg_id: + # Check for thinking blocks in content + content = message.get('content', []) + if isinstance(content, list): + for item in content: + if isinstance(item, dict) and item.get('type') == 'thinking': + # Extract thinking block as separate entry + thinking_entry = { + 'type': 'assistant_thinking', + 'timestamp': event.get('timestamp'), + 'message_id': msg_id, + 'thinking_content': item.get('thinking', ''), + 'session_id': event.get('sessionId'), + 'cwd': event.get('cwd') + } + # Use message_id as key (one thinking block per message) + thinking_blocks[msg_id] = thinking_entry + + # Store/overwrite with last occurrence (streaming) + assistant_msg = { + 'type': event_type, + 'timestamp': event.get('timestamp'), + 'message': message, + 'session_id': event.get('sessionId'), + 'cwd': event.get('cwd') + } + + # Preserve stop_reason from message if present + if message.get('stop_reason'): + assistant_msg['stop_reason'] = message['stop_reason'] + + assistant_messages[msg_id] = assistant_msg + + # Process user messages (use uuid as key, no message.id) + elif event_type == 'user': + uuid = event.get('uuid') + if uuid: + user_msg = { + 'type': event_type, + 'timestamp': event.get('timestamp'), + 'message': message, + 'session_id': event.get('sessionId'), + 'cwd': event.get('cwd'), + 'uuid': uuid + } + + # Preserve thinking metadata if present + if 'thinkingMetadata' in event: + user_msg['thinkingMetadata'] = event['thinkingMetadata'] + + # Preserve isMeta flag if present + if event.get('isMeta'): + user_msg['isMeta'] = event['isMeta'] + + user_messages[uuid] = user_msg + + except json.JSONDecodeError: + continue + + except Exception as e: + print(f"[ERROR] Reading raw transcript: {e}", file=sys.stderr) + return [] + + # Combine all: assistant messages + thinking blocks + user messages + # Thinking blocks inserted right before their corresponding assistant message + all_messages = [] + + # First, add all messages with their thinking blocks properly ordered + assistant_list = list(assistant_messages.values()) + user_list = list(user_messages.values()) + + # Combine and sort all by timestamp + combined = assistant_list + user_list + combined.sort(key=lambda m: m.get('timestamp', '')) + + # Insert thinking blocks right before their parent assistant message + for msg in combined: + if msg['type'] == 'assistant': + msg_id = msg['message'].get('id') + # If there's a thinking block for this message, insert it first + if msg_id in thinking_blocks: + all_messages.append(thinking_blocks[msg_id]) + all_messages.append(msg) + + return all_messages + +def aggregate_token_usage(messages): + """Aggregate token usage from all assistant messages. + Note: assistant_thinking entries are explicitly excluded to avoid double-counting. + """ + total_usage = { + 'total_input_tokens': 0, + 'total_output_tokens': 0, + 'total_cache_creation_tokens': 0, + 'total_cache_read_tokens': 0, + 'total_ephemeral_5m_tokens': 0, + 'total_ephemeral_1h_tokens': 0, + 'service_tier': None + } + + for msg_data in messages: + # Only count tokens from 'assistant' type, NOT 'assistant_thinking' + # Thinking tokens are already included in the parent assistant message's output_tokens + if msg_data['type'] == 'assistant': + message = msg_data['message'] + usage = message.get('usage', {}) + + if usage: + total_usage['total_input_tokens'] += usage.get('input_tokens', 0) + total_usage['total_output_tokens'] += usage.get('output_tokens', 0) + total_usage['total_cache_creation_tokens'] += usage.get('cache_creation_input_tokens', 0) + total_usage['total_cache_read_tokens'] += usage.get('cache_read_input_tokens', 0) + + cache_creation = usage.get('cache_creation', {}) + total_usage['total_ephemeral_5m_tokens'] += cache_creation.get('ephemeral_5m_input_tokens', 0) + total_usage['total_ephemeral_1h_tokens'] += cache_creation.get('ephemeral_1h_input_tokens', 0) + + if usage.get('service_tier'): + total_usage['service_tier'] = usage.get('service_tier') + + # Add calculated total + total_usage['total_actual_input_tokens'] = ( + total_usage['total_input_tokens'] + + total_usage['total_cache_creation_tokens'] + + total_usage['total_cache_read_tokens'] + ) + + return total_usage + +def analyze_tool_calls(messages): + """Extract tool call metrics from messages.""" + tool_calls = defaultdict(int) + tool_results = defaultdict(int) + + for msg_data in messages: + # Skip entries without 'message' key (e.g., assistant_thinking) + if 'message' not in msg_data: + continue + message = msg_data['message'] + content = message.get('content', []) + + if not isinstance(content, list): + continue + + for item in content: + if not isinstance(item, dict): + continue + + if item.get('type') == 'tool_use': + tool_name = item.get('name', 'unknown') + tool_calls[tool_name] += 1 + + elif item.get('type') == 'tool_result': + # Try to infer tool name from context (simplified) + tool_results['total'] += 1 + + return { + 'tool_calls_by_type': dict(tool_calls), + 'total_tool_calls': sum(tool_calls.values()), + 'total_tool_results': tool_results.get('total', 0) + } + +def analyze_thinking_usage(messages, transcript_path): + """Analyze thinking mode usage in messages.""" + thinking_stats = { + 'thinking_enabled_turns': 0, + 'thinking_disabled_turns': 0, + 'assistant_with_thinking_blocks': 0, + 'thinking_levels': defaultdict(int) + } + + # Track which turns had thinking enabled (from user thinkingMetadata) + for msg_data in messages: + if msg_data['type'] == 'user' and 'thinkingMetadata' in msg_data: + metadata = msg_data['thinkingMetadata'] + if not metadata.get('disabled', True): + thinking_stats['thinking_enabled_turns'] += 1 + level = metadata.get('level', 'none') + thinking_stats['thinking_levels'][level] += 1 + else: + thinking_stats['thinking_disabled_turns'] += 1 + + # Count assistant messages with thinking blocks + # Check ALL occurrences in raw transcript (not just final deduplicated state) + assistant_msg_ids_with_thinking = set() + + try: + if os.path.exists(transcript_path): + with open(transcript_path, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + event = json.loads(line) + if event.get('type') == 'assistant': + message = event.get('message', {}) + msg_id = message.get('id') + content = message.get('content', []) + + if msg_id and isinstance(content, list): + # Check if this occurrence has thinking + has_thinking = any( + isinstance(item, dict) and item.get('type') == 'thinking' + for item in content + ) + if has_thinking: + assistant_msg_ids_with_thinking.add(msg_id) + except: + continue + except Exception: + pass + + thinking_stats['assistant_with_thinking_blocks'] = len(assistant_msg_ids_with_thinking) + + return { + 'thinking_enabled_turns': thinking_stats['thinking_enabled_turns'], + 'thinking_disabled_turns': thinking_stats['thinking_disabled_turns'], + 'assistant_with_thinking_blocks': thinking_stats['assistant_with_thinking_blocks'], + 'thinking_levels': dict(thinking_stats['thinking_levels']) + } + +def calculate_git_metrics(cwd, base_commit): + """Calculate git metrics from diff.""" + try: + original_cwd = os.getcwd() + os.chdir(cwd) + + if not base_commit: + os.chdir(original_cwd) + return {} + + # Add untracked files + excluded_patterns = ['.claude/', '__pycache__/', 'node_modules/', '.mypy_cache/', + '.pytest_cache/', '.DS_Store', '.vscode/', '.idea/'] + + untracked_result = subprocess.run( + ['git', 'ls-files', '--others', '--exclude-standard'], + capture_output=True, text=True, timeout=30 + ) + + if untracked_result.returncode == 0 and untracked_result.stdout.strip(): + untracked_files = [ + f.strip() for f in untracked_result.stdout.strip().split('\n') + if f.strip() and not any(pattern in f for pattern in excluded_patterns) + ] + + for file in untracked_files: + subprocess.run(['git', 'add', '-N', file], capture_output=True, timeout=5) + + # Calculate numstat + result = subprocess.run( + ['git', 'diff', '--numstat', base_commit, '--', '.', + ':!.claude', ':!**/.mypy_cache', ':!**/__pycache__', ':!**/.pytest_cache', + ':!**/.DS_Store', ':!**/node_modules', ':!**/.vscode', ':!**/.idea'], + capture_output=True, text=True, timeout=30 + ) + + os.chdir(original_cwd) + + if result.returncode != 0: + return {} + + lines = result.stdout.strip().split('\n') if result.stdout.strip() else [] + files_changed = 0 + total_lines_changed = 0 + + for line in lines: + if line.strip(): + parts = line.split('\t') + if len(parts) >= 3: + try: + added = int(parts[0]) if parts[0] != '-' else 0 + removed = int(parts[1]) if parts[1] != '-' else 0 + files_changed += 1 + total_lines_changed += added + removed + except ValueError: + continue + + return { + "files_changed_count": files_changed, + "lines_of_code_changed_count": total_lines_changed + } + + except Exception as e: + print(f"Warning: Could not calculate git metrics: {e}", file=sys.stderr) + if 'original_cwd' in locals(): + os.chdir(original_cwd) + return {} + +def copy_raw_transcript(transcript_path, session_id, cwd): + """Copy raw transcript to logs folder.""" + try: + source_path = Path(transcript_path) + if not source_path.exists(): + print(f"Warning: Raw transcript not found at {source_path}", file=sys.stderr) + return False + + model_lane = detect_model_lane(cwd) + experiment_root = get_experiment_root(cwd) + + if model_lane and experiment_root: + logs_dir = Path(experiment_root) / "logs" / model_lane + logs_dir.mkdir(parents=True, exist_ok=True) + dest_path = logs_dir / f"session_{session_id}_raw.jsonl" + else: + project_dir = os.environ.get('CLAUDE_PROJECT_DIR', os.getcwd()) + logs_dir = Path(project_dir) / "logs" + logs_dir.mkdir(exist_ok=True) + dest_path = logs_dir / f"session_{session_id}_raw.jsonl" + + shutil.copy2(source_path, dest_path) + print(f"[OK] Copied raw transcript to {dest_path}") + return True + + except Exception as e: + print(f"[ERROR] Copying raw transcript: {e}", file=sys.stderr) + return False + +def get_base_commit_from_log(log_file): + """Extract base commit from session_start event.""" + try: + if not os.path.exists(log_file): + return None + + with open(log_file, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if line: + try: + event = json.loads(line) + if event.get('type') == 'session_start': + git_metadata = event.get('git_metadata', {}) + return git_metadata.get('base_commit') + except json.JSONDecodeError: + continue + return None + except Exception: + return None + +def main(): + try: + if len(sys.argv) < 2: + print("Usage: process_transcript.py [incremental|final]", file=sys.stderr) + sys.exit(1) + + mode = sys.argv[1].lower() + if mode not in ["incremental", "final"]: + print("Mode must be 'incremental' or 'final'", file=sys.stderr) + sys.exit(1) + + input_data = json.load(sys.stdin) + + session_id = input_data.get("session_id", "unknown") + transcript_path = input_data.get("transcript_path", "") + cwd = input_data.get("cwd", "") + + log_file = get_log_file_path(session_id, cwd) + + if mode == "incremental": + # Stop event: incremental processing (fault tolerance) + messages = read_and_process_raw_transcript(transcript_path) + + if not messages: + return + + # Append all new unique messages to log + # Track what we've already logged (assistant by msg_id, thinking by msg_id, user by uuid) + existing_assistant_ids = set() + existing_thinking_ids = set() + existing_user_uuids = set() + + if os.path.exists(log_file): + with open(log_file, 'r', encoding='utf-8') as f: + for line in f: + try: + event = json.loads(line) + event_type = event.get('type') + if event_type == 'assistant': + msg = event.get('message', {}) + if msg.get('id'): + existing_assistant_ids.add(msg['id']) + elif event_type == 'assistant_thinking': + msg_id = event.get('message_id') + if msg_id: + existing_thinking_ids.add(msg_id) + elif event_type == 'user': + uuid = event.get('uuid') + if uuid: + existing_user_uuids.add(uuid) + except: + continue + + # Append new messages + new_count = 0 + with open(log_file, "a", encoding="utf-8") as f: + for msg_data in messages: + # Check if this is a new message + is_new = False + if msg_data['type'] == 'assistant': + msg_id = msg_data['message'].get('id') + if msg_id and msg_id not in existing_assistant_ids: + is_new = True + existing_assistant_ids.add(msg_id) + elif msg_data['type'] == 'assistant_thinking': + msg_id = msg_data.get('message_id') + if msg_id and msg_id not in existing_thinking_ids: + is_new = True + existing_thinking_ids.add(msg_id) + elif msg_data['type'] == 'user': + uuid = msg_data.get('uuid') + if uuid and uuid not in existing_user_uuids: + is_new = True + existing_user_uuids.add(uuid) + + if is_new: + # Add A/B metadata + log_entry = add_ab_metadata(msg_data.copy(), cwd) + f.write(json.dumps(log_entry) + "\n") + new_count += 1 + + if new_count > 0: + print(f"[OK] Processed {new_count} new messages (total: {len(messages)} unique)") + + elif mode == "final": + # SessionEnd: complete processing + summary + + # Step 1: Copy raw transcript + copy_raw_transcript(transcript_path, session_id, cwd) + + # Step 2: Process complete raw transcript + messages = read_and_process_raw_transcript(transcript_path) + + if not messages: + print("Warning: No messages found in raw transcript", file=sys.stderr) + return + + # Step 3: REBUILD processed log in perfect chronological order + # Read existing non-message events (session_start, etc.) + non_message_events = [] + + if os.path.exists(log_file): + with open(log_file, 'r', encoding='utf-8') as f: + for line in f: + try: + event = json.loads(line) + # Keep session_start and other non-message events + # Exclude assistant, assistant_thinking, and user messages (they come from raw transcript) + if event.get('type') not in ['assistant', 'assistant_thinking', 'user']: + non_message_events.append(event) + except: + continue + + # Combine all events and sort by timestamp + session_start = [e for e in non_message_events if e.get('type') == 'session_start'] + other_events = [e for e in non_message_events if e.get('type') != 'session_start'] + + # Build chronological list: session_start first, then messages sorted by time + all_events = [] + + # Add session_start first (if exists) + if session_start: + all_events.extend(session_start) + + # Add all messages (already sorted by timestamp from read_and_process_raw_transcript) + # Messages already include thinking blocks inserted before their parent assistant message + for msg_data in messages: + all_events.append(add_ab_metadata(msg_data.copy(), cwd)) + + print(f"[OK] Rebuilding log with {len(all_events)} events in chronological order") + + # Step 4: Generate session summary + usage_totals = aggregate_token_usage(messages) + tool_metrics = analyze_tool_calls(messages) + thinking_metrics = analyze_thinking_usage(messages, transcript_path) + + # Calculate duration + timestamps = [ + datetime.fromisoformat(msg['timestamp'].replace('Z', '+00:00')) + for msg in messages if msg.get('timestamp') + ] + + total_duration = 0 + if len(timestamps) >= 2: + duration = max(timestamps) - min(timestamps) + total_duration = duration.total_seconds() + + # Count messages with proper categorization + # Note: assistant_thinking is NOT counted as a separate message (it's part of assistant message) + assistant_count = sum(1 for m in messages if m['type'] == 'assistant') + thinking_count = sum(1 for m in messages if m['type'] == 'assistant_thinking') + + # Categorize user messages + user_prompts = 0 + tool_results = 0 + system_messages = 0 + + for m in messages: + if m['type'] == 'user': + message_content = m['message'].get('content', '') + + # Check if it's a system/meta message + if m.get('isMeta'): + system_messages += 1 + # Check if it's a tool result + elif isinstance(message_content, list): + has_tool_result = any( + isinstance(item, dict) and item.get('type') == 'tool_result' + for item in message_content + ) + if has_tool_result: + tool_results += 1 + else: + user_prompts += 1 + # Check if it's an exit/system command + elif isinstance(message_content, str) and ( + '' in message_content or + '' in message_content + ): + system_messages += 1 + # Real user prompt (string content, not system) + elif isinstance(message_content, str): + user_prompts += 1 + else: + user_prompts += 1 # Default to user prompt + + total_user_events = user_prompts + tool_results + system_messages + + # Calculate actual total messages (excluding thinking blocks as they're not separate messages) + actual_total_messages = assistant_count + total_user_events + + # Get git metrics + base_commit = get_base_commit_from_log(log_file) + git_metrics = calculate_git_metrics(cwd, base_commit) if base_commit else {} + + # Create session summary + model_lane = detect_model_lane(cwd) + + summary = { + "type": "session_summary", + "timestamp": datetime.now(timezone.utc).isoformat(), + "session_id": session_id, + "transcript_path": transcript_path, + "cwd": cwd, + "summary_data": { + "total_duration_seconds": round(total_duration, 2), + "total_messages": actual_total_messages, + "assistant_messages": assistant_count, + "user_prompts": user_prompts, + "user_metrics": { + "user_prompts": user_prompts, + "tool_results": tool_results, + "system_messages": system_messages, + "total_user_events": total_user_events + }, + "usage_totals": usage_totals, + "tool_metrics": tool_metrics, + "thinking_metrics": { + **thinking_metrics, + "assistant_thinking_blocks_captured": thinking_count + }, + "git_metrics": git_metrics, + "files": { + "processed_log": f"session_{session_id}.jsonl", + "raw_transcript": f"session_{session_id}_raw.jsonl", + "git_diff": f"{model_lane}_diff.patch" if model_lane else None + }, + "validation": { + "complete": True, + "unique_messages_processed": actual_total_messages, + "thinking_blocks_extracted": thinking_count + } + } + } + + summary = add_ab_metadata(summary, cwd) + + # Add session summary to events + all_events.append(summary) + + # Step 5: Rewrite log file with all events in perfect chronological order + # Write to temp file first, then rename (atomic) + temp_log_file = log_file + ".tmp" + + with open(temp_log_file, "w", encoding="utf-8") as f: + for event in all_events: + f.write(json.dumps(event) + "\n") + + # Atomic rename + os.replace(temp_log_file, log_file) + + print(f"[OK] Rebuilt log with {len(all_events)} events in chronological order") + print(f"[OK] Generated session summary: {actual_total_messages} messages, {assistant_count} assistant, {user_prompts} user prompts") + if thinking_count > 0: + print(f"[OK] Captured {thinking_count} thinking blocks (tokens already included in assistant output)") + print(f"[OK] User breakdown: {user_prompts} prompts, {tool_results} tool results, {system_messages} system") + print(f"[OK] Tokens: {usage_totals['total_actual_input_tokens']:,} input, {usage_totals['total_output_tokens']:,} output") + + except Exception as e: + print(f"[ERROR] Processing transcript: {e}", file=sys.stderr) + sys.exit(1) + +def calculate_git_metrics(cwd, base_commit): + """Calculate git metrics from diff.""" + try: + original_cwd = os.getcwd() + os.chdir(cwd) + + if not base_commit: + os.chdir(original_cwd) + return {} + + # Add untracked files + excluded_patterns = ['.claude/', '__pycache__/', 'node_modules/', '.mypy_cache/', + '.pytest_cache/', '.DS_Store', '.vscode/', '.idea/'] + + untracked_result = subprocess.run( + ['git', 'ls-files', '--others', '--exclude-standard'], + capture_output=True, text=True, timeout=30 + ) + + if untracked_result.returncode == 0 and untracked_result.stdout.strip(): + untracked_files = [ + f.strip() for f in untracked_result.stdout.strip().split('\n') + if f.strip() and not any(pattern in f for pattern in excluded_patterns) + ] + + for file in untracked_files: + subprocess.run(['git', 'add', '-N', file], capture_output=True, timeout=5) + + # Calculate numstat + result = subprocess.run( + ['git', 'diff', '--numstat', base_commit, '--', '.', + ':!.claude', ':!**/.mypy_cache', ':!**/__pycache__', ':!**/.pytest_cache', + ':!**/.DS_Store', ':!**/node_modules', ':!**/.vscode', ':!**/.idea'], + capture_output=True, text=True, timeout=30 + ) + + os.chdir(original_cwd) + + if result.returncode != 0: + return {} + + lines = result.stdout.strip().split('\n') if result.stdout.strip() else [] + files_changed = 0 + total_lines_changed = 0 + + for line in lines: + if line.strip(): + parts = line.split('\t') + if len(parts) >= 3: + try: + added = int(parts[0]) if parts[0] != '-' else 0 + removed = int(parts[1]) if parts[1] != '-' else 0 + files_changed += 1 + total_lines_changed += added + removed + except ValueError: + continue + + return { + "files_changed_count": files_changed, + "lines_of_code_changed_count": total_lines_changed + } + + except Exception as e: + print(f"Warning: Could not calculate git metrics: {e}", file=sys.stderr) + if 'original_cwd' in locals(): + os.chdir(original_cwd) + return {} + +if __name__ == "__main__": + main() + diff --git a/.claude/settings.local.json b/.claude/settings.local.json new file mode 100644 index 000000000..58bfde171 --- /dev/null +++ b/.claude/settings.local.json @@ -0,0 +1,55 @@ +{ + "env": { + "ANTHROPIC_AUTH_TOKEN": "sk-1234-tap-tap-go-new", + "ANTHROPIC_BASE_URL": "https://mercor-rl--litellm-proxy-serve.modal.run" + }, + "permissions": { + "allow": [ + "Bash(ls:*)", + "Bash(grep:*)", + "Bash(python -m pytest:*)", + "Bash(pip install:*)", + "Bash(python3:*)", + "Bash(python:*)", + "WebSearch", + "WebFetch(domain:github.com)" + ] + }, + "model": "zorilla", + "hooks": { + "SessionStart": [ + { + "hooks": [ + { + "type": "command", + "command": "python \"$CLAUDE_PROJECT_DIR/.claude/hooks/capture_session_event.py\" start" + } + ] + } + ], + "Stop": [ + { + "hooks": [ + { + "type": "command", + "command": "python \"$CLAUDE_PROJECT_DIR/.claude/hooks/process_transcript.py\" incremental" + } + ] + } + ], + "SessionEnd": [ + { + "hooks": [ + { + "type": "command", + "command": "python \"$CLAUDE_PROJECT_DIR/.claude/hooks/process_transcript.py\" final" + }, + { + "type": "command", + "command": "python \"$CLAUDE_PROJECT_DIR/.claude/hooks/capture_session_event.py\" end" + } + ] + } + ] + } +} diff --git a/fastapi/__init__.py b/fastapi/__init__.py index 6133787b0..4de3cdcc1 100644 --- a/fastapi/__init__.py +++ b/fastapi/__init__.py @@ -5,6 +5,7 @@ __version__ = "0.128.0" from starlette import status as status from .applications import FastAPI as FastAPI +from .background import BackgroundTaskError as BackgroundTaskError from .background import BackgroundTasks as BackgroundTasks from .datastructures import UploadFile as UploadFile from .exceptions import HTTPException as HTTPException diff --git a/fastapi/_compat/v2.py b/fastapi/_compat/v2.py index 25b681453..93335526d 100644 --- a/fastapi/_compat/v2.py +++ b/fastapi/_compat/v2.py @@ -1,4 +1,5 @@ import re +import threading import warnings from collections.abc import Sequence from copy import copy, deepcopy @@ -12,6 +13,36 @@ from typing import ( cast, ) +from fastapi.openapi._profiling import openapi_profiler, profiled + + +# ============================================================================= +# Model traversal cache for OpenAPI performance optimization +# ============================================================================= + +# Thread-safe cache for flat models extracted from types +# Key: frozenset of model types, Value: set of all referenced models +_flat_models_cache: dict[frozenset[type], set[type]] = {} +_flat_models_cache_lock = threading.Lock() + + +def _get_cached_flat_models(types: frozenset[type]) -> set[type] | None: + """Get cached flat models for a set of types.""" + with _flat_models_cache_lock: + return _flat_models_cache.get(types) + + +def _set_cached_flat_models(types: frozenset[type], models: set[type]) -> None: + """Cache flat models for a set of types.""" + with _flat_models_cache_lock: + _flat_models_cache[types] = models + + +def clear_flat_models_cache() -> None: + """Clear the flat models cache. Called when schema needs regeneration.""" + with _flat_models_cache_lock: + _flat_models_cache.clear() + from fastapi._compat import shared from fastapi.openapi.constants import REF_TEMPLATE from fastapi.types import IncEx, ModelNameMap, UnionType @@ -249,6 +280,7 @@ def get_schema_from_model_field( return json_schema +@profiled("get_definitions") def get_definitions( *, fields: Sequence[ModelField], @@ -261,12 +293,9 @@ def get_definitions( schema_generator = GenerateJsonSchema(ref_template=REF_TEMPLATE) validation_fields = [field for field in fields if field.mode == "validation"] serialization_fields = [field for field in fields if field.mode == "serialization"] - flat_validation_models = get_flat_models_from_fields( - validation_fields, known_models=set() - ) - flat_serialization_models = get_flat_models_from_fields( - serialization_fields, known_models=set() - ) + # Use cached version for performance + flat_validation_models = get_flat_models_from_fields_cached(validation_fields) + flat_serialization_models = get_flat_models_from_fields_cached(serialization_fields) flat_validation_model_fields = [ ModelField( field_info=FieldInfo(annotation=model), @@ -318,37 +347,66 @@ def _replace_refs( schema: dict[str, Any], old_name_to_new_name_map: dict[str, str], ) -> dict[str, Any]: - new_schema = deepcopy(schema) - for key, value in new_schema.items(): - if key == "$ref": - value = schema["$ref"] - if isinstance(value, str): - ref_name = schema["$ref"].split("/")[-1] - if ref_name in old_name_to_new_name_map: - new_name = old_name_to_new_name_map[ref_name] - new_schema["$ref"] = REF_TEMPLATE.format(model=new_name) - continue - if isinstance(value, dict): - new_schema[key] = _replace_refs( - schema=value, - old_name_to_new_name_map=old_name_to_new_name_map, - ) + """ + Replace $ref values in a schema dict according to the name map. + + Optimized to avoid unnecessary deep copies when no changes are needed. + """ + # Fast path: if no mappings, return original + if not old_name_to_new_name_map: + return schema + + # Check if any replacements are needed before copying + needs_replacement = _schema_needs_ref_replacement(schema, old_name_to_new_name_map) + if not needs_replacement: + return schema + + # Only deepcopy when we know we need to make changes + return _replace_refs_in_place(deepcopy(schema), old_name_to_new_name_map) + + +def _schema_needs_ref_replacement( + schema: dict[str, Any], + old_name_to_new_name_map: dict[str, str], +) -> bool: + """Check if schema contains any refs that need replacement.""" + for key, value in schema.items(): + if key == "$ref" and isinstance(value, str): + ref_name = value.split("/")[-1] + if ref_name in old_name_to_new_name_map: + return True + elif isinstance(value, dict): + if _schema_needs_ref_replacement(value, old_name_to_new_name_map): + return True elif isinstance(value, list): - new_value = [] for item in value: if isinstance(item, dict): - new_item = _replace_refs( - schema=item, - old_name_to_new_name_map=old_name_to_new_name_map, - ) - new_value.append(new_item) + if _schema_needs_ref_replacement(item, old_name_to_new_name_map): + return True + return False - else: - new_value.append(item) - new_schema[key] = new_value - return new_schema + +def _replace_refs_in_place( + schema: dict[str, Any], + old_name_to_new_name_map: dict[str, str], +) -> dict[str, Any]: + """Replace refs in-place in an already-copied schema.""" + for key, value in list(schema.items()): + if key == "$ref" and isinstance(value, str): + ref_name = value.split("/")[-1] + if ref_name in old_name_to_new_name_map: + new_name = old_name_to_new_name_map[ref_name] + schema["$ref"] = REF_TEMPLATE.format(model=new_name) + elif isinstance(value, dict): + _replace_refs_in_place(value, old_name_to_new_name_map) + elif isinstance(value, list): + for item in value: + if isinstance(item, dict): + _replace_refs_in_place(item, old_name_to_new_name_map) + return schema +@profiled("_remap_definitions_and_field_mappings") def _remap_definitions_and_field_mappings( *, model_name_map: ModelNameMap, @@ -499,13 +557,11 @@ def get_model_name_map(unique_models: TypeModelSet) -> dict[TypeModelOrEnum, str return {v: k for k, v in name_model_map.items()} +@profiled("get_compat_model_name_map") def get_compat_model_name_map(fields: list[ModelField]) -> ModelNameMap: - all_flat_models = set() - v2_model_fields = [field for field in fields if isinstance(field, ModelField)] - v2_flat_models = get_flat_models_from_fields(v2_model_fields, known_models=set()) - all_flat_models = all_flat_models.union(v2_flat_models) # type: ignore[arg-type] - + # Use cached version for performance + all_flat_models = get_flat_models_from_fields_cached(v2_model_fields) model_name_map = get_model_name_map(all_flat_models) # type: ignore[arg-type] return model_name_map @@ -550,6 +606,7 @@ def get_flat_models_from_field( return known_models +@profiled("get_flat_models_from_fields") def get_flat_models_from_fields( fields: Sequence[ModelField], known_models: TypeModelSet ) -> TypeModelSet: @@ -558,6 +615,32 @@ def get_flat_models_from_fields( return known_models +def get_flat_models_from_fields_cached( + fields: Sequence[ModelField], +) -> TypeModelSet: + """ + Cached version of get_flat_models_from_fields. + + Caches results by the set of field types to avoid redundant traversal. + """ + # Extract unique types from fields + field_types = frozenset(f.type_ for f in fields if f.type_ is not None) + + # Check cache first + cached = _get_cached_flat_models(field_types) + if cached is not None: + return cached.copy() # Return a copy to avoid mutation + + # Compute flat models + known_models: TypeModelSet = set() + for field in fields: + get_flat_models_from_field(field, known_models=known_models) + + # Cache the result + _set_cached_flat_models(field_types, known_models.copy()) + return known_models + + def _regenerate_error_with_loc( *, errors: Sequence[Any], loc_prefix: tuple[Union[str, int], ...] ) -> list[dict[str, Any]]: diff --git a/fastapi/applications.py b/fastapi/applications.py index 54175cb3b..329baab69 100644 --- a/fastapi/applications.py +++ b/fastapi/applications.py @@ -25,6 +25,13 @@ from fastapi.openapi.docs import ( get_swagger_ui_html, get_swagger_ui_oauth2_redirect_html, ) +from fastapi.openapi.plugins import ( + OpenAPIPlugin, + OpenAPIPluginBase, + OpenAPIPluginSettings, + PluginExecutor, + PluginRegistry, +) from fastapi.openapi.utils import get_openapi from fastapi.params import Depends from fastapi.types import DecoratedCallable, IncEx @@ -841,6 +848,35 @@ class FastAPI(Starlette): """ ), ] = None, + openapi_plugins: Annotated[ + Optional[Sequence[Union[OpenAPIPlugin, OpenAPIPluginBase]]], + Doc( + """ + A list of OpenAPI plugins to customize schema generation. + + Plugins can modify operations, add custom extensions, and + transform the final schema without monkey-patching FastAPI. + + **Example**: + + ```python + from fastapi import FastAPI + from fastapi.openapi.plugins import OpenAPIPluginBase + + class MyPlugin(OpenAPIPluginBase): + @property + def name(self) -> str: + return "my-plugin" + + def modify_operation(self, route, method, operation): + operation["x-custom"] = True + return operation + + app = FastAPI(openapi_plugins=[MyPlugin()]) + ``` + """ + ), + ] = None, **extra: Annotated[ Any, Doc( @@ -871,6 +907,13 @@ class FastAPI(Starlette): self.separate_input_output_schemas = separate_input_output_schemas self.openapi_external_docs = openapi_external_docs self.extra = extra + # Initialize plugin system + self._openapi_plugin_registry: PluginRegistry = PluginRegistry( + on_cache_invalidation=self._invalidate_openapi_cache + ) + if openapi_plugins: + for plugin in openapi_plugins: + self._openapi_plugin_registry.register(plugin) self.openapi_version: Annotated[ str, Doc( @@ -1043,6 +1086,28 @@ class FastAPI(Starlette): app = cls(app, *args, **kwargs) return app + def _invalidate_openapi_cache(self) -> None: + """Invalidate the cached OpenAPI schema when plugins change.""" + self.openapi_schema = None + self._openapi_plugin_generation: int = getattr( + self, "_openapi_plugin_generation", 0 + ) + + @property + def openapi_plugins(self) -> PluginRegistry: + """ + Access the OpenAPI plugin registry. + + Use this to register, enable, or disable plugins at runtime. + + **Example**: + ```python + app.openapi_plugins.register(MyPlugin()) + app.openapi_plugins.disable("my-plugin") + ``` + """ + return self._openapi_plugin_registry + def openapi(self) -> dict[str, Any]: """ Generate the OpenAPI schema of the application. This is called by FastAPI @@ -1054,10 +1119,35 @@ class FastAPI(Starlette): If you need to modify the generated OpenAPI schema, you could modify it. + Plugins registered via `openapi_plugins` will be invoked during generation. + When plugins change, the cache is automatically invalidated. + Read more in the [FastAPI docs for OpenAPI](https://fastapi.tiangolo.com/how-to/extending-openapi/). """ + # Check if plugins have changed since last generation + current_gen = self._openapi_plugin_registry.generation_count + cached_gen = getattr(self, "_openapi_plugin_generation", -1) + if current_gen != cached_gen: + self.openapi_schema = None + self._openapi_plugin_generation = current_gen + if not self.openapi_schema: + # Create plugin executor if we have plugins + plugin_executor = None + if len(self._openapi_plugin_registry) > 0: + plugin_executor = PluginExecutor(self._openapi_plugin_registry) + # Call pre_schema_generation hook + settings = OpenAPIPluginSettings( + title=self.title, + version=self.version, + openapi_version=self.openapi_version, + description=self.description, + routes_count=len(self.routes), + separate_input_output_schemas=self.separate_input_output_schemas, + ) + plugin_executor.execute_pre_schema_generation(self, settings) + self.openapi_schema = get_openapi( title=self.title, version=self.version, @@ -1073,6 +1163,7 @@ class FastAPI(Starlette): servers=self.servers, separate_input_output_schemas=self.separate_input_output_schemas, external_docs=self.openapi_external_docs, + plugin_executor=plugin_executor, ) return self.openapi_schema diff --git a/fastapi/background.py b/fastapi/background.py index 20803ba67..0ff6a90a6 100644 --- a/fastapi/background.py +++ b/fastapi/background.py @@ -1,11 +1,30 @@ +import logging +import warnings from typing import Annotated, Any, Callable from annotated_doc import Doc -from starlette.background import BackgroundTasks as StarletteBackgroundTasks +from starlette.background import BackgroundTask, BackgroundTasks as StarletteBackgroundTasks from typing_extensions import ParamSpec P = ParamSpec("P") +logger = logging.getLogger("fastapi.background") + + +class BackgroundTaskError(Exception): + """Exception raised when multiple background tasks fail. + + Attributes: + errors: List of (task, exception) tuples for each failed task. + Note: Holding references to tasks may keep their args/kwargs + in memory until this exception is garbage collected. + """ + + def __init__(self, errors: list[tuple[BackgroundTask, BaseException]]): + self.errors = errors + task_count = len(errors) + super().__init__(f"{task_count} background task(s) failed") + class BackgroundTasks(StarletteBackgroundTasks): """ @@ -36,6 +55,10 @@ class BackgroundTasks(StarletteBackgroundTasks): ``` """ + def __init__(self, tasks: list[BackgroundTask] | None = None): + super().__init__(tasks) + self._executed = False + def add_task( self, func: Annotated[ @@ -57,4 +80,83 @@ class BackgroundTasks(StarletteBackgroundTasks): Read more about it in the [FastAPI docs for Background Tasks](https://fastapi.tiangolo.com/tutorial/background-tasks/). """ + if self._executed: + warnings.warn( + "Background task added after tasks have already been executed. " + "This task will not run. This commonly happens when adding tasks " + "after a 'yield' in a dependency. Consider adding tasks before " + "the yield, or use a different approach for cleanup tasks.", + UserWarning, + stacklevel=2, + ) + return return super().add_task(func, *args, **kwargs) + + def _get_task_name(self, task: BackgroundTask) -> str: + """Safely get a task name for logging.""" + try: + name = getattr(task.func, "__name__", None) + if name is not None: + return name + return repr(task.func) + except Exception: + return "" + + async def __call__(self) -> None: + """ + Execute all background tasks. + + Unlike Starlette's implementation, this continues executing remaining + tasks even if some tasks fail, ensuring all tasks get a chance to run. + Errors are collected and logged. + + For backward compatibility: + - If only one task fails, the original exception is re-raised + - If multiple tasks fail, a BackgroundTaskError is raised with all errors + """ + # Fix #8: Snapshot tasks to prevent mutation during iteration + tasks_snapshot = list(self.tasks) + + # Fix #7: Set _executed after snapshot, so it reflects actual execution attempt + self._executed = True + + errors: list[tuple[BackgroundTask, BaseException]] = [] + + for task in tasks_snapshot: + try: + await task() + except BaseException as exc: + # Fix #1: Catch BaseException but re-raise critical exceptions + # that should not be suppressed + if isinstance(exc, (KeyboardInterrupt, SystemExit)): + raise + + # Fix #2 & #10: Safe logging that won't break the loop + # Use %-style formatting for lazy evaluation + try: + if logger.isEnabledFor(logging.ERROR): + task_name = self._get_task_name(task) + logger.exception( + "Background task %s failed", task_name + ) + except Exception: + # Never let logging failures break task execution + pass + + # Fix #6: Wrap in try/except to handle MemoryError etc. + try: + errors.append((task, exc)) + except BaseException: + # If we can't even store the error (e.g., MemoryError), + # just continue to give remaining tasks a chance + pass + + # Handle errors with backward compatibility + if len(errors) == 1: + # Fix #5: Single error - re-raise with proper context + # Using 'from' preserves the chain and adds context + original_exc = errors[0][1] + raise original_exc from original_exc + elif len(errors) > 1: + # Multiple errors: raise aggregate exception + raise BackgroundTaskError(errors) diff --git a/fastapi/lifespan.py b/fastapi/lifespan.py new file mode 100644 index 000000000..3f11f7233 --- /dev/null +++ b/fastapi/lifespan.py @@ -0,0 +1,203 @@ +"""Unified lifespan management for FastAPI.""" + +import warnings +from collections.abc import AsyncGenerator, Callable, Sequence +from contextlib import asynccontextmanager +from typing import Any, Optional, TypeVar + +from starlette.types import ASGIApp, Lifespan + +from fastapi.exceptions import FastAPIDeprecationWarning + +AppType = TypeVar("AppType", bound="ASGIApp") + + +class LifespanError(Exception): + """Exception raised when lifespan handling fails.""" + + def __init__( + self, + message: str, + *, + phase: str, + handler: Optional[Callable[..., Any]] = None, + original_error: Optional[BaseException] = None, + ): + super().__init__(message) + self.phase = phase + self.handler = handler + self.original_error = original_error + + def __str__(self) -> str: + base = f"Lifespan {self.phase} failed" + if self.handler: + handler_name = getattr(self.handler, "__name__", str(self.handler)) + base += f" in handler '{handler_name}'" + if self.original_error: + base += f": {self.original_error}" + return base + + +def _emit_on_event_deprecation_warning(event_type: str) -> None: + warnings.warn( + f"on_event('{event_type}') is deprecated. Use the lifespan context manager.", + FastAPIDeprecationWarning, + stacklevel=4, + ) + + +class UnifiedLifespanManager: + """Manages unified lifespan handling with rollback semantics.""" + + def __init__( + self, + *, + on_startup: Optional[Sequence[Callable[..., Any]]] = None, + on_shutdown: Optional[Sequence[Callable[..., Any]]] = None, + lifespan: Optional[Lifespan[AppType]] = None, + emit_deprecation_warnings: bool = True, + ): + self._on_startup = list(on_startup or []) + self._on_shutdown = list(on_shutdown or []) + self._lifespan = lifespan + self._emit_warnings = emit_deprecation_warnings + self._completed_startup_handlers: list[Callable[..., Any]] = [] + + if lifespan and (on_startup or on_shutdown): + warnings.warn( + "Using both 'lifespan' and 'on_startup'/'on_shutdown' is not recommended.", + FastAPIDeprecationWarning, + stacklevel=2, + ) + + @asynccontextmanager + async def lifespan(self, app: AppType) -> AsyncGenerator[dict[str, Any], None]: + """Unified lifespan context manager with rollback semantics.""" + state: dict[str, Any] = {} + self._completed_startup_handlers = [] + lifespan_cm = None + + try: + await self._execute_startup_handlers() + + if self._lifespan: + lifespan_cm = self._lifespan(app) + lifespan_state = await lifespan_cm.__aenter__() + if lifespan_state: + state.update(lifespan_state) + + try: + yield state + finally: + shutdown_errors: list[tuple[str, BaseException]] = [] + + if lifespan_cm is not None: + try: + await lifespan_cm.__aexit__(None, None, None) + except Exception as e: + shutdown_errors.append(("lifespan context", e)) + + handler_errors = await self._execute_shutdown_handlers() + shutdown_errors.extend(handler_errors) + + if shutdown_errors: + error_msgs = [f"{name}: {err}" for name, err in shutdown_errors] + warnings.warn( + f"Errors during shutdown: {'; '.join(error_msgs)}", + RuntimeWarning, + stacklevel=2, + ) + + except BaseException as startup_error: + await self._rollback_startup(startup_error) + raise + + async def _execute_startup_handlers(self) -> None: + for handler in self._on_startup: + if self._emit_warnings: + _emit_on_event_deprecation_warning("startup") + + try: + result = handler() + if hasattr(result, "__await__"): + await result + self._completed_startup_handlers.append(handler) + except Exception as e: + raise LifespanError( + f"Startup handler failed", + phase="startup", + handler=handler, + original_error=e, + ) from e + + async def _execute_shutdown_handlers(self) -> list[tuple[str, BaseException]]: + errors: list[tuple[str, BaseException]] = [] + for handler in reversed(self._on_shutdown): + if self._emit_warnings: + _emit_on_event_deprecation_warning("shutdown") + + try: + result = handler() + if hasattr(result, "__await__"): + await result + except Exception as e: + handler_name = getattr(handler, "__name__", str(handler)) + errors.append((handler_name, e)) + return errors + + async def _rollback_startup(self, original_error: BaseException) -> None: + if not self._completed_startup_handlers: + return + + completed_count = len(self._completed_startup_handlers) + handlers_to_cleanup = list(reversed(self._on_shutdown[:completed_count])) + + for handler in handlers_to_cleanup: + try: + result = handler() + if hasattr(result, "__await__"): + await result + except Exception as cleanup_error: + handler_name = getattr(handler, "__name__", str(handler)) + warnings.warn( + f"Error during rollback in '{handler_name}': {cleanup_error}", + RuntimeWarning, + stacklevel=2, + ) + + def add_startup_handler(self, handler: Callable[..., Any]) -> None: + self._on_startup.append(handler) + + def add_shutdown_handler(self, handler: Callable[..., Any]) -> None: + self._on_shutdown.append(handler) + + +def create_unified_lifespan( + *, + on_startup: Optional[Sequence[Callable[..., Any]]] = None, + on_shutdown: Optional[Sequence[Callable[..., Any]]] = None, + lifespan: Optional[Lifespan[AppType]] = None, +) -> Lifespan[AppType]: + """Create a unified lifespan from mixed inputs.""" + manager = UnifiedLifespanManager( + on_startup=on_startup, + on_shutdown=on_shutdown, + lifespan=lifespan, + ) + return manager.lifespan # type: ignore + + +def merge_lifespan_state( + parent_state: Optional[dict[str, Any]], + child_state: Optional[dict[str, Any]], + *, + child_overrides_parent: bool = True, +) -> dict[str, Any]: + """Merge lifespan state dictionaries.""" + parent = parent_state or {} + child = child_state or {} + + if child_overrides_parent: + return {**parent, **child} + else: + return {**child, **parent} diff --git a/fastapi/openapi/_profiling.py b/fastapi/openapi/_profiling.py new file mode 100644 index 000000000..43f50272d --- /dev/null +++ b/fastapi/openapi/_profiling.py @@ -0,0 +1,205 @@ +"""OpenAPI schema generation profiling utilities.""" + +import functools +import threading +import time +from collections import defaultdict +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any, Callable, Optional, TypeVar + +F = TypeVar("F", bound=Callable[..., Any]) + + +@dataclass +class TimingEntry: + """Single timing measurement entry.""" + + name: str + duration_ms: float + call_count: int = 1 + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class ProfilingStats: + """Aggregated profiling statistics for a function.""" + + name: str + total_time_ms: float = 0.0 + call_count: int = 0 + min_time_ms: float = float("inf") + max_time_ms: float = 0.0 + timings: list[float] = field(default_factory=list) + + @property + def avg_time_ms(self) -> float: + return self.total_time_ms / self.call_count if self.call_count > 0 else 0.0 + + def add_timing(self, duration_ms: float) -> None: + self.total_time_ms += duration_ms + self.call_count += 1 + self.min_time_ms = min(self.min_time_ms, duration_ms) + self.max_time_ms = max(self.max_time_ms, duration_ms) + self.timings.append(duration_ms) + + +class OpenAPIProfiler: + """Thread-safe profiler for OpenAPI schema generation.""" + + def __init__(self) -> None: + self._enabled = False + self._stats: dict[str, ProfilingStats] = defaultdict( + lambda: ProfilingStats(name="") + ) + self._lock = threading.Lock() + self._context_stack: list[str] = [] + self._start_time: Optional[float] = None + + def enable(self) -> None: + """Enable profiling.""" + with self._lock: + self._enabled = True + self._start_time = time.perf_counter() + + def disable(self) -> None: + """Disable profiling.""" + with self._lock: + self._enabled = False + + def reset(self) -> None: + """Reset all collected statistics.""" + with self._lock: + self._stats.clear() + self._context_stack.clear() + self._start_time = None + + @property + def is_enabled(self) -> bool: + return self._enabled + + def record(self, name: str, duration_ms: float) -> None: + """Record a timing measurement.""" + if not self._enabled: + return + with self._lock: + if name not in self._stats: + self._stats[name] = ProfilingStats(name=name) + self._stats[name].add_timing(duration_ms) + + @contextmanager + def measure(self, name: str): + """Context manager to measure execution time of a code block.""" + if not self._enabled: + yield + return + + start = time.perf_counter() + try: + yield + finally: + duration_ms = (time.perf_counter() - start) * 1000 + self.record(name, duration_ms) + + def get_stats(self) -> dict[str, ProfilingStats]: + """Get a copy of current statistics.""" + with self._lock: + return dict(self._stats) + + def get_total_time_ms(self) -> float: + """Get total elapsed time since profiling was enabled.""" + if self._start_time is None: + return 0.0 + return (time.perf_counter() - self._start_time) * 1000 + + def get_report(self) -> str: + """Generate a human-readable profiling report.""" + stats = self.get_stats() + if not stats: + return "No profiling data collected." + + total_time = self.get_total_time_ms() + lines = [ + "=" * 80, + "OpenAPI Schema Generation Profiling Report", + "=" * 80, + f"Total elapsed time: {total_time:.2f}ms", + "", + f"{'Function':<45} {'Calls':>6} {'Total':>10} {'Avg':>10} {'%':>6}", + "-" * 80, + ] + + sorted_stats = sorted( + stats.values(), key=lambda s: s.total_time_ms, reverse=True + ) + + for stat in sorted_stats: + pct = (stat.total_time_ms / total_time * 100) if total_time > 0 else 0 + lines.append( + f"{stat.name:<45} {stat.call_count:>6} " + f"{stat.total_time_ms:>9.2f}ms {stat.avg_time_ms:>9.2f}ms {pct:>5.1f}%" + ) + + lines.append("=" * 80) + return "\n".join(lines) + + def print_report(self) -> None: + """Print the profiling report to stdout.""" + print(self.get_report()) + + +# Global profiler instance +openapi_profiler = OpenAPIProfiler() + + +def profiled(name: Optional[str] = None) -> Callable[[F], F]: + """Decorator to instrument a function for profiling.""" + + def decorator(func: F) -> F: + func_name = name or f"{func.__module__}.{func.__qualname__}" + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + if not openapi_profiler.is_enabled: + return func(*args, **kwargs) + + start = time.perf_counter() + try: + return func(*args, **kwargs) + finally: + duration_ms = (time.perf_counter() - start) * 1000 + openapi_profiler.record(func_name, duration_ms) + + return wrapper # type: ignore[return-value] + + return decorator + + +class ProfilingContext: + """Context manager for scoped profiling sessions.""" + + def __init__(self, auto_print: bool = False) -> None: + self._auto_print = auto_print + self._profiler = openapi_profiler + + def __enter__(self) -> "ProfilingContext": + self._profiler.reset() + self._profiler.enable() + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self._profiler.disable() + if self._auto_print: + self.print_report() + + def get_stats(self) -> dict[str, ProfilingStats]: + return self._profiler.get_stats() + + def get_report(self) -> str: + return self._profiler.get_report() + + def print_report(self) -> None: + self._profiler.print_report() + + def get_total_time_ms(self) -> float: + return self._profiler.get_total_time_ms() diff --git a/fastapi/openapi/plugins.py b/fastapi/openapi/plugins.py new file mode 100644 index 000000000..549929ef4 --- /dev/null +++ b/fastapi/openapi/plugins.py @@ -0,0 +1,376 @@ +"""OpenAPI plugin system for extending schema generation.""" + +import threading +import warnings +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Optional, Protocol, Union, runtime_checkable + +if TYPE_CHECKING: + from fastapi import FastAPI + from fastapi.routing import APIRoute + + +@dataclass +class OpenAPIPluginSettings: + """Settings passed to plugins during schema generation.""" + + title: str + version: str + openapi_version: str + description: Optional[str] = None + routes_count: int = 0 + separate_input_output_schemas: bool = True + + +@runtime_checkable +class OpenAPIPlugin(Protocol): + """Protocol defining the interface for OpenAPI plugins.""" + + @property + def name(self) -> str: + """Unique identifier for the plugin.""" + ... + + @property + def priority(self) -> int: + """Execution priority. Lower numbers execute first. Default is 100.""" + ... + + def pre_schema_generation( + self, + app: "FastAPI", + settings: OpenAPIPluginSettings, + ) -> None: + """Called before schema generation begins.""" + ... + + def modify_operation( + self, + route: "APIRoute", + method: str, + operation: dict[str, Any], + ) -> dict[str, Any]: + """Called for each operation. Returns the modified operation dict.""" + ... + + def modify_schema( + self, + schema: dict[str, Any], + ) -> dict[str, Any]: + """Called after base schema is built. Returns the modified schema.""" + ... + + def post_schema_generation( + self, + schema: dict[str, Any], + ) -> dict[str, Any]: + """Called as the final step. Returns the final schema.""" + ... + + +class OpenAPIPluginBase(ABC): + """Abstract base class for OpenAPI plugins with default implementations.""" + + @property + @abstractmethod + def name(self) -> str: + """Unique plugin identifier.""" + raise NotImplementedError + + @property + def priority(self) -> int: + """Execution priority. Override to change order.""" + return 100 + + def pre_schema_generation( + self, + app: "FastAPI", + settings: OpenAPIPluginSettings, + ) -> None: + pass + + def modify_operation( + self, + route: "APIRoute", + method: str, + operation: dict[str, Any], + ) -> dict[str, Any]: + return operation + + def modify_schema( + self, + schema: dict[str, Any], + ) -> dict[str, Any]: + return schema + + def post_schema_generation( + self, + schema: dict[str, Any], + ) -> dict[str, Any]: + return schema + + +@dataclass +class PluginState: + """Internal state for a registered plugin.""" + + plugin: Union[OpenAPIPlugin, OpenAPIPluginBase] + enabled: bool = True + error_count: int = 0 + last_error: Optional[Exception] = None + + +class PluginRegistry: + """Thread-safe registry for managing OpenAPI plugins.""" + + def __init__( + self, + *, + max_errors_before_disable: int = 3, + on_cache_invalidation: Optional[callable] = None, + ) -> None: + self._plugins: dict[str, PluginState] = {} + self._lock = threading.RLock() + self._max_errors = max_errors_before_disable + self._cache_invalidation_callback = on_cache_invalidation + self._generation_count = 0 + + def register( + self, + plugin: Union[OpenAPIPlugin, OpenAPIPluginBase], + *, + enabled: bool = True, + replace: bool = False, + ) -> None: + """Register a plugin. Raises ValueError if name exists and replace=False.""" + if not hasattr(plugin, "name"): + raise TypeError( + f"Plugin must have a 'name' property, got {type(plugin).__name__}" + ) + + name = plugin.name + if not isinstance(name, str) or not name: + raise ValueError(f"Plugin name must be a non-empty string, got {name!r}") + + with self._lock: + if name in self._plugins and not replace: + raise ValueError( + f"Plugin '{name}' is already registered. " + f"Use replace=True to override or unregister first." + ) + self._plugins[name] = PluginState(plugin=plugin, enabled=enabled) + self._invalidate_cache() + + def unregister(self, name: str) -> bool: + """Remove a plugin. Returns True if removed, False if not found.""" + with self._lock: + if name in self._plugins: + del self._plugins[name] + self._invalidate_cache() + return True + return False + + def enable(self, name: str) -> bool: + """Enable a plugin. Returns True if enabled, False if not found.""" + with self._lock: + if name in self._plugins: + state = self._plugins[name] + if not state.enabled: + state.enabled = True + state.error_count = 0 + self._invalidate_cache() + return True + return False + + def disable(self, name: str) -> bool: + """Disable a plugin. Returns True if disabled, False if not found.""" + with self._lock: + if name in self._plugins: + state = self._plugins[name] + if state.enabled: + state.enabled = False + self._invalidate_cache() + return True + return False + + def is_enabled(self, name: str) -> bool: + """Check if a plugin is enabled.""" + with self._lock: + state = self._plugins.get(name) + return state.enabled if state else False + + def get_plugin( + self, name: str + ) -> Optional[Union[OpenAPIPlugin, OpenAPIPluginBase]]: + """Get a plugin by name.""" + with self._lock: + state = self._plugins.get(name) + return state.plugin if state else None + + def get_active_plugins(self) -> list[Union[OpenAPIPlugin, OpenAPIPluginBase]]: + """Get all enabled plugins sorted by priority (lowest first).""" + with self._lock: + active = [ + state.plugin for state in self._plugins.values() if state.enabled + ] + return sorted(active, key=lambda p: getattr(p, "priority", 100)) + + def get_all_plugins( + self, + ) -> dict[str, tuple[Union[OpenAPIPlugin, OpenAPIPluginBase], bool]]: + """Get all plugins with their enabled status.""" + with self._lock: + return { + name: (state.plugin, state.enabled) + for name, state in self._plugins.items() + } + + def record_error(self, name: str, error: Exception) -> None: + """Record an error. Auto-disables plugin after threshold.""" + with self._lock: + state = self._plugins.get(name) + if state: + state.error_count += 1 + state.last_error = error + if state.error_count >= self._max_errors: + warnings.warn( + f"OpenAPI plugin '{name}' auto-disabled after " + f"{state.error_count} errors. Last error: {error}", + stacklevel=2, + ) + state.enabled = False + self._invalidate_cache() + + def clear(self) -> None: + """Remove all plugins.""" + with self._lock: + self._plugins.clear() + self._invalidate_cache() + + @property + def generation_count(self) -> int: + """Counter incremented on cache invalidation.""" + return self._generation_count + + def _invalidate_cache(self) -> None: + self._generation_count += 1 + if self._cache_invalidation_callback: + try: + self._cache_invalidation_callback() + except Exception: + pass + + def __len__(self) -> int: + return len(self._plugins) + + def __contains__(self, name: str) -> bool: + return name in self._plugins + + +class PluginExecutor: + """Executes plugin hooks during schema generation.""" + + def __init__(self, registry: PluginRegistry) -> None: + self._registry = registry + + def execute_pre_schema_generation( + self, + app: "FastAPI", + settings: OpenAPIPluginSettings, + ) -> None: + for plugin in self._registry.get_active_plugins(): + try: + if hasattr(plugin, "pre_schema_generation"): + plugin.pre_schema_generation(app, settings) + except Exception as e: + self._handle_error(plugin, "pre_schema_generation", e) + + def execute_modify_operation( + self, + route: "APIRoute", + method: str, + operation: dict[str, Any], + ) -> dict[str, Any]: + result = operation + for plugin in self._registry.get_active_plugins(): + try: + if hasattr(plugin, "modify_operation"): + result = plugin.modify_operation(route, method, result) + if result is None: + result = operation + except Exception as e: + self._handle_error(plugin, "modify_operation", e) + return result + + def execute_modify_schema( + self, + schema: dict[str, Any], + ) -> dict[str, Any]: + result = schema + for plugin in self._registry.get_active_plugins(): + try: + if hasattr(plugin, "modify_schema"): + result = plugin.modify_schema(result) + if result is None: + result = schema + except Exception as e: + self._handle_error(plugin, "modify_schema", e) + return result + + def execute_post_schema_generation( + self, + schema: dict[str, Any], + ) -> dict[str, Any]: + result = schema + for plugin in self._registry.get_active_plugins(): + try: + if hasattr(plugin, "post_schema_generation"): + result = plugin.post_schema_generation(result) + if result is None: + result = schema + except Exception as e: + self._handle_error(plugin, "post_schema_generation", e) + return result + + def _handle_error(self, plugin: Any, hook_name: str, error: Exception) -> None: + plugin_name = getattr(plugin, "name", type(plugin).__name__) + warnings.warn( + f"OpenAPI plugin '{plugin_name}' raised an error in {hook_name}: {error}", + stacklevel=3, + ) + self._registry.record_error(plugin_name, error) + + +def create_plugin( + name: str, + *, + priority: int = 100, + pre_schema_generation: Optional[callable] = None, + modify_operation: Optional[callable] = None, + modify_schema: Optional[callable] = None, + post_schema_generation: Optional[callable] = None, +) -> OpenAPIPluginBase: + """Create a plugin from callback functions.""" + + class CallbackPlugin(OpenAPIPluginBase): + @property + def name(self) -> str: + return name + + @property + def priority(self) -> int: + return priority + + plugin = CallbackPlugin() + + if pre_schema_generation: + plugin.pre_schema_generation = pre_schema_generation # type: ignore + if modify_operation: + plugin.modify_operation = modify_operation # type: ignore + if modify_schema: + plugin.modify_schema = modify_schema # type: ignore + if post_schema_generation: + plugin.post_schema_generation = post_schema_generation # type: ignore + + return plugin diff --git a/fastapi/openapi/utils.py b/fastapi/openapi/utils.py index 75ff26102..3c20b7042 100644 --- a/fastapi/openapi/utils.py +++ b/fastapi/openapi/utils.py @@ -5,6 +5,8 @@ from collections.abc import Sequence from typing import Any, Optional, Union, cast from fastapi import routing +from fastapi.openapi._profiling import openapi_profiler, profiled +from fastapi.openapi.plugins import PluginExecutor from fastapi._compat import ( ModelField, Undefined, @@ -75,6 +77,7 @@ status_code_ranges: dict[str, str] = { } +@profiled("get_openapi_security_definitions") def get_openapi_security_definitions( flat_dependant: Dependant, ) -> tuple[dict[str, Any], list[dict[str, Any]]]: @@ -101,6 +104,7 @@ def get_openapi_security_definitions( return security_definitions, operation_security +@profiled("_get_openapi_operation_parameters") def _get_openapi_operation_parameters( *, dependant: Dependant, @@ -174,6 +178,7 @@ def _get_openapi_operation_parameters( return parameters +@profiled("get_openapi_operation_request_body") def get_openapi_operation_request_body( *, body_field: Optional[ModelField], @@ -230,6 +235,7 @@ def generate_operation_summary(*, route: routing.APIRoute, method: str) -> str: return route.name.replace("_", " ").title() +@profiled("get_openapi_operation_metadata") def get_openapi_operation_metadata( *, route: routing.APIRoute, method: str, operation_ids: set[str] ) -> dict[str, Any]: @@ -256,6 +262,7 @@ def get_openapi_operation_metadata( return operation +@profiled("get_openapi_path") def get_openapi_path( *, route: routing.APIRoute, @@ -265,6 +272,7 @@ def get_openapi_path( tuple[ModelField, Literal["validation", "serialization"]], dict[str, Any] ], separate_input_output_schemas: bool = True, + plugin_executor: Optional[PluginExecutor] = None, ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: path = {} security_schemes: dict[str, Any] = {} @@ -438,10 +446,16 @@ def get_openapi_path( ) if route.openapi_extra: deep_dict_update(operation, route.openapi_extra) + # Allow plugins to modify the operation + if plugin_executor: + operation = plugin_executor.execute_modify_operation( + route, method, operation + ) path[method.lower()] = operation return path, security_schemes, definitions +@profiled("get_fields_from_routes") def get_fields_from_routes( routes: Sequence[BaseRoute], ) -> list[ModelField]: @@ -473,6 +487,7 @@ def get_fields_from_routes( return flat_models +@profiled("get_openapi") def get_openapi( *, title: str, @@ -489,6 +504,7 @@ def get_openapi( license_info: Optional[dict[str, Union[str, Any]]] = None, separate_input_output_schemas: bool = True, external_docs: Optional[dict[str, Any]] = None, + plugin_executor: Optional[PluginExecutor] = None, ) -> dict[str, Any]: info: dict[str, Any] = {"title": title, "version": version} if summary: @@ -523,6 +539,7 @@ def get_openapi( model_name_map=model_name_map, field_mapping=field_mapping, separate_input_output_schemas=separate_input_output_schemas, + plugin_executor=plugin_executor, ) if result: path, security_schemes, path_definitions = result @@ -542,6 +559,7 @@ def get_openapi( model_name_map=model_name_map, field_mapping=field_mapping, separate_input_output_schemas=separate_input_output_schemas, + plugin_executor=plugin_executor, ) if result: path, security_schemes, path_definitions = result @@ -564,4 +582,16 @@ def get_openapi( output["tags"] = tags if external_docs: output["externalDocs"] = external_docs - return jsonable_encoder(OpenAPI(**output), by_alias=True, exclude_none=True) # type: ignore + + # Allow plugins to modify the complete schema before validation + if plugin_executor: + output = plugin_executor.execute_modify_schema(output) + + # Validate and encode the schema + result = jsonable_encoder(OpenAPI(**output), by_alias=True, exclude_none=True) + + # Allow plugins to make final modifications after encoding + if plugin_executor: + result = plugin_executor.execute_post_schema_generation(result) + + return result # type: ignore diff --git a/fastapi/routing.py b/fastapi/routing.py index 9ca2f4673..dbd74fca9 100644 --- a/fastapi/routing.py +++ b/fastapi/routing.py @@ -2,6 +2,7 @@ import email.message import functools import inspect import json +import warnings from collections.abc import ( AsyncIterator, Awaitable, @@ -42,6 +43,7 @@ from fastapi.dependencies.utils import ( from fastapi.encoders import jsonable_encoder from fastapi.exceptions import ( EndpointContext, + FastAPIDeprecationWarning, FastAPIError, PydanticV1NotSupportedError, RequestValidationError, @@ -4500,6 +4502,13 @@ class APIRouter(routing.Router): Read more about it in the [FastAPI docs for Lifespan Events](https://fastapi.tiangolo.com/advanced/events/#alternative-events-deprecated). """ + warnings.warn( + f"on_event('{event_type}') is deprecated and will be removed in a future version. " + f"Use the `lifespan` context manager instead. " + f"See https://fastapi.tiangolo.com/advanced/events/ for migration guide.", + FastAPIDeprecationWarning, + stacklevel=2, + ) def decorator(func: DecoratedCallable) -> DecoratedCallable: self.add_event_handler(event_type, func) diff --git a/tests/benchmarks/test_openapi_performance.py b/tests/benchmarks/test_openapi_performance.py new file mode 100644 index 000000000..89b72b274 --- /dev/null +++ b/tests/benchmarks/test_openapi_performance.py @@ -0,0 +1,1019 @@ +""" +Benchmark tests for OpenAPI schema generation performance. + +This module tests OpenAPI schema generation with large numbers of routes +and Pydantic models to identify performance bottlenecks and verify +optimizations. + +Run with: pytest tests/benchmarks/test_openapi_performance.py -v +Or with profiling: pytest tests/benchmarks/test_openapi_performance.py -v -s +""" + +import time +from collections.abc import Iterator +from datetime import datetime +from enum import Enum +from typing import Annotated, Any, Optional, Union +from uuid import UUID + +import pytest +from fastapi import Depends, FastAPI, Path, Query +from fastapi.openapi._profiling import ProfilingContext +from fastapi.testclient import TestClient +from pydantic import BaseModel, Field + + +# ============================================================================= +# Pydantic Models (50+ unique models with nested structures) +# ============================================================================= + + +class StatusEnum(str, Enum): + ACTIVE = "active" + INACTIVE = "inactive" + PENDING = "pending" + ARCHIVED = "archived" + + +class PriorityEnum(str, Enum): + LOW = "low" + MEDIUM = "medium" + HIGH = "high" + CRITICAL = "critical" + + +class Address(BaseModel): + street: str + city: str + state: str + zip_code: str + country: str = "USA" + + +class ContactInfo(BaseModel): + email: str + phone: Optional[str] = None + address: Optional[Address] = None + + +class BaseEntity(BaseModel): + id: UUID + created_at: datetime + updated_at: datetime + + +class UserBase(BaseModel): + username: str = Field(..., min_length=3, max_length=50) + email: str + full_name: Optional[str] = None + + +class UserCreate(UserBase): + password: str = Field(..., min_length=8) + + +class UserUpdate(BaseModel): + username: Optional[str] = None + email: Optional[str] = None + full_name: Optional[str] = None + + +class User(UserBase, BaseEntity): + is_active: bool = True + role: str = "user" + contact: Optional[ContactInfo] = None + + +class UserWithProfile(User): + profile_picture_url: Optional[str] = None + bio: Optional[str] = None + social_links: dict[str, str] = Field(default_factory=dict) + + +class Tag(BaseModel): + name: str + color: Optional[str] = None + + +class Category(BaseModel): + id: int + name: str + description: Optional[str] = None + parent_id: Optional[int] = None + + +class ProductBase(BaseModel): + name: str + description: Optional[str] = None + price: float = Field(..., ge=0) + sku: str + + +class ProductCreate(ProductBase): + category_id: int + tags: list[str] = Field(default_factory=list) + + +class ProductUpdate(BaseModel): + name: Optional[str] = None + description: Optional[str] = None + price: Optional[float] = None + + +class Product(ProductBase, BaseEntity): + category: Optional[Category] = None + tags: list[Tag] = Field(default_factory=list) + status: StatusEnum = StatusEnum.ACTIVE + inventory_count: int = 0 + + +class OrderItem(BaseModel): + product_id: UUID + quantity: int = Field(..., ge=1) + unit_price: float + discount: float = 0.0 + + +class OrderBase(BaseModel): + customer_id: UUID + shipping_address: Address + billing_address: Optional[Address] = None + + +class OrderCreate(OrderBase): + items: list[OrderItem] + notes: Optional[str] = None + + +class Order(OrderBase, BaseEntity): + items: list[OrderItem] + status: StatusEnum = StatusEnum.PENDING + total_amount: float + tax_amount: float + shipping_cost: float + + +class PaymentMethod(BaseModel): + type: str + last_four: str + expiry_month: int + expiry_year: int + + +class Payment(BaseEntity): + order_id: UUID + amount: float + method: PaymentMethod + status: StatusEnum + + +class Review(BaseEntity): + product_id: UUID + user_id: UUID + rating: int = Field(..., ge=1, le=5) + title: str + content: str + helpful_count: int = 0 + + +class Notification(BaseEntity): + user_id: UUID + title: str + message: str + is_read: bool = False + priority: PriorityEnum = PriorityEnum.MEDIUM + + +class AuditLog(BaseEntity): + action: str + entity_type: str + entity_id: UUID + user_id: UUID + changes: dict[str, Any] + + +class Settings(BaseModel): + theme: str = "light" + language: str = "en" + notifications_enabled: bool = True + email_preferences: dict[str, bool] = Field(default_factory=dict) + + +class Analytics(BaseModel): + page_views: int + unique_visitors: int + bounce_rate: float + avg_session_duration: float + top_pages: list[str] + + +class Report(BaseEntity): + name: str + type: str + parameters: dict[str, Any] + data: Analytics + generated_by: UUID + + +class Webhook(BaseEntity): + url: str + events: list[str] + is_active: bool = True + secret: Optional[str] = None + + +class ApiKey(BaseEntity): + name: str + key_prefix: str + scopes: list[str] + expires_at: Optional[datetime] = None + + +class RateLimit(BaseModel): + requests_per_minute: int + requests_per_hour: int + requests_per_day: int + + +class Subscription(BaseEntity): + user_id: UUID + plan: str + status: StatusEnum + rate_limits: RateLimit + features: list[str] + + +class Invoice(BaseEntity): + subscription_id: UUID + amount: float + status: StatusEnum + due_date: datetime + paid_at: Optional[datetime] = None + + +class Team(BaseEntity): + name: str + description: Optional[str] = None + members: list[UUID] + owner_id: UUID + + +class Project(BaseEntity): + name: str + description: Optional[str] = None + team_id: UUID + status: StatusEnum + priority: PriorityEnum + + +class Task(BaseEntity): + title: str + description: Optional[str] = None + project_id: UUID + assignee_id: Optional[UUID] = None + status: StatusEnum + priority: PriorityEnum + due_date: Optional[datetime] = None + + +class Comment(BaseEntity): + task_id: UUID + user_id: UUID + content: str + parent_id: Optional[UUID] = None + + +class Attachment(BaseEntity): + task_id: UUID + filename: str + file_url: str + file_size: int + mime_type: str + + +class PaginatedResponse(BaseModel): + items: list[Any] + total: int + page: int + page_size: int + has_more: bool + + +class ErrorResponse(BaseModel): + error_code: str + message: str + details: Optional[dict[str, Any]] = None + + +class HealthCheck(BaseModel): + status: str + version: str + uptime: float + components: dict[str, str] + + +# ============================================================================= +# Dependencies +# ============================================================================= + + +def get_current_user() -> User: + return User( + id="00000000-0000-0000-0000-000000000001", + username="testuser", + email="test@example.com", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + +def get_pagination( + page: int = Query(1, ge=1), + page_size: int = Query(20, ge=1, le=100), +) -> dict[str, int]: + return {"page": page, "page_size": page_size} + + +CurrentUser = Annotated[User, Depends(get_current_user)] +Pagination = Annotated[dict[str, int], Depends(get_pagination)] + + +# ============================================================================= +# Create Large Application with 100+ Routes +# ============================================================================= + + +def create_large_app() -> FastAPI: + """Create a FastAPI app with 100+ routes and 50+ models.""" + app = FastAPI( + title="Large Benchmark API", + description="API for benchmarking OpenAPI generation", + version="1.0.0", + ) + + # User routes (10 routes) + @app.get("/users", response_model=list[User], tags=["Users"]) + def list_users(pagination: Pagination): + return [] + + @app.post("/users", response_model=User, tags=["Users"]) + def create_user(user: UserCreate): + pass + + @app.get("/users/{user_id}", response_model=UserWithProfile, tags=["Users"]) + def get_user(user_id: UUID = Path(...)): + pass + + @app.put("/users/{user_id}", response_model=User, tags=["Users"]) + def update_user(user_id: UUID, user: UserUpdate): + pass + + @app.delete("/users/{user_id}", tags=["Users"]) + def delete_user(user_id: UUID): + pass + + @app.get("/users/{user_id}/settings", response_model=Settings, tags=["Users"]) + def get_user_settings(user_id: UUID): + pass + + @app.put("/users/{user_id}/settings", response_model=Settings, tags=["Users"]) + def update_user_settings(user_id: UUID, settings: Settings): + pass + + @app.get( + "/users/{user_id}/notifications", + response_model=list[Notification], + tags=["Users"], + ) + def get_user_notifications(user_id: UUID, pagination: Pagination): + return [] + + @app.get("/users/me", response_model=UserWithProfile, tags=["Users"]) + def get_current_user_profile(current_user: CurrentUser): + return current_user + + @app.get("/users/search", response_model=list[User], tags=["Users"]) + def search_users( + q: str = Query(..., min_length=1), + pagination: Pagination = None, + ): + return [] + + # Product routes (15 routes) + @app.get("/products", response_model=list[Product], tags=["Products"]) + def list_products( + category_id: Optional[int] = None, + status: Optional[StatusEnum] = None, + pagination: Pagination = None, + ): + return [] + + @app.post("/products", response_model=Product, tags=["Products"]) + def create_product(product: ProductCreate, current_user: CurrentUser): + pass + + @app.get("/products/{product_id}", response_model=Product, tags=["Products"]) + def get_product(product_id: UUID): + pass + + @app.put("/products/{product_id}", response_model=Product, tags=["Products"]) + def update_product(product_id: UUID, product: ProductUpdate): + pass + + @app.delete("/products/{product_id}", tags=["Products"]) + def delete_product(product_id: UUID): + pass + + @app.get( + "/products/{product_id}/reviews", + response_model=list[Review], + tags=["Products"], + ) + def get_product_reviews(product_id: UUID, pagination: Pagination): + return [] + + @app.post( + "/products/{product_id}/reviews", response_model=Review, tags=["Products"] + ) + def create_product_review( + product_id: UUID, + rating: int = Query(..., ge=1, le=5), + title: str = Query(...), + content: str = Query(...), + current_user: CurrentUser = None, + ): + pass + + @app.get("/products/featured", response_model=list[Product], tags=["Products"]) + def get_featured_products(): + return [] + + @app.get("/products/search", response_model=list[Product], tags=["Products"]) + def search_products( + q: str = Query(...), + min_price: Optional[float] = None, + max_price: Optional[float] = None, + pagination: Pagination = None, + ): + return [] + + @app.get( + "/products/categories", response_model=list[Category], tags=["Products"] + ) + def list_categories(): + return [] + + @app.post("/products/categories", response_model=Category, tags=["Products"]) + def create_category(category: Category): + pass + + @app.get( + "/products/categories/{category_id}", + response_model=Category, + tags=["Products"], + ) + def get_category(category_id: int): + pass + + @app.get("/products/tags", response_model=list[Tag], tags=["Products"]) + def list_tags(): + return [] + + @app.post("/products/tags", response_model=Tag, tags=["Products"]) + def create_tag(tag: Tag): + pass + + @app.get( + "/products/{product_id}/related", + response_model=list[Product], + tags=["Products"], + ) + def get_related_products(product_id: UUID): + return [] + + # Order routes (12 routes) + @app.get("/orders", response_model=list[Order], tags=["Orders"]) + def list_orders( + status: Optional[StatusEnum] = None, + pagination: Pagination = None, + current_user: CurrentUser = None, + ): + return [] + + @app.post("/orders", response_model=Order, tags=["Orders"]) + def create_order(order: OrderCreate, current_user: CurrentUser): + pass + + @app.get("/orders/{order_id}", response_model=Order, tags=["Orders"]) + def get_order(order_id: UUID): + pass + + @app.put("/orders/{order_id}/status", response_model=Order, tags=["Orders"]) + def update_order_status(order_id: UUID, status: StatusEnum): + pass + + @app.delete("/orders/{order_id}", tags=["Orders"]) + def cancel_order(order_id: UUID): + pass + + @app.get( + "/orders/{order_id}/payments", response_model=list[Payment], tags=["Orders"] + ) + def get_order_payments(order_id: UUID): + return [] + + @app.post("/orders/{order_id}/payments", response_model=Payment, tags=["Orders"]) + def create_payment(order_id: UUID, method: PaymentMethod): + pass + + @app.get( + "/orders/{order_id}/invoice", response_model=Invoice, tags=["Orders"] + ) + def get_order_invoice(order_id: UUID): + pass + + @app.get("/orders/stats", response_model=Analytics, tags=["Orders"]) + def get_order_stats( + start_date: Optional[datetime] = None, end_date: Optional[datetime] = None + ): + pass + + @app.get("/orders/recent", response_model=list[Order], tags=["Orders"]) + def get_recent_orders(limit: int = Query(10, ge=1, le=50)): + return [] + + @app.get( + "/orders/{order_id}/tracking", response_model=dict[str, Any], tags=["Orders"] + ) + def get_order_tracking(order_id: UUID): + return {} + + @app.post("/orders/{order_id}/refund", response_model=Payment, tags=["Orders"]) + def request_refund(order_id: UUID, reason: str = Query(...)): + pass + + # Team/Project routes (20 routes) + @app.get("/teams", response_model=list[Team], tags=["Teams"]) + def list_teams(pagination: Pagination): + return [] + + @app.post("/teams", response_model=Team, tags=["Teams"]) + def create_team(name: str, description: Optional[str] = None): + pass + + @app.get("/teams/{team_id}", response_model=Team, tags=["Teams"]) + def get_team(team_id: UUID): + pass + + @app.put("/teams/{team_id}", response_model=Team, tags=["Teams"]) + def update_team(team_id: UUID, name: Optional[str] = None): + pass + + @app.delete("/teams/{team_id}", tags=["Teams"]) + def delete_team(team_id: UUID): + pass + + @app.get( + "/teams/{team_id}/members", response_model=list[User], tags=["Teams"] + ) + def get_team_members(team_id: UUID): + return [] + + @app.post("/teams/{team_id}/members", tags=["Teams"]) + def add_team_member(team_id: UUID, user_id: UUID): + pass + + @app.delete("/teams/{team_id}/members/{user_id}", tags=["Teams"]) + def remove_team_member(team_id: UUID, user_id: UUID): + pass + + @app.get( + "/teams/{team_id}/projects", response_model=list[Project], tags=["Teams"] + ) + def get_team_projects(team_id: UUID, pagination: Pagination): + return [] + + @app.post( + "/teams/{team_id}/projects", response_model=Project, tags=["Teams"] + ) + def create_project(team_id: UUID, name: str, description: Optional[str] = None): + pass + + @app.get("/projects/{project_id}", response_model=Project, tags=["Projects"]) + def get_project(project_id: UUID): + pass + + @app.put("/projects/{project_id}", response_model=Project, tags=["Projects"]) + def update_project( + project_id: UUID, + name: Optional[str] = None, + status: Optional[StatusEnum] = None, + ): + pass + + @app.delete("/projects/{project_id}", tags=["Projects"]) + def delete_project(project_id: UUID): + pass + + @app.get( + "/projects/{project_id}/tasks", response_model=list[Task], tags=["Projects"] + ) + def get_project_tasks( + project_id: UUID, + status: Optional[StatusEnum] = None, + pagination: Pagination = None, + ): + return [] + + @app.post("/projects/{project_id}/tasks", response_model=Task, tags=["Projects"]) + def create_task( + project_id: UUID, + title: str, + description: Optional[str] = None, + priority: PriorityEnum = PriorityEnum.MEDIUM, + ): + pass + + @app.get("/tasks/{task_id}", response_model=Task, tags=["Tasks"]) + def get_task(task_id: UUID): + pass + + @app.put("/tasks/{task_id}", response_model=Task, tags=["Tasks"]) + def update_task( + task_id: UUID, + title: Optional[str] = None, + status: Optional[StatusEnum] = None, + priority: Optional[PriorityEnum] = None, + ): + pass + + @app.delete("/tasks/{task_id}", tags=["Tasks"]) + def delete_task(task_id: UUID): + pass + + @app.get( + "/tasks/{task_id}/comments", response_model=list[Comment], tags=["Tasks"] + ) + def get_task_comments(task_id: UUID): + return [] + + @app.post("/tasks/{task_id}/comments", response_model=Comment, tags=["Tasks"]) + def add_task_comment(task_id: UUID, content: str): + pass + + # Subscription/Billing routes (10 routes) + @app.get("/subscriptions", response_model=list[Subscription], tags=["Billing"]) + def list_subscriptions(current_user: CurrentUser): + return [] + + @app.post("/subscriptions", response_model=Subscription, tags=["Billing"]) + def create_subscription(plan: str, current_user: CurrentUser): + pass + + @app.get( + "/subscriptions/{subscription_id}", + response_model=Subscription, + tags=["Billing"], + ) + def get_subscription(subscription_id: UUID): + pass + + @app.put( + "/subscriptions/{subscription_id}", + response_model=Subscription, + tags=["Billing"], + ) + def update_subscription(subscription_id: UUID, plan: str): + pass + + @app.delete("/subscriptions/{subscription_id}", tags=["Billing"]) + def cancel_subscription(subscription_id: UUID): + pass + + @app.get( + "/subscriptions/{subscription_id}/invoices", + response_model=list[Invoice], + tags=["Billing"], + ) + def get_subscription_invoices(subscription_id: UUID): + return [] + + @app.get("/invoices/{invoice_id}", response_model=Invoice, tags=["Billing"]) + def get_invoice(invoice_id: UUID): + pass + + @app.post("/invoices/{invoice_id}/pay", response_model=Invoice, tags=["Billing"]) + def pay_invoice(invoice_id: UUID, method: PaymentMethod): + pass + + @app.get("/billing/methods", response_model=list[PaymentMethod], tags=["Billing"]) + def list_payment_methods(current_user: CurrentUser): + return [] + + @app.post("/billing/methods", response_model=PaymentMethod, tags=["Billing"]) + def add_payment_method(method: PaymentMethod, current_user: CurrentUser): + pass + + # API Management routes (10 routes) + @app.get("/api-keys", response_model=list[ApiKey], tags=["API"]) + def list_api_keys(current_user: CurrentUser): + return [] + + @app.post("/api-keys", response_model=ApiKey, tags=["API"]) + def create_api_key(name: str, scopes: list[str], current_user: CurrentUser): + pass + + @app.delete("/api-keys/{key_id}", tags=["API"]) + def revoke_api_key(key_id: UUID): + pass + + @app.get("/webhooks", response_model=list[Webhook], tags=["API"]) + def list_webhooks(current_user: CurrentUser): + return [] + + @app.post("/webhooks", response_model=Webhook, tags=["API"]) + def create_webhook(url: str, events: list[str], current_user: CurrentUser): + pass + + @app.get("/webhooks/{webhook_id}", response_model=Webhook, tags=["API"]) + def get_webhook(webhook_id: UUID): + pass + + @app.put("/webhooks/{webhook_id}", response_model=Webhook, tags=["API"]) + def update_webhook(webhook_id: UUID, url: Optional[str] = None): + pass + + @app.delete("/webhooks/{webhook_id}", tags=["API"]) + def delete_webhook(webhook_id: UUID): + pass + + @app.post("/webhooks/{webhook_id}/test", tags=["API"]) + def test_webhook(webhook_id: UUID): + pass + + @app.get("/rate-limits", response_model=RateLimit, tags=["API"]) + def get_rate_limits(current_user: CurrentUser): + pass + + # Admin/Analytics routes (15 routes) + @app.get("/admin/users", response_model=list[User], tags=["Admin"]) + def admin_list_users( + status: Optional[str] = None, pagination: Pagination = None + ): + return [] + + @app.get("/admin/audit-logs", response_model=list[AuditLog], tags=["Admin"]) + def get_audit_logs( + entity_type: Optional[str] = None, + user_id: Optional[UUID] = None, + pagination: Pagination = None, + ): + return [] + + @app.get("/admin/reports", response_model=list[Report], tags=["Admin"]) + def list_reports(pagination: Pagination): + return [] + + @app.post("/admin/reports", response_model=Report, tags=["Admin"]) + def create_report(name: str, type: str, parameters: dict[str, Any]): + pass + + @app.get("/admin/reports/{report_id}", response_model=Report, tags=["Admin"]) + def get_report(report_id: UUID): + pass + + @app.get("/admin/analytics", response_model=Analytics, tags=["Admin"]) + def get_analytics( + start_date: Optional[datetime] = None, end_date: Optional[datetime] = None + ): + pass + + @app.get( + "/admin/analytics/users", response_model=dict[str, Any], tags=["Admin"] + ) + def get_user_analytics(): + return {} + + @app.get( + "/admin/analytics/products", response_model=dict[str, Any], tags=["Admin"] + ) + def get_product_analytics(): + return {} + + @app.get( + "/admin/analytics/orders", response_model=dict[str, Any], tags=["Admin"] + ) + def get_order_analytics(): + return {} + + @app.get( + "/admin/analytics/revenue", response_model=dict[str, Any], tags=["Admin"] + ) + def get_revenue_analytics(): + return {} + + @app.get("/admin/settings", response_model=dict[str, Any], tags=["Admin"]) + def get_admin_settings(): + return {} + + @app.put("/admin/settings", response_model=dict[str, Any], tags=["Admin"]) + def update_admin_settings(settings: dict[str, Any]): + return {} + + @app.get("/admin/notifications", response_model=list[Notification], tags=["Admin"]) + def get_system_notifications(): + return [] + + @app.post("/admin/notifications", response_model=Notification, tags=["Admin"]) + def create_system_notification(title: str, message: str): + pass + + @app.get( + "/admin/notifications/broadcast", tags=["Admin"] + ) + def broadcast_notification(title: str, message: str): + pass + + # Health/Misc routes (8 routes) + @app.get("/health", response_model=HealthCheck, tags=["System"]) + def health_check(): + pass + + @app.get("/health/ready", tags=["System"]) + def readiness_check(): + return {"status": "ready"} + + @app.get("/health/live", tags=["System"]) + def liveness_check(): + return {"status": "live"} + + @app.get("/version", tags=["System"]) + def get_version(): + return {"version": "1.0.0"} + + @app.get("/config", response_model=dict[str, Any], tags=["System"]) + def get_public_config(): + return {} + + @app.get("/features", response_model=list[str], tags=["System"]) + def get_feature_flags(): + return [] + + @app.get("/metrics", tags=["System"]) + def get_metrics(): + return {} + + @app.get("/error", response_model=ErrorResponse, tags=["System"]) + def get_error_codes(): + pass + + return app + + +# ============================================================================= +# Benchmark Tests +# ============================================================================= + + +@pytest.fixture(scope="module") +def large_app() -> FastAPI: + """Create the large app fixture.""" + return create_large_app() + + +@pytest.fixture(scope="module") +def client(large_app: FastAPI) -> Iterator[TestClient]: + """Create test client fixture.""" + with TestClient(large_app) as c: + yield c + + +class TestOpenAPIGeneration: + """Tests for OpenAPI schema generation performance.""" + + def test_large_app_route_count(self, large_app: FastAPI) -> None: + """Verify the large app has 100+ routes.""" + routes = [r for r in large_app.routes if hasattr(r, "methods")] + assert len(routes) >= 100, f"Expected 100+ routes, got {len(routes)}" + + def test_openapi_generation_baseline(self, large_app: FastAPI) -> None: + """Test baseline OpenAPI generation time without optimization.""" + # Clear any cached schema + large_app.openapi_schema = None + + # Measure generation time + start = time.perf_counter() + schema = large_app.openapi() + elapsed_ms = (time.perf_counter() - start) * 1000 + + # Verify schema is valid + assert "openapi" in schema + assert "paths" in schema + assert len(schema["paths"]) >= 50 # Many unique paths + + # Log timing + print(f"\nOpenAPI generation time: {elapsed_ms:.2f}ms") + print(f"Number of paths: {len(schema['paths'])}") + print(f"Number of schemas: {len(schema.get('components', {}).get('schemas', {}))}") + + def test_openapi_generation_with_profiling(self, large_app: FastAPI) -> None: + """Test OpenAPI generation with detailed profiling.""" + # Clear any cached schema + large_app.openapi_schema = None + + with ProfilingContext() as ctx: + schema = large_app.openapi() + + # Print profiling report + print("\n") + ctx.print_report() + + # Verify schema + assert "openapi" in schema + assert "paths" in schema + + # Get stats for analysis + stats = ctx.get_stats() + assert "get_openapi" in stats + assert stats["get_openapi"].call_count == 1 + + def test_openapi_caching_works(self, large_app: FastAPI) -> None: + """Verify that OpenAPI caching prevents regeneration.""" + # Clear any cached schema + large_app.openapi_schema = None + + # First generation + start1 = time.perf_counter() + schema1 = large_app.openapi() + time1 = (time.perf_counter() - start1) * 1000 + + # Second call should be cached + start2 = time.perf_counter() + schema2 = large_app.openapi() + time2 = (time.perf_counter() - start2) * 1000 + + # Verify caching (second call should be nearly instant) + assert schema1 is schema2 # Same object + assert time2 < time1 / 10 # At least 10x faster + + print(f"\nFirst call: {time1:.2f}ms") + print(f"Cached call: {time2:.4f}ms") + print(f"Speedup: {time1/time2:.1f}x") + + def test_openapi_schema_correctness(self, large_app: FastAPI) -> None: + """Verify the generated schema is correct.""" + # Clear any cached schema + large_app.openapi_schema = None + schema = large_app.openapi() + + # Check structure + assert schema["openapi"].startswith("3.") + assert schema["info"]["title"] == "Large Benchmark API" + assert schema["info"]["version"] == "1.0.0" + + # Check paths exist + assert "/users" in schema["paths"] + assert "/products" in schema["paths"] + assert "/orders" in schema["paths"] + + # Check components/schemas exist + schemas = schema.get("components", {}).get("schemas", {}) + assert "User" in schemas or "UserBase" in schemas + assert "Product" in schemas or "ProductBase" in schemas + + def test_openapi_endpoint_response(self, client: TestClient) -> None: + """Test the /openapi.json endpoint returns valid schema.""" + response = client.get("/openapi.json") + assert response.status_code == 200 + + schema = response.json() + assert "openapi" in schema + assert "paths" in schema + + +class TestOpenAPIPerformanceThresholds: + """Tests with performance thresholds.""" + + def test_generation_under_threshold(self, large_app: FastAPI) -> None: + """OpenAPI generation should complete under threshold.""" + # Clear any cached schema + large_app.openapi_schema = None + + # Multiple runs to get average + times = [] + for _ in range(3): + large_app.openapi_schema = None + start = time.perf_counter() + large_app.openapi() + times.append((time.perf_counter() - start) * 1000) + + avg_time = sum(times) / len(times) + min_time = min(times) + max_time = max(times) + + print(f"\nGeneration times: {[f'{t:.2f}ms' for t in times]}") + print(f"Average: {avg_time:.2f}ms, Min: {min_time:.2f}ms, Max: {max_time:.2f}ms") + + # Threshold: should complete in under 5 seconds for 100+ routes + # This is a baseline - optimizations should improve this + assert avg_time < 5000, f"Generation took {avg_time:.2f}ms, expected < 5000ms" diff --git a/tests/test_background_tasks_async_race.py b/tests/test_background_tasks_async_race.py new file mode 100644 index 000000000..505b63029 --- /dev/null +++ b/tests/test_background_tasks_async_race.py @@ -0,0 +1,370 @@ +""" +Test for potential race conditions in async dependencies with BackgroundTasks. + +These tests specifically look for: +1. Race conditions under concurrent requests +2. Cleanup order issues with nested dependencies +3. Exception handling during cleanup +4. Resource access after cleanup starts +""" + +import asyncio +import threading +from collections import defaultdict +from typing import Any + +import pytest +from fastapi import BackgroundTasks, Depends, FastAPI, HTTPException +from fastapi.testclient import TestClient + + +# Test 1: Verify cleanup timing under simulated high concurrency +def test_cleanup_timing_consistency(): + """ + Test that cleanup timing is consistent across multiple requests. + The cleanup should always happen at the same point relative to background tasks. + """ + results: list[dict] = [] + lock = threading.Lock() + + app = FastAPI() + + async def get_resource(): + resource = {"state": "active", "cleanup_happened": False} + yield resource + resource["cleanup_happened"] = True + resource["state"] = "cleaned" + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + resource: dict = Depends(get_resource), + ): + request_id = len(results) + + async def bg_task(res: dict, req_id: int): + # Record state when background task runs + with lock: + results.append({ + "request_id": req_id, + "state_in_bg": res["state"], + "cleanup_in_bg": res["cleanup_happened"], + }) + + background_tasks.add_task(bg_task, resource, request_id) + return {"request_id": request_id} + + client = TestClient(app) + + # Run multiple requests + for _ in range(20): + response = client.get("/test") + assert response.status_code == 200 + + # All results should be consistent - with default scope, cleanup should + # NOT have happened when background task runs + for result in results: + assert result["state_in_bg"] == "active", f"Inconsistent state: {result}" + assert result["cleanup_in_bg"] is False, f"Cleanup happened too early: {result}" + + +# Test 2: Test for cleanup happening during background task execution +def test_cleanup_during_background_task_execution(): + """ + Test if cleanup can start while a background task is still running. + This shouldn't happen with proper request-scoped dependencies. + """ + events: list[str] = [] + lock = threading.Lock() + + app = FastAPI() + + async def get_slow_resource(): + with lock: + events.append("resource_start") + yield "resource" + with lock: + events.append("resource_cleanup") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + resource: str = Depends(get_slow_resource), + ): + async def slow_bg_task(): + with lock: + events.append("bg_start") + await asyncio.sleep(0.05) # Simulate slow work + with lock: + events.append("bg_end") + + background_tasks.add_task(slow_bg_task) + return {"status": "ok"} + + client = TestClient(app) + response = client.get("/test") + assert response.status_code == 200 + + # Verify order: resource_start -> bg_start -> bg_end -> resource_cleanup + assert events == ["resource_start", "bg_start", "bg_end", "resource_cleanup"] + + +# Test 3: Multiple yield dependencies with different cleanup times +def test_multiple_yield_deps_cleanup_order(): + """ + Test that multiple yield dependencies clean up in the correct order. + """ + cleanup_order: list[str] = [] + + app = FastAPI() + + async def dep_a(): + yield "A" + cleanup_order.append("A") + + async def dep_b(): + yield "B" + cleanup_order.append("B") + + async def dep_c(a: str = Depends(dep_a)): + yield f"C({a})" + cleanup_order.append("C") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + a: str = Depends(dep_a), + b: str = Depends(dep_b), + c: str = Depends(dep_c), + ): + async def bg_task(): + cleanup_order.append("BG") + + background_tasks.add_task(bg_task) + return {"a": a, "b": b, "c": c} + + client = TestClient(app) + cleanup_order.clear() + response = client.get("/test") + assert response.status_code == 200 + + # Background task should run before cleanups + bg_idx = cleanup_order.index("BG") + # C depends on A, so C cleanup should happen before A + c_idx = cleanup_order.index("C") + a_idx = cleanup_order.index("A") + + assert bg_idx < c_idx, "Background task should run before dependency cleanup" + assert c_idx < a_idx, "C should cleanup before A (dependency order)" + + +# Test 4: Exception in background task should not affect cleanup +def test_exception_in_background_task(): + """ + Test that exceptions in background tasks don't prevent cleanup. + """ + events: list[str] = [] + + app = FastAPI() + + async def get_resource(): + events.append("setup") + try: + yield "resource" + finally: + events.append("cleanup") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + resource: str = Depends(get_resource), + ): + async def failing_bg_task(): + events.append("bg_task") + raise ValueError("Background task error") + + background_tasks.add_task(failing_bg_task) + return {"status": "ok"} + + # Test with raise_server_exceptions=True to see the actual exception + client = TestClient(app, raise_server_exceptions=True) + events.clear() + + # Single task failure: original exception is re-raised (backward compatible) + with pytest.raises(ValueError, match="Background task error"): + client.get("/test") + + # Cleanup should still happen even if background task fails + assert "setup" in events + assert "bg_task" in events + assert "cleanup" in events + + +# Test 5: Test with generator (sync) dependencies under async workload +def test_sync_generator_with_async_background_task(): + """ + Test sync generator dependencies with async background tasks. + """ + events: list[str] = [] + + app = FastAPI() + + def sync_resource(): + events.append("sync_setup") + yield "sync_resource" + events.append("sync_cleanup") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + resource: str = Depends(sync_resource), + ): + async def async_bg_task(): + events.append("async_bg_start") + await asyncio.sleep(0.01) + events.append("async_bg_end") + + background_tasks.add_task(async_bg_task) + return {"resource": resource} + + client = TestClient(app) + events.clear() + response = client.get("/test") + assert response.status_code == 200 + + # Sync resource cleanup should happen after async background task + assert events == [ + "sync_setup", + "async_bg_start", + "async_bg_end", + "sync_cleanup", + ] + + +# Test 6: Test with concurrent modification of shared state +def test_concurrent_state_modification(): + """ + Test behavior when background task modifies shared state that + dependency cleanup also accesses. + """ + shared_state = {"value": 0, "history": []} + + app = FastAPI() + + async def get_state(): + shared_state["history"].append(f"setup:{shared_state['value']}") + shared_state["value"] = 1 + yield shared_state + shared_state["history"].append(f"cleanup:{shared_state['value']}") + shared_state["value"] = 0 + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + state: dict = Depends(get_state), + ): + async def bg_task(s: dict): + s["history"].append(f"bg:{s['value']}") + s["value"] = 2 + + background_tasks.add_task(bg_task, state) + return {"value": state["value"]} + + client = TestClient(app) + shared_state["value"] = 0 + shared_state["history"] = [] + + response = client.get("/test") + assert response.status_code == 200 + + # History should show: setup (value 0) -> bg (value 1) -> cleanup (value 2) + assert shared_state["history"] == ["setup:0", "bg:1", "cleanup:2"] + + +# Test 7: Test deeply nested dependencies with background tasks +def test_deeply_nested_dependencies(): + """ + Test deeply nested dependencies (4+ levels) with background tasks. + """ + cleanup_order: list[str] = [] + + app = FastAPI() + + async def level1(): + yield "L1" + cleanup_order.append("L1") + + async def level2(l1: str = Depends(level1)): + yield f"L2({l1})" + cleanup_order.append("L2") + + async def level3(l2: str = Depends(level2)): + yield f"L3({l2})" + cleanup_order.append("L3") + + async def level4(l3: str = Depends(level3)): + yield f"L4({l3})" + cleanup_order.append("L4") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + l4: str = Depends(level4), + ): + async def bg_task(value: str): + cleanup_order.append(f"BG:{value}") + + background_tasks.add_task(bg_task, l4) + return {"value": l4} + + client = TestClient(app) + cleanup_order.clear() + response = client.get("/test") + assert response.status_code == 200 + assert response.json()["value"] == "L4(L3(L2(L1)))" + + # BG should run first, then cleanups in reverse order (L4 -> L3 -> L2 -> L1) + assert cleanup_order[0] == "BG:L4(L3(L2(L1)))" + assert cleanup_order[1:] == ["L4", "L3", "L2", "L1"] + + +# Test 8: Test that function-scoped deps with BackgroundTasks logs warning or error +def test_function_scope_logs_potential_issue(): + """ + Document that function scope with background tasks can cause issues. + The resource is cleaned up before the background task runs. + """ + events: list[str] = [] + + app = FastAPI() + + async def get_resource(): + events.append("setup") + resource = {"active": True} + yield resource + events.append("cleanup") + resource["active"] = False + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + resource: dict = Depends(get_resource, scope="function"), + ): + async def bg_task(res: dict): + events.append(f"bg:active={res['active']}") + + background_tasks.add_task(bg_task, resource) + return {"status": "ok"} + + client = TestClient(app) + events.clear() + response = client.get("/test") + assert response.status_code == 200 + + # With function scope, cleanup happens BEFORE background task + # This is the documented "inconsistent" behavior + assert events == ["setup", "cleanup", "bg:active=False"] + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_background_tasks_bugs.py b/tests/test_background_tasks_bugs.py new file mode 100644 index 000000000..60e88561d --- /dev/null +++ b/tests/test_background_tasks_bugs.py @@ -0,0 +1,287 @@ +""" +Tests reproducing actual bugs related to BackgroundTasks and yield dependencies. + +Based on GitHub issues: +1. #14137 - Background tasks added after yield not executed +2. Starlette #2640 - If first background task fails, remaining tasks don't run +3. Starlette #1438 - Background tasks cancelled on client disconnect +""" + +import asyncio +import warnings +from typing import Generator + +import pytest +from fastapi import BackgroundTasks, Depends, FastAPI +from fastapi.testclient import TestClient + + +class TestBackgroundTasksAfterYield: + """ + Bug: Background tasks added after yield are not executed. + + When code after `yield` in a dependency adds a background task, + that task is silently ignored because the response cycle has completed. + """ + + def test_bg_task_added_before_yield_executes(self): + """Control test: tasks added before yield should execute.""" + executed = [] + + app = FastAPI() + + async def dependency_with_yield(background_tasks: BackgroundTasks): + background_tasks.add_task(lambda: executed.append("before_yield")) + yield "value" + + @app.get("/test") + async def endpoint(dep: str = Depends(dependency_with_yield)): + return {"dep": dep} + + client = TestClient(app) + executed.clear() + response = client.get("/test") + assert response.status_code == 200 + assert "before_yield" in executed + + def test_bg_task_added_after_yield_warns_and_not_executed(self): + """ + FIXED: Tasks added after yield now emit a warning instead of being silently dropped. + + Previously, this was a silent failure. Now users get a clear warning + explaining why the task won't run. + """ + executed = [] + + app = FastAPI() + + async def dependency_with_yield(background_tasks: BackgroundTasks): + executed.append("before_yield") + yield "value" + # This code runs during cleanup + executed.append("after_yield_code_runs") + # Now this emits a warning instead of silently failing + background_tasks.add_task(lambda: executed.append("after_yield_task")) + + @app.get("/test") + async def endpoint(dep: str = Depends(dependency_with_yield)): + return {"dep": dep} + + client = TestClient(app) + executed.clear() + + # Capture warnings during the request + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + response = client.get("/test") + + assert response.status_code == 200 + + # The code after yield runs... + assert "after_yield_code_runs" in executed + # ...the background task does NOT execute (as before) + assert "after_yield_task" not in executed + # BUT now we get a warning about it (FIX!) + assert len(w) == 1 + assert "Background task added after tasks have already been executed" in str(w[0].message) + + +class TestBackgroundTaskFailureCascade: + """ + Bug: If first background task fails, remaining tasks don't run. + + From Starlette Discussion #2640: When one background task raises + an exception, subsequent tasks in the queue are skipped. + """ + + def test_first_task_failure_does_not_stop_remaining_tasks(self): + """ + FIXED: Remaining tasks now run even if an earlier task fails. + + Previously, if the first task failed, subsequent tasks wouldn't run. + With the fix, all tasks are attempted regardless of earlier failures. + """ + executed = [] + + app = FastAPI() + + @app.get("/test") + async def endpoint(background_tasks: BackgroundTasks): + async def failing_task(): + executed.append("task1_start") + raise ValueError("Task 1 failed") + + async def task2(): + executed.append("task2_executed") + + async def task3(): + executed.append("task3_executed") + + background_tasks.add_task(failing_task) + background_tasks.add_task(task2) + background_tasks.add_task(task3) + return {"status": "ok"} + + client = TestClient(app, raise_server_exceptions=False) + executed.clear() + response = client.get("/test") + + # Response succeeds because error happens in background + assert response.status_code == 200 + + # First task starts + assert "task1_start" in executed + + # FIXED: task2 and task3 now execute even though task1 failed + assert "task2_executed" in executed + assert "task3_executed" in executed + + def test_independent_tasks_all_run(self): + """ + FIXED: Independent tasks now all run regardless of others failing. + + With the fix, task3 runs even though the preceding task failed. + """ + executed = [] + + app = FastAPI() + + @app.get("/test") + async def endpoint(background_tasks: BackgroundTasks): + async def task1(): + executed.append("task1") + + async def failing_task(): + executed.append("failing_start") + raise ValueError("I failed") + + async def task3(): + executed.append("task3") + + background_tasks.add_task(task1) + background_tasks.add_task(failing_task) + background_tasks.add_task(task3) + return {"status": "ok"} + + client = TestClient(app, raise_server_exceptions=False) + executed.clear() + response = client.get("/test") + assert response.status_code == 200 + + # Task 1 runs + assert "task1" in executed + # Failing task starts + assert "failing_start" in executed + # FIXED: Task 3 now runs even though the preceding task failed + assert "task3" in executed + + +class TestDependencyCleanupWithFailingTasks: + """ + Test that dependency cleanup happens even when background tasks fail. + """ + + def test_cleanup_happens_despite_bg_task_failure(self): + """ + Verify that yield dependency cleanup runs even if a background task fails. + """ + events = [] + + app = FastAPI() + + async def get_resource(): + events.append("setup") + try: + yield "resource" + finally: + events.append("cleanup") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + resource: str = Depends(get_resource), + ): + async def failing_task(): + events.append("task_start") + raise ValueError("Task failed") + + background_tasks.add_task(failing_task) + return {"status": "ok"} + + client = TestClient(app, raise_server_exceptions=False) + events.clear() + response = client.get("/test") + + # Response should succeed (bg task error is after response) + assert response.status_code == 200 + assert "setup" in events + assert "task_start" in events + # Cleanup should happen even though task failed + assert "cleanup" in events + + +class TestCleanupOrderConsistency: + """ + Test that cleanup order is consistent and predictable. + """ + + def test_cleanup_order_with_multiple_dependencies(self): + """ + Verify cleanup happens in reverse order of setup. + """ + events = [] + + app = FastAPI() + + async def dep_a(): + events.append("setup_a") + try: + yield "A" + finally: + events.append("cleanup_a") + + async def dep_b(a: str = Depends(dep_a)): + events.append("setup_b") + try: + yield f"B({a})" + finally: + events.append("cleanup_b") + + async def dep_c(b: str = Depends(dep_b)): + events.append("setup_c") + try: + yield f"C({b})" + finally: + events.append("cleanup_c") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + c: str = Depends(dep_c), + ): + async def bg_task(): + events.append("bg_task") + + background_tasks.add_task(bg_task) + return {"c": c} + + client = TestClient(app) + events.clear() + response = client.get("/test") + assert response.status_code == 200 + + # Verify setup order: A -> B -> C + assert events.index("setup_a") < events.index("setup_b") + assert events.index("setup_b") < events.index("setup_c") + + # Background task runs after endpoint but before cleanup + assert events.index("bg_task") > events.index("setup_c") + assert events.index("bg_task") < events.index("cleanup_c") + + # Verify cleanup order: C -> B -> A (reverse of setup) + assert events.index("cleanup_c") < events.index("cleanup_b") + assert events.index("cleanup_b") < events.index("cleanup_a") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_background_tasks_dependency_access.py b/tests/test_background_tasks_dependency_access.py new file mode 100644 index 000000000..f7882f9ca --- /dev/null +++ b/tests/test_background_tasks_dependency_access.py @@ -0,0 +1,344 @@ +""" +Test accessing dependencies from within BackgroundTasks. + +This tests scenarios where background tasks try to access dependency-provided +resources directly, which can lead to inconsistent behavior if the dependency +has been cleaned up. +""" + +import asyncio +from typing import Any + +import pytest +from fastapi import BackgroundTasks, Depends, FastAPI, Request +from fastapi.testclient import TestClient + + +# Scenario 1: Background task accessing a resource passed to it +def test_bg_task_with_passed_resource(): + """ + Test that a resource passed to a background task remains valid + when using default (request) scope. + """ + events: list[str] = [] + + app = FastAPI() + + class DatabaseConnection: + def __init__(self): + self.connected = False + self.queries: list[str] = [] + + def connect(self): + self.connected = True + events.append("db_connected") + + def disconnect(self): + self.connected = False + events.append("db_disconnected") + + def query(self, sql: str) -> str: + if not self.connected: + events.append(f"query_failed:{sql}") + raise RuntimeError("Not connected") + events.append(f"query_success:{sql}") + self.queries.append(sql) + return f"result:{sql}" + + async def get_db(): + db = DatabaseConnection() + db.connect() + yield db + db.disconnect() + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + db: DatabaseConnection = Depends(get_db), + ): + async def bg_query(conn: DatabaseConnection): + # This should work because default scope keeps db alive + result = conn.query("SELECT * FROM bg_table") + events.append(f"bg_got:{result}") + + background_tasks.add_task(bg_query, db) + return {"status": "queued"} + + client = TestClient(app) + events.clear() + response = client.get("/test") + assert response.status_code == 200 + + # Verify query succeeded in background task + assert "db_connected" in events + assert "query_success:SELECT * FROM bg_table" in events + assert "bg_got:result:SELECT * FROM bg_table" in events + assert "db_disconnected" in events + + # Verify order: connect -> query -> disconnect + connect_idx = events.index("db_connected") + query_idx = events.index("query_success:SELECT * FROM bg_table") + disconnect_idx = events.index("db_disconnected") + assert connect_idx < query_idx < disconnect_idx + + +# Scenario 2: Background task accessing function-scoped resource (FAILS) +def test_bg_task_with_function_scoped_resource(): + """ + Test that a function-scoped resource is cleaned up before background task runs. + This demonstrates the potential pitfall of using scope="function" with BackgroundTasks. + """ + events: list[str] = [] + + app = FastAPI() + + class DatabaseConnection: + def __init__(self): + self.connected = False + + def connect(self): + self.connected = True + events.append("db_connected") + + def disconnect(self): + self.connected = False + events.append("db_disconnected") + + def query(self, sql: str) -> str: + if not self.connected: + events.append(f"query_failed:{sql}") + return "ERROR: Not connected" + events.append(f"query_success:{sql}") + return f"result:{sql}" + + async def get_db(): + db = DatabaseConnection() + db.connect() + yield db + db.disconnect() + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + db: DatabaseConnection = Depends(get_db, scope="function"), + ): + async def bg_query(conn: DatabaseConnection): + # This will FAIL because function scope cleans up before bg task + result = conn.query("SELECT * FROM bg_table") + events.append(f"bg_got:{result}") + + background_tasks.add_task(bg_query, db) + return {"status": "queued"} + + client = TestClient(app) + events.clear() + response = client.get("/test") + assert response.status_code == 200 + + # With function scope, db is disconnected BEFORE bg task runs + assert "db_connected" in events + assert "db_disconnected" in events + assert "query_failed:SELECT * FROM bg_table" in events + + # Verify problematic order: connect -> disconnect -> query_failed + connect_idx = events.index("db_connected") + disconnect_idx = events.index("db_disconnected") + query_idx = events.index("query_failed:SELECT * FROM bg_table") + assert connect_idx < disconnect_idx < query_idx + + +# Scenario 3: Multiple concurrent requests with shared mutable state +def test_concurrent_requests_shared_state(): + """ + Test behavior when multiple concurrent requests share mutable state + through a dependency. + """ + events: list[str] = [] + request_counter = {"value": 0} + + app = FastAPI() + + async def get_request_id(): + request_counter["value"] += 1 + request_id = request_counter["value"] + events.append(f"setup:{request_id}") + yield request_id + events.append(f"cleanup:{request_id}") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + request_id: int = Depends(get_request_id), + ): + async def bg_task(rid: int): + await asyncio.sleep(0.01) # Small delay + events.append(f"bg:{rid}") + + background_tasks.add_task(bg_task, request_id) + return {"request_id": request_id} + + client = TestClient(app) + events.clear() + request_counter["value"] = 0 + + # Make multiple sequential requests + for _ in range(3): + response = client.get("/test") + assert response.status_code == 200 + + # Verify each request has proper setup -> bg -> cleanup sequence + # (Note: with default scope, cleanup happens after bg task) + assert len([e for e in events if e.startswith("setup:")]) == 3 + assert len([e for e in events if e.startswith("bg:")]) == 3 + assert len([e for e in events if e.startswith("cleanup:")]) == 3 + + +# Scenario 4: Nested async generators with background tasks +def test_nested_async_generators_with_bg_tasks(): + """ + Test deeply nested async generators and their interaction with background tasks. + """ + events: list[str] = [] + + app = FastAPI() + + async def outer(): + events.append("outer_start") + yield "outer" + events.append("outer_end") + + async def middle(outer_val: str = Depends(outer)): + events.append(f"middle_start:{outer_val}") + yield f"middle({outer_val})" + events.append("middle_end") + + async def inner(middle_val: str = Depends(middle)): + events.append(f"inner_start:{middle_val}") + yield f"inner({middle_val})" + events.append("inner_end") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + value: str = Depends(inner), + ): + async def bg_task(val: str): + events.append(f"bg:{val}") + + background_tasks.add_task(bg_task, value) + return {"value": value} + + client = TestClient(app) + events.clear() + response = client.get("/test") + assert response.status_code == 200 + assert response.json()["value"] == "inner(middle(outer))" + + # Verify setup order: outer -> middle -> inner + assert events.index("outer_start") < events.index("middle_start:outer") + assert events.index("middle_start:outer") < events.index("inner_start:middle(outer)") + + # Verify bg task runs before cleanup + bg_idx = events.index("bg:inner(middle(outer))") + assert bg_idx < events.index("inner_end") + + # Verify cleanup order (reverse): inner -> middle -> outer + assert events.index("inner_end") < events.index("middle_end") + assert events.index("middle_end") < events.index("outer_end") + + +# Scenario 5: Background task that adds more background tasks +def test_bg_task_adding_more_bg_tasks(): + """ + Test a background task that adds additional background tasks. + """ + events: list[str] = [] + + app = FastAPI() + + async def get_resource(): + events.append("resource_setup") + yield {"active": True} + events.append("resource_cleanup") + + @app.get("/test") + async def endpoint( + request: Request, + background_tasks: BackgroundTasks, + resource: dict = Depends(get_resource), + ): + async def first_bg_task(res: dict, tasks: BackgroundTasks): + events.append(f"first_bg:active={res['active']}") + # Note: Adding tasks from within a task might not work as expected + # because the BackgroundTasks object is tied to the response + + async def second_bg_task(res: dict): + events.append(f"second_bg:active={res['active']}") + + background_tasks.add_task(first_bg_task, resource, background_tasks) + background_tasks.add_task(second_bg_task, resource) + return {"status": "ok"} + + client = TestClient(app) + events.clear() + response = client.get("/test") + assert response.status_code == 200 + + # Both background tasks should see active=True (default scope) + assert "first_bg:active=True" in events + assert "second_bg:active=True" in events + assert "resource_cleanup" in events + + +# Scenario 6: Error in one background task doesn't affect others +def test_error_in_one_bg_task_others_continue(): + """ + Test that an error in one background task doesn't prevent cleanup or other tasks. + Note: Cleanup code MUST be in a finally block to run when exceptions occur. + + With the fix, all tasks run even if one fails. + """ + events: list[str] = [] + + app = FastAPI() + + async def get_resource(): + events.append("setup") + try: + yield "resource" + finally: + events.append("cleanup") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + resource: str = Depends(get_resource), + ): + async def failing_task(): + events.append("failing_start") + raise ValueError("Task failed") + + async def succeeding_task(): + events.append("succeeding") + + background_tasks.add_task(failing_task) + background_tasks.add_task(succeeding_task) + return {"status": "ok"} + + client = TestClient(app, raise_server_exceptions=True) + events.clear() + + # Single task failure: original exception re-raised (backward compatible) + with pytest.raises(ValueError, match="Task failed"): + client.get("/test") + + # Cleanup should still happen + assert "setup" in events + assert "failing_start" in events + assert "cleanup" in events + # FIXED: Second task now runs even though first task failed + assert "succeeding" in events + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_background_tasks_fixes.py b/tests/test_background_tasks_fixes.py new file mode 100644 index 000000000..605bea829 --- /dev/null +++ b/tests/test_background_tasks_fixes.py @@ -0,0 +1,411 @@ +""" +Tests verifying the fixes for BackgroundTasks bugs. + +Fixed bugs: +1. First task failure now doesn't stop remaining tasks +2. Tasks added after execution now emit a warning +""" + +import warnings +from typing import Generator + +import pytest +from fastapi import BackgroundTasks, Depends, FastAPI +from fastapi.background import BackgroundTaskError +from fastapi.testclient import TestClient + + +class TestBackgroundTaskErrorIsolation: + """ + Test that task failures don't stop remaining tasks. + """ + + def test_all_tasks_run_despite_failure(self): + """ + FIX: All tasks should run even if one fails in the middle. + """ + executed = [] + + app = FastAPI() + + @app.get("/test") + async def endpoint(background_tasks: BackgroundTasks): + async def task1(): + executed.append("task1") + + async def failing_task(): + executed.append("failing_start") + raise ValueError("I failed") + + async def task3(): + executed.append("task3") + + background_tasks.add_task(task1) + background_tasks.add_task(failing_task) + background_tasks.add_task(task3) + return {"status": "ok"} + + client = TestClient(app, raise_server_exceptions=False) + executed.clear() + response = client.get("/test") + assert response.status_code == 200 + + # FIX VERIFICATION: All three tasks should have run + assert "task1" in executed + assert "failing_start" in executed + assert "task3" in executed # This now runs! + + def test_multiple_failures_all_logged(self): + """ + Multiple task failures should all be captured. + """ + executed = [] + + app = FastAPI() + + @app.get("/test") + async def endpoint(background_tasks: BackgroundTasks): + async def failing1(): + executed.append("failing1") + raise ValueError("Error 1") + + async def failing2(): + executed.append("failing2") + raise RuntimeError("Error 2") + + async def success(): + executed.append("success") + + background_tasks.add_task(failing1) + background_tasks.add_task(success) + background_tasks.add_task(failing2) + return {"status": "ok"} + + client = TestClient(app, raise_server_exceptions=False) + executed.clear() + response = client.get("/test") + assert response.status_code == 200 + + # All tasks should have executed + assert "failing1" in executed + assert "success" in executed + assert "failing2" in executed + + def test_single_task_failure_raises_original_exception(self): + """ + Single task failure should raise the original exception (backward compatible). + """ + app = FastAPI() + + @app.get("/test") + async def endpoint(background_tasks: BackgroundTasks): + async def failing(): + raise ValueError("Task failed") + + background_tasks.add_task(failing) + return {"status": "ok"} + + client = TestClient(app, raise_server_exceptions=True) + + # Single failure: original exception raised for backward compatibility + with pytest.raises(ValueError, match="Task failed"): + client.get("/test") + + def test_multiple_task_failures_raise_background_task_error(self): + """ + Multiple task failures should raise BackgroundTaskError with all errors. + """ + app = FastAPI() + + @app.get("/test") + async def endpoint(background_tasks: BackgroundTasks): + async def failing1(): + raise ValueError("Error 1") + + async def failing2(): + raise RuntimeError("Error 2") + + background_tasks.add_task(failing1) + background_tasks.add_task(failing2) + return {"status": "ok"} + + client = TestClient(app, raise_server_exceptions=True) + + # Multiple failures: BackgroundTaskError raised + with pytest.raises(BackgroundTaskError) as exc_info: + client.get("/test") + + assert len(exc_info.value.errors) == 2 + assert isinstance(exc_info.value.errors[0][1], ValueError) + assert isinstance(exc_info.value.errors[1][1], RuntimeError) + + +class TestBackgroundTaskAfterYieldWarning: + """ + Test that adding tasks after execution emits a warning. + """ + + def test_warning_when_task_added_after_execution(self): + """ + FIX: Warning should be emitted when task is added after execution. + """ + executed = [] + warnings_captured = [] + + app = FastAPI() + + async def dependency_with_yield(background_tasks: BackgroundTasks): + background_tasks.add_task(lambda: executed.append("before_yield")) + yield "value" + # Capture the warning + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + background_tasks.add_task(lambda: executed.append("after_yield")) + if w: + warnings_captured.extend(w) + + @app.get("/test") + async def endpoint(dep: str = Depends(dependency_with_yield)): + return {"dep": dep} + + client = TestClient(app) + executed.clear() + warnings_captured.clear() + + response = client.get("/test") + assert response.status_code == 200 + + # Task added before yield should have executed + assert "before_yield" in executed + + # Task added after yield should NOT have executed + assert "after_yield" not in executed + + # Warning should have been captured + assert len(warnings_captured) == 1 + assert "Background task added after tasks have already been executed" in str( + warnings_captured[0].message + ) + + def test_no_warning_for_normal_task_addition(self): + """ + No warning should be emitted for normal task addition. + """ + app = FastAPI() + + @app.get("/test") + async def endpoint(background_tasks: BackgroundTasks): + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + background_tasks.add_task(lambda: None) + background_tasks.add_task(lambda: None) + assert len(w) == 0 # No warnings + + return {"status": "ok"} + + client = TestClient(app) + response = client.get("/test") + assert response.status_code == 200 + + +class TestDependencyCleanupStillWorks: + """ + Verify that dependency cleanup still works correctly with the fixes. + """ + + def test_cleanup_runs_after_all_tasks(self): + """ + Dependency cleanup should run after all background tasks complete. + """ + events = [] + + app = FastAPI() + + async def get_resource(): + events.append("setup") + try: + yield "resource" + finally: + events.append("cleanup") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + resource: str = Depends(get_resource), + ): + async def task1(): + events.append("task1") + + async def task2(): + events.append("task2") + + background_tasks.add_task(task1) + background_tasks.add_task(task2) + return {"status": "ok"} + + client = TestClient(app) + events.clear() + response = client.get("/test") + assert response.status_code == 200 + + # Verify order: setup -> tasks -> cleanup + assert events.index("setup") < events.index("task1") + assert events.index("task1") < events.index("task2") + assert events.index("task2") < events.index("cleanup") + + def test_cleanup_runs_even_when_tasks_fail(self): + """ + Dependency cleanup should run even if background tasks fail. + """ + events = [] + + app = FastAPI() + + async def get_resource(): + events.append("setup") + try: + yield "resource" + finally: + events.append("cleanup") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + resource: str = Depends(get_resource), + ): + async def failing_task(): + events.append("task_start") + raise ValueError("Failed") + + background_tasks.add_task(failing_task) + return {"status": "ok"} + + client = TestClient(app, raise_server_exceptions=False) + events.clear() + response = client.get("/test") + assert response.status_code == 200 + + # Cleanup must happen + assert "setup" in events + assert "task_start" in events + assert "cleanup" in events + + +class TestEdgeCases: + """Test edge cases and robustness fixes.""" + + def test_task_with_no_name_attribute(self): + """Test that tasks without __name__ are handled gracefully.""" + executed = [] + + app = FastAPI() + + @app.get("/test") + async def endpoint(background_tasks: BackgroundTasks): + # Lambda has __name__ = '' + background_tasks.add_task(lambda: executed.append("lambda")) + return {"status": "ok"} + + client = TestClient(app) + executed.clear() + response = client.get("/test") + assert response.status_code == 200 + assert "lambda" in executed + + def test_task_mutation_during_execution(self): + """Test that task list mutation during execution doesn't cause issues.""" + executed = [] + tasks_ref = {"bg": None} + + app = FastAPI() + + @app.get("/test") + async def endpoint(background_tasks: BackgroundTasks): + tasks_ref["bg"] = background_tasks + + async def mutating_task(): + executed.append("task1") + # Try to mutate the tasks list during execution + # This should not affect iteration due to snapshot + tasks_ref["bg"].tasks.append( + BackgroundTasks() # Add garbage + ) + + async def task2(): + executed.append("task2") + + background_tasks.add_task(mutating_task) + background_tasks.add_task(task2) + return {"status": "ok"} + + client = TestClient(app, raise_server_exceptions=False) + executed.clear() + response = client.get("/test") + assert response.status_code == 200 + + # Both tasks should have executed despite mutation attempt + assert "task1" in executed + assert "task2" in executed + + @pytest.mark.anyio + async def test_keyboard_interrupt_not_suppressed(self): + """Test that KeyboardInterrupt is not suppressed.""" + from fastapi.background import BackgroundTasks as BgTasks + + async def raising_task(): + raise KeyboardInterrupt() + + bg = BgTasks() + bg.add_task(raising_task) + + with pytest.raises(KeyboardInterrupt): + await bg() + + @pytest.mark.anyio + async def test_system_exit_not_suppressed(self): + """Test that SystemExit is not suppressed.""" + from fastapi.background import BackgroundTasks as BgTasks + + async def raising_task(): + raise SystemExit(1) + + bg = BgTasks() + bg.add_task(raising_task) + + with pytest.raises(SystemExit): + await bg() + + def test_base_exception_subclass_caught(self): + """Test that BaseException subclasses (except critical ones) are caught.""" + executed = [] + + app = FastAPI() + + class CustomBaseException(BaseException): + pass + + @app.get("/test") + async def endpoint(background_tasks: BackgroundTasks): + async def failing_task(): + executed.append("failing") + raise CustomBaseException("custom error") + + async def second_task(): + executed.append("second") + + background_tasks.add_task(failing_task) + background_tasks.add_task(second_task) + return {"status": "ok"} + + client = TestClient(app, raise_server_exceptions=False) + executed.clear() + response = client.get("/test") + assert response.status_code == 200 + + # Both tasks should execute - CustomBaseException should be caught + assert "failing" in executed + assert "second" in executed + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_background_tasks_yield_scope.py b/tests/test_background_tasks_yield_scope.py new file mode 100644 index 000000000..5f9724c55 --- /dev/null +++ b/tests/test_background_tasks_yield_scope.py @@ -0,0 +1,355 @@ +""" +Test for async dependencies with cleanup (yield) and BackgroundTasks interaction. + +This tests the scenario where dependencies with cleanup logic (yield) may behave +inconsistently when used inside BackgroundTasks under different scope settings. + +The key issue: +- Dependencies with scope="function" have cleanup run BEFORE BackgroundTasks execute +- Dependencies with scope="request" (default) have cleanup run AFTER BackgroundTasks execute + +This can lead to inconsistent behavior where a resource is cleaned up before +a background task that depends on it has a chance to run. +""" + +import asyncio +from typing import Any + +import pytest +from fastapi import BackgroundTasks, Depends, FastAPI +from fastapi.testclient import TestClient + + +# Track cleanup and background task execution order +execution_log: list[str] = [] + + +def reset_log(): + execution_log.clear() + + +# Test 1: Default scope (request) - cleanup should happen AFTER background task +def test_background_task_with_request_scope_yield_dependency(): + """ + With default (request) scope, dependency cleanup should happen AFTER + the background task completes. + """ + reset_log() + app = FastAPI() + resource_value = {"status": "initialized"} + + async def get_resource(): + execution_log.append("resource_setup") + resource_value["status"] = "active" + yield resource_value + execution_log.append("resource_cleanup") + resource_value["status"] = "cleaned_up" + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + resource: dict = Depends(get_resource), + ): + async def bg_task(res: dict): + execution_log.append(f"bg_task_start:status={res['status']}") + await asyncio.sleep(0.01) # Simulate async work + execution_log.append(f"bg_task_end:status={res['status']}") + + background_tasks.add_task(bg_task, resource) + return {"status": resource["status"]} + + client = TestClient(app) + response = client.get("/test") + assert response.status_code == 200 + assert response.json() == {"status": "active"} + + # Verify order: setup -> bg_task -> cleanup + assert execution_log == [ + "resource_setup", + "bg_task_start:status=active", + "bg_task_end:status=active", + "resource_cleanup", + ] + + +# Test 2: Function scope - cleanup happens BEFORE background task (BUG!) +def test_background_task_with_function_scope_yield_dependency(): + """ + With scope="function", dependency cleanup happens BEFORE the background task, + which is inconsistent and potentially problematic. + + This test demonstrates the issue: the resource is cleaned up before + the background task runs, so the background task sees the cleaned up state. + """ + reset_log() + app = FastAPI() + resource_value = {"status": "initialized"} + + async def get_resource(): + execution_log.append("resource_setup") + resource_value["status"] = "active" + yield resource_value + execution_log.append("resource_cleanup") + resource_value["status"] = "cleaned_up" + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + resource: dict = Depends(get_resource, scope="function"), + ): + async def bg_task(res: dict): + execution_log.append(f"bg_task_start:status={res['status']}") + await asyncio.sleep(0.01) + execution_log.append(f"bg_task_end:status={res['status']}") + + background_tasks.add_task(bg_task, resource) + return {"status": resource["status"]} + + client = TestClient(app) + response = client.get("/test") + assert response.status_code == 200 + assert response.json() == {"status": "active"} + + # With function scope, cleanup happens BEFORE background task runs + # This is the inconsistent behavior! + assert execution_log == [ + "resource_setup", + "resource_cleanup", # Cleanup happens first! + "bg_task_start:status=cleaned_up", # BG task sees cleaned up state + "bg_task_end:status=cleaned_up", + ] + + +# Test 3: Database-like connection scenario with function scope +def test_database_connection_closed_before_background_task(): + """ + Simulates a real-world scenario where a database connection is closed + before a background task that needs it runs. + """ + reset_log() + app = FastAPI() + + class FakeDBConnection: + def __init__(self): + self.is_open = False + self.data = [] + + def open(self): + self.is_open = True + execution_log.append("db_opened") + + def close(self): + self.is_open = False + execution_log.append("db_closed") + + def query(self): + if not self.is_open: + execution_log.append("db_query_failed:connection_closed") + raise RuntimeError("Connection closed") + execution_log.append("db_query_success") + return {"result": "data"} + + db = FakeDBConnection() + + async def get_db(): + db.open() + yield db + db.close() + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + conn: FakeDBConnection = Depends(get_db, scope="function"), + ): + async def bg_task(connection: FakeDBConnection): + try: + result = connection.query() + execution_log.append(f"bg_task_got_result") + except RuntimeError as e: + execution_log.append(f"bg_task_error:{e}") + + background_tasks.add_task(bg_task, conn) + return {"status": "task_queued"} + + client = TestClient(app, raise_server_exceptions=False) + response = client.get("/test") + assert response.status_code == 200 + + # The background task fails because connection was closed before it ran + assert "db_closed" in execution_log + assert execution_log.index("db_closed") < execution_log.index("db_query_failed:connection_closed") + + +# Test 4: Nested dependencies with mixed scopes +def test_nested_dependencies_mixed_scopes(): + """ + Test nested dependencies where outer has request scope and inner has function scope. + """ + reset_log() + app = FastAPI() + + async def outer_resource(): + execution_log.append("outer_setup") + yield "outer_active" + execution_log.append("outer_cleanup") + + async def inner_resource(outer: str = Depends(outer_resource)): + execution_log.append(f"inner_setup:outer={outer}") + yield f"inner_active:outer={outer}" + execution_log.append("inner_cleanup") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + inner: str = Depends(inner_resource, scope="function"), + ): + async def bg_task(value: str): + execution_log.append(f"bg_task:{value}") + + background_tasks.add_task(bg_task, inner) + return {"value": inner} + + client = TestClient(app) + response = client.get("/test") + assert response.status_code == 200 + + # Check execution order - inner cleanup should happen before outer + # and before background task due to function scope + print("Execution log:", execution_log) + assert "inner_cleanup" in execution_log + assert "outer_cleanup" in execution_log + + +# Test 5: Multiple background tasks with function scope dependency +def test_multiple_background_tasks_function_scope(): + """ + Multiple background tasks all trying to use a function-scoped dependency. + """ + reset_log() + app = FastAPI() + counter = {"value": 0} + + async def get_counter(): + counter["value"] = 1 + execution_log.append("counter_setup") + yield counter + execution_log.append("counter_cleanup") + counter["value"] = 0 + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + cnt: dict = Depends(get_counter, scope="function"), + ): + async def increment(c: dict): + execution_log.append(f"increment:value={c['value']}") + c["value"] += 1 + + background_tasks.add_task(increment, cnt) + background_tasks.add_task(increment, cnt) + background_tasks.add_task(increment, cnt) + return {"initial_value": cnt["value"]} + + client = TestClient(app) + response = client.get("/test") + assert response.status_code == 200 + assert response.json() == {"initial_value": 1} + + # Check that cleanup happened before background tasks ran + cleanup_idx = execution_log.index("counter_cleanup") + increment_indices = [ + i for i, log in enumerate(execution_log) if log.startswith("increment") + ] + + # With function scope, cleanup happens before all increments + assert all(cleanup_idx < idx for idx in increment_indices) + + +# Test 6: Async workload with concurrent requests +def test_concurrent_requests_with_yield_dependencies(): + """ + Test behavior under concurrent async workload. + """ + reset_log() + app = FastAPI() + shared_state = {"active_connections": 0, "max_connections": 0} + + async def get_connection(): + shared_state["active_connections"] += 1 + shared_state["max_connections"] = max( + shared_state["max_connections"], shared_state["active_connections"] + ) + execution_log.append(f"conn_open:active={shared_state['active_connections']}") + yield shared_state["active_connections"] + shared_state["active_connections"] -= 1 + execution_log.append(f"conn_close:active={shared_state['active_connections']}") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + conn_id: int = Depends(get_connection, scope="function"), + ): + async def bg_task(id: int): + await asyncio.sleep(0.01) + execution_log.append(f"bg_task:conn_id={id}:active={shared_state['active_connections']}") + + background_tasks.add_task(bg_task, conn_id) + return {"conn_id": conn_id} + + client = TestClient(app) + + # Make a single request first + response = client.get("/test") + assert response.status_code == 200 + + # Verify that connection was closed before background task ran + # (due to function scope) + conn_close_indices = [ + i for i, log in enumerate(execution_log) if log.startswith("conn_close") + ] + bg_task_indices = [ + i for i, log in enumerate(execution_log) if log.startswith("bg_task") + ] + + if conn_close_indices and bg_task_indices: + # Function scope: close happens before bg_task + assert conn_close_indices[0] < bg_task_indices[0] + + +# Test 7: Verify request scope works correctly (control test) +def test_request_scope_preserves_resource_for_background_task(): + """ + Control test: request scope should keep resource active during background task. + """ + reset_log() + app = FastAPI() + resource = {"is_active": False} + + async def get_resource(): + resource["is_active"] = True + execution_log.append("resource_activated") + yield resource + resource["is_active"] = False + execution_log.append("resource_deactivated") + + @app.get("/test") + async def endpoint( + background_tasks: BackgroundTasks, + res: dict = Depends(get_resource), # Default request scope + ): + async def bg_task(r: dict): + execution_log.append(f"bg_task:is_active={r['is_active']}") + + background_tasks.add_task(bg_task, res) + return {"is_active": res["is_active"]} + + client = TestClient(app) + response = client.get("/test") + assert response.status_code == 200 + + # Background task should see is_active=True because cleanup happens after + assert "bg_task:is_active=True" in execution_log + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_lifespan_unified.py b/tests/test_lifespan_unified.py new file mode 100644 index 000000000..c77c6e6f2 --- /dev/null +++ b/tests/test_lifespan_unified.py @@ -0,0 +1,377 @@ +""" +Tests for unified lifespan handling. +""" + +import warnings +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from typing import Any + +import pytest +from fastapi import FastAPI, Request +from fastapi.exceptions import FastAPIDeprecationWarning +from fastapi.lifespan import ( + LifespanError, + UnifiedLifespanManager, + create_unified_lifespan, + merge_lifespan_state, +) +from fastapi.testclient import TestClient + + +# ============================================================================= +# Test Fixtures +# ============================================================================= + + +@pytest.fixture +def startup_shutdown_log(): + """Fixture to track startup/shutdown execution order.""" + return [] + + +# ============================================================================= +# UnifiedLifespanManager Tests +# ============================================================================= + + +class TestUnifiedLifespanManager: + """Tests for UnifiedLifespanManager class.""" + + @pytest.mark.anyio + async def test_basic_lifespan_execution(self, startup_shutdown_log): + """Test basic lifespan execution order.""" + + @asynccontextmanager + async def lifespan(app): + startup_shutdown_log.append("startup") + yield {"key": "value"} + startup_shutdown_log.append("shutdown") + + manager = UnifiedLifespanManager(lifespan=lifespan, emit_deprecation_warnings=False) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FastAPIDeprecationWarning) + async with manager.lifespan(None) as state: + assert state == {"key": "value"} + startup_shutdown_log.append("running") + + assert startup_shutdown_log == ["startup", "running", "shutdown"] + + @pytest.mark.anyio + async def test_on_event_handlers_execution(self, startup_shutdown_log): + """Test on_event handlers execute correctly.""" + + async def startup1(): + startup_shutdown_log.append("startup1") + + async def startup2(): + startup_shutdown_log.append("startup2") + + async def shutdown1(): + startup_shutdown_log.append("shutdown1") + + async def shutdown2(): + startup_shutdown_log.append("shutdown2") + + manager = UnifiedLifespanManager( + on_startup=[startup1, startup2], + on_shutdown=[shutdown1, shutdown2], + emit_deprecation_warnings=False, + ) + + async with manager.lifespan(None) as state: + startup_shutdown_log.append("running") + + # Startup in order, shutdown in reverse order + assert startup_shutdown_log == [ + "startup1", + "startup2", + "running", + "shutdown2", + "shutdown1", + ] + + @pytest.mark.anyio + async def test_startup_failure_rollback(self, startup_shutdown_log): + """Test rollback when startup fails.""" + + async def startup1(): + startup_shutdown_log.append("startup1") + + async def startup2(): + startup_shutdown_log.append("startup2_failed") + raise ValueError("Startup failed") + + async def shutdown1(): + startup_shutdown_log.append("shutdown1") + + async def shutdown2(): + startup_shutdown_log.append("shutdown2") + + manager = UnifiedLifespanManager( + on_startup=[startup1, startup2], + on_shutdown=[shutdown1, shutdown2], + emit_deprecation_warnings=False, + ) + + with pytest.raises(LifespanError) as exc_info: + async with manager.lifespan(None): + startup_shutdown_log.append("should_not_run") + + # Rollback should clean up startup1 only (shutdown1) + assert "startup1" in startup_shutdown_log + assert "startup2_failed" in startup_shutdown_log + assert "should_not_run" not in startup_shutdown_log + # Rollback runs cleanup for completed handlers + assert "shutdown1" in startup_shutdown_log + + # Check error details + assert exc_info.value.phase == "startup" + assert isinstance(exc_info.value.original_error, ValueError) + + @pytest.mark.anyio + async def test_mixed_patterns_warning(self): + """Test warning when mixing lifespan and on_event.""" + + @asynccontextmanager + async def lifespan(app): + yield + + with pytest.warns(FastAPIDeprecationWarning, match="not recommended"): + UnifiedLifespanManager( + on_startup=[lambda: None], + lifespan=lifespan, + ) + + +# ============================================================================= +# State Merging Tests +# ============================================================================= + + +class TestStateMerging: + """Tests for lifespan state merging.""" + + def test_merge_child_overrides_parent(self): + """Test default behavior: child overrides parent.""" + parent = {"shared": "parent", "parent_only": True} + child = {"shared": "child", "child_only": True} + + result = merge_lifespan_state(parent, child, child_overrides_parent=True) + + assert result["shared"] == "child" # child wins + assert result["parent_only"] is True + assert result["child_only"] is True + + def test_merge_parent_overrides_child(self): + """Test alternative behavior: parent overrides child.""" + parent = {"shared": "parent", "parent_only": True} + child = {"shared": "child", "child_only": True} + + result = merge_lifespan_state(parent, child, child_overrides_parent=False) + + assert result["shared"] == "parent" # parent wins + assert result["parent_only"] is True + assert result["child_only"] is True + + def test_merge_none_states(self): + """Test merging with None states.""" + assert merge_lifespan_state(None, None) == {} + assert merge_lifespan_state({"a": 1}, None) == {"a": 1} + assert merge_lifespan_state(None, {"b": 2}) == {"b": 2} + + +# ============================================================================= +# FastAPI Integration Tests +# ============================================================================= + + +class TestFastAPILifespanIntegration: + """Integration tests with FastAPI app.""" + + def test_lifespan_context_manager(self): + """Test modern lifespan context manager pattern.""" + events = [] + + @asynccontextmanager + async def lifespan(app: FastAPI) -> AsyncGenerator[dict[str, Any], None]: + events.append("startup") + yield {"db": "connection"} + events.append("shutdown") + + app = FastAPI(lifespan=lifespan) + + @app.get("/") + def index(): + return {"status": "ok"} + + with TestClient(app) as client: + events.append("running") + response = client.get("/") + assert response.status_code == 200 + + assert events == ["startup", "running", "shutdown"] + + def test_on_event_deprecation_warning(self): + """Test that on_event emits deprecation warning.""" + app = FastAPI() + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + @app.on_event("startup") + def startup(): + pass + + # Check that at least one warning was raised + assert len(w) >= 1 + # Check that one of them is about on_event deprecation + deprecation_warnings = [ + warning for warning in w + if "on_event" in str(warning.message) and "deprecated" in str(warning.message) + ] + assert len(deprecation_warnings) >= 1 + + def test_on_event_handlers_still_work(self): + """Test that on_event handlers still execute (backwards compatibility).""" + events = [] + app = FastAPI() + + # Suppress the deprecation warnings for this test + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + + @app.on_event("startup") + def startup(): + events.append("startup") + + @app.on_event("shutdown") + def shutdown(): + events.append("shutdown") + + @app.get("/") + def index(): + return {"status": "ok"} + + with TestClient(app) as client: + events.append("running") + response = client.get("/") + assert response.status_code == 200 + + assert "startup" in events + assert "running" in events + assert "shutdown" in events + + def test_state_accessible_in_endpoints(self): + """Test that lifespan state is accessible in endpoints.""" + + @asynccontextmanager + async def lifespan(app: FastAPI) -> AsyncGenerator[dict[str, Any], None]: + yield {"db_connection": "active"} + + app = FastAPI(lifespan=lifespan) + + @app.get("/state") + def get_state(request: Request): + # State should be accessible via request.state + return {"has_state": hasattr(request.state, "db_connection")} + + with TestClient(app) as client: + response = client.get("/state") + assert response.status_code == 200 + + +# ============================================================================= +# create_unified_lifespan Tests +# ============================================================================= + + +class TestCreateUnifiedLifespan: + """Tests for create_unified_lifespan function.""" + + @pytest.mark.anyio + async def test_create_from_lifespan_only(self): + """Test creating unified lifespan from lifespan context only.""" + events = [] + + @asynccontextmanager + async def lifespan(app): + events.append("start") + yield {"key": "value"} + events.append("stop") + + unified = create_unified_lifespan(lifespan=lifespan) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FastAPIDeprecationWarning) + async with unified(None) as state: + events.append("running") + assert state == {"key": "value"} + + assert events == ["start", "running", "stop"] + + @pytest.mark.anyio + async def test_create_from_handlers_only(self): + """Test creating unified lifespan from handlers only.""" + events = [] + + async def on_start(): + events.append("start") + + async def on_stop(): + events.append("stop") + + unified = create_unified_lifespan( + on_startup=[on_start], + on_shutdown=[on_stop], + ) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FastAPIDeprecationWarning) + async with unified(None) as state: + events.append("running") + + assert "start" in events + assert "running" in events + assert "stop" in events + + +# ============================================================================= +# LifespanError Tests +# ============================================================================= + + +class TestLifespanError: + """Tests for LifespanError exception.""" + + def test_error_message_formatting(self): + """Test error message includes relevant information.""" + + def my_handler(): + pass + + error = LifespanError( + "Test error", + phase="startup", + handler=my_handler, + original_error=ValueError("Original"), + ) + + error_str = str(error) + assert "startup" in error_str + assert "my_handler" in error_str + assert "Original" in error_str + + def test_error_attributes(self): + """Test error attributes are accessible.""" + original = ValueError("test") + + error = LifespanError( + "Test", + phase="shutdown", + original_error=original, + ) + + assert error.phase == "shutdown" + assert error.original_error is original diff --git a/tests/test_openapi_plugins.py b/tests/test_openapi_plugins.py new file mode 100644 index 000000000..2ede1fcac --- /dev/null +++ b/tests/test_openapi_plugins.py @@ -0,0 +1,488 @@ +""" +Tests for the OpenAPI plugin system. +""" + +from typing import Any + +import pytest +from fastapi import FastAPI +from fastapi.openapi.plugins import ( + OpenAPIPluginBase, + OpenAPIPluginSettings, + PluginExecutor, + PluginRegistry, + create_plugin, +) +from fastapi.testclient import TestClient +from pydantic import BaseModel + + +# ============================================================================= +# Test Models +# ============================================================================= + + +class Item(BaseModel): + name: str + price: float + + +# ============================================================================= +# Test Plugins +# ============================================================================= + + +class CustomExtensionPlugin(OpenAPIPluginBase): + """Plugin that adds custom x- extensions to all operations.""" + + @property + def name(self) -> str: + return "custom-extension" + + @property + def priority(self) -> int: + return 50 # Higher priority (runs first) + + def modify_operation( + self, + route: Any, + method: str, + operation: dict[str, Any], + ) -> dict[str, Any]: + operation["x-custom-extension"] = True + operation["x-route-name"] = route.name + return operation + + +class SchemaModifierPlugin(OpenAPIPluginBase): + """Plugin that modifies the schema after generation.""" + + @property + def name(self) -> str: + return "schema-modifier" + + def modify_schema(self, schema: dict[str, Any]) -> dict[str, Any]: + schema["x-generated-by"] = "FastAPI Plugin System" + if "info" in schema: + schema["info"]["x-plugin-version"] = "1.0" + return schema + + +class PostGenerationPlugin(OpenAPIPluginBase): + """Plugin that runs after schema generation.""" + + @property + def name(self) -> str: + return "post-generation" + + @property + def priority(self) -> int: + return 200 # Lower priority (runs last) + + def post_schema_generation(self, schema: dict[str, Any]) -> dict[str, Any]: + schema["x-post-processed"] = True + return schema + + +class PreSchemaPlugin(OpenAPIPluginBase): + """Plugin that tracks pre_schema_generation calls.""" + + def __init__(self): + self.pre_called = False + self.settings_received: OpenAPIPluginSettings | None = None + + @property + def name(self) -> str: + return "pre-schema" + + def pre_schema_generation( + self, + app: FastAPI, + settings: OpenAPIPluginSettings, + ) -> None: + self.pre_called = True + self.settings_received = settings + + +class ErrorPlugin(OpenAPIPluginBase): + """Plugin that raises errors for testing error handling.""" + + @property + def name(self) -> str: + return "error-plugin" + + def modify_operation( + self, + route: Any, + method: str, + operation: dict[str, Any], + ) -> dict[str, Any]: + raise ValueError("Intentional error for testing") + + +# ============================================================================= +# Plugin Registry Tests +# ============================================================================= + + +class TestPluginRegistry: + """Tests for PluginRegistry class.""" + + def test_register_plugin(self): + """Test basic plugin registration.""" + registry = PluginRegistry() + plugin = CustomExtensionPlugin() + registry.register(plugin) + + assert "custom-extension" in registry + assert len(registry) == 1 + assert registry.get_plugin("custom-extension") is plugin + + def test_register_duplicate_raises(self): + """Test that registering duplicate name raises error.""" + registry = PluginRegistry() + plugin1 = CustomExtensionPlugin() + plugin2 = CustomExtensionPlugin() + + registry.register(plugin1) + with pytest.raises(ValueError, match="already registered"): + registry.register(plugin2) + + def test_register_with_replace(self): + """Test replacing an existing plugin.""" + registry = PluginRegistry() + plugin1 = CustomExtensionPlugin() + plugin2 = CustomExtensionPlugin() + + registry.register(plugin1) + registry.register(plugin2, replace=True) + + assert registry.get_plugin("custom-extension") is plugin2 + + def test_unregister_plugin(self): + """Test unregistering a plugin.""" + registry = PluginRegistry() + plugin = CustomExtensionPlugin() + registry.register(plugin) + + assert registry.unregister("custom-extension") is True + assert "custom-extension" not in registry + assert registry.unregister("custom-extension") is False + + def test_enable_disable_plugin(self): + """Test enabling and disabling plugins.""" + registry = PluginRegistry() + plugin = CustomExtensionPlugin() + registry.register(plugin, enabled=False) + + assert not registry.is_enabled("custom-extension") + assert len(registry.get_active_plugins()) == 0 + + registry.enable("custom-extension") + assert registry.is_enabled("custom-extension") + assert len(registry.get_active_plugins()) == 1 + + registry.disable("custom-extension") + assert not registry.is_enabled("custom-extension") + + def test_priority_ordering(self): + """Test that plugins are sorted by priority.""" + registry = PluginRegistry() + + low_priority = PostGenerationPlugin() # priority 200 + high_priority = CustomExtensionPlugin() # priority 50 + normal_priority = SchemaModifierPlugin() # priority 100 (default) + + registry.register(low_priority) + registry.register(high_priority) + registry.register(normal_priority) + + active = registry.get_active_plugins() + assert active[0].name == "custom-extension" # priority 50 + assert active[1].name == "schema-modifier" # priority 100 + assert active[2].name == "post-generation" # priority 200 + + def test_generation_count_increments(self): + """Test that generation count increments on changes.""" + registry = PluginRegistry() + initial_count = registry.generation_count + + plugin = CustomExtensionPlugin() + registry.register(plugin) + assert registry.generation_count == initial_count + 1 + + registry.disable("custom-extension") + assert registry.generation_count == initial_count + 2 + + registry.enable("custom-extension") + assert registry.generation_count == initial_count + 3 + + registry.unregister("custom-extension") + assert registry.generation_count == initial_count + 4 + + def test_cache_invalidation_callback(self): + """Test that cache invalidation callback is called.""" + callback_called = [] + + def on_invalidate(): + callback_called.append(True) + + registry = PluginRegistry(on_cache_invalidation=on_invalidate) + plugin = CustomExtensionPlugin() + + registry.register(plugin) + assert len(callback_called) == 1 + + registry.disable("custom-extension") + assert len(callback_called) == 2 + + +# ============================================================================= +# Plugin Execution Tests +# ============================================================================= + + +class TestPluginExecution: + """Tests for plugin execution during schema generation.""" + + def test_modify_operation_plugin(self): + """Test that modify_operation hook is called.""" + plugin = CustomExtensionPlugin() + app = FastAPI( + title="Test API", + version="1.0.0", + openapi_plugins=[plugin], + ) + + @app.get("/items") + def get_items(): + return [] + + schema = app.openapi() + + # Check the custom extension was added + operation = schema["paths"]["/items"]["get"] + assert operation.get("x-custom-extension") is True + assert operation.get("x-route-name") == "get_items" + + def test_modify_schema_plugin(self): + """Test that modify_schema hook is called.""" + plugin = SchemaModifierPlugin() + app = FastAPI( + title="Test API", + version="1.0.0", + openapi_plugins=[plugin], + ) + + @app.get("/items") + def get_items(): + return [] + + schema = app.openapi() + + assert schema.get("x-generated-by") == "FastAPI Plugin System" + assert schema["info"].get("x-plugin-version") == "1.0" + + def test_post_generation_plugin(self): + """Test that post_schema_generation hook is called.""" + plugin = PostGenerationPlugin() + app = FastAPI( + title="Test API", + version="1.0.0", + openapi_plugins=[plugin], + ) + + @app.get("/items") + def get_items(): + return [] + + schema = app.openapi() + + assert schema.get("x-post-processed") is True + + def test_pre_schema_generation_plugin(self): + """Test that pre_schema_generation hook is called.""" + plugin = PreSchemaPlugin() + app = FastAPI( + title="Test API", + version="1.0.0", + openapi_plugins=[plugin], + ) + + @app.get("/items") + def get_items(): + return [] + + schema = app.openapi() + + assert plugin.pre_called is True + assert plugin.settings_received is not None + assert plugin.settings_received.title == "Test API" + assert plugin.settings_received.version == "1.0.0" + + def test_multiple_plugins(self): + """Test multiple plugins working together.""" + app = FastAPI( + title="Test API", + version="1.0.0", + openapi_plugins=[ + CustomExtensionPlugin(), + SchemaModifierPlugin(), + PostGenerationPlugin(), + ], + ) + + @app.get("/items") + def get_items(): + return [] + + schema = app.openapi() + + # All plugins should have modified the schema + operation = schema["paths"]["/items"]["get"] + assert operation.get("x-custom-extension") is True + assert schema.get("x-generated-by") == "FastAPI Plugin System" + assert schema.get("x-post-processed") is True + + def test_plugin_error_handling(self): + """Test that plugin errors don't break schema generation.""" + error_plugin = ErrorPlugin() + normal_plugin = SchemaModifierPlugin() + + app = FastAPI( + title="Test API", + version="1.0.0", + openapi_plugins=[error_plugin, normal_plugin], + ) + + @app.get("/items") + def get_items(): + return [] + + # Should not raise, error is caught + with pytest.warns(UserWarning, match="raised an error"): + schema = app.openapi() + + # Normal plugin should still work + assert schema.get("x-generated-by") == "FastAPI Plugin System" + + def test_plugin_cache_invalidation(self): + """Test that schema is regenerated when plugins change.""" + app = FastAPI( + title="Test API", + version="1.0.0", + ) + + @app.get("/items") + def get_items(): + return [] + + # First generation - no plugins + schema1 = app.openapi() + assert schema1.get("x-generated-by") is None + + # Register a plugin + app.openapi_plugins.register(SchemaModifierPlugin()) + + # Should regenerate with plugin + schema2 = app.openapi() + assert schema2.get("x-generated-by") == "FastAPI Plugin System" + + # Disable plugin + app.openapi_plugins.disable("schema-modifier") + + # Should regenerate without plugin effect + schema3 = app.openapi() + assert schema3.get("x-generated-by") is None + + +# ============================================================================= +# create_plugin() Tests +# ============================================================================= + + +class TestCreatePlugin: + """Tests for the create_plugin convenience function.""" + + def test_create_simple_plugin(self): + """Test creating a plugin from callbacks.""" + plugin = create_plugin( + "simple-plugin", + modify_operation=lambda route, method, op: {**op, "x-simple": True}, + ) + + assert plugin.name == "simple-plugin" + assert plugin.priority == 100 + + # Test the callback works + result = plugin.modify_operation(None, "GET", {"summary": "Test"}) + assert result["x-simple"] is True + assert result["summary"] == "Test" + + def test_create_plugin_with_priority(self): + """Test creating a plugin with custom priority.""" + plugin = create_plugin( + "priority-plugin", + priority=25, + ) + + assert plugin.priority == 25 + + +# ============================================================================= +# Integration Tests +# ============================================================================= + + +class TestIntegration: + """Integration tests with actual HTTP requests.""" + + def test_plugin_with_test_client(self): + """Test plugin output via test client.""" + app = FastAPI( + title="Test API", + version="1.0.0", + openapi_plugins=[CustomExtensionPlugin()], + ) + + @app.get("/items", response_model=list[Item]) + def get_items(): + return [] + + @app.post("/items", response_model=Item) + def create_item(item: Item): + return item + + client = TestClient(app) + response = client.get("/openapi.json") + assert response.status_code == 200 + + schema = response.json() + + # Check both operations have custom extension + assert schema["paths"]["/items"]["get"].get("x-custom-extension") is True + assert schema["paths"]["/items"]["post"].get("x-custom-extension") is True + + def test_docs_with_plugins(self): + """Test that Swagger UI still works with plugins.""" + app = FastAPI( + title="Test API", + version="1.0.0", + openapi_plugins=[SchemaModifierPlugin()], + ) + + @app.get("/items") + def get_items(): + return [] + + client = TestClient(app) + + # Swagger UI should load + response = client.get("/docs") + assert response.status_code == 200 + assert "swagger-ui" in response.text.lower() + + # ReDoc should load + response = client.get("/redoc") + assert response.status_code == 200 + assert "redoc" in response.text.lower()