From f2989cd48b1bac8e4af12fd89d8d2c5544ee0195 Mon Sep 17 00:00:00 2001 From: wuyang6 Date: Wed, 13 May 2026 18:15:53 +0800 Subject: [PATCH] Stream Jupyter workspace file downloads --- backend/api/server.py | 71 +++++++++++++++++-- .../assistant-ui/activity-panel.tsx | 2 + src/jupyter_runtime.py | 44 ++++++++++++ 3 files changed, 111 insertions(+), 6 deletions(-) diff --git a/backend/api/server.py b/backend/api/server.py index 00f2fb7..dd0daf8 100644 --- a/backend/api/server.py +++ b/backend/api/server.py @@ -83,6 +83,9 @@ FEISHU_SPREADSHEET_MAX_CELL_CHARS = 5000 FEISHU_SPREADSHEET_WRITE_BATCH_ROWS = 500 FEISHU_ONLINE_DOCS_FILENAME = 'online-docs.json' JUPYTER_WORKSPACE_FILENAME = 'jupyter_workspace.json' +JUPYTER_STREAM_CHUNK_BYTES = 1024 * 1024 +FEISHU_REMOTE_DOC_MAX_BYTES = 50 * 1024 * 1024 +FEISHU_REMOTE_SHEET_MAX_BYTES = 100 * 1024 * 1024 @dataclass @@ -1096,11 +1099,20 @@ def create_app(state: AgentState) -> FastAPI: if runtime is None: raise HTTPException(status_code=404, detail='Jupyter runtime is not connected') try: - data, filename = runtime.read_file_bytes(path) + response, filename = runtime.open_file_stream(path) except JupyterRuntimeError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc - return Response( - data, + + def stream_chunks() -> Any: + try: + for chunk in response.iter_content(chunk_size=JUPYTER_STREAM_CHUNK_BYTES): + if chunk: + yield chunk + finally: + response.close() + + return StreamingResponse( + stream_chunks(), media_type=_content_type_for_filename(filename), headers={ 'content-disposition': ( @@ -3050,9 +3062,27 @@ def _materialize_jupyter_file_for_feishu( if runtime is None: raise HTTPException(status_code=404, detail='Jupyter runtime is not connected') try: - data, filename = runtime.read_file_bytes(remote_path) + info = runtime.file_info(remote_path) + filename = _safe_uploaded_filename(str(info.get('name') or 'remote-file')) + suffix = Path(filename).suffix.lower() + max_bytes = ( + FEISHU_REMOTE_SHEET_MAX_BYTES + if suffix in FEISHU_SPREADSHEET_SUFFIXES + else FEISHU_REMOTE_DOC_MAX_BYTES + ) + size = info.get('size') + if isinstance(size, int) and size > max_bytes: + raise HTTPException( + status_code=413, + detail=( + f'远端文件过大，当前在线文档转换上限为 ' + f'{_format_bytes(max_bytes)}，请先在 Jupyter 侧裁剪或抽样后再转换。' + ), + ) + response, stream_filename = runtime.open_file_stream(remote_path) except JupyterRuntimeError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc + filename = filename or _safe_uploaded_filename(stream_filename) target_dir = ( state.account_paths(account_id)['base'] / 'integrations' @@ -3061,8 +3091,27 @@ def _materialize_jupyter_file_for_feishu( / session_id ) target_dir.mkdir(parents=True, exist_ok=True) - target = target_dir / _safe_uploaded_filename(filename) - target.write_bytes(data) + target = target_dir / filename + written = 0 + try: + with target.open('wb') as handle: + for chunk in response.iter_content(chunk_size=JUPYTER_STREAM_CHUNK_BYTES): + if not chunk: + continue + written += len(chunk) + if written > max_bytes: + handle.close() + target.unlink(missing_ok=True) + raise HTTPException( + status_code=413, + detail=( + f'远端文件过大，当前在线文档转换上限为 ' + f'{_format_bytes(max_bytes)}，请先在 Jupyter 侧裁剪或抽样后再转换。' + ), + ) + handle.write(chunk) + finally: + response.close() return target @@ -3080,6 +3129,16 @@ def _safe_uploaded_filename(filename: str) -> str: return cleaned[:180] or 'remote-file' +def _format_bytes(value: int) -> str: + if value >= 1024 * 1024 * 1024: + return f'{value / (1024 * 1024 * 1024):.1f}GB' + if value >= 1024 * 1024: + return f'{value / (1024 * 1024):.0f}MB' + if value >= 1024: + return f'{value / 1024:.0f}KB' + return f'{value}B' + + def _resolve_account_file(state: AgentState, account_id: str | None, raw_path: str) -> Path: path = Path(raw_path).expanduser().resolve() account_base = state.account_paths(account_id)['base'].resolve() diff --git a/frontend/app/components/assistant-ui/activity-panel.tsx b/frontend/app/components/assistant-ui/activity-panel.tsx index 4c37ca0..1271ede 100644 --- a/frontend/app/components/assistant-ui/activity-panel.tsx +++ b/frontend/app/components/assistant-ui/activity-panel.tsx @@ -282,6 +282,7 @@ type SessionFile = { name: string; path: string; kind: "input" | "output"; + source?: "local" | "jupyter"; size: number; modified_at: string; download_url: string; @@ -593,6 +594,7 @@ function SessionFileRow({ file }: { file: SessionFile }) { )}
+ {file.source === "jupyter" ? "JUPYTER · " : ""} {fileExtension(file.name).toUpperCase() || "FILE"} ·{" "} {formatBytes(file.size)}
diff --git a/src/jupyter_runtime.py b/src/jupyter_runtime.py index f1eb7c1..d5420a5 100644 --- a/src/jupyter_runtime.py +++ b/src/jupyter_runtime.py @@ -647,6 +647,50 @@ print(json.dumps(entries, ensure_ascii=False)) name = str(payload.get('name') or PurePosixPath(target).name or 'remote-file') return data, name + def file_info(self, path: str) -> dict[str, Any]: + target = self.resolve_workspace_path(path) + api_path = api_path_for_absolute_path(target) + url = f'{self.binding.base_url}/api/contents/{quote(api_path, safe="/")}' + response = self._session.get( + url, + params={'content': 0}, + headers={'X-XSRFToken': self._xsrf_token}, + timeout=30, + ) + if response.status_code >= 400: + raise JupyterRuntimeError( + f'Unable to inspect remote file: HTTP {response.status_code} {response.text[:500]}' + ) + payload = response.json() + if payload.get('type') != 'file': + raise JupyterRuntimeError('Remote path is not a file.') + name = str(payload.get('name') or PurePosixPath(target).name or 'remote-file') + size = payload.get('size') + return { + 'name': name, + 'path': target, + 'size': int(size) if isinstance(size, int) and size >= 0 else None, + } + + def open_file_stream(self, path: str) -> tuple[Any, str]: + """打开远端文件流，由调用方负责关闭 response。""" + + target = self.resolve_workspace_path(path) + api_path = api_path_for_absolute_path(target) + url = f'{self.binding.base_url}/files/{quote(api_path, safe="/")}' + response = self._session.get( + url, + headers={'X-XSRFToken': self._xsrf_token}, + stream=True, + timeout=(10, 300), + ) + if response.status_code >= 400: + response.close() + raise JupyterRuntimeError( + f'Unable to stream remote file: HTTP {response.status_code} {response.text[:500]}' + ) + return response, str(PurePosixPath(target).name or 'remote-file') + def read_text( self, path: str,