跳到主要内容

数据截至 (上游 commit 771b20c93db8)

内核桥:Jupyter 消息协议与流式执行(全项目最深)

30 秒导读: 沙箱里真正跑代码的是一个 Jupyter 内核。这一章讲的是夹在「FastAPI 路由」和「内核」之间的那座桥——messaging.py。它把用户的一次 execute(code) 翻译成 Jupyter 规定的 execute_request 消息发进内核,再把内核异步、乱序、分很多条吐回来的消息,重组成一条有序的输出流;还要在客户端悄悄断线时打断内核、在连接掉线时给在途执行补个「结束」信号防止永久卡死。全项目工程含量最高的一章。

这一章接在 02-server-routing.md(FastAPI 路由把请求交给哪个 WebSocket)之后。路由层做的最后一件事,就是把 ws.execute(...) 这个异步生成器塞进流式响应里(template/server/main.py:121-127)。本章讲的就是这个生成器内部的一切。

1. 这是什么(零基础也能懂)​

先建立一个心智模型​

想象你有一个「永远开着的 Python 终端」(其实是 Jupyter 内核)。你不能直接摸它,只能隔着一根网线跟它喊话。这根网线就是一条 WebSocket 长连接,喊话的「语言」是一套叫 Jupyter Wire Protocol(Jupyter 线协议) 的规矩:你发一条格式固定的 JSON 说「请执行这段代码」,内核就会陆陆续续回你好几条 JSON——"我开始忙了"、"这是一行 stdout"、"这是一张图"、"这是最终结果"、"我空闲了"。

messaging.py 干的活,就是把这套一来一回的喊话,包装成一个干净的 Python 异步生成器:

# 示意,非源码:调用方眼中的样子
async for chunk in ws.execute(code, env_vars, token):
    # chunk 已经是结构化的 Stdout / Result / Error / EndOfExecution
    send_to_client(chunk)

调用方完全不用知道 Jupyter 协议、不用管消息乱序、不用管连接掉线——桥全帮你抹平了。

它要解决的三个真问题​

真问题桥怎么接
内核回消息是异步 + 乱序 + 一次执行分很多条后台任务 _receive_message 收所有消息,按 parent_msg_id 分发到对应执行的队列;_wait_for_result 从队列有序取出
客户端可能悄悄断线(网络挂了,没说再见)_wait_for_result 每 5 秒发一个 keepalive 逼服务器写 socket,写失败才发现断了,进而打断内核(否则下一次执行被卡住)
WebSocket 掉线/重连时,在途执行会永远收不到结束收消息循环退出时,给所有在途执行补一个 UnexpectedEndOfExecution,防止生成器永久 hang

一句话直觉: 把它想成一个快递分拣中心。内核是四面八方发货的仓库(消息乱序到达),桥按每个包裹上的「订单号」(parent_msg_id)分拣进对应的传送带(队列),客户那头则从传送带上按顺序一件件取货,直到看见「本单发完」的标签(EndOfExecution)。

2. 顶层全景(它大概怎么转)​

一次执行的生命周期​

先给一句「怎么读这张图」:左边是一次 execute() 调用的主流程(串行发送),右边是永远在后台跑的收消息循环,两者靠中间那条 per-execution 队列解耦。

   execute(code)  ← 来自 FastAPI 路由(main.py:122)
        │
        │ async with _lock  ← 串行化:同一时刻只放一次执行进内核
        ▼
  ┌───────────────┐   注入 env vars(按首行缩进对齐)
  │ 组装 complete │   生成 msg_id + Execution(带自己的 queue)
  │   _code       │
  └──────┬────────┘
         │ _get_execute_request → JSON
         │ ws.send(...)  ← 失败则 reconnect 重发,最多 3 次
         ▼
  ┌──────────────────────┐        ┌───────────────────────────┐
  │ _wait_for_result     │        │  _receive_message (后台)   │
  │  while True:         │◀──队列──│  async for msg in ws:      │
  │   queue.get(timeout) │  put()  │    _process_message(msg)   │
  │   超时→yield keepalive│        │      按 msg_type 映射成     │
  │   拿到→yield 结构化    │        │      Stdout/Result/Error…  │
  │   见 EndOfExec→break  │        │      按 parent_msg_id 分发  │
  └──────────┬───────────┘        └────────────▲──────────────┘
             │                                 │
             ▼                          内核 ws://localhost:8888
     yield 给路由 → 客户端        /api/kernels/{id}/channels
             │
    (客户端断开 → CancelledError → shield(interrupt()) 打断内核)
             │
             ▼
   env_vars 清理:另起一个后台执行请求 _cleanup_env_vars

部件一句话职责​

部件(符号)干什么位置
ContextWebSocket一个 context 一条长连接的门面,握着 ws、收消息任务、执行表messaging.py:60
Execution一次执行的信箱:一条 asyncio.Queue + errored/input_accepted 标志messaging.py:42
connect / reconnect建连并起后台收消息任务messaging.py:75-102
execute主流程:串行化、注入 env、发送带重连、流式 yield、断开打断messaging.py:318
_wait_for_result从某次执行的队列有序取输出,带 keepalive 超时messaging.py:254
_receive_message后台读所有内核消息,掉线时给在途执行补收尾messaging.py:422
_process_message把每种 Jupyter msg_type 映射成结构化对象入队messaging.py:445
_get_execute_request按 Jupyter 协议构造 execute_request JSONmessaging.py:120
interrupt走 Jupyter REST 打断内核messaging.py:104
change_current_directory各语言用不同 magic 切工作目录messaging.py:289

关键设计:每个 context 一条长连接。 __init__ 只是把 url 拼成 ws://localhost:8888/api/kernels/{context_id}/channels(messaging.py:70),真正建连在 connect()——一次建连后 _receive_task 就一直在后台跑,后续所有执行复用同一条连接(messaging.py:90-102)。

3. 核心原理(逐个机制,由浅入深)​

3.1 把一次执行翻成 Jupyter 消息:_get_execute_request​

要解决的小问题: 内核只认 Jupyter 线协议的消息格式,不认「跑这段代码」这种大白话。

思路: Jupyter 每条消息都由 header(谁发的、什么类型、唯一 msg_id)、parent_header、metadata、content(真正的负载)组成。执行请求的 msg_type 固定是 execute_request,content.code 放代码。

真实实现(messaging.py:120-149,_get_execute_request):

# 真实源码节选,messaging.py:124-147
"header": {
    "msg_id": msg_id,              # 唯一 id,后续所有回信的 parent_header 都带它
    "session": self.session_id,
    "msg_type": "execute_request",
    "version": "5.3",
    ...
},
"content": {
    "code": code,
    "silent": background,          # 后台执行(如切目录、清 env)设 True,不产 execute_result
    "store_history": True,
    "stop_on_error": True,
    "allow_stdin": False,
},

关键细节——msg_id 是整条流水线的主键。 发出去时这个 msg_id 就是本次执行的身份证;内核回来的每条消息都会在 parent_header.msg_id 里带上它,_process_message 正是靠它把回信分拣到对的 Execution(messaging.py:470-477)。silent=background 这个开关很妙:切目录、清 env 这类「杂活」用 silent=True 发,内核就不会把它们当成用户代码去产生 execute_result 污染输出。

3.2 注入环境变量,还得「对齐缩进」:_indent_code_with_level​

要解决的小问题: 用户想给这次执行加几个环境变量。最直白的做法是在用户代码前面拼一段 os.environ[...] = ...。但对 Python、YAML 那种缩进敏感的语言,如果用户代码本身是缩进的(比如整段被包在某个块里),你在前面贴一行顶格代码,缩进就乱了。

思路: 先看用户代码第一行非空行的缩进是多少,把注入的 env 代码也缩进到同样层级,再拼上去。

用户代码(假设首行有 4 空格缩进):
    result = compute()
    print(result)

注入前先把 env 代码对齐到 4 空格:
    os.environ['API_KEY'] = 'xxx'   ← _indent_code_with_level 加的缩进
    result = compute()
    print(result)

真实实现: _get_code_indentation 取首行非空行的前导空白(messaging.py:208-218),_indent_code_with_level 给每个非空行补上这段缩进(messaging.py:220-234),execute 里把「全局 env + 调用级 env」拼成一段、缩进对齐、再前置到用户代码(messaging.py:354-358):

# 真实源码,messaging.py:354-358
if global_env_vars_snippet or env_vars_snippet:
    indented_env_code = self._indent_code_with_level(
        f"{global_env_vars_snippet}\n{env_vars_snippet}", code_indent
    )
    complete_code = f"{indented_env_code}\n{complete_code}"

两层 env vars 从哪来:

  • 全局 env:第一次执行时懒加载,向 envd 拉取(get_envs,envs.py:10-20;E2B_LOCAL 时返回一个测试变量)。只拉一次,存 _global_env_vars(messaging.py:347-349)。
  • 调用级 env:本次 execute 传进来的 env_vars,每次都注入(messaging.py:351-352)。

各语言的 set/unset 片段由 _set_env_var_snippet / _delete_env_var_snippet 生成(messaging.py:151-177),覆盖 python/js/ts/r/java/bash 五套写法。

3.3 有序取输出 + keepalive 侦测断线:_wait_for_result​

要解决的小问题(两个):

  1. 内核消息是后台任务异步塞进队列的,主流程怎么有序、阻塞式地一条条取出来?
  2. 如果客户端悄悄断线(网络挂了,TCP 层没立刻报错),服务器怎么尽快发现,好去打断那个还在傻跑的内核?

思路: 用一个 while True 循环 await queue.get();但不是死等,而是加一个 KEEPALIVE_INTERVAL(5 秒)超时。超时了就 yield 一个 keepalive,逼 Starlette 往 socket 写一次——如果客户端断了,这次写会失败,uvicorn 随即抛 http.disconnect,把这个生成器 cancel 掉。这是 issue #213 的核心修复思路。

真实实现(messaging.py:254-287,_wait_for_result):

# 真实源码节选,messaging.py:264-287
while True:
    try:
        output = await asyncio.wait_for(queue.get(), timeout=KEEPALIVE_INTERVAL)
    except asyncio.TimeoutError:
        yield {"type": "keepalive"}   # 逼服务器写 socket → 断线才暴露(#213)
        continue

    if output.type == OutputType.END_OF_EXECUTION:
        break                          # 正常收尾

    if output.type == OutputType.UNEXPECTED_END_OF_EXECUTION:
        yield Error(name="UnexpectedEndOfExecution",
                    value="Connection ... closed before ... finished", traceback="")
        break                          # 连接掉了,给个明确错误再退出

为什么 keepalive 能侦测断线? 关键不在「发心跳」本身,而在它强制了一次真正的 socket 写。没有写操作,一个 await queue.get() 会永远静静地等——即便客户端早就没了,内核还在空跑,谁也不知道。周期性写一下,操作系统/uvicorn 才有机会在写失败时上报断连。

3.4 客户端断开就打断内核:asyncio.shield(interrupt())​

要解决的小问题: 3.3 侦测到客户端断了,_wait_for_result 会被 cancel。但内核还在跑那段可能很重的代码——不打断它,下一次执行会被这个 context 的串行锁卡住(见 3.5)。

思路: 在 execute 里包住流式循环,捕获 CancelledError / GeneratorExit,然后调 interrupt()。但有个坑:此刻整个任务正在被 cancel,你新发的「打断内核」HTTP 请求也会被顺手 cancel 掉。所以用 asyncio.shield 把它护住,让它跑完。

真实实现(messaging.py:398-411):

# 真实源码节选,messaging.py:401-411
except (asyncio.CancelledError, GeneratorExit):
    logger.warning(f"Client disconnected during execution ({message_id}), interrupting kernel")
    try:
        await asyncio.shield(self.interrupt())   # 护住,别被外层 cancel 一起带走
    except asyncio.CancelledError:
        pass
    raise                                        # 打断完仍要把 cancel 传播出去

interrupt() 走的是 Jupyter REST API(不是 WebSocket):POST {JUPYTER_BASE_URL}/api/kernels/{id}/interrupt(messaging.py:104-118)。注意最后那个 raise——打断内核只是善后,取消信号本身必须继续往上传,否则 asyncio 会抱怨「cancel 被吞了」。

3.5 串行化 + 发送重连:_lock 与 MAX_RECONNECT_RETRIES​

要解决的小问题: 一条连接、一个内核,同时塞两次执行会乱套(输出交织、状态错乱)。而且发送那一刻连接可能刚好掉了。

串行化: 整个 execute 主体包在 async with self._lock 里(messaging.py:327),保证同一 context 同一时刻只有一次执行在内核里。这条锁还顺带简化了掉线善后的推理(见 3.6:因为有锁,掉线时在途执行「要么没有,要么只有一个」)。

先等上次 env 清理做完: 进锁后第一件事,是等上一次执行留下的 _cleanup_task 跑完(messaging.py:329-336)——因为清理 env 也是一次真正发进内核的执行(见 3.7),不等它,新执行会和清理撞车。

发送带重连重发: 发送放在一个 for i in range(1 + MAX_RECONNECT_RETRIES) 里(共 4 次机会),send 抛 ConnectionClosedError 就 reconnect() 再试;for...else 兜底——四次都失败,直接给队列塞一个 WebSocketError + UnexpectedEndOfExecution,让 _wait_for_result 干净退出而不是 hang(messaging.py:366-393)。

# 真实源码节选,messaging.py:366-393
for i in range(1 + MAX_RECONNECT_RETRIES):
    try:
        request = self._get_execute_request(message_id, complete_code, False)
        await self._ws.send(request)
        break
    except (ConnectionClosedError, WebSocketException) as e:
        if i < MAX_RECONNECT_RETRIES:
            await self.reconnect()
else:
    # 四次都没发出去 → 塞错误 + 结束信号,别让 _wait_for_result 卡死
    await execution.queue.put(Error(name="WebSocketError", ...))
    await execution.queue.put(UnexpectedEndOfExecution())

3.6 掉线不 hang:给在途执行补 UnexpectedEndOfExecution​

要解决的小问题: 后台收消息循环 async for message in self._ws 一旦因连接断开而退出,内核之后的收尾消息(那个决定 EndOfExecution 的 idle 状态)就再也不会来了。_wait_for_result 会永远等下去。

思路: 在收消息循环的 finally 里,遍历所有还在 _executions 里的执行,给每个都塞一个 WebSocketError + UnexpectedEndOfExecution,强制它们结束(messaging.py:432-443):

# 真实源码节选,messaging.py:432-443
finally:
    # 重连中可能丢结果 → 取消所有在途执行,防止无限 hang
    # 因为有锁,在途执行要么没有,要么只有一个
    for key, execution in self._executions.items():
        await execution.queue.put(Error(name="WebSocketError",
            value="The connections was lost, rerun the code to get the results", traceback=""))
        await execution.queue.put(UnexpectedEndOfExecution())

那句注释「因为有锁,要么没有,要么只有一个」把 3.5 的串行锁和这里的善后逻辑串了起来——正是锁让这段补救可以简单地「一把撸完」而不必担心并发的多个在途执行。

3.7 env 清理:一次独立的后台执行请求​

要解决的小问题: 本次注入的 env vars 不能污染后续执行。执行结束后要把它们恢复到全局值或删掉。

思路: 清理不是在同一段代码里做,而是另起一次执行请求(silent=True),在主执行流之后异步发出。

主流程末尾建一个 _cleanup_task(messaging.py:417-420);_cleanup_env_vars 自己造一个新 msg_id + 后台 Execution,发送 _reset_env_vars_code 生成的恢复代码,并用 _wait_for_result 等它跑完(messaging.py:236-252)。

_reset_env_vars_code 的规则很清楚(messaging.py:189-206):某个 key 若在全局 env 里存在,就恢复成全局值;否则删掉。这样调用级临时变量退场、全局变量保持原样。下一次 execute 进锁时会先等这个清理任务(3.5),保证时序正确。

4. 深入实现:_process_message 的消息映射表​

这是桥的「翻译词典」——内核吐回的每种 msg_type 被翻成哪个结构化对象、塞进队列(messaging.py:445-588)。调用方后面拿到的每个 chunk,都是这里造出来的。

Jupyter msg_type / 状态映射成说明行
status + restartingError(ContextRestarting) + EndOfExecution内核重启,给所有在途执行发结束454-468
errorError(ename, evalue, traceback)运行时异常;errored 去重防重复480-495
stream name=stdoutStdout(text, timestamp)标准输出497-506
stream name=stderrStderr(text, timestamp)标准错误508-516
display_dataResult(is_main_result=False, ...)富显示(图/HTML),非主结果518-523
execute_resultResult(is_main_result=True, ...)cell 的主返回值525-530
status busy(且后台)置 input_accepted=True后台执行没有 execute_input,靠 busy 标记已开始532-535
status idleEndOfExecution(若已 accepted)执行真正结束的信号537-540
status errorError + EndOfExecution542-551
execute_reply status=error/abort/okError / Error(Aborted)+End / 跳过回执553-579
execute_inputNumberOfExecutions(execution_count) + input_accepted=True内核已接收输入,回带执行序号581-586
其它只 logger.warning未处理类型587-588

几个要点:

  • is_main_result 区分主结果与显示数据。 execute_result 是 cell 的那个「返回值」(主结果,is_main_result=True);display_data 是过程中 display() 之类调出来的富内容(可以有很多个,is_main_result=False)。二者都用同一个 Result 模型,只靠这个布尔位区分(messaging.py:518-530)。富结果内部怎么从 data 的 MIME 字典抽出 png/html/chart,见 04-rich-results.md。

  • idle 才是「结束」的权威信号。 不是 execute_reply、不是 error,而是内核状态回到 idle 且此前 input_accepted 为真,才 put EndOfExecution(messaging.py:537-540)。这也是为什么后台执行要靠 busy 来补 input_accepted——它们收不到 execute_input。

  • 分拣靠 parent_header.msg_id。 每条消息先取 parent_msg_ig(源码里的拼写),查不到对应 Execution 就直接丢弃(messaging.py:470-477)——这就是「多次执行共用一条连接却不串味」的机制。

Result 之外的输出类型定义在 api/models/output.py:OutputType 枚举、EndOfExecution、UnexpectedEndOfExecution、NumberOfExecutions(output.py:7-31);Stdout/Stderr 在 api/models/logs.py,Error 在 api/models/error.py。

change_current_directory:各语言的切目录 magic​

context 创建时要切到指定 cwd(由 contexts.py:65 调用)。不同语言机制不同(messaging.py:289-316):

语言用什么备注
python%cd {path}IPython magic
javascript / typescriptprocess.chdir('{path}')
rsetwd('{path}')
javaSystem.setProperty("user.dir", ...)并不真的切目录,只设属性(源码有注释,messaging.py:302)
bashcd '{path}'

都以 silent=True 后台发送,并用 _wait_for_result 等结果;出错则抛 ExecutionError(errors.py:1-2,messaging.py:314-316)。

5. 巧妙之处(可借鉴的技术)​

  • 用「强制 socket 写」侦测静默断线。 阻塞在 queue.get() 上永远发现不了客户端消失;_wait_for_result 用带超时的 wait_for + yield keepalive,把「侦测断线」这个网络层难题,转化为「周期性触发一次写」——写失败自然暴露断连(messaging.py:264-273,#213)。这套「靠 keepalive 逼写来暴露死连接」的思路,任何长连接流式接口都能借。

  • asyncio.shield 让善后逻辑在 cancel 洪流中活下来。 任务被 cancel 时,你在 except CancelledError 里发的补救请求也会被 cancel。shield(interrupt()) 护住它跑完,再 raise 把 cancel 传出去——「打断内核」和「尊重取消语义」两不误(messaging.py:408-411)。

  • 每条掉线路径都保证发「结束信号」。 hang 的根源是消费者永远等不到结束。这里三处都堵死了:发送四次失败(else 分支)、收消息循环退出(finally)、内核重启(restarting)——统统 put 一个终结信号(UnexpectedEndOfExecution 或 EndOfExecution),_wait_for_result 一定能退出。

  • 缩进感知的 env 注入。 为了在缩进敏感语言里安全地「前置」env 代码,先探测用户代码首行缩进再对齐注入(_get_code_indentation + _indent_code_with_level,messaging.py:208-234)。一个容易被忽略、但不做就出错的细节。

  • env 清理作为独立后台执行 + 下次执行前等待。 把「恢复环境」从主执行里拆出去成一次独立的 silent 执行,再用 _lock 里的「等 _cleanup_task」保证时序(messaging.py:329-336,417-420)。职责干净、互不污染。

6. 边界与局限​

  • 一 context 一内核,执行严格串行。 _lock 保证同一 context 不并发执行——想并行得开多个 context。这是刻意的简化,也换来了掉线善后逻辑的简单(3.6)。

  • Java 的「切目录」是假的。 只设 user.dir 系统属性,不真的改进程工作目录(源码明确注释,messaging.py:302)。

  • 重连不重放。 掉线重连后,在途执行的结果已经丢了——桥不会重放,只会给它一个 UnexpectedEndOfExecution 让你「重跑代码拿结果」(错误文案就是这么写的,messaging.py:439)。

  • env 注入是「拼字符串进代码」。 value 直接内插进 os.environ['k']='v' 这类片段(messaging.py:151-163),未见转义处理——含引号的 value 可能破坏语法(inferred,基于片段构造方式)。

  • silent=True 的后台执行不产 execute_input。 所以要靠 busy 状态补 input_accepted,否则永远等不到 idle 判定的结束(messaging.py:532-535)。

7. 横向对比​

同 shelf 的 cloudflare sandbox-sdk 等沙箱执行方案,大多也要解决「流式取输出 + 断线善后」这对孪生难题,但 E2B 这一支的独特取舍是:直接复用 Jupyter 内核与线协议,于是「富结果(图/DataFrame/HTML)」几乎白得(内核已按 MIME 吐好),代价是绑定 Jupyter 那套 msg_type/状态机语义。本章的重组逻辑,本质是把 Jupyter 面向 notebook 的异步消息模型,压平成一条面向 HTTP 流的线性输出。上层如何把这条流经 HTTP 发给客户端、客户端 SDK 又如何解析,见 01-sdk-client.md;内核本身怎么被装进沙箱镜像,见 05-template-runtime.md。

8. 代码地图(导航索引)​

主题文件符号
WebSocket 门面 / 长连接template/server/messaging.pyContextWebSocket, connect, reconnect
一次执行的信箱template/server/messaging.pyExecution
主执行流程(串行/注入/重连/断开打断)template/server/messaging.pyexecute
有序取输出 + keepalivetemplate/server/messaging.py_wait_for_result, KEEPALIVE_INTERVAL
后台收消息 / 掉线善后template/server/messaging.py_receive_message
Jupyter 消息 → 结构化对象template/server/messaging.py_process_message
构造 execute_requesttemplate/server/messaging.py_get_execute_request
打断内核(REST)template/server/messaging.pyinterrupt
env 注入 / 清理 / 缩进对齐template/server/messaging.py_set_env_vars_code, _reset_env_vars_code, _indent_code_with_level, _cleanup_env_vars, _get_code_indentation
切工作目录(各语言 magic)template/server/messaging.pychange_current_directory
重连重发次数上限template/server/messaging.pyMAX_RECONNECT_RETRIES
输出/结束/序号类型template/server/api/models/output.pyOutputType, EndOfExecution, UnexpectedEndOfExecution, NumberOfExecutions
执行错误异常template/server/errors.pyExecutionError
拉取全局 envtemplate/server/envs.pyget_envs
内核 REST 基址template/server/consts.pyJUPYTER_BASE_URL
路由把 execute() 塞进流式响应template/server/main.pyStreamingListJsonResponse, ws.execute (main.py:121-127)