/*
 * Copyright 2008-2026 Async-IO.org
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package org.atmosphere.ai;

import org.atmosphere.ai.approval.ApprovalRegistry;
import org.atmosphere.ai.approval.ApprovalStrategy;
import org.atmosphere.ai.tool.ToolRegistry;
import org.atmosphere.cpr.AtmosphereResource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

/**
 * A {@link StreamingSession} wrapper that adds {@link #stream(String)} support.
 * Created by the {@code @AiEndpoint} infrastructure to enable the simple pattern:
 *
 * <pre>{@code
 * @Prompt
 * public void onPrompt(String message, StreamingSession session) {
 *     session.stream(message);
 * }
 * }</pre>
 *
 * <p>When {@code stream(message)} is called, this session:</p>
 * <ol>
 *   <li>Loads conversation history from {@link AiConversationMemory} (if enabled)</li>
 *   <li>Builds an {@link AiRequest} from the message + stored system prompt + model + history</li>
 *   <li>Runs {@link AiInterceptor#preProcess} in FIFO order</li>
 *   <li>Wraps the delegate in a {@link MemoryCapturingSession} to capture the response</li>
 *   <li>Delegates to {@link AgentRuntime#stream(AiRequest, StreamingSession)}</li>
 *   <li>Runs {@link AiInterceptor#postProcess} in LIFO order</li>
 * </ol>
 */
public class AiStreamingSession implements StreamingSession {

    private static final Logger logger = LoggerFactory.getLogger(AiStreamingSession.class);

    /**
     * Active sessions keyed by {@link AtmosphereResource#uuid()}. The value is a
     * {@link java.util.concurrent.CopyOnWriteArrayList} because Atmosphere
     * explicitly allows overlapping prompts on a single resource (AG-UI,
     * streaming-chat): two concurrent {@code @Prompt} invocations register two
     * sessions under the same UUID, and registering the second must not clobber
     * the first. Identity-checked removal on complete/error ensures prompt A
     * finishing after prompt B started does not remove B's mapping. Approval
     * routing walks every session in the list for the target UUID and
     * short-circuits only on {@link ApprovalRegistry.ResolveResult#RESOLVED}.
     */
    private static final ConcurrentHashMap<String, java.util.concurrent.CopyOnWriteArrayList<AiStreamingSession>> ACTIVE_SESSIONS =
            new ConcurrentHashMap<>();

    /**
     * Request attribute key under which the active {@link StreamingSession} is stored,
     * allowing {@link AiInterceptor#postProcess} to send metadata to the client.
     */
    public static final String STREAMING_SESSION_ATTR = "ai.streaming.session";

    private final StreamingSession delegate;
    private final AgentRuntime runtime;
    private final String systemPrompt;
    private final String model;
    private final List<AiInterceptor> interceptors;
    private final AtmosphereResource resource;
    private final AiConversationMemory memory;
    private final ToolRegistry toolRegistry;
    private final List<AiGuardrail> guardrails;
    private final List<ContextProvider> contextProviders;
    private final AiMetrics metrics;
    private final Class<?> responseType;
    private final AtomicReference<PostPromptHook> preStreamHook = new AtomicReference<>();
    private final AtomicBoolean handoffInProgress = new AtomicBoolean();
    private final ApprovalRegistry approvalRegistry = new ApprovalRegistry();
    /** Endpoint-configured prompt cache policy from {@code @AiEndpoint.promptCache()}. */
    private volatile org.atmosphere.ai.llm.CacheHint.CachePolicy cachePolicy;
    /** Endpoint-configured per-request retry policy from {@code @AiEndpoint.retry()}. */
    private volatile org.atmosphere.ai.RetryPolicy endpointRetryPolicy;
    /**
     * Framework-scoped injectables stashed by the endpoint handler and
     * threaded into {@code @AiTool} dispatch so tool methods can declare
     * {@code AgentFleet}, {@code AgentIdentity}, etc. directly as parameters.
     * Copied defensively by {@link #setInjectables(java.util.Map)}.
     */
    private volatile java.util.Map<Class<?>, Object> injectables = java.util.Map.of();

    /**
     * The runId assigned by {@code RunRegistry} when the endpoint handler
     * registered this run; surfaced via {@link #runId()} so the client's
     * first response carries the id and subsequent reconnects can reattach
     * via {@code X-Atmosphere-Run-Id}.
     */
    private volatile String runId;

    /**
     * Handle returned by the runtime for the currently in-flight prompt.
     * Captured inside {@link #stream(String)} so {@link #cancelInflight()}
     * (called from the disconnect cleanup path) can abort the upstream LLM
     * HTTP request and refund any tokens that would otherwise have been
     * burned for a client who is no longer listening. Volatile because the
     * disconnect path runs on a different thread (Atmosphere's lifecycle
     * dispatcher) than the @Prompt VT.
     */
    private volatile ExecutionHandle currentHandle;

    /**
     * Latches a {@link #cancelInflight()} call that arrived before
     * {@link #currentHandle} was published. Runtimes that publish their
     * handle from a side channel (e.g. by completing a {@code CompletableFuture}
     * inside {@code executeWithHandle}) can unblock the disconnect path
     * before the @Prompt VT has assigned {@code this.currentHandle}. Without
     * this latch the cancellation would be silently dropped and the VT would
     * park on {@code whenDone().join()} forever. After assigning
     * {@code currentHandle}, {@link #stream(String)} re-reads this flag and
     * fires the cancel if it was set during the race window.
     */
    private final java.util.concurrent.atomic.AtomicBoolean cancelPending =
            new java.util.concurrent.atomic.AtomicBoolean();

    /**
     * @param delegate     the underlying streaming session
     * @param runtime    the resolved AI support implementation
     * @param systemPrompt the system prompt from {@code @AiEndpoint}
     * @param model        the model name (may be null for provider default)
     * @param interceptors the interceptor chain
     * @param resource     the atmosphere resource for this client
     */
    public AiStreamingSession(StreamingSession delegate, AgentRuntime runtime,
                              String systemPrompt, String model,
                              List<AiInterceptor> interceptors,
                              AtmosphereResource resource) {
        this(delegate, runtime, systemPrompt, model, interceptors, resource, null,
                null, List.of(), List.of(), AiMetrics.NOOP, null);
    }

    /**
     * @param delegate     the underlying streaming session
     * @param runtime    the resolved AI support implementation
     * @param systemPrompt the system prompt from {@code @AiEndpoint}
     * @param model        the model name (may be null for provider default)
     * @param interceptors the interceptor chain
     * @param resource     the atmosphere resource for this client
     * @param memory       conversation memory (may be null if disabled)
     */
    public AiStreamingSession(StreamingSession delegate, AgentRuntime runtime,
                              String systemPrompt, String model,
                              List<AiInterceptor> interceptors,
                              AtmosphereResource resource,
                              AiConversationMemory memory) {
        this(delegate, runtime, systemPrompt, model, interceptors, resource, memory,
                null, List.of(), List.of(), AiMetrics.NOOP, null);
    }

    /**
     * Full constructor with tools, guardrails, context providers, and metrics.
     */
    public AiStreamingSession(StreamingSession delegate, AgentRuntime runtime,
                              String systemPrompt, String model,
                              List<AiInterceptor> interceptors,
                              AtmosphereResource resource,
                              AiConversationMemory memory,
                              ToolRegistry toolRegistry,
                              List<AiGuardrail> guardrails,
                              List<ContextProvider> contextProviders) {
        this(delegate, runtime, systemPrompt, model, interceptors, resource, memory,
                toolRegistry, guardrails, contextProviders, AiMetrics.NOOP, null);
    }

    /**
     * Full constructor with metrics and structured output.
     */
    public AiStreamingSession(StreamingSession delegate, AgentRuntime runtime,
                              String systemPrompt, String model,
                              List<AiInterceptor> interceptors,
                              AtmosphereResource resource,
                              AiConversationMemory memory,
                              ToolRegistry toolRegistry,
                              List<AiGuardrail> guardrails,
                              List<ContextProvider> contextProviders,
                              AiMetrics metrics,
                              Class<?> responseType) {
        this.delegate = delegate;
        this.runtime = runtime;
        this.systemPrompt = systemPrompt != null ? systemPrompt : "";
        this.model = model;
        this.interceptors = interceptors != null ? interceptors : List.of();
        this.resource = resource;
        this.memory = memory;
        this.toolRegistry = toolRegistry;
        this.guardrails = guardrails != null ? guardrails : List.of();
        this.contextProviders = contextProviders != null ? contextProviders : List.of();
        this.metrics = org.atmosphere.ai.jfr.CompositeAiMetrics.withJfr(
                metrics != null ? metrics : AiMetrics.NOOP);
        this.responseType = responseType;
    }

    /**
     * Register a session as an active session for its resource. Called by the
     * handler after construction so approval responses can be routed to it.
     * Multiple sessions may be registered under the same UUID (overlapping
     * prompts on one resource); each registration appends to the list.
     *
     * @param session the session to register
     */
    public static void registerActive(AiStreamingSession session) {
        var uuid = session.resource.uuid();
        if (uuid != null) {
            ACTIVE_SESSIONS
                    .computeIfAbsent(uuid, k -> new java.util.concurrent.CopyOnWriteArrayList<>())
                    .add(session);
        }
    }

    /**
     * Try to resolve an approval message against every session currently
     * registered for the given resource UUID. Walks the list in registration
     * order and short-circuits only on
     * {@link ApprovalRegistry.ResolveResult#RESOLVED} so a newer session whose
     * registry does not own the pending ID cannot swallow the message before
     * an older overlapping session has a chance to claim it.
     *
     * @param resourceUuid the atmosphere resource UUID
     * @param message      the incoming approval message
     * @return {@code true} iff some session's registry owned the pending ID and
     *     completed the approval future
     */
    public static boolean tryResolveApprovalForResource(String resourceUuid, String message) {
        if (resourceUuid == null) {
            return false;
        }
        var sessions = ACTIVE_SESSIONS.get(resourceUuid);
        if (sessions == null) {
            return false;
        }
        for (var session : sessions) {
            if (session.approvalRegistry.resolve(message) == ApprovalRegistry.ResolveResult.RESOLVED) {
                return true;
            }
        }
        return false;
    }

    /**
     * Reconnect fallback: try to resolve an approval message against every
     * active session across every resource UUID. Used when the resource UUID
     * changed due to transport reconnect and the primary lookup returns null.
     * Short-circuits only on {@link ApprovalRegistry.ResolveResult#RESOLVED};
     * an {@link ApprovalRegistry.ResolveResult#UNKNOWN_ID} from one session is
     * not allowed to consume the message — the next session may own it.
     *
     * @param message the incoming approval message
     * @return {@code true} iff some session's registry owned the pending ID and
     *     completed the approval future
     */
    public static boolean tryResolveAnySession(String message) {
        for (var sessions : ACTIVE_SESSIONS.values()) {
            for (var session : sessions) {
                if (session.approvalRegistry.resolve(message) == ApprovalRegistry.ResolveResult.RESOLVED) {
                    return true;
                }
            }
        }
        return false;
    }

    /**
     * Remove every session registered for a given resource UUID. Called by
     * {@code AiEndpointHandler#handleDisconnect} when the underlying socket
     * is going away and all prompts on it are done.
     *
     * <p>Each removed session has {@link #cancelInflight()} fired before the
     * map entry is dropped so any in-flight LLM call is aborted (refunding
     * tokens for a client that is no longer listening) and any virtual
     * thread parked on a {@code @RequiresApproval} future unblocks instead
     * of waiting out its 5-minute timeout. Closes Correctness Invariant #2
     * (Terminal Path Completeness) for the disconnect-while-streaming
     * scenario.</p>
     *
     * @param resourceUuid the atmosphere resource UUID (may be null)
     */
    public static void removeAllForResource(String resourceUuid) {
        if (resourceUuid == null) {
            return;
        }
        var sessions = ACTIVE_SESSIONS.remove(resourceUuid);
        if (sessions == null) {
            return;
        }
        for (var session : sessions) {
            try {
                session.cancelInflight();
            } catch (RuntimeException e) {
                // Best-effort cleanup. Swallowing here would violate the
                // project's no-swallow rule, so log at TRACE — the disconnect
                // path is already past the point where errors can be returned
                // to the client.
                logger.trace("cancelInflight threw during disconnect cleanup for {}: {}",
                        resourceUuid, e.getMessage(), e);
            }
        }
    }

    /**
     * Abort the in-flight LLM call (if any) and release every pending tool
     * approval registered on this session. Called from the disconnect cleanup
     * path so the upstream HTTP request to the model provider is cancelled
     * — refunding tokens for a client who has gone away — and any virtual
     * thread parked on {@link ApprovalRegistry#awaitApproval} unblocks
     * immediately with a denial outcome instead of waiting out its
     * per-approval timeout.
     *
     * <p>Idempotent: subsequent calls are no-ops once {@link ExecutionHandle}
     * is done. Safe to invoke from any thread.</p>
     */
    public void cancelInflight() {
        // Set the latch BEFORE reading currentHandle so that a concurrent
        // stream() VT publishing its handle after this read still observes
        // the cancel intent and self-cancels post-publish.
        cancelPending.set(true);
        var handle = this.currentHandle;
        if (handle != null && !handle.isDone()) {
            handle.cancel();
        }
        int cancelledApprovals = approvalRegistry.cancelAllPending();
        if (cancelledApprovals > 0 && logger.isDebugEnabled()) {
            logger.debug("Disconnect cancelled {} pending approval(s) for resource {}",
                    cancelledApprovals, resource != null ? resource.uuid() : "null");
        }
    }

    /**
     * @param resourceUuid the atmosphere resource UUID
     * @return {@code true} if at least one {@link AiStreamingSession} is
     *     registered for that uuid; {@code false} otherwise. Used by the admin
     *     {@code cancel-inflight} endpoint to distinguish "no such session"
     *     (404) from "session present, cancellation fired" (200).
     */
    public static boolean resourceHasActiveSessions(String resourceUuid) {
        if (resourceUuid == null) {
            return false;
        }
        var sessions = ACTIVE_SESSIONS.get(resourceUuid);
        return sessions != null && !sessions.isEmpty();
    }

    /**
     * Cancel the in-flight LLM call on every session registered for
     * {@code resourceUuid} WITHOUT removing them from the active map. Used by
     * the admin {@code POST /sessions/{uuid}/cancel-inflight} endpoint when
     * an operator wants to abort the upstream call (refunding tokens,
     * releasing approval futures) but keep the socket open for the next
     * turn.
     *
     * <p>Contrast with {@link #removeAllForResource} which is fired from the
     * full disconnect path — that variant cancels AND drops the map entry
     * because the transport itself is going away.</p>
     */
    public static void cancelInflightForResource(String resourceUuid) {
        if (resourceUuid == null) {
            return;
        }
        var sessions = ACTIVE_SESSIONS.get(resourceUuid);
        if (sessions == null) {
            return;
        }
        for (var session : sessions) {
            try {
                session.cancelInflight();
            } catch (RuntimeException e) {
                logger.trace("cancelInflight threw during admin cancel-inflight for {}: {}",
                        resourceUuid, e.getMessage(), e);
            }
        }
    }

    /**
     * Remove a specific session from the active map on completion/error. Uses
     * identity-checked removal so prompt A finishing after prompt B started on
     * the same resource does not clobber B's mapping. When the list for a UUID
     * becomes empty, the entry is removed from the outer map.
     *
     * @param session the session to remove (identity-checked)
     */
    public static void removeActiveSession(AiStreamingSession session) {
        if (session == null || session.resource == null) {
            return;
        }
        var uuid = session.resource.uuid();
        if (uuid == null) {
            return;
        }
        var sessions = ACTIVE_SESSIONS.get(uuid);
        if (sessions == null) {
            return;
        }
        sessions.remove(session);
        if (sessions.isEmpty()) {
            // Only remove the outer entry if the list is still empty — another
            // thread may have appended between the isEmpty check and the remove.
            ACTIVE_SESSIONS.remove(uuid, sessions);
        }
    }

    /** Sets a hook to fire once at the start of the first {@code stream()} call. */
    public void setPreStreamHook(PostPromptHook hook) {
        this.preStreamHook.set(hook);
    }

    /**
     * Set the endpoint-scoped prompt cache policy from {@code @AiEndpoint.promptCache()}.
     * Every request built by this session attaches a {@link org.atmosphere.ai.llm.CacheHint}
     * with the given policy to the context metadata.
     */
    public void setCachePolicy(org.atmosphere.ai.llm.CacheHint.CachePolicy policy) {
        this.cachePolicy = policy;
    }

    /**
     * Set the endpoint-scoped per-request retry policy from {@code @AiEndpoint.retry()}.
     * Every context built by this session uses this policy instead of the client-level default.
     */
    public void setRetryPolicy(org.atmosphere.ai.RetryPolicy policy) {
        this.endpointRetryPolicy = policy;
    }

    /**
     * Publish the framework-scoped injectables the endpoint handler collected
     * (live {@code AgentFleet}, {@code AgentIdentity}, {@code AgentState},
     * ...). Threaded into the tool-call loop so {@code @AiTool} methods can
     * declare these types directly — no {@link ThreadLocal} shim. Copied
     * defensively so the handler can mutate its source map without affecting
     * the session.
     */
    public void setInjectables(java.util.Map<Class<?>, Object> source) {
        this.injectables = source == null || source.isEmpty()
                ? java.util.Map.of()
                : java.util.Map.copyOf(source);
    }

    /**
     * Replace (or add) a single injectable entry, e.g. so {@code @Prompt} can
     * swap a fleet wrapped in a {@code StreamingActivityListener} before tools
     * dispatch. The internal map stays immutable-at-publish via copy-on-write
     * so concurrent tool invocations see a consistent snapshot.
     */
    public void putInjectable(Class<?> key, Object value) {
        var next = new java.util.LinkedHashMap<Class<?>, Object>(this.injectables);
        next.put(key, value);
        this.injectables = java.util.Map.copyOf(next);
    }

    @Override
    public java.util.Map<Class<?>, Object> injectables() {
        return injectables;
    }

    /** Publish the runId the {@code RunRegistry} assigned on registration. */
    public void setRunId(String runId) {
        this.runId = runId;
    }

    @Override
    public java.util.Optional<String> runId() {
        return java.util.Optional.ofNullable(runId);
    }

    @Override
    public void stream(String message) {
        // Fire pre-stream hook (e.g., journal emit) before LLM call starts.
        // By the time stream() is called, all coordinator/agent work is done.
        var hook = preStreamHook.getAndSet(null);
        if (hook != null) {
            hook.afterPrompt(this);
        }

        // Load conversation history if memory is enabled
        var history = memory != null
                ? memory.getHistory(resource.uuid())
                : List.<org.atmosphere.ai.llm.ChatMessage>of();

        // Populate identity fields from the AtmosphereResource
        var userId = extractAttribute(resource, "ai.userId");
        var sessionId = resource.uuid();
        var agentId = extractAttribute(resource, "ai.agentId");
        var conversationId = extractAttribute(resource, "ai.conversationId");
        if (conversationId == null) {
            conversationId = sessionId;
        }

        var request = new AiRequest(message, systemPrompt, model,
                userId, sessionId, agentId, conversationId, Map.of(), history);

        // Attach available tools to the request
        if (toolRegistry != null && !toolRegistry.allTools().isEmpty()) {
            request = request.withTools(toolRegistry.allTools());
        }

        // Guardrails: inspect request (pre)
        for (var guardrail : guardrails) {
            try {
                var result = guardrail.inspectRequest(request);
                switch (result) {
                    case AiGuardrail.GuardrailResult.Block block -> {
                        logger.warn("Request blocked by guardrail {}: {}",
                                guardrail.getClass().getSimpleName(), block.reason());
                        delegate.error(new SecurityException("Request blocked: " + block.reason()));
                        return;
                    }
                    case AiGuardrail.GuardrailResult.Modify modify ->
                            request = modify.modifiedRequest();
                    case AiGuardrail.GuardrailResult.Pass ignored -> { }
                }
            } catch (Exception e) {
                logger.error("AiGuardrail.inspectRequest failed: {}",
                        guardrail.getClass().getName(), e);
                delegate.error(e);
                return;
            }
        }

        // Context providers: RAG augmentation with query transform + filter/rerank/post-process
        if (!contextProviders.isEmpty()) {
            var contextBuilder = new StringBuilder();
            for (var provider : contextProviders) {
                try {
                    var query = ContextProvider.normalizeQuery(provider.transformQuery(request.message()));
                    if (!ContextProvider.shouldRetrieve(query)) {
                        continue;
                    }
                    var docs = provider.retrieve(query, 5);
                    docs = provider.filter(query, docs);
                    docs = provider.rerank(query, docs);
                    docs = provider.postProcess(query, docs);
                    for (var doc : docs) {
                        contextBuilder.append("\n---\nSource: ").append(ContextProvider.formatCitation(doc))
                                .append("\n").append(doc.content());
                    }
                } catch (Exception e) {
                    logger.error("ContextProvider.retrieve failed: {}",
                            provider.getClass().getName(), e);
                    delegate.sendMetadata("rag.error",
                            provider.getClass().getSimpleName() + ": " + e.getMessage());
                }
            }
            if (!contextBuilder.isEmpty()) {
                var augmented = request.message()
                        + "\n\nRelevant context:" + contextBuilder;
                request = request.withMessage(augmented);
            }
        }

        // Pre-process: FIFO order
        for (var interceptor : interceptors) {
            try {
                request = interceptor.preProcess(request, resource);
            } catch (Exception e) {
                logger.error("AiInterceptor.preProcess failed: {}", interceptor.getClass().getName(), e);
                delegate.error(e);
                return;
            }
        }

        // Per-request ScopePolicy install — interceptors may have set a
        // ScopeConfig under ScopePolicy.REQUEST_SCOPE_METADATA_KEY. Consume
        // it now (mutating a copy of the metadata map) so the governance-
        // internal key does not leak to the provider, then run pre-admission
        // evaluate + system-prompt hardening, and append a PolicyAsGuardrail
        // to the capturing list so the streamed response is measured against
        // the narrowed scope when @AgentScope.postResponseCheck is enabled.
        List<AiGuardrail> effectiveGuardrails = guardrails;
        var mutableMeta = new java.util.HashMap<String, Object>(
                request.metadata() != null ? request.metadata() : Map.of());
        var requestScopePolicy =
                org.atmosphere.ai.governance.scope.ScopePolicyInstaller.extract(mutableMeta);
        if (requestScopePolicy != null) {
            request = new AiRequest(request.message(), request.systemPrompt(),
                    request.model(), request.userId(), request.sessionId(),
                    request.agentId(), request.conversationId(),
                    Map.copyOf(mutableMeta), request.history());
            var ctx = org.atmosphere.ai.governance.PolicyContext.preAdmission(request);
            var tracer = org.atmosphere.ai.governance.GovernanceTracer.start(requestScopePolicy, ctx);
            var startNs = System.nanoTime();
            try {
                var decision = requestScopePolicy.evaluate(ctx);
                var evalMs = (System.nanoTime() - startNs) / 1_000_000.0;
                switch (decision) {
                    case org.atmosphere.ai.governance.PolicyDecision.Deny deny -> {
                        logger.warn("Request denied by per-request scope {}: {}",
                                requestScopePolicy.name(), deny.reason());
                        org.atmosphere.ai.governance.GovernanceDecisionLog.installed().record(
                                org.atmosphere.ai.governance.GovernanceDecisionLog.entry(
                                        requestScopePolicy, ctx, "deny", deny.reason(), evalMs));
                        tracer.end("deny", deny.reason());
                        delegate.error(new SecurityException("Request denied by per-request scope "
                                + requestScopePolicy.name() + ": " + deny.reason()));
                        return;
                    }
                    case org.atmosphere.ai.governance.PolicyDecision.Transform transform -> {
                        org.atmosphere.ai.governance.GovernanceDecisionLog.installed().record(
                                org.atmosphere.ai.governance.GovernanceDecisionLog.entry(
                                        requestScopePolicy, ctx, "transform",
                                        "request rewritten", evalMs));
                        tracer.end("transform", "request rewritten");
                        request = transform.modifiedRequest();
                    }
                    case org.atmosphere.ai.governance.PolicyDecision.Admit ignored -> {
                        org.atmosphere.ai.governance.GovernanceDecisionLog.installed().record(
                                org.atmosphere.ai.governance.GovernanceDecisionLog.entry(
                                        requestScopePolicy, ctx, "admit", "", evalMs));
                        tracer.end("admit", "");
                    }
                }
            } catch (RuntimeException e) {
                var evalMs = (System.nanoTime() - startNs) / 1_000_000.0;
                logger.error("Per-request scope {} evaluate failed — fail-closed",
                        requestScopePolicy.name(), e);
                org.atmosphere.ai.governance.GovernanceDecisionLog.installed().record(
                        org.atmosphere.ai.governance.GovernanceDecisionLog.entry(
                                requestScopePolicy, ctx, "error",
                                "evaluate threw: " + e.getMessage(), evalMs));
                tracer.end("error", e.getMessage(), e);
                delegate.error(new SecurityException("Per-request scope " + requestScopePolicy.name()
                        + " evaluation failed: " + e.getMessage(), e));
                return;
            }
            // Prepend the confinement preamble so the LLM call is fronted
            // by the narrowed scope even when the interceptor overwrote
            // the endpoint's @AiEndpoint.systemPromptResource.
            var hardened = org.atmosphere.ai.governance.scope.ScopePolicyInstaller
                    .hardenSystemPrompt(request.systemPrompt(), List.of(requestScopePolicy));
            request = request.withSystemPrompt(hardened);
            var augmented = new java.util.ArrayList<AiGuardrail>(guardrails.size() + 1);
            augmented.addAll(guardrails);
            augmented.add(new org.atmosphere.ai.governance.PolicyAsGuardrail(requestScopePolicy));
            effectiveGuardrails = List.copyOf(augmented);
        }

        // The capturing-session decorator chain starts at `this` rather than
        // `delegate` so every delegating method — critically, `injectables()`
        // — reaches AiStreamingSession's live map. The delegate chain
        // (DefaultStreamingSession etc.) returns the empty default, which
        // would leave {@code @AiTool} methods without their typed parameter
        // injections (AgentFleet, AgentIdentity, ...) and force the
        // ThreadLocal shim back. Existing capturing sessions forward all
        // StreamingSession methods to their delegate, so starting the chain
        // at `this` doesn't change observable behaviour on any other method.
        StreamingSession target = this;
        if (memory != null) {
            target = new MemoryCapturingSession(target, memory, resource.uuid(), message);
        }

        // LineageCapturingSession ties the prompt → tool calls → RAG → cost
        // chain into a single audit row per @Prompt invocation. NOOP recorder
        // is the framework default so unmounted code paths cost nothing; admin
        // / Spring / Quarkus auto-config installs a real recorder via
        // LineageRecorderHolder.install. Wrapped after MemoryCapturingSession
        // so the lineage entry's terminal classification reflects what the
        // memory capture sees (memory persists only on success).
        var lineageRecorder = org.atmosphere.ai.lineage.LineageRecorderHolder.get();
        if (lineageRecorder != org.atmosphere.ai.lineage.LineageRecorder.NOOP) {
            target = new org.atmosphere.ai.lineage.LineageCapturingSession(
                    target, lineageRecorder, userId, agentId, conversationId, message);
        }

        // Wrap in MetricsCapturingSession for latency/streaming text tracking
        if (metrics != AiMetrics.NOOP) {
            target = new MetricsCapturingSession(target, metrics, model);
        }

        // Wrap in CostAccountingSession so TokenUsage events feed whatever
        // CostAccountant is installed in the holder (default NOOP). This is
        // the observability → enforcement wire: runtimes report usage, the
        // accountant pushes cost into CostCeilingGuardrail.addCost, the next
        // request-side guardrail inspection blocks the tenant once the
        // ceiling trips. Skip wrapping when no accountant is installed so
        // the decorator chain stays flat in the common case.
        var accountant = org.atmosphere.ai.cost.CostAccountantHolder.get();
        if (accountant != org.atmosphere.ai.cost.CostAccountant.NOOP) {
            target = new org.atmosphere.ai.cost.CostAccountingSession(target, accountant);
        }

        // Wrap in GuardrailCapturingSession for post-LLM response inspection.
        // `effectiveGuardrails` adds the per-request ScopePolicy (wrapped as a
        // PolicyAsGuardrail) when one was installed via request metadata so
        // post-response drift detection honours the narrowed scope.
        if (!effectiveGuardrails.isEmpty()) {
            target = new GuardrailCapturingSession(target, effectiveGuardrails);
        }

        // Snapshot the system prompt as the developer (and ScopePolicy) shaped
        // it, before structured-output augmentation appends schema text. The
        // per-stage telemetry below attributes the input bill to the stage
        // that introduced it instead of lumping schema into "system".
        var baseSystemPrompt = request.systemPrompt();
        String structuredSchemaText = null;

        // Wrap in StructuredOutputCapturingSession for typed response parsing
        var effectiveResponseType = responseType != null ? responseType : request.responseType();
        if (effectiveResponseType != null && effectiveResponseType != Void.class) {
            var parser = StructuredOutputParser.resolve();
            target = new StructuredOutputCapturingSession(target, parser, effectiveResponseType);
            structuredSchemaText = parser.schemaInstructions(effectiveResponseType);
            request = request.withSystemPrompt(
                    request.systemPrompt() + "\n\n" + structuredSchemaText);
        }

        // Expose the session to interceptors via request attribute
        resource.getRequest().setAttribute(STREAMING_SESSION_ATTR, target);

        // Build execution context and delegate to runtime.
        // The session-scoped ApprovalStrategy is carried on the context so every
        // runtime bridge routes tool execution through ToolExecutionHelper.executeWithApproval().
        // Tools come from the (potentially guardrail-modified) request, NOT from
        // re-reading the registry, so a guardrail that narrows the tool set via
        // {@code AiRequest#withTools} actually takes effect — without this,
        // guardrails could only block the whole request, not drop specific tools.
        var finalRequest = request;
        var tools = finalRequest.tools() != null
                ? finalRequest.tools()
                : java.util.List.<org.atmosphere.ai.tool.ToolDefinition>of();
        var strategy = ApprovalStrategy.virtualThread(approvalRegistry);
        // Merge endpoint-scoped CacheHint into request metadata if configured.
        java.util.Map<String, Object> effectiveMetadata = request.metadata();
        var policy = this.cachePolicy;
        if (policy != null && policy != org.atmosphere.ai.llm.CacheHint.CachePolicy.NONE) {
            var merged = new java.util.HashMap<String, Object>(
                    effectiveMetadata != null ? effectiveMetadata : java.util.Map.of());
            merged.putIfAbsent(org.atmosphere.ai.llm.CacheHint.METADATA_KEY,
                    new org.atmosphere.ai.llm.CacheHint(policy,
                            java.util.Optional.empty(), java.util.Optional.empty()));
            effectiveMetadata = merged;
        }
        var context = new AgentExecutionContext(
                request.message(), request.systemPrompt(), request.model(),
                request.agentId(), request.sessionId(), request.userId(),
                request.conversationId(),
                java.util.List.copyOf(tools), null, memory,
                contextProviders, effectiveMetadata, request.history(),
                effectiveResponseType, strategy);
        // Apply endpoint-scoped retry policy if configured.
        var endpointRetry = this.endpointRetryPolicy;
        if (endpointRetry != null) {
            context = context.withRetryPolicy(endpointRetry);
        }
        // Per-stage input breakdown — emitted on the websocket @AiEndpoint
        // path to match AiPipeline's channel-bridge instrumentation. Mode
        // parity (Invariant #7): runtime turn-cost telemetry must look the
        // same regardless of which entry the request came in on. Confidence
        // elicitation only ships through the pipeline path today, so it is
        // intentionally absent here.
        InputAssemblyTelemetry.emit(metrics, request.model(),
                baseSystemPrompt, structuredSchemaText, null,
                tools, request.history(), request.message());

        var streamingTarget = target;
        var turnEvent = new org.atmosphere.ai.jfr.AgentTurnEvent();
        turnEvent.runtime = runtime != null ? runtime.name() : "unknown";
        turnEvent.model = request.model() != null ? request.model() : (model != null ? model : "unknown");
        turnEvent.clientId = resource != null ? resource.uuid() : null;
        turnEvent.cacheHit = false;
        turnEvent.begin();
        try {
            // executeWithHandle returns a handle the disconnect path can
            // use to abort the upstream LLM call. For runtimes that haven't
            // overridden it (default impl) execute() runs synchronously and
            // a COMPLETED handle is returned — whenDone().join() is a no-op
            // and the cancel knob is harmless. Runtimes that override it
            // (BuiltInAgentRuntime, etc.) return a real handle backed by
            // their native cancel primitive, so cancelInflight() actually
            // closes the HTTP stream and unwinds the tool loop.
            var handle = runtime.executeWithHandle(context, streamingTarget);
            this.currentHandle = handle;
            // Close the publish-then-cancel race: a disconnect that fired
            // between executeWithHandle returning and currentHandle being
            // assigned would have read null and dropped the cancel. Re-check
            // the latch here so we catch up on that intent before parking.
            if (cancelPending.get() && !handle.isDone()) {
                handle.cancel();
            }
            try {
                handle.whenDone().join();
                turnEvent.status = streamingTarget.hasErrored() ? "error" : "success";
            } catch (java.util.concurrent.CompletionException ce) {
                // Runtime completed exceptionally via Settable.completeExceptionally.
                // The streamingTarget.error(...) path was already invoked by the
                // runtime, so the client has been notified — no further action.
                turnEvent.status = "error";
                var cause = ce.getCause();
                turnEvent.errorType = cause != null ? cause.getClass().getSimpleName() : ce.getClass().getSimpleName();
                if (logger.isTraceEnabled()) {
                    logger.trace("Execution completed exceptionally for {}: {}",
                            resource != null ? resource.uuid() : "null",
                            cause != null ? cause.getMessage() : ce.getMessage());
                }
            }
        } catch (Exception e) {
            turnEvent.status = "error";
            turnEvent.errorType = e.getClass().getSimpleName();
            metrics.recordError(model != null ? model : "unknown", "stream_error");
            logger.error("Streaming error", e);
            streamingTarget.error(e);
        } finally {
            turnEvent.commit();
        }
        try {
            // Post-process follows execution
        } finally {
            // Post-process: LIFO order (matching AtmosphereInterceptor convention)
            for (int i = interceptors.size() - 1; i >= 0; i--) {
                try {
                    interceptors.get(i).postProcess(finalRequest, resource);
                } catch (Exception e) {
                    logger.error("AiInterceptor.postProcess failed: {}",
                            interceptors.get(i).getClass().getName(), e);
                }
            }
        }
    }

    @Override
    public void handoff(String agentName, String message) {
        if (!handoffInProgress.compareAndSet(false, true)) {
            throw new IllegalStateException("Nested handoffs not supported");
        }

        logger.info("Handoff from {} to agent '{}'", resource.uuid(), agentName);
        emit(new AiEvent.Handoff(
                extractAttribute(resource, "ai.agentId"),
                agentName,
                "User request routed to " + agentName));

        // Copy conversation history to the target agent's memory
        if (memory != null) {
            var fromId = resource.uuid();
            var toId = agentName + ":" + fromId;
            memory.copyTo(fromId, toId);
        }

        // Dispatch to the target agent via its handler. We look up the
        // AtmosphereHandlerWrapper and invoke the handler's onStateChange
        // with the message, so the target's @Prompt method runs — including
        // demo mode, tool registration, and all pipeline hooks.
        var framework = resource.getAtmosphereConfig().framework();
        var targetPath = "/atmosphere/agent/" + agentName;
        var handlerWrapper = framework.getAtmosphereHandlers().get(targetPath);
        if (handlerWrapper == null) {
            delegate.error(new IllegalArgumentException(
                    "Agent '" + agentName + "' not found at " + targetPath));
            handoffInProgress.set(false);
            return;
        }

        try {
            // Use the handler's public API: get the handler via the wrapper's
            // accessor and call onStateChange with a crafted event.
            var handler = handlerWrapper.atmosphereHandler();
            var event = new org.atmosphere.cpr.AtmosphereResourceEventImpl(
                    (org.atmosphere.cpr.AtmosphereResourceImpl) resource);
            event.setMessage(message);
            handler.onStateChange(event);
        } catch (Exception e) {
            logger.error("Handoff to '{}' failed", agentName, e);
            delegate.error(e);
        } finally {
            handoffInProgress.set(false);
        }
    }

    /**
     * Check if an incoming message is an approval response and route it
     * to the approval registry. Called by the handler before dispatching to @Prompt.
     *
     * <p>Only returns {@code true} when the registry matched and resolved a
     * pending approval in <em>this</em> session. Stale / unknown IDs and
     * non-approval messages fall through so the handler dispatches them to
     * @Prompt as normal user input.
     *
     * @param message the incoming message
     * @return true if the message resolved a pending approval in this session
     */
    public boolean tryResolveApproval(String message) {
        return approvalRegistry.resolve(message) == ApprovalRegistry.ResolveResult.RESOLVED;
    }

    /** The session-scoped approval registry, exposed for the endpoint handler to route responses. */
    public ApprovalRegistry approvalRegistry() {
        return approvalRegistry;
    }

    // -- Delegate all StreamingSession methods --

    @Override
    public String sessionId() {
        return delegate.sessionId();
    }

    @Override
    public void send(String text) {
        delegate.send(text);
    }

    @Override
    public void sendContent(Content content) {
        // Forward multi-modal content to the wrapped leaf session. Without
        // this override AiStreamingSession inherits the interface default,
        // which drops binary parts to a metadata breadcrumb — that is the
        // wrong shape for an AiEndpoint-attached session where the delegate
        // is typically a DefaultStreamingSession that can emit the binary
        // frame on the wire.
        delegate.sendContent(content);
    }

    @Override
    public void sendMetadata(String key, Object value) {
        delegate.sendMetadata(key, value);
    }

    @Override
    public void progress(String message) {
        delegate.progress(message);
    }

    @Override
    public void complete() {
        removeActiveSession(this);
        delegate.complete();
    }

    @Override
    public void complete(String summary) {
        removeActiveSession(this);
        delegate.complete(summary);
    }

    @Override
    public void error(Throwable t) {
        removeActiveSession(this);
        delegate.error(t);
    }

    @Override
    public void emit(AiEvent event) {
        delegate.emit(event);
    }

    @Override
    public boolean isClosed() {
        return delegate.isClosed();
    }

    @Override
    public boolean hasErrored() {
        return delegate.hasErrored();
    }

    // visible for testing
    AgentRuntime runtime() {
        return runtime;
    }

    List<AiInterceptor> interceptors() {
        return interceptors;
    }

    String systemPrompt() {
        return systemPrompt;
    }

    private static String extractAttribute(org.atmosphere.cpr.AtmosphereResource resource, String key) {
        var attr = resource.getRequest().getAttribute(key);
        return attr != null ? attr.toString() : null;
    }
}
