interactive-story/backend/app/contextwindow.py

"""M11: what the inference server will *actually* accept, as opposed to what we budgeted.

M8 found the failure this module exists to prevent. The application budgets a
prompt up to `Settings.context_token_budget` — 16,384 by default — while Ollama
enforces a window of its own, and on a machine with no VRAM that window defaults
to **4,096**. The request still returns HTTP 200. Nothing warns anybody. What
actually happens is worse than an error: `llama.cpp` drops the **oldest** tokens,
and the oldest tokens in this application are the system block — the narrator
rules and the campaign canon. The symptom is a narrator that forgets canon deep
into a long session, with nothing on screen explaining why, and every acceptance
test that reads a returned 200 as success passing throughout.

The invariant M11 requires:

    The application must not silently budget more narrator input
    than the configured Ollama runtime will actually accept.

Note the word *silently*. There are two honest outcomes and this module produces
both: either the window is **verified**, in which case the budget is capped to it
so the prompt physically cannot overflow; or it is **unverified**, in which case
the assembly says so, in the context report, on the connection test, and in the
turn's stored provenance. What must not happen is the third thing — assembling
16,384 tokens against a 4,096-token server and calling the result a turn.

## Why this is not solved by sending `num_ctx`

It was tried, and it is documented in `DEVELOPMENT.md`. Ollama's
OpenAI-compatible endpoint accepts `num_ctx` — nested in `options` or at the top
level — returns 200, and ignores it. Worse, it reloads the model at its own
default, so priming the server through the native API first does not help
either: the next request resets the window. The window is a property of how the
model is loaded, not of the request, so the only things that change it are a
model with `num_ctx` baked in (`/api/create`) or `OLLAMA_CONTEXT_LENGTH` on the
server. Both are operator actions. This module's job is not to change the
window; it is to find out what it is and refuse to lie about it.

## How the window is found

Ollama's native API sits beside the OpenAI-compatible one on the same host, so
this asks the server the application is already talking to, and nothing else. No
new destination, the same endpoint policy, the same TLS trust store.

    /api/ps    a loaded model reports `context_length`: the window the runtime
               is enforcing *right now*. This is the truth when it is available.
    /api/show  an unloaded model may carry `num_ctx` in its baked parameters,
               which is the window it will load with; `model_info` carries the
               architecture's own ceiling, which caps everything else.

`/api/ps` is asked first because a model that is loaded has already settled the
question. `/api/show` answers it for a model that is not loaded yet, which is the
ordinary case at the start of a session.

## What it deliberately does not do

It does not hard-code 4,096, which would cripple a correctly configured
deployment; it does not raise the budget, which is the operator's decision; it
does not fall back to a cloud probe, a bundled table of model sizes, or a guess
from the model's name. An unknown window is reported as unknown.

## The server that cannot be asked

Discovery above is Ollama's native API. Nothing restricts `endpoint_url` to
Ollama — any allowed address serving an OpenAI-compatible `/v1` is accepted —
and on vLLM, llama.cpp's own server, or anything else, `/api/ps` and `/api/show`
are simply not there. Discovery then fails exactly as designed and the window is
reported unknown, which is honest but leaves the invariant at the top of this
file unenforced: the budget stands at whatever is configured, and if that server
enforces a smaller window it drops the oldest tokens again.

`context_window_override` is the operator's answer to that. It is a number the
operator states because they know how the server was launched, and it is used
**only when the server could not be asked**:

    verified window     -> always wins; a declaration cannot raise it
    no verified window  -> the declaration becomes the ceiling, source DECLARED
    neither             -> unknown, exactly as before

This does not weaken what `verified` claims. `verified` still means the server
itself answered, so `window_verified` in a turn's provenance keeps the meaning
the M11 report gives it, and a declared window is identifiable as a declaration
wherever it appears. What the declaration buys is enforcement: the prompt is
capped, so the failure mode is a shorter prompt rather than a silently truncated
one.
"""

from __future__ import annotations

import logging
import math
import re
import time
from dataclasses import dataclass

import httpx

from . import endpoints, tlstrust

log = logging.getLogger(__name__)

#: Short, because this sits in the turn path. A server that does not answer in
#: two seconds has told us what we need to know: we cannot verify the window
#: right now, and the turn should proceed unverified rather than stall.
PROBE_TIMEOUT = 2.0
CONNECT_TIMEOUT = 1.5

#: A verified window is stable — it changes when an operator reloads a model —
#: so it is worth keeping. A failure is cached too, and for much less time,
#: because the commonest cause is a server that is starting up.
POSITIVE_TTL = 600.0
NEGATIVE_TTL = 60.0

#: Sources, in the order of how much they prove.
LOADED = "loaded"        # /api/ps: what the runtime is enforcing now
PARAMETERS = "parameters"  # /api/show: what the model will load with
DECLARED = "declared"    # the operator said so; the server could not be asked
UNKNOWN = "unknown"

#: Sources that mean *the server answered*, as opposed to somebody asserting.
FROM_SERVER = (LOADED, PARAMETERS)


@dataclass(frozen=True)
class Window:
    """What was learned about the server's input window, and how."""

    #: The total context in tokens — input *and* output share it — or None when
    #: it could not be determined.
    tokens: int | None
    #: One of LOADED, PARAMETERS, UNKNOWN.
    source: str
    #: The architecture's own ceiling, when the server reported one. Useful to a
    #: reader deciding whether raising the window is even possible.
    model_max: int | None = None
    #: Why the window is unknown, or how it was found. Shown to the user.
    detail: str = ""
    #: v1.1: the server answered a discovery request at all, whatever it said.
    #: A server that answered but could not report a window may simply not have
    #: the model loaded yet, which `ensure_window` can fix; one that did not
    #: answer cannot be helped by asking it to load anything.
    reachable: bool = False

    @property
    def verified(self) -> bool:
        """The **server** answered. An operator's declaration is not this.

        Kept narrow on purpose. `window_verified` travels in every turn's stored
        provenance and the M11 report counts on it meaning one thing: that the
        runtime was asked and replied. A declaration is a person's claim about a
        server, which is worth acting on and is not the same evidence.
        """
        return self.tokens is not None and self.source in FROM_SERVER

    @property
    def enforceable(self) -> bool:
        """There is a number to cap the prompt to, whoever supplied it."""
        return self.tokens is not None


UNVERIFIED = Window(tokens=None, source=UNKNOWN, detail="not checked")

_cache: dict[tuple[str, str], tuple[float, Window]] = {}


def native_base(endpoint_url: str) -> str:
    """The Ollama-native base beside an OpenAI-compatible endpoint.

    `https://host:1234/v1` -> `https://host:1234`. Anything else is used as
    given, because an endpoint that is not shaped like Ollama's is one this
    cannot interrogate and should not guess about.
    """
    trimmed = (endpoint_url or "").rstrip("/")
    return re.sub(r"/v1$", "", trimmed)


def effective_budget(configured: int, window: Window | int | None) -> int:
    """The budget the prompt may actually use.

    The whole enforcement, in one line: a known window is a ceiling — whether
    the server reported it or the operator declared it. The configured budget
    still wins when it is *smaller*, because a reader who has deliberately asked
    for a shorter prompt should get one.
    """
    tokens = window.tokens if isinstance(window, Window) else window
    if tokens is None or tokens <= 0:
        return configured
    return min(configured, tokens)


#: v1.1 WP-A1: the tokens kept free below the effective window, beyond the reply.
#:
#: The builder counts with `cl100k_base`; the narrator counts with its own
#: tokenizer. The v1 evidence put the largest prompts 23-42 real tokens from the
#: edge of a 16,384 window, and Ollama does not refuse a prompt past the edge —
#: measured on Ollama 0.33 at a 4,096 window, a 6,316-token prompt came back 200
#: with `prompt_tokens` 2,050. So the reserve is deliberate and sized to the
#: window: the larger of a floor and a share, **rounded up to a whole token**.
#:
#:     4,096 -> 256     8,192 -> 410     16,384 -> 820
#:
#: A fixed, documented tolerance, owner-chosen for v1.1. It is not a setting and
#: it is not calibrated per model.
SAFETY_RESERVE_FLOOR = 256
SAFETY_RESERVE_PERCENT = 5


def safety_reserve(effective_window: int) -> int:
    """`max(256, ceil(5% of the effective window))`, in tokens.

    The effective window is the budget the prompt is actually built to — the
    verified or declared window when there is one, the configured budget
    otherwise — so a 16,384 setting against a 4,096 server reserves 256, not 820.
    Integer arithmetic, so the rounding is exact rather than a float's.
    """
    share = math.ceil(max(0, effective_window) * SAFETY_RESERVE_PERCENT / 100)
    return max(SAFETY_RESERVE_FLOOR, share)


#: v1.1 WP-A1: what the server's own count says about a turn that was sent.
FITS = "fits"
EXCEEDED = "exceeded"
TRUNCATION_SUSPECTED = "truncation_suspected"
#: `UNKNOWN` above: the server reported no usable count.


def classify_usage(usage: dict | None, *, estimate: int, budget: int,
                   max_output_tokens: int, window_verified: bool) -> dict:
    """Sets the server's reported prompt count against what the application sent.

    The order of the checks is the order of what they prove:

    ``unknown``
        No positive integer `prompt_tokens`. Nothing can be said, and nothing
        is claimed: an absent count is never read as a prompt that fitted.
    ``truncation_suspected``
        The server read fewer tokens than were sent by more than the safety
        reserve. A tokenizer thriftier than `cl100k_base` may honestly count a
        little less; a shortfall larger than the tolerance the application keeps
        for drift is the signature of a server that cut the prompt — the real
        shape was 6,316 sent and 2,050 read.
    ``exceeded``
        The server's count plus the reply allocation is more than the window
        the prompt was built for. The drift was larger than the whole reserve,
        so the reply may have been cut short.
    ``fits``
        Otherwise.

    `observed_margin` is what was left beside the reply by the server's count:
    `budget - max_output_tokens - server_prompt_tokens`. The safety reserve is
    the tolerance, so a margin between 0 and the reserve is still `fits`.

    A discrepancy is recorded, never acted on: the reply has already streamed
    to the reader and is accepted story.
    """
    prompt = usage.get("prompt_tokens") if isinstance(usage, dict) else None
    reserve = safety_reserve(budget)
    verified_note = "" if window_verified else (
        " The window itself was not verified for this turn.")
    record = {
        "status": UNKNOWN,
        "server_prompt_tokens": None,
        "estimate": estimate,
        "difference": None,
        "budget": budget,
        "max_output_tokens": max_output_tokens,
        "safety_reserve": reserve,
        "observed_margin": None,
        "window_verified": bool(window_verified),
        "detail": "",
    }
    if type(prompt) is not int or prompt <= 0:
        record["detail"] = ("The server reported no prompt token count, so nothing "
                            "confirms the whole prompt was read." + verified_note)
        return record

    record["server_prompt_tokens"] = prompt
    record["difference"] = prompt - estimate
    record["observed_margin"] = budget - max_output_tokens - prompt
    if prompt + reserve < estimate:
        record["status"] = TRUNCATION_SUSPECTED
        record["detail"] = (
            f"The server read {prompt:,} prompt tokens of the {estimate:,} sent, a "
            f"shortfall larger than the {reserve:,}-token safety reserve. A server "
            "that cuts an over-window prompt reports exactly this, and what it cuts "
            "is the start: the narrator's rules and the canon." + verified_note)
    elif prompt + max_output_tokens > budget:
        record["status"] = EXCEEDED
        record["detail"] = (
            f"The server counted {prompt:,} prompt tokens; with {max_output_tokens:,} "
            f"for the reply that is more than the {budget:,}-token window the prompt "
            "was built for, so the reply may have been cut short." + verified_note)
    else:
        record["status"] = FITS
        record["detail"] = (
            f"The server read {prompt:,} prompt tokens, leaving "
            f"{record['observed_margin']:,} beside the reply." + verified_note)
    return record


def cache_clear() -> None:
    """Forgets what was learned. Called when the endpoint or model changes."""
    _cache.clear()


async def probe(endpoint_url: str, model: str, *,
                declared: int | None = None, use_cache: bool = True) -> Window:
    """What window `model` gets, asked of the server and only then declared.

    Returns `UNVERIFIED` for every discovery failure — refused endpoint,
    unreachable server, TLS failure, a server with no Ollama-native API, an
    unparseable answer — unless `declared` supplies a number to fall back on.
    The caller cannot act differently on those failures and the reader is told
    the same thing either way: the window could not be checked.

    `declared` is `Settings.context_window_override`. It never overrides a
    verified answer, so an operator cannot talk the application into a bigger
    prompt than the runtime will read; it only fills a gap discovery left.
    """
    if not endpoint_url or not model:
        return _declared_or(declared,
                            Window(None, UNKNOWN,
                                   detail="no endpoint or model configured"))

    discovered = await _discover(endpoint_url, model, use_cache=use_cache)
    return _declared_or(declared, discovered)


def _declared_or(declared: int | None, discovered: Window) -> Window:
    """The operator's number, but only where the server left a hole.

    A verified window always wins. That ordering is the whole safety property:
    a declaration can lower an unknown ceiling into existence, never raise a
    known one.
    """
    if discovered.verified:
        return discovered
    if not declared or declared <= 0:
        return discovered
    return Window(
        declared, DECLARED, discovered.model_max,
        f"{declared:,} tokens, declared in settings — the server was not able "
        f"to say ({discovered.detail})",
        reachable=discovered.reachable,
    )


async def _discover(endpoint_url: str, model: str, *,
                    use_cache: bool = True) -> Window:
    """The server's own answer, cached. Knows nothing about declarations.

    The cache holds only what was discovered, so changing the declared override
    takes effect on the next turn without having to clear anything: the
    declaration is layered on afterwards, in `_declared_or`.
    """
    key = (endpoint_url, model)
    now = time.monotonic()
    if use_cache:
        hit = _cache.get(key)
        if hit is not None and hit[0] > now:
            return hit[1]

    window = await _ask(endpoint_url, model)
    ttl = POSITIVE_TTL if window.verified else NEGATIVE_TTL
    _cache[key] = (now + ttl, window)
    return window


async def _ask(endpoint_url: str, model: str) -> Window:
    # The same policy the turn itself is held to. A window probe must not be a
    # way to reach an address inference may not (ADR 011, H12).
    reason = endpoints.rejection_reason(endpoint_url)
    if reason is not None:
        return Window(None, UNKNOWN, detail=f"endpoint not allowed — {reason}")

    base = native_base(endpoint_url)
    try:
        async with httpx.AsyncClient(
            timeout=httpx.Timeout(PROBE_TIMEOUT, connect=CONNECT_TIMEOUT),
            verify=tlstrust.ssl_context(),
        ) as client:
            loaded = await _loaded_window(client, base, model)
            if loaded is not None:
                tokens, ceiling = loaded
                return Window(
                    tokens, LOADED, ceiling,
                    f"{tokens:,} tokens, reported by the running model",
                    reachable=True,
                )
            return await _declared_window(client, base, model)
    except (httpx.HTTPError, ValueError, TypeError, KeyError) as exc:
        log.debug("context window probe failed for %s: %s", base, exc)
        return Window(None, UNKNOWN, detail=f"could not ask the server ({type(exc).__name__})")


async def _loaded_window(client, base: str, model: str):
    """`/api/ps`: the window a resident model is actually being served with."""
    resp = await client.get(f"{base}/api/ps")
    if resp.status_code != 200:
        return None
    for entry in (resp.json() or {}).get("models") or []:
        if entry.get("name") == model or entry.get("model") == model:
            tokens = entry.get("context_length")
            if isinstance(tokens, int) and tokens > 0:
                return tokens, None
    return None


async def _declared_window(client, base: str, model: str) -> Window:
    """`/api/show`: what the model will load with, and its architectural cap."""
    resp = await client.post(f"{base}/api/show", json={"model": model})
    if resp.status_code != 200:
        return Window(
            None, UNKNOWN,
            detail=f"the server did not describe the model (HTTP {resp.status_code})",
            reachable=True,
        )
    body = resp.json() or {}
    ceiling = _architecture_ceiling(body.get("model_info") or {})
    declared = _num_ctx(body.get("parameters"))
    if declared is None:
        return Window(
            None, UNKNOWN, ceiling,
            detail=(
                "the model sets no num_ctx, so the server will load it at its own "
                "default — which is 4,096 where there is no VRAM"
            ),
            reachable=True,
        )
    tokens = min(declared, ceiling) if ceiling else declared
    return Window(
        tokens, PARAMETERS, ceiling,
        f"{tokens:,} tokens, from the model's own num_ctx",
        reachable=True,
    )


#: v1.1 WP-A1 corrective: loading the configured model so its window can be read.
#:
#: The first real turn of the A1 evidence found a cold model: `/api/ps` knew
#: nothing, `/api/show` found no `num_ctx`, so the window was unverified and the
#: prompt was built to the configured 16,384. Ollama loaded the model at its own
#: 4,096 default, kept 2,050 of 13,875 tokens and answered 200. That case is
#: preventable, because the window becomes readable the moment the model is
#: resident. Ollama's native `POST /api/generate` with a model and **no prompt**
#: loads the model and generates nothing — measured on Ollama 0.33: HTTP 200,
#: `"response": ""`, `"done_reason": "load"`, and `/api/ps` then reported the
#: window. The OpenAI-compatible request that followed did not reload it.
WARM_PATH = "/api/generate"


async def warm(endpoint_url: str, model: str, *, timeout: float) -> tuple[bool, str]:
    """Asks the configured server to load `model`. One request, no story text.

    Held to the same endpoint policy and TLS trust as inference and the probe, and
    sent to the same host the probe asks. The body names the model and nothing
    else: no prompt, so nothing is generated, and no `options` or `keep_alive`, so
    the model loads the way the server would load it for the turn itself.

    Returns `(loaded, detail)`. Every failure is `(False, why)` and never raises:
    a server that will not load the model on request will fail the turn's own
    call the ordinary way, which is where that failure belongs.
    """
    reason = endpoints.rejection_reason(endpoint_url)
    if reason is not None:
        return False, f"endpoint not allowed — {reason}"
    base = native_base(endpoint_url)
    try:
        async with httpx.AsyncClient(
            timeout=httpx.Timeout(timeout, connect=CONNECT_TIMEOUT),
            verify=tlstrust.ssl_context(),
        ) as client:
            resp = await client.post(f"{base}{WARM_PATH}", json={"model": model})
    except httpx.HTTPError as exc:
        log.debug("model warm-up failed for %s: %s", base, exc)
        return False, f"could not ask the server to load the model ({type(exc).__name__})"
    if resp.status_code != 200:
        return False, f"the server did not load the model (HTTP {resp.status_code})"
    try:
        body = resp.json() or {}
    except ValueError:
        return False, "the server answered the load request with something that was not JSON"
    return True, f"the server loaded the model ({body.get('done_reason') or 'done'})"


async def ensure_window(endpoint_url: str, model: str, *, declared: int | None = None,
                        warm_timeout: float = 300.0) -> tuple[Window, dict]:
    """The window for a turn about to be generated, loading the model once if that is what it takes.

    1. Probe as before.
    2. If the window is not verified, the server answered, and there is a model to
       load: one bounded `warm` request.
    3. If the model loaded, probe again, bypassing the cache that still holds the
       unverified answer.

    Whatever the second probe says is the answer. There is no retry loop, no
    guessed window, and no hard-coded 4,096: a window still unverified leaves the
    configured budget standing, exactly as before, and the turn's accounting
    still catches a server that cut the prompt.

    Returns the window and a `preflight` record for the turn's provenance.
    Not used by the context dry run: loading a model is a side effect, and
    opening a panel should not cause one.
    """
    window = await probe(endpoint_url, model, declared=declared)
    preflight = {"attempted": False, "loaded": None, "verified_before": window.verified,
                 "verified_after": window.verified, "detail": ""}
    if window.verified:
        preflight["detail"] = "the window was already verified"
        return window, preflight
    if not (endpoint_url and model):
        preflight["detail"] = "no endpoint or model configured"
        return window, preflight
    if not window.reachable:
        preflight["detail"] = "the server did not answer, so no model was loaded"
        return window, preflight
    loaded, detail = await warm(endpoint_url, model, timeout=warm_timeout)
    preflight.update(attempted=True, loaded=loaded, detail=detail)
    if loaded:
        window = await probe(endpoint_url, model, declared=declared, use_cache=False)
    preflight["verified_after"] = window.verified
    return window, preflight


def _num_ctx(parameters) -> int | None:
    """Reads `num_ctx` out of the plain-text parameter block Ollama returns."""
    if not isinstance(parameters, str):
        return None
    match = re.search(r"^\s*num_ctx\s+(\d+)\s*$", parameters, re.MULTILINE)
    return int(match.group(1)) if match else None


def _architecture_ceiling(model_info: dict) -> int | None:
    """`<arch>.context_length` — the largest window this model can have."""
    for key, value in model_info.items():
        if key.endswith(".context_length") and isinstance(value, int) and value > 0:
            return value
    return None