Advanced Problem-Solving Scenarios and Runbooks

flowchart TD
    A["User submits complex<br/>comparison query"] --> B["Step Functions Express<br/>Workflow starts"]
    B --> C["ReAct iteration 1<br/>Thought + Action + Observe<br/>~800ms"]
    C --> D["ReAct iteration 2<br/>~750ms cumulative: 1550ms"]
    D --> E["ReAct iteration 3<br/>~700ms cumulative: 2250ms"]
    E --> F["ReAct iteration 4<br/>~800ms cumulative: 3050ms"]
    F --> G{"Budget check:<br/>remaining < 500ms?"}
    G -->|"No: 1950ms remaining"| H["ReAct iteration 5<br/>Bedrock Sonnet slow: 1200ms"]
    H --> I["Observation step<br/>Bedrock Haiku: 400ms"]
    I --> J["UpdateContext Lambda: 100ms<br/>cumulative: 4750ms"]
    J --> K["ReAct iteration 6 starts...<br/>but only 250ms left"]
    K --> L["Bedrock Sonnet call<br/>takes 900ms"]
    L --> M["ExecutionTimedOut<br/>at 5000ms"]
    M --> N["CloudWatch alarm fires:<br/>ExecutionTimedOut count > 0"]
    N --> O["PagerDuty alert to<br/>on-call engineer"]

    style M fill:#ffcdd2
    style N fill:#fff3e0
    style O fill:#fff3e0

"""
Resolution: Adaptive budget manager with hard cutoff and Sonnet-to-Haiku fallback.

Fixes:
1. Enforce hard per-iteration cutoff that accounts for Bedrock P99 latency
2. Dynamically switch from Sonnet to Haiku mid-loop when budget is tight
3. Reserve a fixed 600ms window for best-effort synthesis
4. Increase Express Workflow timeout to 8 seconds but maintain 3s user SLA
   via WebSocket streaming (send partial results as they arrive)
"""

import time
import logging
from dataclasses import dataclass

logger = logging.getLogger("react_budget_fix")

SONNET_P99_LATENCY_MS = 1500
HAIKU_P99_LATENCY_MS = 500
TOOL_P99_LATENCY_MS = 400
SYNTHESIS_RESERVE_MS = 600
NETWORK_OVERHEAD_MS = 200


@dataclass
class AdaptiveBudgetManager:
    """
    Fixed budget manager that prevents execution timeout.

    Key changes from the broken version:
    - Uses P99 latency estimates instead of P50
    - Reserves 600ms for best-effort synthesis (non-negotiable)
    - Switches to Haiku automatically when Sonnet is unaffordable
    - Calculates max affordable iterations upfront
    """
    total_budget_ms: float = 4800.0  # 5000ms timeout - 200ms overhead
    start_time: float = 0.0
    synthesis_reserve_ms: float = SYNTHESIS_RESERVE_MS

    def __post_init__(self):
        if self.start_time == 0.0:
            self.start_time = time.monotonic()

    @property
    def elapsed_ms(self) -> float:
        return (time.monotonic() - self.start_time) * 1000

    @property
    def remaining_ms(self) -> float:
        return self.total_budget_ms - self.elapsed_ms

    @property
    def actionable_remaining_ms(self) -> float:
        """Time available for reasoning (excludes synthesis reserve)."""
        return max(0, self.remaining_ms - self.synthesis_reserve_ms)

    def can_start_iteration(self, iteration: int) -> bool:
        """
        Check if we can safely start another ReAct iteration.

        A full iteration needs:
        - Thought step: Haiku P99 = 500ms, Sonnet P99 = 1500ms
        - Action step: Tool P99 = 400ms
        - Observation step: Haiku P99 = 500ms
        Minimum per iteration (all Haiku): ~1400ms
        """
        min_iteration_cost = HAIKU_P99_LATENCY_MS + TOOL_P99_LATENCY_MS + HAIKU_P99_LATENCY_MS
        return self.actionable_remaining_ms >= min_iteration_cost

    def select_thought_model(self, iteration: int) -> str:
        """
        Select model for the Thought step based on budget.

        Use Sonnet only if:
        1. It is the first iteration (most impactful reasoning step)
        2. We have enough budget for Sonnet + remaining iteration steps
        """
        sonnet_iteration_cost = SONNET_P99_LATENCY_MS + TOOL_P99_LATENCY_MS + HAIKU_P99_LATENCY_MS
        haiku_iteration_cost = HAIKU_P99_LATENCY_MS + TOOL_P99_LATENCY_MS + HAIKU_P99_LATENCY_MS

        if iteration == 0 and self.actionable_remaining_ms >= sonnet_iteration_cost + haiku_iteration_cost:
            # Can afford Sonnet for first iteration AND at least one more Haiku iteration
            return "anthropic.claude-3-sonnet-20240229-v1:0"

        return "anthropic.claude-3-haiku-20240307-v1:0"

    def max_affordable_iterations(self) -> int:
        """Calculate the maximum number of iterations we can afford."""
        min_cost = HAIKU_P99_LATENCY_MS + TOOL_P99_LATENCY_MS + HAIKU_P99_LATENCY_MS
        return max(0, int(self.actionable_remaining_ms / min_cost))

    def log_budget_status(self, iteration: int):
        """Log current budget status for debugging."""
        logger.info(
            "Budget status: iteration=%d elapsed=%.0fms remaining=%.0fms "
            "actionable=%.0fms max_iterations=%d",
            iteration,
            self.elapsed_ms,
            self.remaining_ms,
            self.actionable_remaining_ms,
            self.max_affordable_iterations(),
        )


def handler_thought_with_budget(event, context):
    """
    Fixed Thought step handler with adaptive budget management.

    This replaces the original handler_thought that did not enforce
    hard budget limits and used optimistic latency estimates.
    """
    import json
    import boto3

    bedrock = boto3.client("bedrock-runtime")
    ctx = event["reasoning_context"]
    iteration = ctx["current_iteration"]

    # Reconstruct budget from context
    budget = AdaptiveBudgetManager(
        total_budget_ms=ctx.get("total_budget_ms", 4800),
        start_time=time.monotonic() - (ctx.get("elapsed_ms", 0) / 1000),
    )

    # Hard check: can we afford another iteration?
    if not budget.can_start_iteration(iteration):
        logger.warning(
            "Budget exhausted at iteration %d -- forcing synthesis",
            iteration,
        )
        return {
            "reasoning": "Budget exhausted, synthesizing from available observations",
            "decision": "FINAL_ANSWER",
            "confidence": 0.5,
            "next_action": None,
            "action_params": None,
            "final_answer": None,  # Will be filled by BestEffortSynthesis state
            "force_synthesis": True,
        }

    # Select model based on budget
    model_id = budget.select_thought_model(iteration)
    budget.log_budget_status(iteration)

    obs_summary = _format_observations_compact(ctx.get("observations", []))

    prompt = f"""You are MangaAssist's reasoning engine (iteration {iteration + 1}).
Budget: {budget.actionable_remaining_ms:.0f}ms remaining. Be efficient.

USER QUERY: {ctx['user_query']}

OBSERVATIONS SO FAR:
{obs_summary}

Available tools: opensearch_vector_search, dynamodb_product_lookup,
elasticache_rating_cache, genre_theme_classifier, inventory_check

JSON response:
{{"reasoning": "...", "decision": "CONTINUE|FINAL_ANSWER|CLARIFY",
  "confidence": 0.0-1.0, "next_action": "tool_name|null",
  "action_params": {{}}, "final_answer": "answer|null"}}"""

    response = bedrock.invoke_model(
        modelId=model_id,
        contentType="application/json",
        accept="application/json",
        body=json.dumps({
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": 500,  # Reduced from 600 to save output time
            "temperature": 0.1,
            "messages": [{"role": "user", "content": prompt}],
        }),
    )

    result = json.loads(response["body"].read())
    text = result["content"][0]["text"]
    parsed = json.loads(text)
    parsed["model_used"] = model_id
    parsed["budget_remaining_ms"] = budget.actionable_remaining_ms

    return parsed


def _format_observations_compact(observations: list) -> str:
    """Compact observation formatting to reduce input tokens."""
    if not observations:
        return "None yet."
    return "; ".join(
        f"[{i+1}] {obs.get('summary', 'N/A')[:100]}"
        for i, obs in enumerate(observations)
    )

flowchart TD
    A["Prevention: Execution Timeout"] --> B["1. Adaptive Budget Manager"]
    A --> C["2. Express Workflow Config"]
    A --> D["3. Monitoring and Alerts"]
    A --> E["4. Architecture Changes"]

    B --> B1["Use P99 latency estimates<br/>not P50 or averages"]
    B --> B2["Reserve 600ms for synthesis<br/>(non-negotiable buffer)"]
    B --> B3["Switch Sonnet to Haiku<br/>mid-loop when budget tight"]
    B --> B4["Calculate max iterations<br/>upfront at loop start"]

    C --> C1["Increase Express Workflow<br/>timeout to 8 seconds"]
    C --> C2["Set per-state TimeoutSeconds<br/>to match P99 model latency"]
    C --> C3["Add heartbeat monitoring<br/>for long-running states"]

    D --> D1["CloudWatch alarm on<br/>ExecutionTimedOut > 1/min"]
    D --> D2["Dashboard: P99 latency<br/>per iteration per model"]
    D --> D3["Alert on Bedrock latency<br/>spike > 2x baseline"]

    E --> E1["Stream partial results<br/>via WebSocket as each<br/>iteration completes"]
    E --> E2["Cache reasoning traces<br/>for similar queries to<br/>skip re-computation"]
    E --> E3["Pre-warm Bedrock inference<br/>with periodic keepalive calls"]

    style A fill:#e8f5e9
    style B fill:#e3f2fd
    style C fill:#e3f2fd
    style D fill:#e3f2fd
    style E fill:#e3f2fd

flowchart TD
    A["CoT Response Generated"] --> B["Parse STEP 1...N<br/>and ANSWER"]
    B --> C["Contradiction Detector<br/>(Bedrock Haiku)"]
    C --> D{"Contradiction<br/>score > 0.6?"}
    D -->|"No"| E["Pass through to user"]
    D -->|"Yes"| F["Flag for review"]
    F --> G["Log contradiction trace<br/>to DynamoDB"]
    G --> H["CloudWatch metric:<br/>ContradictionDetected"]
    H --> I{"Contradiction<br/>rate > 5%?"}
    I -->|"No"| J["Auto-resolution:<br/>regenerate with<br/>contradiction prompt"]
    I -->|"Yes"| K["Alert: Prompt template<br/>may need revision"]
    J --> L["Re-invoke with<br/>contradiction-aware prompt"]
    L --> M["Return corrected answer"]

    style D fill:#fff3e0
    style I fill:#fff3e0
    style K fill:#ffcdd2
    style M fill:#c8e6c9

"""
Resolution: Contradiction detection and self-correction in CoT reasoning.

Fixes:
1. Add a cross-step consistency check after CoT generation
2. If contradiction detected, re-invoke with a contradiction-aware prompt
3. Add few-shot exemplars that demonstrate conflict resolution
4. Implement an ObservationConsistencyChecker in the ReAct loop
"""

import json
import logging
from typing import Optional
from dataclasses import dataclass

import boto3

logger = logging.getLogger("cot_consistency")

HAIKU_MODEL_ID = "anthropic.claude-3-haiku-20240307-v1:0"


@dataclass
class ConsistencyCheckResult:
    """Result of checking CoT steps for contradictions."""
    has_contradiction: bool
    contradiction_details: str
    conflicting_steps: list[int]  # 1-indexed step numbers
    severity: float  # 0.0-1.0
    suggested_resolution: str


class CoTConsistencyChecker:
    """
    Checks chain-of-thought reasoning steps for internal contradictions.

    Uses a lightweight Haiku call to evaluate cross-step consistency.
    Runs after CoT generation but before returning the answer to the user.

    Cost: ~200 tokens input + 150 tokens output per check
    At 1M messages/day: ~$0.24/day (negligible)
    Latency: ~200ms (Haiku)
    """

    def __init__(self, bedrock_client=None):
        self.bedrock = bedrock_client or boto3.client("bedrock-runtime")

    async def check_consistency(
        self,
        query: str,
        steps: list[str],
        answer: str,
    ) -> ConsistencyCheckResult:
        """
        Check if the CoT reasoning steps are internally consistent.

        Returns a ConsistencyCheckResult indicating whether contradictions
        exist and how to resolve them.
        """
        steps_text = "\n".join(f"{s}" for s in steps)

        prompt = f"""Analyze these reasoning steps for internal contradictions.

QUERY: {query}

REASONING STEPS:
{steps_text}

FINAL ANSWER: {answer}

Check for:
1. Steps that directly contradict each other
2. Answer that contradicts evidence in the steps
3. Logical inconsistencies (e.g., using genre label to override content analysis)

Respond in JSON:
{{
  "has_contradiction": true/false,
  "contradiction_details": "description of the contradiction or 'none'",
  "conflicting_steps": [step_numbers],
  "severity": 0.0-1.0,
  "suggested_resolution": "how to resolve the contradiction"
}}"""

        response = self.bedrock.invoke_model(
            modelId=HAIKU_MODEL_ID,
            contentType="application/json",
            accept="application/json",
            body=json.dumps({
                "anthropic_version": "bedrock-2023-05-31",
                "max_tokens": 300,
                "temperature": 0.0,
                "messages": [{"role": "user", "content": prompt}],
            }),
        )

        result = json.loads(response["body"].read())
        parsed = json.loads(result["content"][0]["text"])

        return ConsistencyCheckResult(
            has_contradiction=parsed.get("has_contradiction", False),
            contradiction_details=parsed.get("contradiction_details", "none"),
            conflicting_steps=parsed.get("conflicting_steps", []),
            severity=parsed.get("severity", 0.0),
            suggested_resolution=parsed.get("suggested_resolution", ""),
        )

    async def resolve_contradiction(
        self,
        query: str,
        original_steps: list[str],
        contradiction: ConsistencyCheckResult,
    ) -> dict:
        """
        Re-invoke CoT with explicit contradiction awareness.

        Includes the detected contradiction in the prompt so the model
        can reason about it explicitly rather than glossing over it.
        """
        steps_text = "\n".join(f"{s}" for s in original_steps)

        prompt = f"""You are MangaAssist. Your previous reasoning had a contradiction that needs correction.

QUERY: {query}

PREVIOUS REASONING:
{steps_text}

DETECTED CONTRADICTION:
{contradiction.contradiction_details}
Conflicting steps: {contradiction.conflicting_steps}
Suggested resolution: {contradiction.suggested_resolution}

Please re-reason from the beginning, explicitly addressing the conflict.
When genre metadata conflicts with actual content, prioritize content analysis.
When multiple valid perspectives exist, present them as a nuanced answer.

STEP 1: [corrected reasoning]
...
ANSWER: [corrected answer]
CONFIDENCE: [0.0-1.0]"""

        response = self.bedrock.invoke_model(
            modelId="anthropic.claude-3-sonnet-20240229-v1:0",  # Use Sonnet for correction
            contentType="application/json",
            accept="application/json",
            body=json.dumps({
                "anthropic_version": "bedrock-2023-05-31",
                "max_tokens": 800,
                "temperature": 0.1,
                "messages": [{"role": "user", "content": prompt}],
            }),
        )

        result = json.loads(response["body"].read())
        return {
            "corrected_text": result["content"][0]["text"],
            "model_used": "sonnet",
            "correction_type": "contradiction_resolution",
        }


class ObservationConsistencyChecker:
    """
    Checks observation history for cross-step consistency in the ReAct loop.

    Runs after each Observation step to detect when new observations
    contradict earlier ones. If inconsistency is found, the next Thought
    step receives a flag to address it explicitly.
    """

    def __init__(self, bedrock_client=None):
        self.bedrock = bedrock_client or boto3.client("bedrock-runtime")

    async def check_observation_consistency(
        self,
        new_observation: dict,
        observation_history: list[dict],
        user_query: str,
    ) -> dict:
        """
        Verify the new observation is consistent with the history.

        Returns a consistency report for the next Thought step.
        """
        if not observation_history:
            return {"consistent": True}

        history_summary = "; ".join(
            obs.get("summary", "")[:80] for obs in observation_history
        )

        prompt = f"""Check if this new observation is consistent with previous observations.

QUERY: {user_query}
PREVIOUS OBSERVATIONS: {history_summary}
NEW OBSERVATION: {new_observation.get('summary', '')}

JSON response:
{{"consistent": true/false, "conflict": "description or null",
  "resolution_hint": "how the next reasoning step should handle this"}}"""

        response = self.bedrock.invoke_model(
            modelId=HAIKU_MODEL_ID,
            contentType="application/json",
            accept="application/json",
            body=json.dumps({
                "anthropic_version": "bedrock-2023-05-31",
                "max_tokens": 200,
                "temperature": 0.0,
                "messages": [{"role": "user", "content": prompt}],
            }),
        )

        result = json.loads(response["body"].read())
        return json.loads(result["content"][0]["text"])


# ──────────────────────────────────────────────
# Few-shot exemplar for contradiction handling
# ──────────────────────────────────────────────

CONTRADICTION_EXEMPLAR = {
    "query": "Is Chainsaw Man appropriate for a 14-year-old?",
    "steps": [
        "STEP 1: Chainsaw Man is published in Shounen Jump (later moved to Jump+), "
        "targeting teen male readers. The genre classification suggests teen suitability.",
        "STEP 2: However, the actual content includes extreme violence with graphic "
        "dismemberment, body horror elements, sexual themes and innuendo, and mature "
        "psychological themes including depression and existential nihilism.",
        "STEP 3: CONFLICT DETECTED -- genre label (shounen/teen) conflicts with content "
        "maturity. In such cases, content analysis takes precedence over genre labels "
        "because publishers sometimes push boundaries within genre classifications.",
        "STEP 4: For a 14-year-old specifically: the violence exceeds typical shounen "
        "levels (e.g., compared to Naruto or My Hero Academia). The thematic content "
        "is more appropriate for older teens (16+) or adults.",
    ],
    "answer": "While Chainsaw Man is technically a shounen manga, its content is more "
    "intense than typical shounen series. It contains graphic violence, mature themes, "
    "and some sexual content. I would recommend it for older teens (16+) rather than a "
    "14-year-old. For a similar energy with more age-appropriate content, consider "
    "Jujutsu Kaisen or My Hero Academia.",
}

flowchart TD
    A["Prevention: CoT Contradictions"] --> B["1. Prompt Engineering"]
    A --> C["2. Runtime Checks"]
    A --> D["3. Exemplar Management"]
    A --> E["4. Monitoring"]

    B --> B1["Add explicit instruction:<br/>'If evidence contradicts<br/>a label, prioritize evidence'"]
    B --> B2["Include CONFLICT CHECK<br/>as a mandatory reasoning step"]
    B --> B3["Instruct model to flag<br/>low-confidence conclusions"]

    C --> C1["CoTConsistencyChecker<br/>after every CoT generation"]
    C --> C2["ObservationConsistencyChecker<br/>after every ReAct observation"]
    C --> C3["Auto-correction with<br/>contradiction-aware re-prompt"]

    D --> D1["Add contradiction-handling<br/>exemplars to few-shot bank"]
    D --> D2["Include examples where<br/>metadata conflicts with content"]
    D --> D3["Review exemplars quarterly<br/>for new contradiction patterns"]

    E --> E1["Track contradiction rate<br/>as CloudWatch metric"]
    E --> E2["Alert if rate exceeds<br/>5% of CoT queries"]
    E --> E3["Log all contradictions<br/>for human review pipeline"]

    style A fill:#e8f5e9
    style B fill:#e3f2fd
    style C fill:#e3f2fd
    style D fill:#e3f2fd
    style E fill:#e3f2fd

flowchart TD
    A["Popular manga release day<br/>traffic spike: 3x normal"] --> B["Bedrock request rate<br/>exceeds provisioned throughput"]
    B --> C["ThrottlingException<br/>returned to Lambda"]
    C --> D["Step Functions retry policy<br/>MaxAttempts: 10, BackoffRate: 1.5"]
    D --> E["Retry 1: wait 1s"]
    E --> F["Retry 2: wait 1.5s"]
    F --> G["Retry 3: wait 2.25s"]
    G --> H["...continues to retry 10<br/>total wait: ~57 seconds"]
    H --> I["Execution duration<br/>spikes to 60+ seconds"]
    I --> J["Express Workflow timeout<br/>(5 min max) not hit but<br/>user SLA (3s) massively violated"]

    J --> K["CloudWatch Alarms Fire"]
    K --> K1["Alarm: Bedrock ThrottlingException<br/>rate > 10/min"]
    K --> K2["Alarm: Avg execution duration<br/>> 5 seconds"]
    K --> K3["Alarm: API Gateway 5xx<br/>error rate > 1%"]

    K1 --> L["Runbook triggered"]
    K2 --> L
    K3 --> L

    style C fill:#ffcdd2
    style H fill:#ffcdd2
    style I fill:#ffcdd2
    style L fill:#fff3e0

"""
Resolution: Circuit breaker + intelligent retry + graceful degradation.

Fixes:
1. Reduce MaxAttempts to 2 with aggressive backoff
2. Implement application-level circuit breaker
3. Degrade from Sonnet to Haiku (different model = different throughput pool)
4. Cache recent reasoning traces to serve similar queries without Bedrock
5. Return cached/template answers when circuit is open
"""

import json
import time
import logging
import threading
from enum import Enum
from dataclasses import dataclass
from typing import Optional

import boto3

logger = logging.getLogger("bedrock_circuit_breaker")


class CircuitState(str, Enum):
    CLOSED = "CLOSED"        # Normal operation
    OPEN = "OPEN"            # Failing, reject requests immediately
    HALF_OPEN = "HALF_OPEN"  # Testing if service recovered


@dataclass
class CircuitBreakerConfig:
    """Configuration for the Bedrock circuit breaker."""
    failure_threshold: int = 5       # Consecutive failures before opening
    recovery_timeout_s: float = 10.0 # Time in OPEN state before testing
    half_open_max_calls: int = 2     # Test calls in HALF_OPEN state
    success_threshold: int = 2       # Successes needed to close again


class BedrockCircuitBreaker:
    """
    Circuit breaker for Bedrock API calls.

    States:
    - CLOSED: Normal operation. Track consecutive failures.
    - OPEN: Bedrock is throttling. Reject calls immediately, return
      fallback. After recovery_timeout, transition to HALF_OPEN.
    - HALF_OPEN: Allow limited test calls. If they succeed, close
      the circuit. If they fail, reopen.

    This prevents the retry storm that overwhelms Bedrock during spikes.
    """

    def __init__(self, config: CircuitBreakerConfig = None):
        self.config = config or CircuitBreakerConfig()
        self._state = CircuitState.CLOSED
        self._failure_count = 0
        self._success_count = 0
        self._last_failure_time = 0.0
        self._half_open_calls = 0
        self._lock = threading.Lock()

    @property
    def state(self) -> CircuitState:
        with self._lock:
            if self._state == CircuitState.OPEN:
                # Check if recovery timeout has elapsed
                if time.monotonic() - self._last_failure_time >= self.config.recovery_timeout_s:
                    self._state = CircuitState.HALF_OPEN
                    self._half_open_calls = 0
                    self._success_count = 0
                    logger.info("Circuit breaker transitioning to HALF_OPEN")
            return self._state

    def allow_request(self) -> bool:
        """Check if a request should be allowed through."""
        state = self.state

        if state == CircuitState.CLOSED:
            return True
        elif state == CircuitState.OPEN:
            return False
        elif state == CircuitState.HALF_OPEN:
            with self._lock:
                if self._half_open_calls < self.config.half_open_max_calls:
                    self._half_open_calls += 1
                    return True
                return False
        return False

    def record_success(self):
        """Record a successful Bedrock call."""
        with self._lock:
            if self._state == CircuitState.HALF_OPEN:
                self._success_count += 1
                if self._success_count >= self.config.success_threshold:
                    self._state = CircuitState.CLOSED
                    self._failure_count = 0
                    logger.info("Circuit breaker CLOSED -- Bedrock recovered")
            else:
                self._failure_count = 0

    def record_failure(self):
        """Record a failed Bedrock call (throttling or error)."""
        with self._lock:
            self._failure_count += 1
            self._last_failure_time = time.monotonic()

            if self._state == CircuitState.HALF_OPEN:
                # Failed during recovery test -- reopen
                self._state = CircuitState.OPEN
                logger.warning("Circuit breaker re-OPENED -- Bedrock still failing")
            elif self._failure_count >= self.config.failure_threshold:
                self._state = CircuitState.OPEN
                logger.warning(
                    "Circuit breaker OPENED after %d consecutive failures",
                    self._failure_count,
                )


# Global circuit breaker instances (one per model)
_circuit_breakers = {
    "sonnet": BedrockCircuitBreaker(CircuitBreakerConfig(
        failure_threshold=3,
        recovery_timeout_s=15.0,
    )),
    "haiku": BedrockCircuitBreaker(CircuitBreakerConfig(
        failure_threshold=5,
        recovery_timeout_s=10.0,
    )),
}


def invoke_bedrock_with_circuit_breaker(
    bedrock_client,
    model_id: str,
    prompt: str,
    max_tokens: int = 500,
    temperature: float = 0.1,
) -> Optional[dict]:
    """
    Invoke Bedrock with circuit breaker protection.

    Fallback chain:
    1. Try requested model (e.g., Sonnet)
    2. If circuit open, try Haiku (different throughput pool)
    3. If both circuits open, return None (caller uses cached response)
    """
    model_key = "sonnet" if "sonnet" in model_id else "haiku"
    breaker = _circuit_breakers[model_key]

    if breaker.allow_request():
        try:
            response = bedrock_client.invoke_model(
                modelId=model_id,
                contentType="application/json",
                accept="application/json",
                body=json.dumps({
                    "anthropic_version": "bedrock-2023-05-31",
                    "max_tokens": max_tokens,
                    "temperature": temperature,
                    "messages": [{"role": "user", "content": prompt}],
                }),
            )
            breaker.record_success()
            result = json.loads(response["body"].read())
            return result
        except Exception as e:
            if "ThrottlingException" in str(type(e).__name__) or "Throttl" in str(e):
                breaker.record_failure()
                logger.warning("Bedrock throttled for model %s", model_id)
            else:
                raise

    # Sonnet circuit open -- try Haiku as fallback
    if model_key == "sonnet":
        haiku_breaker = _circuit_breakers["haiku"]
        haiku_model = "anthropic.claude-3-haiku-20240307-v1:0"
        if haiku_breaker.allow_request():
            try:
                response = bedrock_client.invoke_model(
                    modelId=haiku_model,
                    contentType="application/json",
                    accept="application/json",
                    body=json.dumps({
                        "anthropic_version": "bedrock-2023-05-31",
                        "max_tokens": max_tokens,
                        "temperature": temperature,
                        "messages": [{"role": "user", "content": prompt}],
                    }),
                )
                haiku_breaker.record_success()
                result = json.loads(response["body"].read())
                result["_fallback_model"] = True
                return result
            except Exception as e:
                if "Throttl" in str(e):
                    haiku_breaker.record_failure()

    # Both circuits open -- return None to signal cached/template response
    logger.error("All Bedrock circuits open -- returning None for cached fallback")
    return None


# ──────────────────────────────────────────────
# Updated Step Functions ASL with conservative retry
# ──────────────────────────────────────────────

FIXED_THOUGHT_STATE_RETRY = {
    "Retry": [
        {
            "ErrorEquals": ["Bedrock.ThrottlingException", "Lambda.TooManyRequestsException"],
            "IntervalSeconds": 1,
            "MaxAttempts": 2,       # Down from 10
            "BackoffRate": 2.0,     # Faster backoff
            "JitterStrategy": "FULL"  # Add jitter to prevent thundering herd
        }
    ],
    "Catch": [
        {
            "ErrorEquals": ["Bedrock.ThrottlingException"],
            "ResultPath": "$.throttle_error",
            "Next": "DegradedResponse"  # New state: return best-effort without Bedrock
        },
        {
            "ErrorEquals": ["States.Timeout"],
            "ResultPath": "$.error",
            "Next": "FallbackResponse"
        }
    ],
}

flowchart TD
    A["Prevention: Infinite Retry Loop"] --> B["1. Retry Policy"]
    A --> C["2. Circuit Breaker"]
    A --> D["3. Capacity Planning"]
    A --> E["4. Load Shedding"]

    B --> B1["MaxAttempts: 2<br/>(not 10)"]
    B --> B2["BackoffRate: 2.0<br/>with FULL jitter"]
    B --> B3["Total retry window<br/>< 5 seconds"]
    B --> B4["Catch ThrottlingException<br/>-> DegradedResponse state"]

    C --> C1["Per-model circuit breaker<br/>Sonnet and Haiku independent"]
    C --> C2["Failure threshold: 3-5<br/>consecutive failures"]
    C --> C3["Recovery test after<br/>10-15 seconds"]
    C --> C4["Fallback chain:<br/>Sonnet -> Haiku -> cached"]

    D --> D1["Request Bedrock<br/>provisioned throughput<br/>for Sonnet before launch events"]
    D --> D2["Monitor on-demand<br/>token consumption rate<br/>against account limits"]
    D --> D3["Pre-scale for known events<br/>(manga release days)"]

    E --> E1["API Gateway throttling<br/>at 1500 req/sec"]
    E --> E2["SQS queue for overflow<br/>with async response"]
    E --> E3["Priority queue: returning<br/>customers > anonymous"]
    E --> E4["Serve cached answers<br/>for popular queries"]

    style A fill:#e8f5e9
    style B fill:#e3f2fd
    style C fill:#e3f2fd
    style D fill:#e3f2fd
    style E fill:#e3f2fd

flowchart TD
    A["User submits query with<br/>embedded prompt injection"] --> B["ToT engine starts<br/>max_branching=3, max_depth=3"]
    B --> C["Depth 0: root node"]
    C --> D["Depth 1: expand_node<br/>LLM returns 10 children<br/>(ignoring max_branching=3)"]
    D --> E["Depth 2: expand 10 nodes<br/>each returns 10 children<br/>= 100 nodes"]
    E --> F["Depth 3: expand 100 nodes<br/>each returns 10 children<br/>= 1000 nodes"]
    F --> G["Total: 1111 Bedrock calls<br/>instead of 39"]

    G --> H["CloudWatch Metrics Spike"]
    H --> H1["Metric: ToTNodesExplored<br/>value: 1111 (baseline: 39)"]
    H --> H2["Metric: BedrockTokenCost<br/>value: $0.22 (baseline: $0.008)"]
    H --> H3["Metric: ToTLatency<br/>value: 30s+ (baseline: 2s)"]

    H1 --> I["Cost Anomaly Detection<br/>alert fires"]
    H2 --> I
    H3 --> J["Execution timeout<br/>or SLA breach alert"]

    I --> K["On-call engineer<br/>investigates"]
    J --> K

    style D fill:#ffcdd2
    style E fill:#ffcdd2
    style F fill:#ffcdd2
    style G fill:#ffcdd2

"""
Resolution: Hard enforcement of branching limits + prompt injection defense.

Fixes:
1. Server-side hard cap on children per node (not reliant on LLM output)
2. Token budget enforcement at the tree level
3. Input sanitization to strip prompt injection patterns
4. Per-query cost cap that terminates ToT if exceeded
"""

import json
import re
import time
import logging
import asyncio
from dataclasses import dataclass, field
from typing import Optional

import boto3

logger = logging.getLogger("tot_hardened")

HAIKU_MODEL_ID = "anthropic.claude-3-haiku-20240307-v1:0"
COST_PER_1K_INPUT_TOKENS_HAIKU = 0.00025
COST_PER_1K_OUTPUT_TOKENS_HAIKU = 0.00125


@dataclass
class ToTCostGuard:
    """
    Guards against cost explosions in Tree-of-Thought reasoning.

    Tracks cumulative token usage and enforces a hard per-query cost cap.
    """
    max_cost_per_query_usd: float = 0.02  # Hard cap: $0.02 per query
    max_nodes_per_query: int = 40         # Hard cap: 40 nodes maximum
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_nodes: int = 0

    @property
    def estimated_cost_usd(self) -> float:
        input_cost = (self.total_input_tokens / 1000) * COST_PER_1K_INPUT_TOKENS_HAIKU
        output_cost = (self.total_output_tokens / 1000) * COST_PER_1K_OUTPUT_TOKENS_HAIKU
        return input_cost + output_cost

    def can_expand(self) -> bool:
        """Check if we can afford to expand more nodes."""
        if self.total_nodes >= self.max_nodes_per_query:
            logger.warning("Node limit reached: %d/%d", self.total_nodes, self.max_nodes_per_query)
            return False
        if self.estimated_cost_usd >= self.max_cost_per_query_usd:
            logger.warning(
                "Cost limit reached: $%.4f/$%.4f",
                self.estimated_cost_usd, self.max_cost_per_query_usd,
            )
            return False
        return True

    def record_expansion(self, input_tokens: int, output_tokens: int, num_children: int):
        """Record token usage from a node expansion."""
        self.total_input_tokens += input_tokens
        self.total_output_tokens += output_tokens
        self.total_nodes += num_children


def sanitize_user_query(query: str) -> str:
    """
    Remove prompt injection patterns from user queries.

    Strips common injection vectors:
    - [system: ...] directives
    - <system>...</system> tags
    - "ignore previous instructions" patterns
    - Numeric overrides like "generate 10 options"
    """
    # Remove bracketed system directives
    sanitized = re.sub(r'\[system[:\s].*?\]', '', query, flags=re.IGNORECASE)
    # Remove XML-style system tags
    sanitized = re.sub(r'<system>.*?</system>', '', sanitized, flags=re.IGNORECASE | re.DOTALL)
    # Remove "ignore previous/above instructions" patterns
    sanitized = re.sub(
        r'ignore\s+(previous|above|prior|all)\s+instructions?',
        '', sanitized, flags=re.IGNORECASE,
    )
    # Remove numeric override attempts like "generate 10 options"
    sanitized = re.sub(
        r'generate\s+\d+\s+(options|branches|paths|alternatives)',
        '', sanitized, flags=re.IGNORECASE,
    )
    return sanitized.strip()


class HardenedTreeOfThought:
    """
    Hardened ToT engine with cost guards and injection defense.

    Key differences from original:
    1. sanitize_user_query strips injection patterns before prompting
    2. ToTCostGuard enforces hard caps on nodes and cost
    3. Server-side enforcement: children list is ALWAYS truncated
    4. If LLM returns invalid JSON, node is treated as terminal (no children)
    """

    def __init__(
        self,
        bedrock_client=None,
        max_branching_factor: int = 3,
        max_depth: int = 3,
        pruning_threshold: float = 0.4,
        max_cost_per_query: float = 0.02,
        max_nodes: int = 40,
    ):
        self.bedrock = bedrock_client or boto3.client("bedrock-runtime")
        self.max_branching_factor = max_branching_factor
        self.max_depth = max_depth
        self.pruning_threshold = pruning_threshold
        self.cost_guard = ToTCostGuard(
            max_cost_per_query_usd=max_cost_per_query,
            max_nodes_per_query=max_nodes,
        )

    async def expand_node_safe(
        self,
        query: str,
        context: str,
        parent_thought: str,
        depth: int,
    ) -> list[dict]:
        """
        Safely expand a node with hard enforcement of limits.

        CRITICAL DIFFERENCES from original:
        1. Query is sanitized before insertion into prompt
        2. Children list is ALWAYS truncated to max_branching_factor
        3. JSON parse failure = empty children (no retry)
        4. Cost guard checked before expansion
        """
        # Check cost guard
        if not self.cost_guard.can_expand():
            return []

        # Sanitize query
        safe_query = sanitize_user_query(query)

        prompt = f"""Generate exactly {self.max_branching_factor} different reasoning branches.

IMPORTANT: Return EXACTLY {self.max_branching_factor} items. No more, no less.

QUERY: {safe_query}
CURRENT REASONING: {parent_thought}
DEPTH: {depth}/{self.max_depth}

JSON array of exactly {self.max_branching_factor} items:
[{{"thought": "...", "score": 0.0-1.0, "is_terminal": true/false, "answer": "..."}}]"""

        try:
            response = self.bedrock.invoke_model(
                modelId=HAIKU_MODEL_ID,
                contentType="application/json",
                accept="application/json",
                body=json.dumps({
                    "anthropic_version": "bedrock-2023-05-31",
                    "max_tokens": 600,
                    "temperature": 0.7,
                    "messages": [{"role": "user", "content": prompt}],
                }),
            )

            result = json.loads(response["body"].read())
            usage = result.get("usage", {})
            text = result["content"][0]["text"]

            try:
                children = json.loads(text)
            except json.JSONDecodeError:
                logger.warning("Invalid JSON from ToT expansion -- treating as terminal")
                return []

            if not isinstance(children, list):
                logger.warning("ToT expansion returned non-list -- treating as terminal")
                return []

            # HARD ENFORCEMENT: truncate to max_branching_factor
            children = children[:self.max_branching_factor]

            # Record usage
            self.cost_guard.record_expansion(
                input_tokens=usage.get("input_tokens", 0),
                output_tokens=usage.get("output_tokens", 0),
                num_children=len(children),
            )

            return children

        except Exception as e:
            logger.error("ToT expansion failed: %s", str(e))
            return []

flowchart TD
    A["Prevention: Branching Explosion"] --> B["1. Input Defense"]
    A --> C["2. Hard Limits"]
    A --> D["3. Cost Monitoring"]
    A --> E["4. Architecture"]

    B --> B1["Sanitize all user input<br/>before LLM prompts"]
    B --> B2["Strip system directives<br/>bracketed tags, overrides"]
    B --> B3["Use input validation<br/>Lambda at API Gateway"]

    C --> C1["Server-side truncation<br/>of children lists ALWAYS"]
    C --> C2["Max nodes per query: 40<br/>(hard cap, not configurable)"]
    C --> C3["Per-query cost cap: $0.02<br/>terminate if exceeded"]
    C --> C4["Token budget per query:<br/>10,000 tokens maximum"]

    D --> D1["CloudWatch alarm:<br/>ToTNodesExplored > 50"]
    D --> D2["AWS Cost Anomaly Detection<br/>on Bedrock service"]
    D --> D3["Per-query cost metric<br/>with P99 dashboard"]
    D --> D4["Daily cost report<br/>with anomaly flagging"]

    E --> E1["Gate ToT behind query<br/>classifier -- only 5%<br/>of queries use it"]
    E --> E2["Consider replacing ToT<br/>with Self-Consistency CoT<br/>for most use cases"]
    E --> E3["Rate-limit ToT executions<br/>to 100/minute max"]

    style A fill:#e8f5e9
    style B fill:#e3f2fd
    style C fill:#e3f2fd
    style D fill:#e3f2fd
    style E fill:#e3f2fd

flowchart TD
    A["ReAct iteration 1<br/>context size: 2 KB"] --> B["ReAct iteration 2<br/>context size: 18 KB"]
    B --> C["ReAct iteration 3<br/>context size: 65 KB"]
    C --> D["ReAct iteration 4<br/>context size: 180 KB"]
    D --> E["ReAct iteration 5<br/>UpdateContext Lambda"]
    E --> F{"Context size<br/>> 256 KB?"}
    F -->|"Yes"| G["States.DataLimitExceeded<br/>Execution terminated"]
    F -->|"No"| H["Continue to next<br/>Thought step"]
    G --> I["CloudWatch alarm:<br/>DataLimitExceeded count > 0"]
    I --> J["Inspect execution history<br/>for payload sizes"]
    J --> K["Identify: observations contain<br/>full synopses + reviews"]

    style G fill:#ffcdd2
    style I fill:#fff3e0

"""
Resolution: Context compression + external state storage.

Fixes:
1. Summarize observations before appending to context
2. Store full observation data in DynamoDB, pass only references
3. Enforce a per-observation size limit
4. Implement sliding window: keep only last N observations in context,
   archive older ones to DynamoDB
"""

import json
import time
import sys
import logging
from typing import Optional
from dataclasses import dataclass

import boto3

logger = logging.getLogger("context_compression")

# Step Functions payload limit
SF_PAYLOAD_LIMIT_BYTES = 256 * 1024  # 256 KB
# Safety margin: stay under 200 KB to leave room for state machine metadata
SF_SAFE_LIMIT_BYTES = 200 * 1024     # 200 KB
# Max characters per observation summary
MAX_OBSERVATION_SUMMARY_CHARS = 300
# Max facts per observation
MAX_FACTS_PER_OBSERVATION = 5
# Max characters per fact
MAX_CHARS_PER_FACT = 150
# Sliding window: keep last N observations in context
OBSERVATION_WINDOW_SIZE = 3

HAIKU_MODEL_ID = "anthropic.claude-3-haiku-20240307-v1:0"


@dataclass
class CompactObservation:
    """A size-limited observation for inclusion in Step Functions payloads."""
    iteration: int
    summary: str         # Max 300 chars
    quality_score: float
    key_facts: list[str] # Max 5 facts, each max 150 chars
    full_data_ref: str   # DynamoDB key for full observation data


class ContextCompressor:
    """
    Compresses reasoning context to fit within Step Functions payload limits.

    Strategy:
    1. Store full observation data in DynamoDB (unlimited size)
    2. Keep only compact summaries in the Step Functions payload
    3. Use a sliding window to limit the number of observations in context
    4. Monitor context size and warn when approaching the limit
    """

    def __init__(self, dynamodb_resource=None, bedrock_client=None):
        self.dynamodb = dynamodb_resource or boto3.resource("dynamodb")
        self.bedrock = bedrock_client or boto3.client("bedrock-runtime")
        self.context_table = self.dynamodb.Table("manga-reasoning-context")

    def compress_observation(
        self,
        observation: dict,
        session_id: str,
        iteration: int,
    ) -> CompactObservation:
        """
        Compress a full observation into a compact form.

        Steps:
        1. Store full observation in DynamoDB
        2. Truncate summary to MAX_OBSERVATION_SUMMARY_CHARS
        3. Truncate each fact to MAX_CHARS_PER_FACT
        4. Keep only top MAX_FACTS_PER_OBSERVATION facts
        """
        # Store full data in DynamoDB
        ref_key = f"{session_id}:obs:{iteration}"
        self.context_table.put_item(Item={
            "context_key": ref_key,
            "data": json.dumps(observation, default=str),
            "ttl": int(time.time()) + 3600,  # 1-hour TTL
        })

        # Compress summary
        summary = observation.get("summary", "")[:MAX_OBSERVATION_SUMMARY_CHARS]

        # Compress facts
        facts = observation.get("relevant_facts", [])
        key_facts = [
            fact[:MAX_CHARS_PER_FACT]
            for fact in facts[:MAX_FACTS_PER_OBSERVATION]
        ]

        return CompactObservation(
            iteration=iteration,
            summary=summary,
            quality_score=observation.get("quality_score", 0.0),
            key_facts=key_facts,
            full_data_ref=ref_key,
        )

    def apply_sliding_window(
        self,
        context: dict,
        window_size: int = OBSERVATION_WINDOW_SIZE,
    ) -> dict:
        """
        Apply a sliding window to keep only the most recent observations.

        Older observations are already stored in DynamoDB via their
        full_data_ref, so we keep only the most recent N in the payload.
        """
        observations = context.get("observations", [])

        if len(observations) > window_size:
            # Archive older observations (already in DynamoDB)
            archived = observations[:-window_size]
            context["archived_observation_count"] = context.get(
                "archived_observation_count", 0
            ) + len(archived)
            context["observations"] = observations[-window_size:]

            logger.info(
                "Sliding window applied: kept %d observations, archived %d",
                window_size, len(archived),
            )

        return context

    def check_context_size(self, context: dict) -> tuple[dict, dict]:
        """
        Check the serialized size of the context and compress if needed.

        Returns the context (possibly compressed) and a size report.
        """
        serialized = json.dumps(context, default=str)
        size_bytes = len(serialized.encode("utf-8"))

        report = {
            "size_bytes": size_bytes,
            "size_kb": round(size_bytes / 1024, 1),
            "limit_kb": round(SF_SAFE_LIMIT_BYTES / 1024, 1),
            "utilization_pct": round((size_bytes / SF_SAFE_LIMIT_BYTES) * 100, 1),
            "compressed": False,
        }

        if size_bytes > SF_SAFE_LIMIT_BYTES:
            logger.warning(
                "Context size %.1f KB exceeds safe limit %.1f KB -- compressing",
                report["size_kb"], report["limit_kb"],
            )
            context = self._emergency_compress(context)
            report["compressed"] = True
            new_size = len(json.dumps(context, default=str).encode("utf-8"))
            report["compressed_size_kb"] = round(new_size / 1024, 1)

        return context, report

    def _emergency_compress(self, context: dict) -> dict:
        """
        Emergency compression when context is near the limit.

        Aggressive strategy:
        1. Keep only last 2 observations
        2. Truncate all facts to 80 chars
        3. Remove thoughts and actions from context (keep only observations)
        4. Remove user preferences if present
        """
        # Keep only last 2 observations
        context["observations"] = context.get("observations", [])[-2:]

        # Truncate facts aggressively
        for obs in context.get("observations", []):
            if isinstance(obs, dict):
                obs["relevant_facts"] = [
                    f[:80] for f in obs.get("relevant_facts", [])[:3]
                ]
                obs["summary"] = obs.get("summary", "")[:200]
                obs.pop("gaps", None)

        # Remove bulk data
        context.pop("thoughts", None)
        context.pop("actions", None)
        context.pop("user_preferences", None)

        return context

    def retrieve_full_observation(self, ref_key: str) -> Optional[dict]:
        """Retrieve the full observation data from DynamoDB when needed."""
        try:
            response = self.context_table.get_item(Key={"context_key": ref_key})
            item = response.get("Item", {})
            if item:
                return json.loads(item.get("data", "{}"))
        except Exception as e:
            logger.error("Failed to retrieve observation %s: %s", ref_key, str(e))
        return None


# ──────────────────────────────────────────────
# Updated handler: UpdateReasoningContext with compression
# ──────────────────────────────────────────────

def handler_update_context_compressed(event, context):
    """
    Fixed UpdateReasoningContext handler with context compression.

    Key changes:
    1. Observations are compressed before appending
    2. Sliding window limits the number of in-context observations
    3. Context size is checked before returning to Step Functions
    4. Emergency compression triggers if near the payload limit
    """
    ctx = event["reasoning_context"]
    thought = event["thought"]
    action = event["action"]
    observation = event["observation"]

    compressor = ContextCompressor()

    # Compress observation before appending
    compact_obs = compressor.compress_observation(
        observation=observation,
        session_id=ctx["session_id"],
        iteration=ctx["current_iteration"],
    )

    # Append compact observation (not the full one)
    ctx["observations"].append({
        "iteration": compact_obs.iteration,
        "summary": compact_obs.summary,
        "quality_score": compact_obs.quality_score,
        "key_facts": compact_obs.key_facts,
        "full_data_ref": compact_obs.full_data_ref,
    })

    # Apply sliding window
    ctx = compressor.apply_sliding_window(ctx)

    # Increment iteration
    ctx["current_iteration"] += 1

    # Accumulate token usage (compact)
    if "tokens" in thought:
        ctx["token_usage"]["input"] += thought["tokens"].get("input", 0)
        ctx["token_usage"]["output"] += thought["tokens"].get("output", 0)

    # Check size and compress if needed
    ctx, size_report = compressor.check_context_size(ctx)

    logger.info(
        "Context update: iteration=%d size=%.1fKB utilization=%.1f%%",
        ctx["current_iteration"],
        size_report["size_kb"],
        size_report["utilization_pct"],
    )

    return ctx

flowchart TD
    A["Prevention: Payload Size Limit"] --> B["1. Context Compression"]
    A --> C["2. External Storage"]
    A --> D["3. Size Monitoring"]
    A --> E["4. Architecture"]

    B --> B1["Summarize observations<br/>to 300 chars max"]
    B --> B2["Limit facts to 5 per<br/>observation, 150 chars each"]
    B --> B3["Sliding window: keep<br/>only last 3 observations"]
    B --> B4["Emergency compression<br/>if context > 200 KB"]

    C --> C1["Store full observations<br/>in DynamoDB"]
    C --> C2["Pass only reference keys<br/>in Step Functions payload"]
    C --> C3["Retrieve full data<br/>only when needed for<br/>final synthesis"]

    D --> D1["Check context size<br/>after every UpdateContext"]
    D --> D2["CloudWatch metric:<br/>ContextSizeBytes per iteration"]
    D --> D3["Alert if context exceeds<br/>200 KB (80% of limit)"]

    E --> E1["Consider Step Functions<br/>Distributed Map for<br/>large-context workflows"]
    E --> E2["Use S3 for intermediate<br/>results instead of<br/>passing through states"]
    E --> E3["Evaluate moving to<br/>ECS-based orchestration<br/>for large-context queries"]

    style A fill:#e8f5e9
    style B fill:#e3f2fd
    style C fill:#e3f2fd
    style D fill:#e3f2fd
    style E fill:#e3f2fd

flowchart TD
    START["Reasoning System<br/>Issue Detected"] --> Q1{"Error type?"}

    Q1 -->|"ExecutionTimedOut"| Q2{"Which state<br/>timed out?"}
    Q1 -->|"ThrottlingException"| Q3{"Single spike<br/>or sustained?"}
    Q1 -->|"DataLimitExceeded"| S5["Scenario 5:<br/>Payload Size Limit"]
    Q1 -->|"Incorrect Answer"| Q4{"Error pattern?"}
    Q1 -->|"Cost Anomaly"| Q5{"Which component<br/>cost spike?"}

    Q2 -->|"ThoughtStep<br/>(Bedrock slow)"| S1A["Scenario 1:<br/>Switch to Haiku +<br/>reduce iterations"]
    Q2 -->|"ActionStep<br/>(Tool slow)"| S1B["Check tool health<br/>OpenSearch/DynamoDB/Redis"]
    Q2 -->|"Overall execution"| S1C["Scenario 1:<br/>Adaptive budget manager"]

    Q3 -->|"Single spike"| S3A["Retry will resolve<br/>Monitor for recurrence"]
    Q3 -->|"Sustained (>1 min)"| S3B["Scenario 3:<br/>Circuit breaker +<br/>Haiku fallback"]

    Q4 -->|"Contradictory<br/>reasoning"| S2["Scenario 2:<br/>Consistency checker +<br/>re-prompting"]
    Q4 -->|"Hallucinated<br/>manga titles"| S2B["Not Scenario 2:<br/>Check RAG pipeline<br/>and OpenSearch index"]
    Q4 -->|"Incorrect<br/>metadata"| S2C["Not Scenario 2:<br/>Check DynamoDB data<br/>quality"]

    Q5 -->|"Bedrock tokens"| Q6{"Normal query<br/>volume?"}
    Q5 -->|"Step Functions"| S5B["Check for infinite<br/>loops or excessive<br/>Map iterations"]
    Q5 -->|"Lambda"| S5C["Check for recursive<br/>invocations or<br/>timeout loops"]

    Q6 -->|"Yes, normal volume"| S4["Scenario 4:<br/>Branching explosion +<br/>input sanitization"]
    Q6 -->|"No, traffic spike"| S3B

    style START fill:#e1f5fe
    style S1A fill:#fff3e0
    style S1B fill:#fff3e0
    style S1C fill:#fff3e0
    style S2 fill:#fff3e0
    style S3A fill:#c8e6c9
    style S3B fill:#fff3e0
    style S4 fill:#fff3e0
    style S5 fill:#fff3e0