# Guide: Adding Browser-Use Tools to Your Agent

Add individual browser actions to your existing agent's tool set. Your agent stays in control and drives the browser action by action.

## Table of Contents
- [When to Use This Pattern](#when-to-use-this-pattern)
- [Pick Your Integration](#pick-your-integration)
- [Shell Command Agents (CLI)](#shell-command-agents-cli)
- [TypeScript/JS: CDP + Playwright](#typescriptjs-cdp--playwright)
- [MCP-Native Agents](#mcp-native-agents)
- [Existing Playwright/Puppeteer/Selenium](#existing-playwrightpuppeteerselenium)
- [Decision Summary](#decision-summary)

---

## When to Use This Pattern

Your agent already has tools (search, code execution, file I/O, etc.) and its own reasoning loop. You want to add browser capabilities — navigate, click, type, extract — as tools your agent can call. You don't want to hand off to browser-use's Agent; your agent makes the decisions.

**Use tools integration when:**
- Your agent needs action-by-action browser control
- You want browser actions alongside your other tools
- Your agent's reasoning should drive what gets clicked/typed

**Use [subagent](subagent.md) instead when:**
- You want to delegate an entire web task as a black box
- You don't need control over individual browser actions

## Pick Your Integration

| Your agent type | Best approach | Control level |
|----------------|---------------|--------------|
| CLI coding agent in sandbox | [CLI commands](#shell-command-agents-cli) | Per-command |
| TypeScript/JS | [CDP + Playwright](#typescriptjs-cdp--playwright) | Playwright API |
| MCP client (Claude Desktop, Cursor) | [Local MCP server](#mcp-native-agents) | MCP tools |
| Existing Playwright/Puppeteer/Selenium | [CDP WebSocket (stealth)](#existing-playwrightpuppeteerselenium) | Your existing API |
| HTTP only / any language | Cloud REST: `POST /browsers` → CDP URL | CDP |

---

## Shell Command Agents (CLI)

**For:** Claude Code, Codex, OpenCode, Cline, Windsurf, Cursor background agents, Hermes, OpenClaw — any coding agent running in a VM/container with terminal access.

**Setup:** Install the CLI and load the browser-use SKILL.md into the agent's context. The agent calls browser commands as shell tool invocations.

```bash
uv pip install 'browser-use[cli]'
```

**Core workflow** — the agent calls these commands one at a time, reading output between each:

```bash
# 1. Navigate
browser-use open https://example.com

# 2. Observe — ALWAYS run state first to get element indices
browser-use state
# Output: URL, title, list of clickable elements with indices
# e.g. [0] <input type="search" placeholder="Search...">
#      [1] <button>Submit</button>
#      [2] <a href="/about">About</a>

# 3. Interact — use indices from state
browser-use input 0 "search query"    # Type into element 0
browser-use click 1                   # Click element 1

# 4. Verify — re-run state to see result
browser-use state

# 5. Extract data
browser-use get text 3               # Get element text
browser-use get html --selector "h1" # Get scoped HTML
browser-use eval "document.title"    # Execute JavaScript
browser-use screenshot result.png    # Capture visual state

# 6. Wait for dynamic content
browser-use wait selector ".results" # Wait for element
browser-use wait text "Success"      # Wait for text

# 7. Cleanup
browser-use close
```

**Key details:**
- Background daemon keeps browser alive between commands (~50ms latency per call)
- Agent's reasoning loop decides which command to call next
- `state` output is the agent's "eyes" — it reads element indices and decides what to click
- Commands can be chained with `&&` when intermediate output isn't needed
- `--json` flag for machine-readable output
- `--headed` for visible browser (debugging)
- `--profile "Default"` for authenticated browsing with saved Chrome logins

---

## TypeScript/JS: CDP + Playwright

**For:** TypeScript agents that need browser primitives. Connect Playwright to a cloud stealth browser.

```typescript
import { chromium } from "playwright";

// Connect to cloud stealth browser (no local Chrome needed)
const browser = await chromium.connectOverCDP(
  "wss://connect.browser-use.com?apiKey=YOUR_KEY&proxyCountryCode=us"
);
const page = browser.contexts()[0].pages()[0];

// Your agent calls these as tools:
await page.goto("https://example.com");
await page.fill("#search", "query");
await page.click("button[type=submit]");
const text = await page.textContent(".result");
const screenshot = await page.screenshot();

await browser.close();
// Browser auto-stops when WebSocket disconnects
```

For local browser (no cloud):
```typescript
import { chromium } from "playwright";

const browser = await chromium.launch();
const page = await browser.newPage();
// ... same Playwright API
await browser.close();
```

---

## MCP-Native Agents

**For:** Claude Desktop, Cursor with MCP, any MCP client that discovers tools via protocol.

Start the local MCP server:
```bash
uvx --from 'browser-use[cli]' browser-use --mcp
```

The agent gets individual browser tools:
- `browser_navigate(url)` — go to URL
- `browser_click(index)` — click element by index
- `browser_type(index, text)` — type into element
- `browser_get_state(include_screenshot)` — get page state with element indices
- `browser_extract_content(query)` — LLM-powered extraction
- `browser_screenshot(full_page)` — capture page
- `browser_scroll(direction)` — scroll up/down
- `browser_go_back()` — browser back
- `browser_list_tabs()`, `browser_switch_tab(id)`, `browser_close_tab(id)` — tab management

The agent calls these one at a time, using its own reasoning to decide the next action.

---

## Existing Playwright/Puppeteer/Selenium

**For:** You already have browser automation scripts and want to run them on stealth infrastructure (anti-fingerprinting, CAPTCHA handling, residential proxies).

Zero code changes — just change the connection URL:

### Playwright
```python
# Before: local browser
browser = await playwright.chromium.launch()

# After: cloud stealth browser
browser = await playwright.chromium.connect_over_cdp(
    "wss://connect.browser-use.com?apiKey=KEY&proxyCountryCode=us"
)
# Rest of your code stays exactly the same
```

### Puppeteer
```javascript
// Before
const browser = await puppeteer.launch();

// After
const browser = await puppeteer.connect({
  browserWSEndpoint: "wss://connect.browser-use.com?apiKey=KEY&proxyCountryCode=us"
});
```

Browser auto-starts on connect, auto-stops on disconnect. Pricing: $0.05/hour.

---

## Decision Summary

| Condition | Best option |
|-----------|------------|
| Agent has terminal access (sandbox/VM) | CLI commands |
| TypeScript/JS | CDP WebSocket + Playwright |
| MCP client (Claude Desktop, Cursor) | Local MCP server |
| HTTP only / any language | Cloud REST: `POST /browsers` → CDP URL |
| Existing Playwright/Puppeteer scripts | CDP WebSocket (stealth cloud browser) |

> **Note:** For Python agents that want fine-grained browser control via direct imports (Actor API, Tools Registry, MCPClient), see the **open-source** skill's reference docs.


import * as React from "react"
import { cva, type VariantProps } from "class-variance-authority"

import { cn } from "#/lib/utils.ts"

const alertVariants = cva(
  "group/alert relative grid w-full gap-0.5 rounded-lg border px-2 py-1.5 text-left text-xs/relaxed has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-1.5 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-3.5",
  {
    variants: {
      variant: {
        default: "bg-card text-card-foreground",
        destructive:
          "bg-card text-destructive *:data-[slot=alert-description]:text-destructive/90 *:[svg]:text-current",
      },
    },
    defaultVariants: {
      variant: "default",
    },
  }
)

function Alert({
  className,
  variant,
  ...props
}: React.ComponentProps<"div"> & VariantProps<typeof alertVariants>) {
  return (
    <div
      data-slot="alert"
      role="alert"
      className={cn(alertVariants({ variant }), className)}
      {...props}
    />
  )
}

function AlertTitle({ className, ...props }: React.ComponentProps<"div">) {
  return (
    <div
      data-slot="alert-title"
      className={cn(
        "font-medium group-has-[>svg]/alert:col-start-2 [&_a]:underline [&_a]:underline-offset-3 [&_a]:hover:text-foreground",
        className
      )}
      {...props}
    />
  )
}

function AlertDescription({
  className,
  ...props
}: React.ComponentProps<"div">) {
  return (
    <div
      data-slot="alert-description"
      className={cn(
        "text-xs/relaxed text-balance text-muted-foreground md:text-pretty [&_a]:underline [&_a]:underline-offset-3 [&_a]:hover:text-foreground [&_p:not(:last-child)]:mb-4",
        className
      )}
      {...props}
    />
  )
}

function AlertAction({ className, ...props }: React.ComponentProps<"div">) {
  return (
    <div
      data-slot="alert-action"
      className={cn("absolute top-1.5 right-2", className)}
      {...props}
    />
  )
}

export { Alert, AlertTitle, AlertDescription, AlertAction }


import argparse
import asyncio
import logging
import os
import subprocess
import sys
from datetime import datetime
from pathlib import Path

from browser_use.utils import create_task_with_error_handling

def setup_environment(debug: bool):
	if not debug:
		os.environ['BROWSER_USE_SETUP_LOGGING'] = 'false'
		os.environ['BROWSER_USE_LOGGING_LEVEL'] = 'critical'
		logging.getLogger().setLevel(logging.CRITICAL)
	else:
		os.environ['BROWSER_USE_SETUP_LOGGING'] = 'true'
		os.environ['BROWSER_USE_LOGGING_LEVEL'] = 'info'

parser = argparse.ArgumentParser(description='Generate ads from landing pages using browser-use + 🍌')
parser.add_argument('--url', nargs='?', help='Landing page URL to analyze')
parser.add_argument('--debug', action='store_true', default=False, help='Enable debug mode (show browser, verbose logs)')
parser.add_argument('--count', type=int, default=1, help='Number of ads to generate in parallel (default: 1)')
group = parser.add_mutually_exclusive_group()
group.add_argument('--instagram', action='store_true', default=False, help='Generate Instagram image ad (default)')
group.add_argument('--tiktok', action='store_true', default=False, help='Generate TikTok video ad using Veo3')
args = parser.parse_args()
if not args.instagram and not args.tiktok:
	args.instagram = True
setup_environment(args.debug)

from typing import Any, cast

import aiofiles
from google import genai
from PIL import Image

from browser_use import Agent, BrowserSession
from browser_use.llm.google import ChatGoogle

GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')

class LandingPageAnalyzer:
	def __init__(self, debug: bool = False):
		self.debug = debug
		self.llm = ChatGoogle(model='gemini-2.0-flash-exp', api_key=GOOGLE_API_KEY)
		self.output_dir = Path('output')
		self.output_dir.mkdir(exist_ok=True)

	async def analyze_landing_page(self, url: str, mode: str = 'instagram') -> dict:
		browser_session = BrowserSession(
			headless=not self.debug,
		)

		agent = Agent(
			task=f"""Go to {url} and quickly extract key brand information for Instagram ad creation.

Steps:
1. Navigate to the website
2. From the initial view, extract ONLY these essentials:
   - Brand/Product name
   - Main tagline or value proposition (one sentence)
   - Primary call-to-action text
   - Any visible pricing or special offer
3. Scroll down half a page, twice (0.5 pages each) to check for any key info
4. Done - keep it simple and focused on the brand

Return ONLY the key brand info, not page structure details.""",
			llm=self.llm,
			browser_session=browser_session,
			max_actions_per_step=2,
			step_timeout=30,
			use_thinking=False,
			vision_detail_level='high',
		)

		screenshot_path = None
		timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')

		async def screenshot_callback(agent_instance):
			nonlocal screenshot_path
			await asyncio.sleep(4)
			screenshot_path = self.output_dir / f'landing_page_{timestamp}.png'
			await agent_instance.browser_session.take_screenshot(path=str(screenshot_path), full_page=False)

		screenshot_task = create_task_with_error_handling(
			screenshot_callback(agent), name='screenshot_callback', suppress_exceptions=True
		)
		history = await agent.run()
		try:
			await screenshot_task
		except Exception as e:
			print(f'Screenshot task failed: {e}')

		analysis = history.final_result() or 'No analysis content extracted'
		return {'url': url, 'analysis': analysis, 'screenshot_path': screenshot_path, 'timestamp': timestamp}

class AdGenerator:
	def __init__(self, api_key: str | None = GOOGLE_API_KEY, mode: str = 'instagram'):
		if not api_key:
			raise ValueError('GOOGLE_API_KEY is missing or empty – set the environment variable or pass api_key explicitly')

		self.client = genai.Client(api_key=api_key)
		self.output_dir = Path('output')
		self.output_dir.mkdir(exist_ok=True)
		self.mode = mode

	async def create_video_concept(self, browser_analysis: str, ad_id: int) -> str:
		"""Generate a unique creative concept for each video ad"""
		if self.mode != 'tiktok':
			return ''

		concept_prompt = f"""Based on this brand analysis:
{browser_analysis}

Create a UNIQUE and SPECIFIC TikTok video concept #{ad_id}.

Be creative and different! Consider various approaches like:
- Different visual metaphors and storytelling angles
- Various trending TikTok formats (transitions, reveals, transformations)
- Different emotional appeals (funny, inspiring, surprising, relatable)
- Unique visual styles (neon, retro, minimalist, maximalist, surreal)
- Different perspectives (first-person, aerial, macro, time-lapse)

Return a 2-3 sentence description of a specific, unique video concept that would work for this brand.
Make it visually interesting and different from typical ads. Be specific about visual elements, transitions, and mood."""

		response = self.client.models.generate_content(model='gemini-2.0-flash-exp', contents=concept_prompt)
		return response.text if response and response.text else ''

	def create_ad_prompt(self, browser_analysis: str, video_concept: str = '') -> str:
		if self.mode == 'instagram':
			prompt = f"""Create an Instagram ad for this brand:

{browser_analysis}

Create a vibrant, eye-catching Instagram ad image with:
- Try to use the colors and style of the logo or brand, else:
- Bold, modern gradient background with bright colors
- Large, playful sans-serif text with the product/service name from the analysis
- Trendy design elements: geometric shapes, sparkles, emojis
- Fun bubbles or badges for any pricing or special offers mentioned
- Call-to-action button with text from the analysis
- Emphasizes the key value proposition from the analysis
- Uses visual elements that match the brand personality
- Square format (1:1 ratio)
- Use color psychology to drive action

Style: Modern Instagram advertisement, (1:1), scroll-stopping, professional but playful, conversion-focused"""
		else:  # tiktok
			if video_concept:
				prompt = f"""Create a TikTok video ad based on this specific concept:

{video_concept}

Brand context: {browser_analysis}

Requirements:
- Vertical 9:16 format
- High quality, professional execution
- Bring the concept to life exactly as described
- No text overlays, pure visual storytelling"""
			else:
				prompt = f"""Create a viral TikTok video ad for this brand:

{browser_analysis}

Create a dynamic, engaging vertical video with:
- Quick hook opening that grabs attention immediately
- Minimal text overlays (focus on visual storytelling)
- Fast-paced but not overwhelming editing
- Authentic, relatable energy that appeals to Gen Z
- Vertical 9:16 format optimized for mobile
- High energy but professional execution

Style: Modern TikTok advertisement, viral potential, authentic energy, minimal text, maximum visual impact"""
		return prompt

	async def generate_ad_image(self, prompt: str, screenshot_path: Path | None = None) -> bytes | None:
		"""Generate ad image bytes using Gemini. Returns None on failure."""
		try:
			from typing import Any

			contents: list[Any] = [prompt]

			if screenshot_path and screenshot_path.exists():
				img = Image.open(screenshot_path)
				w, h = img.size
				side = min(w, h)
				img = img.crop(((w - side) // 2, (h - side) // 2, (w + side) // 2, (h + side) // 2))
				contents = [prompt + '\n\nHere is the actual landing page screenshot to reference for design inspiration:', img]

			response = await self.client.aio.models.generate_content(
				model='gemini-2.5-flash-image-preview',
				contents=contents,
			)

			cand = getattr(response, 'candidates', None)
			if cand:
				for part in getattr(cand[0].content, 'parts', []):
					inline = getattr(part, 'inline_data', None)
					if inline:
						return inline.data
		except Exception as e:
			print(f'❌ Image generation failed: {e}')
		return None

	async def generate_ad_video(self, prompt: str, screenshot_path: Path | None = None, ad_id: int = 1) -> bytes:
		"""Generate ad video using Veo3."""
		sync_client = genai.Client(api_key=GOOGLE_API_KEY)

		# Commented out image input for now - it was using the screenshot as first frame
		# if screenshot_path and screenshot_path.exists():
		# 	import base64
		# 	import io

		# 	img = Image.open(screenshot_path)
		# 	img_buffer = io.BytesIO()
		# 	img.save(img_buffer, format='PNG')
		# 	img_bytes = img_buffer.getvalue()

		# 	operation = sync_client.models.generate_videos(
		# 		model='veo-3.0-generate-001',
		# 		prompt=prompt,
		# 		image=cast(Any, {
		# 			'imageBytes': base64.b64encode(img_bytes).decode('utf-8'),
		# 			'mimeType': 'image/png'
		# 		}),
		# 		config=cast(Any, {'aspectRatio': '9:16', 'resolution': '720p'}),
		# 	)
		# else:
		operation = sync_client.models.generate_videos(
			model='veo-3.0-generate-001',
			prompt=prompt,
			config=cast(Any, {'aspectRatio': '9:16', 'resolution': '720p'}),
		)

		while not operation.done:
			await asyncio.sleep(10)
			operation = sync_client.operations.get(operation)

		if not operation.response or not operation.response.generated_videos:
			raise RuntimeError('No videos generated')
		videos = operation.response.generated_videos
		video = videos[0]
		video_file = getattr(video, 'video', None)
		if not video_file:
			raise RuntimeError('No video file in response')
		sync_client.files.download(file=video_file)
		video_bytes = getattr(video_file, 'video_bytes', None)
		if not video_bytes:
			raise RuntimeError('No video bytes in response')
		return video_bytes

	async def save_results(self, ad_content: bytes, prompt: str, analysis: str, url: str, timestamp: str) -> str:
		if self.mode == 'instagram':
			content_path = self.output_dir / f'ad_{timestamp}.png'
		else:  # tiktok
			content_path = self.output_dir / f'ad_{timestamp}.mp4'

		async with aiofiles.open(content_path, 'wb') as f:
			await f.write(ad_content)

		analysis_path = self.output_dir / f'analysis_{timestamp}.txt'
		async with aiofiles.open(analysis_path, 'w', encoding='utf-8') as f:
			await f.write(f'URL: {url}\n\n')
			await f.write('BROWSER-USE ANALYSIS:\n')
			await f.write(analysis)
			await f.write('\n\nGENERATED PROMPT:\n')
			await f.write(prompt)

		return str(content_path)

def open_file(file_path: str):
	"""Open file with default system viewer"""
	try:
		if sys.platform.startswith('darwin'):
			subprocess.run(['open', file_path], check=True)
		elif sys.platform.startswith('win'):
			subprocess.run(['cmd', '/c', 'start', '', file_path], check=True)
		else:
			subprocess.run(['xdg-open', file_path], check=True)
	except Exception as e:
		print(f'❌ Could not open file: {e}')

async def create_ad_from_landing_page(url: str, debug: bool = False, mode: str = 'instagram', ad_id: int = 1):
	analyzer = LandingPageAnalyzer(debug=debug)

	try:
		if ad_id == 1:
			print(f'🚀 Analyzing {url} for {mode.capitalize()} ad...')
			page_data = await analyzer.analyze_landing_page(url, mode=mode)
		else:
			analyzer_temp = LandingPageAnalyzer(debug=debug)
			page_data = await analyzer_temp.analyze_landing_page(url, mode=mode)

		generator = AdGenerator(mode=mode)

		if mode == 'instagram':
			prompt = generator.create_ad_prompt(page_data['analysis'])
			ad_content = await generator.generate_ad_image(prompt, page_data.get('screenshot_path'))
			if ad_content is None:
				raise RuntimeError(f'Ad image generation failed for ad #{ad_id}')
		else:  # tiktok
			video_concept = await generator.create_video_concept(page_data['analysis'], ad_id)
			prompt = generator.create_ad_prompt(page_data['analysis'], video_concept)
			ad_content = await generator.generate_ad_video(prompt, page_data.get('screenshot_path'), ad_id)

		result_path = await generator.save_results(ad_content, prompt, page_data['analysis'], url, page_data['timestamp'])

		if mode == 'instagram':
			print(f'🎨 Generated image ad #{ad_id}: {result_path}')
		else:
			print(f'🎬 Generated video ad #{ad_id}: {result_path}')

		open_file(result_path)

		return result_path

	except Exception as e:
		print(f'❌ Error for ad #{ad_id}: {e}')
		raise
	finally:
		if ad_id == 1 and page_data.get('screenshot_path'):
			print(f'📸 Page screenshot: {page_data["screenshot_path"]}')

async def generate_single_ad(page_data: dict, mode: str, ad_id: int):
	"""Generate a single ad using pre-analyzed page data"""
	generator = AdGenerator(mode=mode)

	try:
		if mode == 'instagram':
			prompt = generator.create_ad_prompt(page_data['analysis'])
			ad_content = await generator.generate_ad_image(prompt, page_data.get('screenshot_path'))
			if ad_content is None:
				raise RuntimeError(f'Ad image generation failed for ad #{ad_id}')
		else:  # tiktok
			video_concept = await generator.create_video_concept(page_data['analysis'], ad_id)
			prompt = generator.create_ad_prompt(page_data['analysis'], video_concept)
			ad_content = await generator.generate_ad_video(prompt, page_data.get('screenshot_path'), ad_id)

		# Create unique timestamp for each ad
		timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') + f'_{ad_id}'
		result_path = await generator.save_results(ad_content, prompt, page_data['analysis'], page_data['url'], timestamp)

		if mode == 'instagram':
			print(f'🎨 Generated image ad #{ad_id}: {result_path}')
		else:
			print(f'🎬 Generated video ad #{ad_id}: {result_path}')

		return result_path

	except Exception as e:
		print(f'❌ Error for ad #{ad_id}: {e}')
		raise

async def create_multiple_ads(url: str, debug: bool = False, mode: str = 'instagram', count: int = 1):
	"""Generate multiple ads in parallel using asyncio concurrency"""
	if count == 1:
		return await create_ad_from_landing_page(url, debug, mode, 1)

	print(f'🚀 Analyzing {url} for {count} {mode} ads...')

	analyzer = LandingPageAnalyzer(debug=debug)
	page_data = await analyzer.analyze_landing_page(url, mode=mode)

	print(f'🎯 Generating {count} {mode} ads in parallel...')

	tasks = []
	for i in range(count):
		task = create_task_with_error_handling(generate_single_ad(page_data, mode, i + 1), name=f'generate_ad_{i + 1}')
		tasks.append(task)

	results = await asyncio.gather(*tasks, return_exceptions=True)

	successful = []
	failed = []

	for i, result in enumerate(results):
		if isinstance(result, Exception):
			failed.append(i + 1)
		else:
			successful.append(result)

	print(f'\n✅ Successfully generated {len(successful)}/{count} ads')
	if failed:
		print(f'❌ Failed ads: {failed}')

	if page_data.get('screenshot_path'):
		print(f'📸 Page screenshot: {page_data["screenshot_path"]}')

	for ad_path in successful:
		open_file(ad_path)

	return successful

if __name__ == '__main__':
	url = args.url
	if not url:
		url = input('🔗 Enter URL: ').strip() or 'https://www.apple.com/iphone-17-pro/'

	if args.tiktok:
		mode = 'tiktok'
	else:
		mode = 'instagram'

	asyncio.run(create_multiple_ads(url, debug=args.debug, mode=mode, count=args.count))


import * as nodeFs from 'node:fs';
import * as path from 'node:path';
import { parse as parseYaml, stringify as stringifyYaml } from 'yaml';
import { z } from 'zod';
import {
  folderStyleNameProblem,
  isKebabId,
  KEBAB_ID_DESCRIPTION,
  KEBAB_ID_FIX,
} from '../id.js';

import { getGlobalDataDir } from '../global-config.js';
import { FileSystemUtils } from '../../utils/file-system.js';
import {
  acquireFileLock,
  isNodeErrorCode,
  makeLockErrorFactory,
  pathIsDirectory,
  pathIsFile,
  releaseFileLock,
  writeFileAtomically,
} from '../file-state.js';
import { formatZodIssues } from '../zod-issues.js';
import { StoreError } from './errors.js';

const fs = nodeFs.promises;

export const STORE_METADATA_DIR_NAME = '.openspec-store';
export const STORE_METADATA_FILE_NAME = 'store.yaml';
export const STORES_DIR_NAME = 'stores';
export const STORE_REGISTRY_FILE_NAME = 'registry.yaml';

export interface StorePathOptions {
  globalDataDir?: string;
}

export interface StoreGitBackendConfig {
  type: 'git';
  local_path: string;
  remote?: string;
  branch?: string;
}

export type StoreBackendConfig = StoreGitBackendConfig;

export interface StoreRegistryEntryState {
  backend: StoreBackendConfig;
}

export interface StoreRegistryState {
  version: 1;
  stores: Record<string, StoreRegistryEntryState>;
}

export interface StoreRegistryEntry {
  id: string;
  backend: StoreBackendConfig;
}

export interface StoreMetadataState {
  version: 1;
  id: string;
  /** Canonical clone source, team-authored. Optional (slice 3.3). */
  remote?: string;
}

export interface ResolveGitStoreBackendInput {
  localPath: string;
  remote?: string;
  branch?: string;
}

function joinStorePath(basePath: string, ...segments: string[]): string {
  return FileSystemUtils.joinPath(basePath, ...segments);
}

export function getStoresDir(options: StorePathOptions = {}): string {
  return joinStorePath(options.globalDataDir ?? getGlobalDataDir(), STORES_DIR_NAME);
}

export function getStoreRegistryPath(options: StorePathOptions = {}): string {
  return joinStorePath(getStoresDir(options), STORE_REGISTRY_FILE_NAME);
}

export function getStoreMetadataDir(storeRoot: string): string {
  return joinStorePath(storeRoot, STORE_METADATA_DIR_NAME);
}

export function getStoreMetadataPath(storeRoot: string): string {
  return joinStorePath(
    getStoreMetadataDir(storeRoot),
    STORE_METADATA_FILE_NAME
  );
}

export function validateStoreId(id: string): string {
  const folderProblem = folderStyleNameProblem(id, 'Store id');
  if (folderProblem !== null) {
    throw new StoreError(folderProblem, 'invalid_store_id', {
      target: 'store.id',
      fix: KEBAB_ID_FIX,
    });
  }

  if (!isKebabId(id)) {
    throw new StoreError(
      `Store id ${KEBAB_ID_DESCRIPTION}`,
      'invalid_store_id',
      {
        target: 'store.id',
        fix: KEBAB_ID_FIX,
      }
    );
  }

  return id;
}

export function isValidStoreId(id: string): boolean {
  try {
    validateStoreId(id);
    return true;
  } catch {
    return false;
  }
}

function isFileNotFoundError(error: unknown): boolean {
  return isNodeErrorCode(error, 'ENOENT');
}

function normalizeExistingPathForStorage(existingPath: string): string {
  return FileSystemUtils.canonicalizeExistingPath(existingPath);
}

function nonEmptyOptionalString() {
  return z.string().min(1).optional();
}

const GitBackendConfigSchema = z.object({
  type: z.literal('git'),
  local_path: z.string().min(1),
  remote: nonEmptyOptionalString(),
  branch: nonEmptyOptionalString(),
}).strict();

const RegistryEntrySchema = z.object({
  backend: GitBackendConfigSchema,
}).strict();

const RegistryStateSchema = z.object({
  version: z.literal(1),
  stores: z.record(z.string(), RegistryEntrySchema),
  // Legacy code-checkout map data is tolerated on read and dropped on
  // the next write.
  repos: z.unknown().optional(),
}).strict();

const MetadataStateSchema = z.object({
  version: z.literal(1),
  id: z.string(),
  remote: nonEmptyOptionalString(),
}).strict();

function storeStateDiagnostic(label: string): {
  code: string;
  target: string;
  fix: string;
} {
  if (label.includes('metadata')) {
    return {
      code: 'invalid_store_metadata',
      target: 'store.metadata',
      fix: 'Repair .openspec-store/store.yaml.',
    };
  }

  return {
    code: 'invalid_store_registry',
    target: 'store.registry',
    fix: `Repair or remove ${getStoreRegistryPath({})}.`,
  };
}

function invalidStoreStateError(label: string, message: string): StoreError {
  const diagnostic = storeStateDiagnostic(label);
  return new StoreError(`Invalid ${label}: ${message}`, diagnostic.code, {
    target: diagnostic.target,
    fix: diagnostic.fix,
  });
}

function parseYamlObject(content: string, label: string): unknown {
  try {
    return parseYaml(content);
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    throw invalidStoreStateError(label, message);
  }
}

function assertValidStoreIds(ids: string[], label: string): void {
  for (const id of ids) {
    if (!isKebabId(id)) {
      throw invalidStoreStateError(
        label,
        `'${id}': ${KEBAB_ID_DESCRIPTION}`
      );
    }
  }
}

export function parseStoreRegistryState(content: string): StoreRegistryState {
  const raw = parseYamlObject(content, 'store registry state');
  const result = RegistryStateSchema.safeParse(raw);

  if (!result.success) {
    throw invalidStoreStateError(
      'store registry state',
      formatZodIssues(result.error)
    );
  }

  assertValidStoreIds(Object.keys(result.data.stores), 'store id');

  return {
    version: 1,
    stores: result.data.stores,
  };
}

export function parseStoreMetadataState(content: string): StoreMetadataState {
  const raw = parseYamlObject(content, 'store metadata state');
  const result = MetadataStateSchema.safeParse(raw);

  if (!result.success) {
    throw invalidStoreStateError(
      'store metadata state',
      formatZodIssues(result.error)
    );
  }

  validateStoreId(result.data.id);

  return {
    version: 1,
    id: result.data.id,
    ...(result.data.remote !== undefined ? { remote: result.data.remote } : {}),
  };
}

export function serializeStoreRegistryState(state: StoreRegistryState): string {
  const result = RegistryStateSchema.safeParse(state);

  if (!result.success) {
    throw invalidStoreStateError(
      'store registry state',
      formatZodIssues(result.error)
    );
  }

  assertValidStoreIds(Object.keys(result.data.stores), 'store id');

  return stringifyYaml({
    version: 1,
    stores: result.data.stores,
  });
}

export function serializeStoreMetadataState(state: StoreMetadataState): string {
  const result = MetadataStateSchema.safeParse(state);

  if (!result.success) {
    throw invalidStoreStateError(
      'store metadata state',
      formatZodIssues(result.error)
    );
  }

  validateStoreId(result.data.id);

  return stringifyYaml({
    version: 1,
    id: result.data.id,
    ...(result.data.remote !== undefined ? { remote: result.data.remote } : {}),
  });
}

export function listStoreRegistryEntries(
  registry: StoreRegistryState
): StoreRegistryEntry[] {
  return Object.entries(registry.stores)
    .map(([id, store]) => ({ id, backend: store.backend }))
    .sort((a, b) => a.id.localeCompare(b.id));
}

export async function isStoreRoot(candidateRoot: string): Promise<boolean> {
  return pathIsFile(getStoreMetadataPath(candidateRoot));
}

export async function readStoreRegistryState(
  options: StorePathOptions = {}
): Promise<StoreRegistryState | null> {
  const registryPath = getStoreRegistryPath(options);

  if (!(await pathIsFile(registryPath))) {
    return null;
  }

  return parseStoreRegistryState(await fs.readFile(registryPath, 'utf-8'));
}

export async function writeStoreRegistryState(
  state: StoreRegistryState,
  options: StorePathOptions = {}
): Promise<void> {
  await writeFileAtomically(
    getStoreRegistryPath(options),
    serializeStoreRegistryState(state)
  );
}

const storeRegistryLockError = makeLockErrorFactory({
  createSubject: 'the registry lock file',
  busyMessage: 'Store registry is busy.',
  code: 'store_registry_busy',
  target: 'store.registry',
});

export async function updateStoreRegistryState(
  updater: (
    state: StoreRegistryState | null
  ) => StoreRegistryState | Promise<StoreRegistryState>,
  options: StorePathOptions = {}
): Promise<StoreRegistryState> {
  const registryPath = getStoreRegistryPath(options);
  const lockPath = `${registryPath}.lock`;
  const lock = await acquireFileLock({
    lockPath,
    errorFor: storeRegistryLockError,
  });

  try {
    const next = await updater(await readStoreRegistryState(options));
    await writeStoreRegistryState(next, options);
    return next;
  } finally {
    await releaseFileLock(lock, lockPath);
  }
}

export async function readStoreMetadataState(
  storeRoot: string
): Promise<StoreMetadataState> {
  return parseStoreMetadataState(
    await fs.readFile(getStoreMetadataPath(storeRoot), 'utf-8')
  );
}

export async function readOptionalStoreMetadataState(
  storeRoot: string
): Promise<StoreMetadataState | null> {
  try {
    return await readStoreMetadataState(storeRoot);
  } catch (error) {
    if (isFileNotFoundError(error)) {
      return null;
    }

    throw error;
  }
}

export async function writeStoreMetadataState(
  storeRoot: string,
  state: StoreMetadataState
): Promise<void> {
  await FileSystemUtils.writeFile(
    getStoreMetadataPath(storeRoot),
    serializeStoreMetadataState(state)
  );
}

export async function resolveGitStoreBackendConfig(
  input: ResolveGitStoreBackendInput,
  cwd = process.cwd()
): Promise<StoreGitBackendConfig> {
  if (input.localPath.length === 0) {
    throw new Error('Store local path must not be empty.');
  }

  const resolvedPath = path.isAbsolute(input.localPath)
    ? path.resolve(input.localPath)
    : path.resolve(cwd, input.localPath);

  if (!(await pathIsDirectory(resolvedPath))) {
    throw new Error(`Store local path does not exist: ${input.localPath}`);
  }

  if (input.remote !== undefined && input.remote.length === 0) {
    throw new Error('Store backend remote must not be empty when provided.');
  }

  if (input.branch !== undefined && input.branch.length === 0) {
    throw new Error('Store branch must not be empty when provided.');
  }

  return {
    type: 'git',
    local_path: normalizeExistingPathForStorage(resolvedPath),
    ...(input.remote ? { remote: input.remote } : {}),
    ...(input.branch ? { branch: input.branch } : {}),
  };
}


# 12. Common Pitfalls 🔴

> **What you'll learn:**
> - 9 common async Rust bugs and how to fix each one
> - Why blocking the executor is the #1 mistake (and how `spawn_blocking` fixes it)
> - Cancellation hazards: what happens when a future is dropped mid-await
> - Debugging: `tokio-console`, `tracing`, `#[instrument]`
> - Testing: `#[tokio::test]`, `time::pause()`, trait-based mocking

## Blocking the Executor

The #1 mistake in async Rust: running blocking code on the async executor thread. This starves other tasks.

```rust
// ❌ WRONG: Blocks the entire executor thread
async fn bad_handler() -> String {
    let data = std::fs::read_to_string("big_file.txt").unwrap(); // BLOCKS!
    process(&data)
}

// ✅ CORRECT: Offload blocking work to a dedicated thread pool
async fn good_handler() -> String {
    let data = tokio::task::spawn_blocking(|| {
        std::fs::read_to_string("big_file.txt").unwrap()
    }).await.unwrap();
    process(&data)
}

// ✅ ALSO CORRECT: Use tokio's async fs
async fn also_good_handler() -> String {
    let data = tokio::fs::read_to_string("big_file.txt").await.unwrap();
    process(&data)
}
```

```mermaid
graph TB
    subgraph "❌ Blocking Call on Executor"
        T1_BAD["Thread 1: std::fs::read()<br/>🔴 BLOCKED for 500ms"]
        T2_BAD["Thread 2: handling requests<br/>🟢 Working alone"]
        TASKS_BAD["100 pending tasks<br/>⏳ Starved"]
        T1_BAD -->|"can't poll"| TASKS_BAD
    end

    subgraph "✅ spawn_blocking"
        T1_GOOD["Thread 1: polling futures<br/>🟢 Available"]
        T2_GOOD["Thread 2: polling futures<br/>🟢 Available"]
        BT["Blocking pool thread:<br/>std::fs::read()<br/>🔵 Separate pool"]
        TASKS_GOOD["100 tasks<br/>✅ All making progress"]
        T1_GOOD -->|"polls"| TASKS_GOOD
        T2_GOOD -->|"polls"| TASKS_GOOD
    end
```

### std::thread::sleep vs tokio::time::sleep

```rust
// ❌ WRONG: Blocks the executor thread for 5 seconds
async fn bad_delay() {
    std::thread::sleep(Duration::from_secs(5)); // Thread can't poll anything else!
}

// ✅ CORRECT: Yields to the executor, other tasks can run
async fn good_delay() {
    tokio::time::sleep(Duration::from_secs(5)).await; // Non-blocking!
}
```

### Holding MutexGuard Across .await

```rust
use std::sync::Mutex; // std Mutex — NOT async-aware

// ⚠ RISKY: MutexGuard held across .await
async fn bad_mutex(data: &Mutex<Vec<String>>) {
    let mut guard = data.lock().unwrap();
    guard.push("item".into());
    some_io().await; // Guard is held here — blocks other threads from locking!
    guard.push("another".into());
}
// NOTE: This compiles! std::sync::MutexGuard is !Send, but the compiler only
// enforces Send on the Future when you pass it to something that requires it
// (e.g., tokio::spawn). Calling bad_mutex(...).await directly compiles fine.
// However, tokio::spawn(bad_mutex(data)) will fail with a Send bound error.
```

**Why this is usually a problem** — but not always:

Holding a `std::sync::Mutex` across `.await` blocks the **OS thread** for the
duration of the I/O, preventing the executor from polling other tasks on that
thread. For short critical sections this is wasteful; for long I/O it's a
performance trap.

**However**, there are legitimate cases where you *must* hold a lock across an
`.await` — the same way a database transaction holds a lock between read and
commit. Dropping and re-acquiring the lock introduces a **TOCTOU (time-of-check
to time-of-use) race**: another task can modify the data between your two
critical sections. The right fix depends on the use case:

```rust
// OPTION 1: Scope the guard — works when operations are independent
async fn scoped_mutex(data: &Mutex<Vec<String>>) {
    {
        let mut guard = data.lock().unwrap();
        guard.push("item".into());
    } // Guard dropped here
    some_io().await; // Lock is released — other tasks can proceed
    {
        let mut guard = data.lock().unwrap();
        guard.push("another".into());
    }
}
// ⚠ Careful: another task can lock + modify the Vec between the two sections.
//    This is fine if the two pushes are independent, but wrong if "another"
//    depends on state set by "item".

// OPTION 2: Use tokio::sync::Mutex — holds lock across .await without
//           blocking the OS thread. Best when you need transactional
//           read-modify-write across an await point.
use tokio::sync::Mutex as AsyncMutex;

async fn async_mutex(data: &AsyncMutex<Vec<String>>) {
    let mut guard = data.lock().await; // Async lock — doesn't block the thread
    guard.push("item".into());
    some_io().await; // OK — tokio Mutex guard is Send
    guard.push("another".into());
    // Guard held the whole time — no TOCTOU race, no thread blocked.
}
```

> **When to use which Mutex**:
> - `std::sync::Mutex`: Short critical sections with no `.await` inside
> - `tokio::sync::Mutex`: When you need to hold the lock across `.await` points
>   (transactional semantics, TOCTOU avoidance)
> - `parking_lot::Mutex`: Drop-in `std` replacement, faster, smaller, still no `.await`
>
> **Rule of thumb**: Don't blindly split a critical section around an `.await`.
> Ask whether the two halves are truly independent. If they aren't — if the
> second half depends on state from the first — use `tokio::sync::Mutex` or
> redesign the data flow.

### Cancellation Hazards

Dropping a future cancels it — but this can leave things in an inconsistent state:

```rust
// ❌ DANGEROUS: Resource leak on cancellation
async fn transfer(from: &Account, to: &Account, amount: u64) {
    from.debit(amount).await;  // If cancelled HERE...
    to.credit(amount).await;   // ...money vanishes!
}

// ✅ SAFE: Make operations atomic or use compensation
async fn safe_transfer(from: &Account, to: &Account, amount: u64) -> Result<(), Error> {
    // Use a database transaction (all-or-nothing)
    let tx = db.begin_transaction().await?;
    tx.debit(from, amount).await?;
    tx.credit(to, amount).await?;
    tx.commit().await?; // Only commits if everything succeeded
    Ok(())
}

// ✅ ALSO SAFE: Use tokio::select! with cancellation awareness
tokio::select! {
    result = transfer(from, to, amount) => {
        // Transfer completed
    }
    _ = shutdown_signal() => {
        // Don't cancel mid-transfer — let it finish
        // Or: roll back explicitly
    }
}
```

### No Async Drop

Rust's `Drop` trait is synchronous — you **cannot** `.await` inside `drop()`. This is a frequent source of confusion:

```rust
struct DbConnection { /* ... */ }

impl Drop for DbConnection {
    fn drop(&mut self) {
        // ❌ Can't do this — drop() is sync!
        // self.connection.shutdown().await;

        // ✅ Workaround 1: Spawn a cleanup task (fire-and-forget)
        let conn = self.connection.take();
        tokio::spawn(async move {
            let _ = conn.shutdown().await;
        });

        // ✅ Workaround 2: Use a synchronous close
        // self.connection.blocking_close();
    }
}
```

**Best practice**: Provide an explicit `async fn close(self)` method and document that callers should use it. Rely on `Drop` only as a safety net, not the primary cleanup path.

### select! Fairness and Starvation

```rust
use tokio::sync::mpsc;

// ❌ UNFAIR: busy_stream always wins, slow_stream starves
async fn unfair(mut fast: mpsc::Receiver<i32>, mut slow: mpsc::Receiver<i32>) {
    loop {
        tokio::select! {
            Some(v) = fast.recv() => println!("fast: {v}"),
            Some(v) = slow.recv() => println!("slow: {v}"),
            // If both are ready, tokio randomly picks one.
            // But if `fast` is ALWAYS ready, `slow` rarely gets polled.
        }
    }
}

// ✅ FAIR: Use biased select or drain in batches
async fn fair(mut fast: mpsc::Receiver<i32>, mut slow: mpsc::Receiver<i32>) {
    loop {
        tokio::select! {
            biased; // Always check in order — explicit priority

            Some(v) = slow.recv() => println!("slow: {v}"),  // Priority!
            Some(v) = fast.recv() => println!("fast: {v}"),
        }
    }
}
```

### Accidental Sequential Execution

```rust
// ❌ SEQUENTIAL: Takes 2 seconds total
async fn slow() {
    let a = fetch("url_a").await; // 1 second
    let b = fetch("url_b").await; // 1 second (waits for a to finish first!)
}

// ✅ CONCURRENT: Takes 1 second total
async fn fast() {
    let (a, b) = tokio::join!(
        fetch("url_a"), // Both start immediately
        fetch("url_b"),
    );
}

// ✅ ALSO CONCURRENT: Using let + join
async fn also_fast() {
    let fut_a = fetch("url_a"); // Create future (lazy — not started yet)
    let fut_b = fetch("url_b"); // Create future
    let (a, b) = tokio::join!(fut_a, fut_b); // NOW both run concurrently
}
```

> **Trap**: `let a = fetch(url).await; let b = fetch(url).await;` is sequential!
> The second `.await` doesn't start until the first finishes. Use `join!` or
> `spawn` for concurrency.

## Case Study: Debugging a Hung Production Service

A real-world scenario: a service handles requests fine for 10 minutes, then stops responding. No errors in logs. CPU at 0%.

**Diagnosis steps:**

1. **Attach `tokio-console`** — reveals 200+ tasks stuck in `Pending` state
2. **Check task details** — all waiting on the same `Mutex::lock().await`
3. **Root cause** — one task held a `std::sync::MutexGuard` across an `.await` and panicked, poisoning the mutex. All other tasks now fail on `lock().unwrap()`

**The fix:**

| Before (broken) | After (fixed) |
|-----------------|---------------|
| `std::sync::Mutex` | `tokio::sync::Mutex` |
| `.lock().unwrap()` across `.await` | Scope lock before `.await` |
| No timeout on lock acquisition | `tokio::time::timeout(dur, mutex.lock())` |
| No recovery on poisoned mutex | `tokio::sync::Mutex` doesn't poison |

**Prevention checklist:**
- [ ] Use `tokio::sync::Mutex` if the guard crosses any `.await`
- [ ] Add `#[tracing::instrument]` to async functions for span tracking
- [ ] Run `tokio-console` in staging to catch hung tasks early
- [ ] Add health check endpoints that verify task responsiveness

🏋 Exercise: Spot the Bugs (click to expand)

**Challenge**: Find all the async pitfalls in this code and fix them.

```rust
use std::sync::Mutex;

async fn process_requests(urls: Vec<String>) -> Vec<String> {
    let results = Mutex::new(Vec::new());

    for url in &urls {
        let response = reqwest::get(url).await.unwrap().text().await.unwrap();
        std::thread::sleep(std::time::Duration::from_millis(100)); // Rate limit
        let mut guard = results.lock().unwrap();
        guard.push(response);
        expensive_parse(&guard).await; // Parse all results so far
    }

    results.into_inner().unwrap()
}
```

🔑 Solution

**Bugs found:**

1. **Sequential fetches** — URLs are fetched one at a time instead of concurrently
2. **`std::thread::sleep`** — Blocks the executor thread
3. **MutexGuard held across `.await`** — `guard` is alive when `expensive_parse` is awaited
4. **No concurrency** — Should use `join!` or `FuturesUnordered`

```rust
use tokio::sync::Mutex;
use std::sync::Arc;
use futures::stream::{self, StreamExt};

async fn process_requests(urls: Vec<String>) -> Vec<String> {
    // Fix 4: Process URLs concurrently with buffer_unordered
    let results: Vec<String> = stream::iter(urls)
        .map(|url| async move {
            let response = reqwest::get(&url).await.unwrap().text().await.unwrap();
            // Fix 2: Use tokio::time::sleep instead of std::thread::sleep
            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
            response
        })
        .buffer_unordered(10) // Up to 10 concurrent requests
        .collect()
        .await;

    // Fix 3: Parse after collecting — no mutex needed at all!
    for result in &results {
        expensive_parse(result).await;
    }

    results
}
```

**Key takeaway**: Often you can restructure async code to eliminate mutexes entirely. Collect results with streams/join, then process. Simpler, faster, no deadlock risk.

---

### Debugging Async Code

Async stack traces are notoriously cryptic — they show the executor's poll loop rather than your logical call chain. Here are the essential debugging tools.

#### tokio-console: Real-Time Task Inspector

[tokio-console](https://github.com/tokio-rs/console) gives you an `htop`-like view of every spawned task: its state, poll duration, waker activity, and resource usage.

```toml
# Cargo.toml
[dependencies]
console-subscriber = "0.4"
tokio = { version = "1", features = ["full", "tracing"] }
```

```rust
#[tokio::main]
async fn main() {
    console_subscriber::init(); // Replaces the default tracing subscriber
    // ... rest of your application
}
```

Then in another terminal:

```bash
$ RUSTFLAGS="--cfg tokio_unstable" cargo run   # Required compile-time flag
$ tokio-console                                # Connects to 127.0.0.1:6669
```

#### tracing + #[instrument]: Structured Logging for Async

The [`tracing`](https://docs.rs/tracing) crate understands `Future` lifetimes. Spans stay open across `.await` points, giving you a logical call stack even when the OS thread has moved on:

```rust
use tracing::{info, instrument};

#[instrument(skip(db_pool), fields(user_id = %user_id))]
async fn handle_request(user_id: u64, db_pool: &Pool) -> Result<Response> {
    info!("looking up user");
    let user = db_pool.get_user(user_id).await?;  // span stays open across .await
    info!(email = %user.email, "found user");
    let orders = fetch_orders(user_id).await?;     // still the same span
    Ok(build_response(user, orders))
}
```

Output (with `tracing_subscriber::fmt::json()`):

```json
{"timestamp":"...","level":"INFO","span":{"name":"handle_request","user_id":"42"},"message":"looking up user"}
{"timestamp":"...","level":"INFO","span":{"name":"handle_request","user_id":"42"},"fields":{"email":"a@b.com"},"message":"found user"}
```

#### Debugging Checklist

| Symptom | Likely Cause | Tool |
|---------|-------------|------|
| Task hangs forever | Missing `.await` or deadlocked `Mutex` | `tokio-console` task view |
| Low throughput | Blocking call on async thread | `tokio-console` poll-time histogram |
| `Future is not Send` | Non-Send type held across `.await` | Compiler error + `#[instrument]` to locate |
| Mysterious cancellation | Parent `select!` dropped a branch | `tracing` span lifecycle events |

> **Tip**: Enable `RUSTFLAGS="--cfg tokio_unstable"` to get task-level metrics
> in tokio-console. This is a compile-time flag, not a runtime one.

### Testing Async Code

Async code introduces unique testing challenges — you need a runtime, time control, and strategies for testing concurrent behavior.

**Basic async tests** with `#[tokio::test]`:

```rust
// Cargo.toml
// [dev-dependencies]
// tokio = { version = "1", features = ["full", "test-util"] }

#[tokio::test]
async fn test_basic_async() {
    let result = fetch_data().await;
    assert_eq!(result, "expected");
}

// Single-threaded test (useful for !Send types):
#[tokio::test(flavor = "current_thread")]
async fn test_single_threaded() {
    let rc = std::rc::Rc::new(42);
    let val = async { *rc }.await;
    assert_eq!(val, 42);
}

// Multi-threaded with explicit worker count:
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_concurrent_behavior() {
    // Tests race conditions with real concurrency
    let counter = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
    let c1 = counter.clone();
    let c2 = counter.clone();
    let (a, b) = tokio::join!(
        tokio::spawn(async move { c1.fetch_add(1, std::sync::atomic::Ordering::SeqCst) }),
        tokio::spawn(async move { c2.fetch_add(1, std::sync::atomic::Ordering::SeqCst) }),
    );
    a.unwrap();
    b.unwrap();
    assert_eq!(counter.load(std::sync::atomic::Ordering::SeqCst), 2);
}
```

**Time manipulation** — test timeouts without actually waiting:

```rust
use tokio::time::{self, Duration, Instant};

#[tokio::test]
async fn test_timeout_behavior() {
    // Pause time — sleep() advances instantly, no real wall-clock delay
    time::pause();

    let start = Instant::now();
    time::sleep(Duration::from_secs(3600)).await; // "waits" 1 hour — takes 0ms
    assert!(start.elapsed() >= Duration::from_secs(3600));
    // Test ran in milliseconds, not an hour!
}

#[tokio::test]
async fn test_retry_timing() {
    time::pause();

    // Test that our retry logic waits the expected durations
    let start = Instant::now();
    let result = retry_with_backoff(|| async {
        Err::<(), _>("simulated failure")
    }, 3, Duration::from_secs(1))
    .await;

    assert!(result.is_err());
    // 1s + 2s + 4s = 7s of backoff (exponential)
    assert!(start.elapsed() >= Duration::from_secs(7));
}

#[tokio::test]
async fn test_deadline_exceeded() {
    time::pause();

    let result = tokio::time::timeout(
        Duration::from_secs(5),
        async {
            // Simulate slow operation
            time::sleep(Duration::from_secs(10)).await;
            "done"
        }
    ).await;

    assert!(result.is_err()); // Timed out
}
```

**Mocking async dependencies** — use trait objects or generics:

```rust
// Define a trait for the dependency:
trait Storage {
    async fn get(&self, key: &str) -> Option<String>;
    async fn set(&self, key: &str, value: String);
}

// Production implementation:
struct RedisStorage { /* ... */ }
impl Storage for RedisStorage {
    async fn get(&self, key: &str) -> Option<String> {
        // Real Redis call
        todo!()
    }
    async fn set(&self, key: &str, value: String) {
        todo!()
    }
}

// Test mock:
struct MockStorage {
    data: std::sync::Mutex<std::collections::HashMap<String, String>>,
}

impl MockStorage {
    fn new() -> Self {
        MockStorage { data: std::sync::Mutex::new(std::collections::HashMap::new()) }
    }
}

impl Storage for MockStorage {
    async fn get(&self, key: &str) -> Option<String> {
        self.data.lock().unwrap().get(key).cloned()
    }
    async fn set(&self, key: &str, value: String) {
        self.data.lock().unwrap().insert(key.to_string(), value);
    }
}

// Tested function is generic over Storage:
async fn cache_lookup<S: Storage>(store: &S, key: &str) -> String {
    match store.get(key).await {
        Some(val) => val,
        None => {
            let val = "computed".to_string();
            store.set(key, val.clone()).await;
            val
        }
    }
}

#[tokio::test]
async fn test_cache_miss_then_hit() {
    let mock = MockStorage::new();

    // First call: miss → computes and stores
    let val = cache_lookup(&mock, "key1").await;
    assert_eq!(val, "computed");

    // Second call: hit → returns stored value
    let val = cache_lookup(&mock, "key1").await;
    assert_eq!(val, "computed");
    assert!(mock.data.lock().unwrap().contains_key("key1"));
}
```

**Testing channels and task communication**:

```rust
#[tokio::test]
async fn test_producer_consumer() {
    let (tx, mut rx) = tokio::sync::mpsc::channel(10);

    tokio::spawn(async move {
        for i in 0..5 {
            tx.send(i).await.unwrap();
        }
        // tx dropped here — channel closes
    });

    let mut received = Vec::new();
    while let Some(val) = rx.recv().await {
        received.push(val);
    }

    assert_eq!(received, vec![0, 1, 2, 3, 4]);
}
```

| Test Pattern | When to Use | Key Tool |
|-------------|-------------|----------|
| `#[tokio::test]` | All async tests | `tokio = { features = ["macros", "rt"] }` |
| `time::pause()` | Testing timeouts, retries, periodic tasks | `tokio::time::pause()` |
| Trait mocking | Testing business logic without I/O | Generic `<S: Storage>` |
| `current_thread` flavor | Testing `!Send` types or deterministic scheduling | `#[tokio::test(flavor = "current_thread")]` |
| `multi_thread` flavor | Testing race conditions | `#[tokio::test(flavor = "multi_thread")]` |

> **Key Takeaways — Common Pitfalls**
> - Never block the executor — use `spawn_blocking` for CPU/sync work
> - Never hold a `MutexGuard` across `.await` — scope locks tightly or use `tokio::sync::Mutex`
> - Cancellation drops the future instantly — use "cancel-safe" patterns for partial operations
> - Use `tokio-console` and `#[tracing::instrument]` for debugging async code
> - Test async code with `#[tokio::test]` and `time::pause()` for deterministic timing

> **See also:** [Ch 8 — Tokio Deep Dive](ch08-tokio-deep-dive.md) for sync primitives, [Ch 13 — Production Patterns](ch13-production-patterns.md) for graceful shutdown and structured concurrency

***


import pytest
import pytest_httpbin

from scrapling import DynamicFetcher

DynamicFetcher.adaptive = True

@pytest_httpbin.use_class_based_httpbin
class TestDynamicFetcherAsync:
    @pytest.fixture
    def fetcher(self):
        return DynamicFetcher

    @pytest.fixture
    def urls(self, httpbin):
        return {
            "status_200": f"{httpbin.url}/status/200",
            "status_404": f"{httpbin.url}/status/404",
            "status_501": f"{httpbin.url}/status/501",
            "basic_url": f"{httpbin.url}/get",
            "html_url": f"{httpbin.url}/html",
            "delayed_url": f"{httpbin.url}/delay/10",
            "cookies_url": f"{httpbin.url}/cookies/set/test/value",
        }

    @pytest.mark.asyncio
    async def test_basic_fetch(self, fetcher, urls):
        """Test doing a basic fetch request with multiple statuses"""
        response = await fetcher.async_fetch(urls["status_200"])
        assert response.status == 200

    @pytest.mark.asyncio
    async def test_cookies_loading(self, fetcher, urls):
        """Test if cookies are set after the request"""
        response = await fetcher.async_fetch(urls["cookies_url"])
        cookies = {response.cookies[0]['name']: response.cookies[0]['value']}
        assert cookies == {"test": "value"}

    @pytest.mark.asyncio
    async def test_automation(self, fetcher, urls):
        """Test if automation breaks the code or not"""

        async def scroll_page(page):
            await page.mouse.wheel(10, 0)
            await page.mouse.move(100, 400)
            await page.mouse.up()
            return page

        response = await fetcher.async_fetch(urls["html_url"], page_action=scroll_page)
        assert response.status == 200

    @pytest.mark.parametrize(
        "kwargs",
        [
            {"real_chrome": True, "disable_resources": True},
            {"wait_selector": "h1", "wait_selector_state": "attached"},
            {"wait_selector": "h1", "wait_selector_state": "visible"},
            {
                "google_search": True,
                "real_chrome": True,
                "wait": 10,
                "locale": "en-US",
                "extra_headers": {"ayo": ""},
                "useragent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:131.0) Gecko/20100101 Firefox/131.0",
                "cookies": [{"name": "test", "value": "123", "domain": "example.com", "path": "/"}],
                "network_idle": True,
                "selector_config": {"keep_comments": False, "keep_cdata": False},
            },
        ],
    )
    @pytest.mark.asyncio
    async def test_properties(self, fetcher, urls, kwargs):
        """Test if different arguments break the code or not"""
        response = await fetcher.async_fetch(urls["html_url"], **kwargs)
        assert response.status == 200

    @pytest.mark.asyncio
    async def test_cdp_url_invalid(self, fetcher, urls):
        """Test if invalid CDP URLs raise appropriate exceptions"""
        with pytest.raises(TypeError):
            await fetcher.async_fetch(urls["html_url"], cdp_url="blahblah")

        with pytest.raises(TypeError):
            await fetcher.async_fetch(
                urls["html_url"], cdp_url="blahblah"
            )

        with pytest.raises(Exception):
            await fetcher.async_fetch(urls["html_url"], cdp_url="ws://blahblah")


# SPDX-FileCopyrightText: 2024-present Adam Fourney <adamfo@microsoft.com>
#
# SPDX-License-Identifier: MIT
import argparse
import sys
import codecs
from typing import Any, Dict
from textwrap import dedent
from importlib.metadata import entry_points
from .__about__ import __version__
from ._markitdown import MarkItDown, StreamInfo, DocumentConverterResult

def main():
    parser = argparse.ArgumentParser(
        description="Convert various file formats to markdown.",
        prog="markitdown",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        usage=dedent(
            """
            SYNTAX:

                markitdown <OPTIONAL: FILENAME>
                If FILENAME is empty, markitdown reads from stdin.

            EXAMPLE:

                markitdown example.pdf

                OR

                cat example.pdf | markitdown

                OR

                markitdown < example.pdf

                OR to save to a file use

                markitdown example.pdf -o example.md

                OR

                markitdown example.pdf > example.md
            """
        ).strip(),
    )

    parser.add_argument(
        "-v",
        "--version",
        action="version",
        version=f"%(prog)s {__version__}",
        help="show the version number and exit",
    )

    parser.add_argument(
        "-o",
        "--output",
        help="Output file name. If not provided, output is written to stdout.",
    )

    parser.add_argument(
        "-x",
        "--extension",
        help="Provide a hint about the file extension (e.g., when reading from stdin).",
    )

    parser.add_argument(
        "-m",
        "--mime-type",
        help="Provide a hint about the file's MIME type.",
    )

    parser.add_argument(
        "-c",
        "--charset",
        help="Provide a hint about the file's charset (e.g., UTF-8).",
    )

    cloud_group = parser.add_mutually_exclusive_group()
    cloud_group.add_argument(
        "-d",
        "--use-docintel",
        action="store_true",
        help="Use Document Intelligence to extract text instead of offline conversion. Requires a valid Document Intelligence Endpoint.",
    )

    cloud_group.add_argument(
        "--use-cu",
        "--use-content-understanding",
        action="store_true",
        dest="use_cu",
        help="Use Azure Content Understanding to extract text. Requires --cu-endpoint.",
    )

    parser.add_argument(
        "-e",
        "--endpoint",
        type=str,
        help="Document Intelligence Endpoint. Required if using Document Intelligence.",
    )

    parser.add_argument(
        "--cu-endpoint",
        type=str,
        help="Content Understanding Endpoint. Required if using --use-cu.",
    )

    parser.add_argument(
        "--cu-analyzer",
        type=str,
        help="Content Understanding analyzer ID. If not specified, auto-selects by file type.",
    )

    parser.add_argument(
        "--cu-file-types",
        type=str,
        help="Comma-separated list of file types to route to Content Understanding (e.g., pdf,jpeg,mp4). If omitted, all supported types are routed.",
    )

    parser.add_argument(
        "-p",
        "--use-plugins",
        action="store_true",
        help="Use 3rd-party plugins to convert files. Use --list-plugins to see installed plugins.",
    )

    parser.add_argument(
        "--list-plugins",
        action="store_true",
        help="List installed 3rd-party plugins. Plugins are loaded when using the -p or --use-plugin option.",
    )

    parser.add_argument(
        "--keep-data-uris",
        action="store_true",
        help="Keep data URIs (like base64-encoded images) in the output. By default, data URIs are truncated.",
    )

    parser.add_argument("filename", nargs="?")
    args = parser.parse_args()

    # Parse the extension hint
    extension_hint = args.extension
    if extension_hint is not None:
        extension_hint = extension_hint.strip().lower()
        if len(extension_hint) > 0:
            if not extension_hint.startswith("."):
                extension_hint = "." + extension_hint
        else:
            extension_hint = None

    # Parse the mime type
    mime_type_hint = args.mime_type
    if mime_type_hint is not None:
        mime_type_hint = mime_type_hint.strip()
        if len(mime_type_hint) > 0:
            if mime_type_hint.count("/") != 1:
                _exit_with_error(f"Invalid MIME type: {mime_type_hint}")
        else:
            mime_type_hint = None

    # Parse the charset
    charset_hint = args.charset
    if charset_hint is not None:
        charset_hint = charset_hint.strip()
        if len(charset_hint) > 0:
            try:
                charset_hint = codecs.lookup(charset_hint).name
            except LookupError:
                _exit_with_error(f"Invalid charset: {charset_hint}")
        else:
            charset_hint = None

    stream_info = None
    if (
        extension_hint is not None
        or mime_type_hint is not None
        or charset_hint is not None
    ):
        stream_info = StreamInfo(
            extension=extension_hint, mimetype=mime_type_hint, charset=charset_hint
        )

    if args.list_plugins:
        # List installed plugins, then exit
        print("Installed MarkItDown 3rd-party Plugins:\n")
        plugin_entry_points = list(entry_points(group="markitdown.plugin"))
        if len(plugin_entry_points) == 0:
            print("  * No 3rd-party plugins installed.")
            print(
                "\nFind plugins by searching for the hashtag #markitdown-plugin on GitHub.\n"
            )
        else:
            for entry_point in plugin_entry_points:
                print(f"  * {entry_point.name:<16}\t(package: {entry_point.value})")
            print(
                "\nUse the -p (or --use-plugins) option to enable 3rd-party plugins.\n"
            )
        sys.exit(0)

    if args.use_docintel:
        if args.endpoint is None:
            _exit_with_error(
                "Document Intelligence Endpoint is required when using Document Intelligence."
            )
        elif args.filename is None:
            _exit_with_error("Filename is required when using Document Intelligence.")

        markitdown = MarkItDown(
            enable_plugins=args.use_plugins, docintel_endpoint=args.endpoint
        )
    elif args.use_cu:
        if args.cu_endpoint is None:
            _exit_with_error(
                "Content Understanding Endpoint (--cu-endpoint) is required when using --use-cu."
            )
        elif args.filename is None:
            _exit_with_error("Filename is required when using Content Understanding.")

        cu_kwargs: Dict[str, Any] = {
            "cu_endpoint": args.cu_endpoint,
        }
        if args.cu_analyzer is not None:
            cu_kwargs["cu_analyzer_id"] = args.cu_analyzer
        if args.cu_file_types is not None:
            # Parse comma-separated file types into ContentUnderstandingFileType list
            from .converters import ContentUnderstandingFileType

            type_names = [
                t.strip().lower() for t in args.cu_file_types.split(",") if t.strip()
            ]
            cu_types = []
            for name in type_names:
                # Try matching by value (e.g., "pdf", "jpeg", "mp4")
                try:
                    cu_types.append(ContentUnderstandingFileType(name))
                except ValueError:
                    _exit_with_error(f"Unknown file type: {name}")
            cu_kwargs["cu_file_types"] = cu_types

        markitdown = MarkItDown(enable_plugins=args.use_plugins, **cu_kwargs)
    else:
        markitdown = MarkItDown(enable_plugins=args.use_plugins)

    if args.filename is None:
        result = markitdown.convert_stream(
            sys.stdin.buffer,
            stream_info=stream_info,
            keep_data_uris=args.keep_data_uris,
        )
    else:
        result = markitdown.convert(
            args.filename, stream_info=stream_info, keep_data_uris=args.keep_data_uris
        )

    _handle_output(args, result)

def _handle_output(args, result: DocumentConverterResult):
    """Handle output to stdout or file"""
    if args.output:
        with open(args.output, "w", encoding="utf-8") as f:
            f.write(result.markdown)
    else:
        # Handle stdout encoding errors more gracefully
        print(
            result.markdown.encode(sys.stdout.encoding, errors="replace").decode(
                sys.stdout.encoding
            )
        )

def _exit_with_error(message: str):
    print(message)
    sys.exit(1)

if __name__ == "__main__":
    main()


//! Provider integrations included in `rig-core`.
//!
//! - Anthropic
//! - Azure OpenAI
//! - ChatGPT and GitHub Copilot auth-backed clients
//! - Cohere
//! - DeepSeek
//! - Gemini
//! - Groq
//! - Hugging Face
//! - Hyperbolic
//! - Llamafile
//! - MiniMax
//! - Mira
//! - Mistral
//! - Moonshot
//! - Ollama
//! - OpenAI
//! - OpenRouter
//! - Perplexity
//! - Together
//! - Voyage AI
//! - xAI
//! - Xiaomi MiMo
//! - Z.ai
//!
//! Each provider module defines a `Client` type and model types for the
//! capabilities it supports. Capability traits such as
//! [`CompletionClient`](crate::client::CompletionClient) and
//! [`EmbeddingsClient`](crate::client::EmbeddingsClient) are implemented only
//! when the provider declares that capability.
//!
//! # Provider implementation checklist
//!
//! When adding or changing a provider, verify that the integration includes:
//!
//! - for OpenAI-chat-compatible APIs: completions driven by
//!   [`GenericCompletionModel`](crate::providers::openai::completion::GenericCompletionModel)
//!   via an
//!   [`OpenAICompatibleProvider`](crate::providers::openai::completion::OpenAICompatibleProvider)
//!   impl on the provider extension (never a hand-rolled completion model,
//!   request struct, or message conversion — dialect differences go in the
//!   trait's hooks);
//! - public `Client` and `ClientBuilder` aliases with the correct generics,
//!   including a `ClientBuilder` API-key generic matching `ProviderBuilder::ApiKey`;
//! - the `Provider`, `ProviderBuilder`, `Capabilities`, and `ProviderClient`
//!   implementations;
//! - explicit API-key marker/auth types with redacted debug behavior for
//!   credential-bearing values;
//! - model constants where they are useful and current;
//! - request conversion from Rig request types, such as
//!   [`CompletionRequest`](crate::completion::CompletionRequest), without
//!   inventing unsupported provider API fields;
//! - response conversion into Rig response types, including usage and tool or
//!   multimodal content where applicable;
//! - streaming support when the provider supports streaming;
//! - provider-response error preservation plus `ProviderResponseExt` and
//!   telemetry fields consistent with nearby providers where applicable;
//! - unit, cassette, or live-test coverage appropriate to the changed behavior;
//! - root facade feature/docs updates for companion provider crates; and
//! - examples and documentation that match the actual API, feature flags, and
//!   credential requirements.
//!
//! # Example
//! ```no_run
//! use rig_core::{
//!     client::{CompletionClient, ProviderClient},
//!     completion::{AssistantContent, CompletionModel},
//!     providers::openai,
//! };
//!
//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
//! // Initialize the OpenAI client
//! let openai = openai::Client::from_env()?;
//!
//! // Create a model and send a low-level completion request.
//! let model = openai.completion_model(openai::GPT_5_2);
//! let request = model
//!     .completion_request("Discuss the fate of Middle Earth.")
//!     .preamble("\
//!         You are Gandalf the white and you will be conversing with other \
//!         powerful beings to discuss the fate of Middle Earth.\
//!     ".to_string())
//!     .build();
//! let response = model.completion(request).await?;
//! for item in response.choice {
//!     if let AssistantContent::Text(text) = item {
//!         println!("{}", text.text);
//!     }
//! }
//! # Ok(())
//! # }
//! ```
pub mod anthropic;
pub mod azure;
pub mod chatgpt;
pub mod cohere;
pub mod copilot;
pub mod deepseek;
pub mod doubleword;
pub mod gemini;
pub mod groq;
pub mod huggingface;
pub mod hyperbolic;
pub(crate) mod internal;
pub mod llamafile;
pub mod minimax;
pub mod mira;
pub mod mistral;
pub mod moonshot;
pub mod ollama;
pub mod openai;
pub mod openrouter;
pub mod perplexity;
pub mod together;
pub mod voyageai;
pub mod xai;
pub mod xiaomimimo;
pub mod zai;


"""
OCR Service Layer for MarkItDown
Provides LLM Vision-based image text extraction.
"""

import base64
from typing import Any, BinaryIO
from dataclasses import dataclass

from markitdown import StreamInfo

@dataclass
class OCRResult:
    """Result from OCR extraction."""

    text: str
    confidence: float | None = None
    backend_used: str | None = None
    error: str | None = None

class LLMVisionOCRService:
    """OCR service using LLM vision models (OpenAI-compatible)."""

    def __init__(
        self,
        client: Any,
        model: str,
        default_prompt: str | None = None,
    ) -> None:
        """
        Initialize LLM Vision OCR service.

        Args:
            client: OpenAI-compatible client
            model: Model name (e.g., 'gpt-4o', 'gemini-2.0-flash')
            default_prompt: Default prompt for OCR extraction
        """
        self.client = client
        self.model = model
        self.default_prompt = default_prompt or (
            "Extract all text from this image. "
            "Return ONLY the extracted text, maintaining the original "
            "layout and order. Do not add any commentary or description."
        )

    def extract_text(
        self,
        image_stream: BinaryIO,
        prompt: str | None = None,
        stream_info: StreamInfo | None = None,
        **kwargs: Any,
    ) -> OCRResult:
        """Extract text using LLM vision."""
        if self.client is None:
            return OCRResult(
                text="",
                backend_used="llm_vision",
                error="LLM client not configured",
            )

        try:
            image_stream.seek(0)

            content_type: str | None = None
            if stream_info:
                content_type = stream_info.mimetype

            if not content_type:
                try:
                    from PIL import Image

                    image_stream.seek(0)
                    img = Image.open(image_stream)
                    fmt = img.format.lower() if img.format else "png"
                    content_type = f"image/{fmt}"
                except Exception:
                    content_type = "image/png"

            image_stream.seek(0)
            base64_image = base64.b64encode(image_stream.read()).decode("utf-8")
            data_uri = f"data:{content_type};base64,{base64_image}"

            actual_prompt = prompt or self.default_prompt
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": actual_prompt},
                            {
                                "type": "image_url",
                                "image_url": {"url": data_uri},
                            },
                        ],
                    }
                ],
            )

            text = response.choices[0].message.content
            return OCRResult(
                text=text.strip() if text else "",
                backend_used="llm_vision",
            )
        except Exception as e:
            return OCRResult(text="", backend_used="llm_vision", error=str(e))
        finally:
            image_stream.seek(0)