chore: import upstream snapshot with attribution

python/sglang/test/scripted_runtime/background_http_poster.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+from concurrent.futures import Future
+from typing import Any, Coroutine, Optional
+
+import aiohttp
+
+logger = logging.getLogger(__name__)
+
+JOIN_TIMEOUT_S: float = 10.0
+
+
+class BackgroundHttpPoster:
+
+    def __init__(self) -> None:
+        self._session: Optional[aiohttp.ClientSession] = None
+        self._loop = asyncio.new_event_loop()
+        self._thread = threading.Thread(
+            target=self._run_loop, name="scripted-runtime-async", daemon=True
+        )
+        self._thread.start()
+
+    def _run_loop(self) -> None:
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+
+    def submit_coro(self, coro: Coroutine) -> None:
+        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
+        future.add_done_callback(self._log_coro_exception)
+
+    @staticmethod
+    def _log_coro_exception(future: Future) -> None:
+        try:
+            future.result()
+        except asyncio.CancelledError:
+            pass
+        except Exception:
+            logger.exception("scripted_runtime: background async coroutine failed")
+
+    async def post(self, url: str, json: Any) -> None:
+        session = self._ensure_session()
+        async with session.post(url, json=json) as resp:
+            await resp.read()
+
+    def _ensure_session(self) -> aiohttp.ClientSession:
+        if self._session is None or self._session.closed:
+            self._session = aiohttp.ClientSession(
+                connector=aiohttp.TCPConnector(limit=0)
+            )
+        return self._session
+
+    def close(self) -> None:
+        try:
+            if self._session is not None:
+                future = asyncio.run_coroutine_threadsafe(
+                    self._session.close(), self._loop
+                )
+                future.result(timeout=JOIN_TIMEOUT_S)
+        except Exception:
+            logger.exception("scripted_runtime: failed to close aiohttp session")
+        try:
+            self._loop.call_soon_threadsafe(self._loop.stop)
+            self._thread.join(timeout=JOIN_TIMEOUT_S)
+        except Exception:
+            logger.exception("scripted_runtime: failed to stop background async loop")

python/sglang/test/scripted_runtime/context/__init__.py

@@ -0,0 +1,3 @@
+from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+__all__ = ["ScriptedContext"]

python/sglang/test/scripted_runtime/context/api.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional
+
+from sglang.test.scripted_runtime.context import (
+    engine,
+    lifecycle,
+    queries,
+    radix,
+)
+from sglang.test.scripted_runtime.context.kv_pool_exhauster import (
+    ScriptedKvPoolExhauster,
+)
+from sglang.test.scripted_runtime.context.lock_ref_exhauster import (
+    ScriptedLockRefExhauster,
+)
+from sglang.test.scripted_runtime.context.req_starter import ScriptedContextReqStarter
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.test.scripted_runtime.background_http_poster import BackgroundHttpPoster
+    from sglang.test.scripted_runtime.req_handle import ScriptedReqHandle
+    from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook
+    from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
+        ScriptedTokenizerRecvProxy,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+class ScriptedContext:
+
+    def __init__(
+        self,
+        *,
+        scheduler_hook: ScriptedSchedulerHook,
+        tokenizer_recv_proxy: Optional[ScriptedTokenizerRecvProxy],
+        http_poster: BackgroundHttpPoster,
+    ) -> None:
+        assert (
+            scheduler_hook._is_driver
+        ), "ScriptedContext only exists on the driver rank"
+        self.scheduler = scheduler_hook.scheduler
+        self._scheduler_hook = scheduler_hook
+        self._tokenizer_recv_proxy = tokenizer_recv_proxy
+        self._http_poster = http_poster
+
+        self._seen_rids: set[str] = set()
+        self._kv_exhauster = ScriptedKvPoolExhauster(self.scheduler)
+        self._lock_ref_exhauster = ScriptedLockRefExhauster(self.scheduler)
+        self._req_starter = ScriptedContextReqStarter(self)
+
+    def start_req(
+        self,
+        *,
+        prompt_len: int,
+        max_new_tokens: int = 8,
+        rid: Optional[str] = None,
+        ignore_eos: bool = False,
+        priority: Optional[int] = None,
+        dp_rank: Optional[int] = None,
+        prompt_token: int = 1,
+        return_logprob: bool = False,
+        logprob_start_len: Optional[int] = None,
+        top_logprobs_num: Optional[int] = None,
+        stop_token_ids: Optional[List[int]] = None,
+        temperature: Optional[float] = None,
+        lora_path: Optional[str] = None,
+    ) -> ScriptedReqHandle:
+        return self._req_starter.start_req(
+            prompt_len=prompt_len,
+            max_new_tokens=max_new_tokens,
+            rid=rid,
+            ignore_eos=ignore_eos,
+            priority=priority,
+            dp_rank=dp_rank,
+            prompt_token=prompt_token,
+            return_logprob=return_logprob,
+            logprob_start_len=logprob_start_len,
+            top_logprobs_num=top_logprobs_num,
+            stop_token_ids=stop_token_ids,
+            temperature=temperature,
+            lora_path=lora_path,
+        )
+
+    def pause_generation(self, *, mode: Literal["retract", "in_place"]) -> None:
+        return lifecycle.pause_generation(self, mode=mode)
+
+    def continue_generation(self, *, torch_empty_cache: bool = False) -> None:
+        return lifecycle.continue_generation(self, torch_empty_cache=torch_empty_cache)
+
+    def abort_all(self) -> None:
+        return lifecycle.abort_all(self)
+
+    def abort(self, handle: ScriptedReqHandle, *, await_arrival: bool = True) -> None:
+        return lifecycle.abort(self, rid=handle.rid, await_arrival=await_arrival)
+
+    def flush_cache(self) -> None:
+        return lifecycle.flush_cache(self)
+
+    def evict_radix(self, *, prefix_tokens: Optional[List[int]]) -> None:
+        assert (
+            prefix_tokens is None
+        ), "evict_radix currently supports only full eviction (prefix_tokens=None)"
+        return lifecycle.flush_cache(self)
+
+    def exhaust_kv(self, *, leave_pages: int) -> None:
+        return self._kv_exhauster.exhaust(leave_pages=leave_pages)
+
+    def exhaust_lock_refs(self, *, leave_refs: int) -> None:
+        return self._lock_ref_exhauster.exhaust(leave_refs=leave_refs)
+
+    def _release_exhausted_pools(self) -> None:
+        self._kv_exhauster.release()
+        self._lock_ref_exhauster.release()
+
+    def get_all_node_hit_counts(self) -> Dict[int, int]:
+        return radix.get_all_node_hit_counts(self)
+
+    def get_all_node_lock_refs(self) -> Dict[int, int]:
+        return radix.get_all_node_lock_refs(self)
+
+    @property
+    def is_idle(self) -> bool:
+        return queries.is_idle(self)
+
+    @property
+    def is_fully_idle(self) -> bool:
+        return queries.is_fully_idle(self)
+
+    @property
+    def last_batch_forward_mode(self) -> Optional[str]:
+        return queries.last_batch_forward_mode(self)
+
+    def find_req_by_rid(self, rid: str) -> Optional[Req]:
+        return queries.find_req_by_rid(self, rid)
+
+    def is_finished(self, rid: str) -> bool:
+        return queries.is_finished(self, rid)
+
+    def is_chunking(self, rid: str) -> bool:
+        return queries.is_chunking(self, rid)
+
+    def status(self, rid: str) -> str:
+        return queries.status(self, rid)
+
+    def remaining_prompt_tokens(self, rid: str) -> int:
+        return queries.remaining_prompt_tokens(self, rid)
+
+    def list_active_reqs(self) -> List[Req]:
+        return queries.list_active_reqs(self)
+
+    def chunks_done(self, rid: str) -> int:
+        return queries.chunks_done(self, rid)
+
+    def chunked_parks(self, rid: str) -> int:
+        return queries.chunked_parks(self, rid)
+
+    def batch_composition(self) -> Dict[str, List[str]]:
+        return queries.batch_composition(self)
+
+    def engine_stats(self) -> Dict[str, int]:
+        return engine.engine_stats(self)

python/sglang/test/scripted_runtime/context/engine.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+def engine_stats(ctx: ScriptedContext) -> Dict[str, int]:
+    s = ctx.scheduler
+    return {
+        "kv_pool_free": s.token_to_kv_pool_allocator.available_size(),
+        "req_pool_free": s.req_to_token_pool.available_size(),
+        "req_pool_total": s.req_to_token_pool.size,
+        "page_size": s.page_size,
+    }

python/sglang/test/scripted_runtime/context/http_post.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+logger = logging.getLogger(__name__)
+
+RECV_MSG_ARRIVAL_TIMEOUT_S: float = 60.0
+
+
+def _http_post_and_await_recv_msg(
+    ctx: ScriptedContext,
+    *,
+    path: str,
+    json: Optional[Dict[str, Any]],
+    predicate: Callable[[Any], bool],
+    description: str,
+    timeout_s: float = RECV_MSG_ARRIVAL_TIMEOUT_S,
+) -> None:
+    _submit_post(ctx, path=path, json=json)
+    ctx._tokenizer_recv_proxy.wait_until_arrived(
+        predicate,
+        timeout_s=timeout_s,
+        description=description,
+    )
+
+
+def _http_post_fire_and_forget(
+    ctx: ScriptedContext,
+    *,
+    path: str,
+    json: Optional[Dict[str, Any]],
+) -> None:
+    _submit_post(ctx, path=path, json=json)
+
+
+def _submit_post(
+    ctx: ScriptedContext,
+    *,
+    path: str,
+    json: Optional[Dict[str, Any]],
+) -> None:
+    server_args = ctx.scheduler.server_args
+    url = f"http://{server_args.host}:{server_args.port}{path}"
+
+    async def _post() -> None:
+        await ctx._http_poster.post(url, json)
+
+    ctx._http_poster.submit_coro(_post())

python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+
+if TYPE_CHECKING:
+    import torch
+
+    from sglang.srt.managers.scheduler import Scheduler
+
+
+class ScriptedKvPoolExhauster:
+
+    def __init__(self, scheduler: Scheduler) -> None:
+        self.scheduler = scheduler
+        self._held: List[torch.Tensor] = []
+
+    def exhaust(self, *, leave_pages: int) -> None:
+        allocator = self.scheduler.token_to_kv_pool_allocator
+
+        leave_tokens = leave_pages * self.scheduler.page_size
+        need = allocator.available_size() - leave_tokens
+        if need <= 0:
+            return
+
+        held = allocator.alloc(need)
+        assert (
+            held is not None
+        ), f"exhaust_kv: allocator could not grab {need} tokens to create pressure"
+        self._held.append(held)
+
+    def release(self) -> None:
+        for held in self._held:
+            self.scheduler.token_to_kv_pool_allocator.free(held)
+        self._held.clear()

python/sglang/test/scripted_runtime/context/lifecycle.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from sglang.srt.managers.io_struct import (
+    AbortReq,
+    ContinueGenerationReqInput,
+    FlushCacheReqInput,
+    PauseGenerationReqInput,
+)
+from sglang.test.scripted_runtime.context.http_post import (
+    _http_post_and_await_recv_msg,
+    _http_post_fire_and_forget,
+)
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+def _await_control(
+    ctx: ScriptedContext,
+    *,
+    path: str,
+    json,
+    expect_type: type,
+    await_arrival: bool = True,
+) -> None:
+    if not await_arrival:
+        _http_post_fire_and_forget(ctx, path=path, json=json)
+        return
+    _http_post_and_await_recv_msg(
+        ctx,
+        path=path,
+        json=json,
+        predicate=lambda obj: isinstance(obj, expect_type),
+        description=expect_type.__name__,
+    )
+
+
+def pause_generation(
+    ctx: ScriptedContext, *, mode: Literal["retract", "in_place"]
+) -> None:
+    _await_control(
+        ctx,
+        path="/pause_generation",
+        json={"mode": mode},
+        expect_type=PauseGenerationReqInput,
+    )
+
+
+def continue_generation(ctx: ScriptedContext, *, torch_empty_cache: bool) -> None:
+    _await_control(
+        ctx,
+        path="/continue_generation",
+        json={"torch_empty_cache": torch_empty_cache},
+        expect_type=ContinueGenerationReqInput,
+    )
+
+
+def abort_all(ctx: ScriptedContext) -> None:
+    _await_control(
+        ctx,
+        path="/abort_request",
+        json={"rid": "", "abort_all": True},
+        expect_type=AbortReq,
+    )
+
+
+def abort(ctx: ScriptedContext, *, rid: str, await_arrival: bool = True) -> None:
+    _await_control(
+        ctx,
+        path="/abort_request",
+        json={"rid": rid, "abort_all": False},
+        expect_type=AbortReq,
+        await_arrival=await_arrival,
+    )
+
+
+def flush_cache(ctx: ScriptedContext) -> None:
+    _await_control(
+        ctx,
+        path="/flush_cache",
+        json=None,
+        expect_type=FlushCacheReqInput,
+    )

python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List
+
+from sglang.test.scripted_runtime.context.radix import _node_lock_ref
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import Scheduler
+
+
+class ScriptedLockRefExhauster:
+
+    def __init__(self, scheduler: Scheduler) -> None:
+        self.scheduler = scheduler
+        self._locked: List[Any] = []
+
+    def exhaust(self, *, leave_refs: int) -> None:
+        tree_cache = self.scheduler.tree_cache
+        if tree_cache.disable:
+            return
+
+        while True:
+            evictable = self._evictable_nodes()
+            if len(evictable) <= leave_refs:
+                return
+
+            target = evictable[0]
+            tree_cache.inc_lock_ref(target)
+
+            newly_locked = [node for node in evictable if _node_lock_ref(node) > 0]
+            if not newly_locked:
+                return
+            self._locked.append(target)
+
+    def release(self) -> None:
+        tree_cache = self.scheduler.tree_cache
+        for node in self._locked:
+            tree_cache.dec_lock_ref(node)
+        self._locked.clear()
+
+    def _evictable_nodes(self) -> List[Any]:
+        evictable: List[Any] = []
+        stack = list(self.scheduler.tree_cache.root_node.children.values())
+        while stack:
+            node = stack.pop()
+            if _node_lock_ref(node) == 0:
+                evictable.append(node)
+            stack.extend(node.children.values())
+        return evictable

python/sglang/test/scripted_runtime/context/queries.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Iterator, List, Optional
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+def _get_all_reqs(ctx: ScriptedContext) -> Iterator[Req]:
+    s = ctx.scheduler
+    if s.chunked_req is not None:
+        yield s.chunked_req
+    yield from s.waiting_queue
+    if s.ps.pp_size > 1:
+        for mb in (*s.mbs, *s.last_mbs, *s.running_mbs):
+            if mb is not None:
+                yield from mb.reqs
+    else:
+        if s.running_batch is not None:
+            yield from s.running_batch.reqs
+        if s.last_batch is not None:
+            yield from s.last_batch.reqs
+
+
+def list_active_reqs(ctx: ScriptedContext) -> List[Req]:
+    return list(set(_get_all_reqs(ctx)))
+
+
+def batch_composition(ctx: ScriptedContext) -> Dict[str, List[str]]:
+    s = ctx.scheduler
+    chunked_rid = s.chunked_req.rid if s.chunked_req is not None else None
+    chunked = [chunked_rid] if chunked_rid is not None else []
+    running = (
+        [r.rid for r in s.running_batch.reqs] if s.running_batch is not None else []
+    )
+
+    prefill: List[str] = []
+    decode: List[str] = []
+    batch = s.last_batch
+    if batch is not None and not batch.is_empty() and batch.forward_mode is not None:
+        bucket = prefill if batch.forward_mode.is_extend() else decode
+        bucket.extend(r.rid for r in batch.reqs if r.rid != chunked_rid)
+
+    return {
+        "prefill": prefill,
+        "decode": decode,
+        "chunked": chunked,
+        "running": running,
+    }
+
+
+def is_idle(ctx: ScriptedContext) -> bool:
+    s = ctx.scheduler
+    return (
+        s.chunked_req is None
+        and len(s.waiting_queue) == 0
+        and (s.running_batch is None or s.running_batch.is_empty())
+    )
+
+
+def is_fully_idle(ctx: ScriptedContext) -> bool:
+    s = ctx.scheduler
+    return is_idle(ctx) and (s.last_batch is None or s.last_batch.is_empty())
+
+
+def last_batch_forward_mode(ctx: ScriptedContext) -> Optional[str]:
+    s = ctx.scheduler
+    if s.last_batch is not None and s.last_batch.forward_mode is not None:
+        return s.last_batch.forward_mode.name
+    return None
+
+
+def find_req_by_rid(ctx: ScriptedContext, rid: str) -> Optional[Req]:
+    req = next((r for r in _get_all_reqs(ctx) if r.rid == rid), None)
+    if req is not None:
+        ctx._seen_rids.add(rid)
+    return req
+
+
+def is_finished(ctx: ScriptedContext, rid: str) -> bool:
+    req = find_req_by_rid(ctx, rid)
+    if req is not None:
+        return req.finished()
+    if rid in ctx._seen_rids:
+        return True
+    # Fallback: if the req ran in a forward batch (recorded in _batch_log) but
+    # is now absent from all active scheduler sets, it must have finished.
+    # This catches requests that completed without ever being observed via
+    # find_req_by_rid (e.g. when Python short-circuit evaluation prevents the
+    # query while another request is still running).
+    log = ctx._scheduler_hook._batch_log
+    if any(rid in record.rids for record in log):
+        ctx._seen_rids.add(rid)
+        return True
+    return False
+
+
+def is_chunking(ctx: ScriptedContext, rid: str) -> bool:
+    s = ctx.scheduler
+    return s.chunked_req is not None and s.chunked_req.rid == rid
+
+
+def status(ctx: ScriptedContext, rid: str) -> str:
+    s = ctx.scheduler
+    if rid in {r.rid for r in s.waiting_queue}:
+        return "waiting"
+    req = find_req_by_rid(ctx, rid)
+    if req is not None:
+        return "finished" if req.finished() else "running"
+    if rid in ctx._seen_rids:
+        return "finished"
+    return "unknown"
+
+
+def remaining_prompt_tokens(ctx: ScriptedContext, rid: str) -> int:
+    req = find_req_by_rid(ctx, rid)
+    if req is None:
+        return 0
+    return max(0, len(req.origin_input_ids) - req.kv_committed_len)
+
+
+def chunks_done(ctx: ScriptedContext, rid: str) -> int:
+    log = ctx._scheduler_hook._batch_log
+    held = sum(1 for record in log if record.chunked_rid == rid and rid in record.rids)
+    if held == 0:
+        return 0
+    completed = any(
+        rid in record.extend_rids and record.chunked_rid != rid for record in log
+    )
+    return held + (1 if completed else 0)
+
+
+def chunked_parks(ctx: ScriptedContext, rid: str) -> int:
+    return sum(
+        1
+        for record in ctx._scheduler_hook._batch_log
+        if record.chunked_rid == rid and rid not in record.rids
+    )

python/sglang/test/scripted_runtime/context/radix.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Dict
+
+from sglang.srt.mem_cache.swa_radix_cache import TreeNode as SWATreeNode
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+def get_all_node_hit_counts(ctx: ScriptedContext) -> Dict[int, int]:
+    return _collect_node_attr(ctx, lambda node: node.hit_count)
+
+
+def get_all_node_lock_refs(ctx: ScriptedContext) -> Dict[int, int]:
+    return _collect_node_attr(ctx, _node_lock_ref)
+
+
+def _node_lock_ref(node: Any) -> int:
+    if isinstance(node, SWATreeNode):
+        return node.full_lock_ref + node.swa_lock_ref
+    return node.lock_ref
+
+
+def _collect_node_attr(
+    ctx: ScriptedContext, get_value: Callable[[Any], int]
+) -> Dict[int, int]:
+    values: Dict[int, int] = {}
+    stack = list(ctx.scheduler.tree_cache.root_node.children.values())
+    while stack:
+        node = stack.pop()
+        values[node.id] = get_value(node)
+        stack.extend(node.children.values())
+    return values

python/sglang/test/scripted_runtime/context/req_starter.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import uuid
+from typing import TYPE_CHECKING, List, Optional
+
+from sglang.test.scripted_runtime.context.http_post import (
+    _http_post_and_await_recv_msg,
+)
+from sglang.test.scripted_runtime.req_handle import ScriptedReqHandle
+
+if TYPE_CHECKING:
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+class ScriptedContextReqStarter:
+    def __init__(self, ctx: ScriptedContext) -> None:
+        self._ctx = ctx
+        self._req_counter = 0
+
+    def start_req(
+        self,
+        *,
+        prompt_len: int,
+        max_new_tokens: int,
+        rid: Optional[str],
+        ignore_eos: bool,
+        priority: Optional[int],
+        dp_rank: Optional[int],
+        prompt_token: int = 1,
+        return_logprob: bool = False,
+        logprob_start_len: Optional[int] = None,
+        top_logprobs_num: Optional[int] = None,
+        stop_token_ids: Optional[List[int]] = None,
+        temperature: Optional[float] = None,
+        lora_path: Optional[str] = None,
+    ) -> ScriptedReqHandle:
+        ctx = self._ctx
+
+        if rid is None:
+            rid = f"scripted-{self._req_counter}-{uuid.uuid4().hex}"
+            self._req_counter += 1
+
+        sampling_params = {"max_new_tokens": max_new_tokens, "ignore_eos": ignore_eos}
+        if stop_token_ids is not None:
+            sampling_params["stop_token_ids"] = stop_token_ids
+        if temperature is not None:
+            sampling_params["temperature"] = temperature
+        payload = {
+            "input_ids": [prompt_token] * prompt_len,
+            "sampling_params": sampling_params,
+            "rid": rid,
+            "stream": True,
+        }
+        payload["return_logprob"] = return_logprob
+        if logprob_start_len is not None:
+            payload["logprob_start_len"] = logprob_start_len
+        if top_logprobs_num is not None:
+            payload["top_logprobs_num"] = top_logprobs_num
+        if priority is not None:
+            payload["priority"] = priority
+        if dp_rank is not None:
+            payload["routed_dp_rank"] = dp_rank
+        if lora_path is not None:
+            payload["lora_path"] = lora_path
+        _http_post_and_await_recv_msg(
+            ctx,
+            path="/generate",
+            json=payload,
+            predicate=lambda obj: getattr(obj, "rid", None) == rid,
+            description=f"request with rid {rid!r}",
+        )
+
+        return ScriptedReqHandle(rid=rid, context=ctx)

python/sglang/test/scripted_runtime/http_server.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+import multiprocessing as mp
+import os
+import tempfile
+import time
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Tuple
+
+import requests
+import zmq
+
+from sglang.srt.entrypoints.http_server import launch_server
+from sglang.srt.environ import envs
+from sglang.srt.managers.io_struct import sock_recv, sock_send, wrap_as_pickle
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils.network import get_free_port, get_zmq_socket_on_host
+from sglang.test.scripted_runtime.io_struct import (
+    HookReady,
+    OutOfBandError,
+    RunScript,
+    ScriptFailed,
+    ScriptSucceeded,
+    Shutdown,
+)
+from sglang.test.scripted_runtime.utils import close_zmq_socket
+
+DEFAULT_RUN_TIMEOUT_S: float = 120.0
+SHUTDOWN_JOIN_TIMEOUT_S: float = 60.0
+LISTENER_ACCEPT_TIMEOUT_S: float = 300.0
+HTTP_READY_TIMEOUT_S: float = 300.0
+HTTP_READY_POLL_INTERVAL_S: float = 0.5
+SERVER_HOST: str = "127.0.0.1"
+
+
+class ScriptedHttpServer:
+
+    def __init__(
+        self,
+        *,
+        ctx: zmq.Context,
+        socket: zmq.Socket,
+        server_process: mp.process.BaseProcess,
+        out_of_band_error_path: Path,
+        http_port: int,
+    ) -> None:
+        self._ctx = ctx
+        self._socket = socket
+        self._server_process = server_process
+        self._out_of_band_error_path = out_of_band_error_path
+        self._base_url = f"http://{SERVER_HOST}:{http_port}"
+        self._shutdown_done = False
+        self._dirty: Optional[str] = None
+
+    @classmethod
+    def start(cls, **engine_kwargs: Any) -> ScriptedHttpServer:
+        out_of_band_error_path = _create_oob_error_file()
+
+        ctx = zmq.Context()
+        dispatch_port, socket = get_zmq_socket_on_host(ctx, zmq.PAIR, host=SERVER_HOST)
+        server_process, http_port = _spawn_server_process(
+            endpoint=f"tcp://{SERVER_HOST}:{dispatch_port}",
+            out_of_band_error_path=out_of_band_error_path,
+            engine_kwargs=engine_kwargs,
+        )
+
+        self = cls(
+            ctx=ctx,
+            socket=socket,
+            server_process=server_process,
+            out_of_band_error_path=out_of_band_error_path,
+            http_port=http_port,
+        )
+        try:
+            self._await_handshake()
+            self._await_http_ready()
+        except BaseException:
+            self._teardown()
+            raise
+        return self
+
+    def execute_script(
+        self,
+        script_fn: Callable,
+        *,
+        args: Tuple[Any, ...] = (),
+        timeout_s: float = DEFAULT_RUN_TIMEOUT_S,
+    ) -> None:
+        if self._dirty:
+            raise RuntimeError(f"ScriptedHttpServer is dirty: {self._dirty}")
+
+        fn_path = f"{script_fn.__module__}:{script_fn.__qualname__}"
+        sock_send(self._socket, wrap_as_pickle(RunScript(fn_path=fn_path, args=args)))
+
+        if not self._socket.poll(int(timeout_s * 1000)):
+            if not self._server_process.is_alive():
+                self._dirty = f"server process died before responding to {fn_path!r}"
+                raise RuntimeError(self._dirty)
+            self._dirty = f"script {fn_path!r} timed out after {timeout_s}s"
+            raise TimeoutError(self._dirty)
+
+        reply = sock_recv(self._socket)
+        match reply:
+            case ScriptFailed(traceback=tb):
+                raise AssertionError(f"scripted-runtime script failed:\n{tb}")
+            case ScriptSucceeded():
+                return
+            case _:
+                raise RuntimeError(
+                    f"scheduler hook replied with unexpected message {reply!r}"
+                )
+
+    def shutdown(self) -> None:
+        if self._shutdown_done:
+            return
+
+        fatal_error: Optional[OutOfBandError] = None
+        try:
+            try:
+                sock_send(self._socket, wrap_as_pickle(Shutdown()))
+            except zmq.ZMQError:
+                pass
+
+            self._server_process.join(timeout=SHUTDOWN_JOIN_TIMEOUT_S)
+            self._terminate_process(self._server_process)
+            fatal_error = self._read_out_of_band_error()
+        finally:
+            self._teardown()
+            self._shutdown_done = True
+
+        if fatal_error:
+            raise AssertionError(
+                f"scripted-runtime server failed:\n{fatal_error.traceback}"
+            )
+
+    def _await_handshake(self) -> None:
+        if not self._socket.poll(int(LISTENER_ACCEPT_TIMEOUT_S * 1000)):
+            raise TimeoutError(
+                f"ScriptedHttpServer: HTTP server did not connect within "
+                f"{LISTENER_ACCEPT_TIMEOUT_S}s"
+            )
+
+        ready = sock_recv(self._socket)
+        if not isinstance(ready, HookReady):
+            raise RuntimeError(
+                f"ScriptedHttpServer: expected HookReady handshake, got {ready!r}"
+            )
+
+    def _await_http_ready(self) -> None:
+        # HookReady only means the scheduler dispatch loop started; the uvicorn
+        # entrypoint may not be bound yet. The first script runs
+        # _reset_engine_state, which POSTs to this server's own HTTP port, so
+        # block until the port is bound and routing before any script can run.
+        #
+        # Poll /model_info rather than /health: once the server reports Up,
+        # /health runs the generation-based health check (a real probe request
+        # through the scheduler), which cannot complete while the scheduler
+        # waits for scripts between RunScript commands. Any /model_info
+        # response proves the socket is bound and routes are registered, which
+        # is all the control POSTs need.
+        url = f"{self._base_url}/model_info"
+        deadline = time.monotonic() + HTTP_READY_TIMEOUT_S
+        while time.monotonic() < deadline:
+            if not self._server_process.is_alive():
+                raise RuntimeError(
+                    "ScriptedHttpServer: server process died during HTTP startup"
+                )
+            try:
+                requests.get(url, timeout=2.0)
+                return
+            except requests.RequestException:
+                pass
+            time.sleep(HTTP_READY_POLL_INTERVAL_S)
+
+        raise TimeoutError(
+            f"ScriptedHttpServer: HTTP endpoint {url} not bound within "
+            f"{HTTP_READY_TIMEOUT_S}s"
+        )
+
+    def _teardown(self) -> None:
+        try:
+            close_zmq_socket(self._socket, self._ctx)
+        except Exception:  # noqa: BLE001 — best-effort cleanup
+            pass
+        self._terminate_process(self._server_process)
+        self._cleanup_files()
+
+    def _read_out_of_band_error(self) -> Optional[OutOfBandError]:
+        try:
+            text = self._out_of_band_error_path.read_text()
+        except OSError:
+            return None
+        text = text.strip()
+        return OutOfBandError.from_json(text) if text else None
+
+    def _cleanup_files(self) -> None:
+        try:
+            self._out_of_band_error_path.unlink()
+        except OSError:
+            pass
+
+    @staticmethod
+    def _terminate_process(process: mp.process.BaseProcess) -> None:
+        if not process.is_alive():
+            return
+        process.terminate()
+        process.join(timeout=10.0)
+        if process.is_alive():
+            process.kill()
+            process.join(timeout=10.0)
+
+
+def _create_oob_error_file() -> Path:
+    err_fd, err_path = tempfile.mkstemp(
+        prefix="scripted_runtime_oob_error_", suffix=".json"
+    )
+    os.close(err_fd)
+    return Path(err_path)
+
+
+def _spawn_server_process(
+    *,
+    endpoint: str,
+    out_of_band_error_path: Path,
+    engine_kwargs: Dict[str, Any],
+) -> Tuple[mp.process.BaseProcess, int]:
+    mp_ctx = mp.get_context("spawn")
+    launch_kwargs: Dict[str, Any] = dict(
+        host=SERVER_HOST,
+        port=get_free_port(),
+        kv_canary="raise",
+        kv_canary_real_data="partial",
+        kv_canary_sweep_interval=100,
+        disable_prefill_cuda_graph=True,
+    )
+    launch_kwargs.update(engine_kwargs)
+    http_port = launch_kwargs["port"]
+    server_process = mp_ctx.Process(
+        target=_launch_scripted_http_server,
+        kwargs=launch_kwargs,
+        name="scripted-runtime-http-server",
+        daemon=False,
+    )
+
+    sys_path_entry = str(Path(__file__).resolve().parent)
+    with (
+        envs.SGLANG_TEST_SCRIPTED_RUNTIME.override(True),
+        envs.SGLANG_TEST_SCRIPTED_RUNTIME_IPC_ADDR.override(endpoint),
+        envs.SGLANG_TEST_SCRIPTED_RUNTIME_OUT_OF_BAND_ERROR_PATH.override(
+            str(out_of_band_error_path)
+        ),
+        envs.SGLANG_TEST_SCRIPTED_RUNTIME_SYS_PATH_ENTRY.override(sys_path_entry),
+    ):
+        server_process.start()
+
+    return server_process, http_port
+
+
+def _launch_scripted_http_server(**engine_kwargs: Any) -> None:
+    launch_server(ServerArgs(**engine_kwargs))

python/sglang/test/scripted_runtime/io_struct.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import dataclasses
+import json
+from dataclasses import dataclass
+from typing import Any, Tuple, Union
+
+
+@dataclass(frozen=True, slots=True)
+class RunScript:
+
+    fn_path: str
+    args: Tuple[Any, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class Shutdown:
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class HookReady:
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class ScriptSucceeded:
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class ScriptFailed:
+
+    traceback: str
+
+
+@dataclass(frozen=True, slots=True)
+class OutOfBandError:
+
+    traceback: str
+
+    def to_json(self) -> str:
+        return json.dumps(dataclasses.asdict(self))
+
+    @classmethod
+    def from_json(cls, text: str) -> OutOfBandError:
+        return cls(**json.loads(text))
+
+
+ScriptedCommand = Union[RunScript, Shutdown]
+ScriptedReply = Union[HookReady, ScriptSucceeded, ScriptFailed]

python/sglang/test/scripted_runtime/req_handle.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+from sglang.test.scripted_runtime.context.radix import _node_lock_ref
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.test.scripted_runtime.context.api import ScriptedContext
+
+
+@dataclass(frozen=True, slots=True)
+class ScriptedReqHandle:
+    rid: str
+    context: ScriptedContext
+
+    @property
+    def req(self) -> Optional[Req]:
+        return self.context.find_req_by_rid(self.rid)
+
+    @property
+    def finished(self) -> bool:
+        return self.context.is_finished(self.rid)
+
+    @property
+    def is_chunking(self) -> bool:
+        return self.context.is_chunking(self.rid)
+
+    @property
+    def chunks_done(self) -> int:
+        return self.context.chunks_done(self.rid)
+
+    @property
+    def status(self) -> str:
+        return self.context.status(self.rid)
+
+    @property
+    def remaining_prompt_tokens(self) -> int:
+        return self.context.remaining_prompt_tokens(self.rid)
+
+    @property
+    def kv_pages(self) -> int:
+        req = self.req
+        if req is None or req.req_pool_idx is None:
+            return 0
+        page_size = self.context.scheduler.page_size
+        return (req.kv_allocated_len + page_size - 1) // page_size
+
+    @property
+    def lock_refs(self) -> int:
+        req = self.req
+        if req is None:
+            return 0
+        node = req.last_node
+        if node is None:
+            return 0
+        return _node_lock_ref(node)

python/sglang/test/scripted_runtime/scheduler_hook.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import logging
+import sys
+import time
+import traceback
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Generator, List, Optional, Tuple
+
+import zmq
+
+from sglang.srt.environ import envs
+from sglang.srt.managers.io_struct import sock_recv, sock_send, wrap_as_pickle
+from sglang.srt.utils.network import get_zmq_socket
+from sglang.test.scripted_runtime.background_http_poster import BackgroundHttpPoster
+from sglang.test.scripted_runtime.context import ScriptedContext
+from sglang.test.scripted_runtime.io_struct import (
+    HookReady,
+    OutOfBandError,
+    RunScript,
+    ScriptFailed,
+    ScriptSucceeded,
+    Shutdown,
+)
+from sglang.test.scripted_runtime.utils import (
+    close_zmq_socket,
+    ensure_script_importable,
+    resolve_fn,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import Scheduler
+    from sglang.test.scripted_runtime.tokenizer_recv_proxy import (
+        ScriptedTokenizerRecvProxy,
+    )
+
+logger = logging.getLogger(__name__)
+
+RESET_DRAIN_MAX_STEPS: int = 200
+# Below the test-side LISTENER_ACCEPT_TIMEOUT_S so a stuck warmup surfaces as
+# this specific error instead of a generic handshake timeout.
+WARMUP_DRIVE_TIMEOUT_S: float = 120.0
+
+
+@dataclass(frozen=True, slots=True)
+class ScriptedBatchRecord:
+    forward_iter: int
+    mode: Optional[str]
+    rids: Tuple[str, ...]
+    extend_rids: Tuple[str, ...]
+    chunked_rid: Optional[str]
+
+
+def _drive_engine_through_warmup(ctx: ScriptedContext) -> Generator:
+    """Run the engine until the server warmup request has been received and
+    fully processed, so scripts never observe foreign warmup traffic."""
+    scheduler = ctx.scheduler
+    server_args = scheduler.server_args
+    if server_args.skip_server_warmup:
+        logger.info("scripted_runtime: skip_server_warmup set, not driving warmup")
+        return
+
+    logger.info("scripted_runtime: driving engine until server warmup completes")
+    start_time = time.monotonic()
+
+    # is_fully_idle() can transiently report idle while a PP microbatch result
+    # is still in flight, so require it to hold for two full microbatch
+    # rotations after the warmup request was observed on the recv socket.
+    quiesce_iters = 2 * (server_args.pp_size + server_args.pp_async_batch_depth)
+    proxy = ctx._tokenizer_recv_proxy
+    deadline = start_time + WARMUP_DRIVE_TIMEOUT_S
+
+    idle_streak = 0
+    while idle_streak < quiesce_iters:
+        if time.monotonic() >= deadline:
+            raise RuntimeError(
+                "scripted_runtime: server warmup did not complete within "
+                f"{WARMUP_DRIVE_TIMEOUT_S}s "
+                f"(work_reqs_seen={proxy.work_reqs_seen}, "
+                f"idle_streak={idle_streak})"
+            )
+        yield
+        if proxy.work_reqs_seen > 0 and scheduler.is_fully_idle():
+            idle_streak += 1
+        else:
+            idle_streak = 0
+
+    logger.info(
+        "scripted_runtime: server warmup drained in %.1fs (work_reqs_seen=%d)",
+        time.monotonic() - start_time,
+        proxy.work_reqs_seen,
+    )
+
+
+def _reset_engine_state(ctx: ScriptedContext) -> Generator:
+    scheduler = ctx.scheduler
+
+    if scheduler._engine_paused:
+        ctx.continue_generation()
+
+    ctx._release_exhausted_pools()
+    ctx.abort_all()
+    for _ in range(RESET_DRAIN_MAX_STEPS):
+        yield
+        if scheduler.is_fully_idle():
+            break
+    else:
+        raise RuntimeError(
+            "scripted_runtime reset: scheduler did not become fully idle "
+            f"within {RESET_DRAIN_MAX_STEPS} steps"
+        )
+
+    ctx.flush_cache()
+    yield
+
+
+class ScriptedSchedulerHook:
+
+    def __init__(
+        self,
+        *,
+        scheduler: Scheduler,
+        tokenizer_recv_proxy: Optional[ScriptedTokenizerRecvProxy],
+    ) -> None:
+        self.scheduler = scheduler
+        self._is_driver = (
+            scheduler.ps.pp_rank == 0
+            and scheduler.ps.tp_rank == 0
+            and scheduler.ps.attn_cp_rank == 0
+        )
+        self._batch_log: List[ScriptedBatchRecord] = []
+
+        if self._is_driver:
+            ensure_script_importable(
+                envs.SGLANG_TEST_SCRIPTED_RUNTIME_SYS_PATH_ENTRY.get()
+            )
+            self._http_poster: Optional[BackgroundHttpPoster] = BackgroundHttpPoster()
+            self._context: Optional[ScriptedContext] = ScriptedContext(
+                scheduler_hook=self,
+                tokenizer_recv_proxy=tokenizer_recv_proxy,
+                http_poster=self._http_poster,
+            )
+            self._script_fn_generator: Optional[Generator] = self._run_dispatch_loop()
+        else:
+            self._http_poster = None
+            self._context = None
+            self._script_fn_generator = None
+
+    def _run_dispatch_loop(self) -> Generator:
+        endpoint = envs.SGLANG_TEST_SCRIPTED_RUNTIME_IPC_ADDR.get()
+        ctx_zmq = zmq.Context()
+        socket = get_zmq_socket(ctx_zmq, zmq.PAIR, endpoint, bind=False)
+        try:
+            yield from _drive_engine_through_warmup(self._context)
+            sock_send(socket, wrap_as_pickle(HookReady()))
+            while True:
+                msg = sock_recv(socket)
+                match msg:
+                    case Shutdown():
+                        return
+                    case RunScript(fn_path=fn_path, args=args):
+                        fn = resolve_fn(fn_path)
+                        ctx = self._context
+                        yield from _reset_engine_state(ctx)
+                        self._batch_log.clear()
+                        sub_gen = fn(ctx, *args)
+                        try:
+                            yield from sub_gen
+                        except Exception:
+                            sock_send(
+                                socket,
+                                wrap_as_pickle(
+                                    ScriptFailed(traceback=traceback.format_exc())
+                                ),
+                            )
+                        else:
+                            sock_send(socket, wrap_as_pickle(ScriptSucceeded()))
+                    case _:
+                        raise ValueError(f"dispatch loop: unknown command {msg!r}")
+        finally:
+            close_zmq_socket(socket, ctx_zmq)
+            self._http_poster.close()
+
+    def on_run_batch(self, batch) -> None:
+        if not self._is_driver:
+            return
+        chunked = self.scheduler.chunked_req
+        self._batch_log.append(
+            ScriptedBatchRecord(
+                forward_iter=batch.forward_iter,
+                mode=(
+                    batch.forward_mode.name.lower()
+                    if batch.forward_mode is not None
+                    else None
+                ),
+                rids=tuple(r.rid for r in batch.reqs),
+                extend_rids=(
+                    tuple(r.rid for r in batch.reqs)
+                    if batch.forward_mode is not None and batch.forward_mode.is_extend()
+                    else ()
+                ),
+                chunked_rid=chunked.rid if chunked is not None else None,
+            )
+        )
+
+    def step(self) -> None:
+        if not self._is_driver:
+            return
+
+        done, exc_tb = _advance_generator(self._script_fn_generator)
+        if not done:
+            return
+
+        if exc_tb is not None:
+            _write_out_of_band_error(exc_tb)
+        sys.exit(0 if exc_tb is None else 1)
+
+
+def _write_out_of_band_error(exc_tb: str) -> None:
+    path = envs.SGLANG_TEST_SCRIPTED_RUNTIME_OUT_OF_BAND_ERROR_PATH.get()
+    if not path:
+        return
+    error = OutOfBandError(traceback=exc_tb or "<no traceback>")
+    try:
+        Path(path).write_text(error.to_json())
+    except OSError:
+        logger.exception(
+            "Failed to write scripted_runtime out-of-band error to %s", path
+        )
+
+
+def _advance_generator(generator: Generator) -> Tuple[bool, Optional[str]]:
+    try:
+        next(generator)
+        return False, None
+    except StopIteration:
+        return True, None
+    except Exception:
+        logger.exception("Failed to advance generator")
+        return True, traceback.format_exc()

python/sglang/test/scripted_runtime/test_case.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from typing import Any, ClassVar, Dict
+
+from sglang.test.scripted_runtime.http_server import ScriptedHttpServer
+from sglang.test.test_utils import CustomTestCase
+
+
+class ScriptedTestCase(CustomTestCase):
+
+    ENGINE_KWARGS: ClassVar[Dict[str, Any]] = {}
+
+    server: ClassVar[ScriptedHttpServer]
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        if not cls.ENGINE_KWARGS:
+            raise AssertionError(
+                f"{cls.__name__} must set ENGINE_KWARGS to a non-empty dict"
+            )
+        cls.server = ScriptedHttpServer.start(**cls.ENGINE_KWARGS)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        try:
+            cls.server.shutdown()
+        finally:
+            super().tearDownClass()

python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import time
+from collections import deque
+from typing import Any, Callable
+
+import zmq
+
+from sglang.srt.managers.io_struct import (
+    BatchTokenizedEmbeddingReqInput,
+    BatchTokenizedGenerateReqInput,
+    TokenizedEmbeddingReqInput,
+    TokenizedGenerateReqInput,
+    sock_recv,
+)
+
+_WORK_REQ_TYPES = (
+    TokenizedGenerateReqInput,
+    TokenizedEmbeddingReqInput,
+    BatchTokenizedGenerateReqInput,
+    BatchTokenizedEmbeddingReqInput,
+)
+
+
+class ScriptedTokenizerRecvProxy:
+
+    def __init__(self, *, underlying: zmq.Socket) -> None:
+        self._underlying = underlying
+        self._buffer: deque = deque()
+        self.work_reqs_seen: int = 0
+
+    def recv_pyobj(self, flags: int = 0) -> Any:
+        self._drain_underlying()
+
+        if self._buffer:
+            return self._buffer.popleft()
+
+        if flags & zmq.NOBLOCK:
+            raise zmq.ZMQError(zmq.EAGAIN, "Resource temporarily unavailable")
+        raise RuntimeError(
+            "ScriptedTokenizerRecvProxy.recv_pyobj: blocking recv is not supported"
+        )
+
+    def recv(self, flags: int = 0) -> bytes:
+        raise NotImplementedError(
+            "TODO: support ScriptedTokenizerRecvProxy.recv for msgpack IPC"
+        )
+
+    def wait_until_arrived(
+        self,
+        predicate: Callable[[Any], bool],
+        *,
+        timeout_s: float,
+        description: str = "matching object",
+    ) -> None:
+        start_len = len(self._buffer)
+        deadline = time.monotonic() + timeout_s
+        while True:
+            self._drain_underlying()
+            for i, obj in enumerate(self._buffer):
+                if i >= start_len and predicate(obj):
+                    return
+            if time.monotonic() >= deadline:
+                raise TimeoutError(
+                    f"ScriptedTokenizerRecvProxy: no {description} arrived on the "
+                    f"recv_from_tokenizer socket within {timeout_s}s"
+                )
+            time.sleep(0.005)
+
+    def _drain_underlying(self) -> None:
+        while True:
+            try:
+                req = sock_recv(self._underlying, zmq.NOBLOCK)
+            except zmq.ZMQError:
+                break
+            if isinstance(req, _WORK_REQ_TYPES):
+                self.work_reqs_seen += 1
+            self._buffer.append(req)

python/sglang/test/scripted_runtime/utils.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import importlib
+import sys
+from typing import Callable, Optional
+
+import zmq
+
+
+def close_zmq_socket(socket: zmq.Socket, ctx: zmq.Context) -> None:
+    socket.setsockopt(zmq.LINGER, 0)
+    socket.close()
+    ctx.term()
+
+
+def ensure_script_importable(sys_path_entry: Optional[str]) -> None:
+    if sys_path_entry and sys_path_entry not in sys.path:
+        sys.path.insert(0, sys_path_entry)
+
+
+def resolve_fn(qualified: str) -> Callable:
+    module_name, sep, fn_name = qualified.partition(":")
+    if not sep or not module_name or not fn_name:
+        raise ValueError(
+            f"scripted-runtime fn path must be 'module.path:function_name', "
+            f"got {qualified!r}"
+        )
+    obj = importlib.import_module(module_name)
+    for part in fn_name.split("."):
+        obj = getattr(obj, part)
+    if not callable(obj):
+        raise TypeError(f"resolved object is not callable: {qualified!r} -> {obj!r}")
+    return obj