Jupyter Notebook API调用WebSocket报错：socket已关闭的原因与解决

问题场景 #

通过Jupyter Notebook API创建会话并执行代码时，调用ws.recv()抛出错误：

raise WebSocketConnectionClosedException("socket is already closed.")

执行代码片段：

url = base + '/api/sessions'
params = '{"path":"%s","type":"notebook","name":"","kernel":{"id":null,"name":"env37"}}' % file_name
response = requests.post(url, headers=headers, data=params)
session = json.loads(response.text)
kernel = session["kernel"]

# 读取notebook文件，并获取每个Cell里的Code
url = base + '/api/contents' + notebook_path
response = requests.get(url, headers=headers)
file = json.loads(response.text)
code = [c['source'] for c in file['content']['cells'] if len(c['source']) > 0]
ws = create_connection("ws://127.0.0.1:8888/api/kernels/" + kernel["id"] + "/channels?session_id" + session["id"],
                       header=headers)
for c in code:
    ws.send(json.dumps(send_execute_request(c)))

# 我们只拿Code执行完的消息结果，其他消息将被忽略
for i in range(0, len(code)):
    try:
        msg_type = ''
        while True:
            rsp = json.loads(ws.recv())
            msg_type = rsp["msg_type"]

Jupyter命令行日志：

[I 18:04:05.904 NotebookApp] Kernel started: 275c3afd-cc10-4a69-8597-9f0d7f3e3a91, name: env37
[W 18:04:05.913 NotebookApp] Notebook example2.ipynb is not trusted
[W 18:04:05.917 NotebookApp] No session ID specified
[W 18:04:07.473 NotebookApp] No channel specified, assuming shell: {'header': {'msg_id': '9f4ce706980c11eebfe64ed5776c682d', 'username': 'test', 'session': '9f4cfa28980c11ee92b64ed5776c682d', 'data': '2023-12-11T18:04:07.471569', 'msg_type': 'execute_request', 'version': '5.0'}, 'parent_header': {'msg_id': '9f4ce706980c11eebfe64ed5776c682d', 'username': 'test', 'session': '9f4cfa28980c11ee92b64ed5776c682d', 'data': '2023-12-11T18:04:07.471569', 'msg_type': 'execute_request', 'version': '5.0'}, 'metadata': {}, 'content': {'code': 'from resync import resync', 'silent': False}}
[W 18:04:07.474 NotebookApp] No channel specified, assuming shell: {'header': {'msg_id': '9f4cfa29980c11eea45e4ed5776c682d', 'username': 'test', 'session': '9f4cfa2a980c11eebef84ed5776c682d', 'data': '2023-12-11T18:04:07.471569', 'msg_type': 'execute_request', 'version': '5.0'}, 'parent_header': {'msg_id': '9f4cfa29980c11eea45e4ed5776c682d', 'username': 'test', 'session': '9f4cfa2a980c11eebef84ed5776c682d', 'data': '2023-12-11T18:04:07.471569', 'msg_type': 'execute_request', 'version': '5.0'}, 'metadata': {}, 'content': {'code': '...', 'silent': False}}
[I 18:04:07.481 NotebookApp] Starting buffering for 275c3afd-cc10-4a69-8597-9f0d7f3e3a91:016c7619-9a09e6bff5dcdcab49729795

关闭内核时日志：

[I 18:04:51.838 NotebookApp] Discarding 10 buffered messages for 275c3afd-cc10-4a69-8597-9f0d7f3e3a91:016c7619-9a09e6bff5dcdcab49729795
[I 18:04:51.838 NotebookApp] Kernel shutdown: 275c3afd-cc10-4a69-8597-9f0d7f3e3a91

问题根源分析 #

从日志和代码可定位核心问题：

1. WebSocket连接URL的session_id参数格式错误，导致服务端无法识别会话ID，触发No session ID specified警告，服务端可能提前关闭连接。
2. 发送的execute_request消息未指定channel，服务端默认使用shell通道，但消息结构不规范导致消息被缓冲，客户端无法接收回复，最终连接超时关闭。
3. 消息接收逻辑不合理，一次性发送所有代码后批量接收，可能导致内核处理消息时连接已因超时关闭。

解决步骤 #

1. 修复WebSocket连接的URL参数 #

原代码中session_id参数拼接错误，需用?session_id={session_id}的键值对格式，修改后代码：

ws_url = f"ws://127.0.0.1:8888/api/kernels/{kernel['id']}/channels?session_id={session['id']}"
ws = create_connection(ws_url, header=headers)

2. 规范execute_request消息结构 #

确保消息包含正确的session ID、channel标识，符合Jupyter协议规范。示例send_execute_request函数：

import uuid

def send_execute_request(code, session_id):
    msg_id = str(uuid.uuid4())
    return {
        "header": {
            "msg_id": msg_id,
            "msg_type": "execute_request",
            "username": "test",
            "session": session_id,
            "version": "5.0"
        },
        "parent_header": {},
        "metadata": {},
        "content": {
            "code": code,
            "silent": False,
            "store_history": True,
            "user_expressions": {},
            "allow_stdin": False
        },
        "channel": "shell"
    }

调用时传入正确的session ID：

for c in code:
    ws.send(json.dumps(send_execute_request(c, session["id"])))

3. 优化消息接收逻辑 #

改为发送一个代码块后，等待该代码块的执行回复（execute_reply）再发送下一个，同时处理其他类型的消息避免缓冲：

for c in code:
    ws.send(json.dumps(send_execute_request(c, session["id"])))
    # 等待当前代码块的执行结果
    while True:
        rsp = json.loads(ws.recv())
        msg_type = rsp["msg_type"]
        if msg_type == "execute_reply":
            # 处理执行结果
            print(f"代码块执行结果: {rsp['content']}")
            break
        # 处理其他输出消息
        elif msg_type == "stream":
            print(f"输出: {rsp['content']['text']}")
        elif msg_type in ["display_data", "execute_result"]:
            print(f"返回结果: {rsp['content']}")

4. 标记Notebook为信任 #

解决日志中Notebook example2.ipynb is not trusted的警告，通过API标记信任：

url = base + '/api/contents' + notebook_path
requests.put(
    url,
    headers=headers,
    json={
        "content": file["content"],
        "type": "notebook",
        "trusted": True
    }
)

5. 设置WebSocket超时与心跳 #

避免长时间无交互导致连接关闭，创建连接时设置超时，并在接收循环中处理超时发送心跳：

from websocket import create_connection, WebSocketTimeoutException

ws = create_connection(ws_url, header=headers, timeout=60)  # 设置60秒超时

try:
    for c in code:
        ws.send(json.dumps(send_execute_request(c, session["id"])))
        while True:
            try:
                rsp = json.loads(ws.recv())
                # 处理消息逻辑
                if rsp["msg_type"] == "execute_reply":
                    break
            except WebSocketTimeoutException:
                # 发送心跳保持连接
                ws.send(json.dumps({"msg_type": "ping"}))
finally:
    ws.close()

内容的提问来源于stack exchange，提问作者simon