teng-lin--notebooklm-py/src/notebooklm/_kernel.py

"""Concrete transport kernel for NotebookLM session operations."""

from __future__ import annotations

from collections.abc import Callable, Mapping
from typing import Any

import httpx

from ._request_types import PostBody
from ._streaming_post import stream_post_with_size_cap
from .auth import AuthTokens, build_cookie_jar
from .types import ConnectionLimits


class Kernel:
    """Own the live HTTP transport and cookie jar.

    Client lifecycle code decides when to open and close. The kernel owns the
    concrete ``httpx.AsyncClient`` instance, its cookie jar, raw POST execution,
    and shielded teardown target.
    """

    def __init__(
        self,
        *,
        async_client_factory: Callable[..., httpx.AsyncClient] = httpx.AsyncClient,
    ) -> None:
        self._async_client_factory = async_client_factory
        self._http_client: httpx.AsyncClient | None = None
        self._timeout: float | None = None
        self._connect_timeout: float | None = None

    @property
    def http_client(self) -> httpx.AsyncClient | None:
        """Return the live HTTP client, or ``None`` when closed.

        The property is read-only by design. Production code mutates the
        underlying client only through :meth:`open` (which constructs the
        live client via the injected ``async_client_factory``) and
        :meth:`aclose` (which nulls it on teardown). Tests that need to
        substitute the live transport at construction time should inject
        an ``async_client_factory`` into the client composition helper
        (the factory is forwarded into this kernel's ``__init__``); tests
        that need to swap the live client AFTER ``open()`` should use the
        dedicated test helper at
        ``tests/_fixtures/kernel_test_helpers.py``.
        """
        return self._http_client

    @property
    def cookies(self) -> httpx.Cookies:
        """Return the live HTTP client's cookie jar.

        Raises ``RuntimeError`` if called before :meth:`open`.
        """
        return self.get_http_client().cookies

    def get_http_client(self) -> httpx.AsyncClient:
        """Return the live HTTP client or raise the legacy not-open error."""
        if self._http_client is None:
            raise RuntimeError("Client not initialized. Use 'async with' context.")
        return self._http_client

    async def open(
        self,
        *,
        auth: AuthTokens,
        timeout: float,
        connect_timeout: float,
        limits: ConnectionLimits,
        capture_cookie_snapshot: Callable[[httpx.Cookies], object],
    ) -> None:
        """Build the HTTP client and capture its normalized cookie baseline."""
        # ClientLifecycle owns the primary idempotency guard. Keep this
        # secondary guard so direct Kernel callers also preserve the live client.
        if self._http_client is not None:
            return

        http_timeout = httpx.Timeout(
            connect=connect_timeout,
            read=timeout,
            write=timeout,
            pool=timeout,
        )
        self._timeout = timeout
        self._connect_timeout = connect_timeout
        cookies = (
            auth.cookie_jar
            if auth.cookie_jar is not None
            else build_cookie_jar(
                cookies=auth.cookies,
                storage_path=auth.storage_path,
            )
        )

        self._http_client = self._async_client_factory(
            headers={
                "Content-Type": "application/x-www-form-urlencoded;charset=UTF-8",
            },
            cookies=cookies,
            timeout=http_timeout,
            follow_redirects=True,
            limits=limits.to_httpx_limits(),
        )
        # If the snapshot raises, __aenter__ has effectively failed, so Python
        # never calls __aexit__ and the freshly built client would leak its
        # connection pool. Close it and reset state before re-raising so a
        # partial open cannot orphan a live client.
        try:
            capture_cookie_snapshot(self._http_client.cookies)
        except BaseException:
            await self.aclose()
            raise

    async def post(
        self,
        url: str,
        headers: Mapping[str, str] | None,
        body: PostBody,
        *,
        read_timeout: float | None = None,
        max_response_bytes: int | None = None,
    ) -> httpx.Response:
        """Issue a raw buffered POST through the live HTTP client."""
        timeout_override: httpx.Timeout | None = None
        if read_timeout is not None:
            # Widen ONLY the read (inactivity) slot; keep connect/write/pool at
            # the stored base values verbatim — including ``None``, which httpx
            # treats as "no timeout", so an explicit infinite-timeout client
            # keeps that intent on its other slots instead of inheriting the
            # chat read window.
            timeout_override = httpx.Timeout(
                connect=self._connect_timeout,
                read=read_timeout,
                write=self._timeout,
                pool=self._timeout,
            )
        headers_arg = dict(headers) if headers is not None else None
        stream_kwargs: dict[str, Any] = {}
        if max_response_bytes is not None:
            stream_kwargs["max_bytes"] = max_response_bytes
        return await stream_post_with_size_cap(
            self.get_http_client(),
            url,
            body=body,
            headers=headers_arg,
            timeout=timeout_override,
            **stream_kwargs,
        )

    async def aclose(self) -> None:
        """Close the live HTTP client and mark the kernel closed."""
        client = self._http_client
        if client is None:
            return
        try:
            await client.aclose()
        finally:
            self._http_client = None
            self._timeout = None
            self._connect_timeout = None


__all__ = ["Kernel"]