From 2274cf52a64c11af6e35dc79e3fe35cb9e5416d5 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Fri, 10 Jul 2026 10:18:09 -0500 Subject: [PATCH 1/6] [Frontend] Abort engine request on client disconnect (free leaked KV) Non-streaming API handlers were not cancelled when the client hung up (Starlette only cancels StreamingResponse, not plain handlers), so the engine kept generating and the sequence's KV blocks leaked until it hit max_tokens. Add a client-disconnect abort path: - engine: EngineCoreMgr.abort_request broadcasts an "abort_request" utility command; _handle_abort_request marks the seq aborted; the scheduler finishes it at the next step via the normal stop path (frees KV, emits a finished RequestOutput). Adds Sequence.aborted. - api_server: _run_nonstream_with_disconnect runs generate_async in a task and polls request.is_disconnected(); on disconnect it cancels the task, whose teardown aborts + pops the request. Wired into /v1/chat/completions and /v1/completions (return HTTP 499 on disconnect). generate_async / generate_async_multimodal / generate_async_fanout / cleanup_streaming_request abort + pop on early exit to avoid leaks. Co-Authored-By: Claude Opus 4.8 (1M context) --- atom/entrypoints/openai/api_server.py | 227 ++++++++++++++++++++------ atom/model_engine/engine_core_mgr.py | 12 ++ atom/model_engine/engine_utility.py | 14 ++ atom/model_engine/scheduler.py | 5 + atom/model_engine/sequence.py | 3 + 5 files changed, 209 insertions(+), 52 deletions(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 6ff565249c..e9aec50035 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -440,17 +440,36 @@ def do_preprocess(): raise engine.core_mgr.add_request([seq]) - while True: - item = await token_queue.get() - token_ids = item.get("token_ids") or [] - if token_ids: - if first_token_at is None: - first_token_at = item.get("ts", time.time()) - last_token_at = item.get("ts", time.time()) - all_token_ids.extend(token_ids) - if item.get("finished", False): - finish_reason = item.get("finish_reason") - break + _finished_ok = False + try: + while True: + item = await token_queue.get() + token_ids = item.get("token_ids") or [] + if token_ids: + if first_token_at is None: + first_token_at = item.get("ts", time.time()) + last_token_at = item.get("ts", time.time()) + all_token_ids.extend(token_ids) + if item.get("finished", False): + finish_reason = item.get("finish_reason") + _finished_ok = True + break + finally: + # Two responsibilities, on EVERY exit path: + # 1) If we didn't finish (client disconnected / cancelled), tell the + # engine to stop so the seq doesn't run to max_tokens and burn GPU. + # 2) Always drop the seq from io_processor.requests. The engine frees + # its own KV on finish, but this dict is only cleaned up here for + # non-stream requests -- without an unconditional pop, every + # completed non-stream request leaks a Sequence (pending grows + # forever). Streaming pops via cleanup_streaming_request instead. + if seq is not None: + if not _finished_ok: + try: + engine.core_mgr.abort_request(seq.id) + except Exception: + pass + engine.io_processor.requests.pop(seq.id, None) text = tokenizer.decode(all_token_ids, skip_special_tokens=True) num_tokens_input = ( @@ -531,17 +550,29 @@ def do_preprocess(): raise engine.core_mgr.add_request([seq]) - while True: - item = await token_queue.get() - token_ids_out = item.get("token_ids") or [] - if token_ids_out: - if first_token_at is None: - first_token_at = item.get("ts", time.time()) - last_token_at = item.get("ts", time.time()) - all_token_ids.extend(token_ids_out) - if item.get("finished", False): - finish_reason = item.get("finish_reason") - break + _finished_ok = False + try: + while True: + item = await token_queue.get() + token_ids_out = item.get("token_ids") or [] + if token_ids_out: + if first_token_at is None: + first_token_at = item.get("ts", time.time()) + last_token_at = item.get("ts", time.time()) + all_token_ids.extend(token_ids_out) + if item.get("finished", False): + finish_reason = item.get("finish_reason") + _finished_ok = True + break + finally: + # See generate_async: abort on early exit, always pop to avoid leak. + if seq is not None: + if not _finished_ok: + try: + engine.core_mgr.abort_request(seq.id) + except Exception: + pass + engine.io_processor.requests.pop(seq.id, None) text = tokenizer.decode(all_token_ids, skip_special_tokens=True) num_tokens_output = len(all_token_ids) @@ -647,19 +678,31 @@ def do_preprocess(): engine.core_mgr.add_request(seqs) num_tokens_input = seqs[0].num_prompt_tokens - while not all(finished): - idx, item = await shared_queue.get() - if finished[idx]: - continue - tokens = item.get("token_ids") or [] - if tokens: - if per_first_token_at[idx] is None: - per_first_token_at[idx] = item.get("ts", time.time()) - per_last_token_at[idx] = item.get("ts", time.time()) - per_tokens[idx].extend(tokens) - if item.get("finished", False): - per_finish_reason[idx] = item.get("finish_reason") - finished[idx] = True + _all_finished = False + try: + while not all(finished): + idx, item = await shared_queue.get() + if finished[idx]: + continue + tokens = item.get("token_ids") or [] + if tokens: + if per_first_token_at[idx] is None: + per_first_token_at[idx] = item.get("ts", time.time()) + per_last_token_at[idx] = item.get("ts", time.time()) + per_tokens[idx].extend(tokens) + if item.get("finished", False): + per_finish_reason[idx] = item.get("finish_reason") + finished[idx] = True + _all_finished = True + finally: + # Abort any sibling still running on early exit; always pop all seqs. + for _seq in seqs: + if not _all_finished: + try: + engine.core_mgr.abort_request(_seq.id) + except Exception: + pass + engine.io_processor.requests.pop(_seq.id, None) finished_at = time.time() outputs: List[Dict[str, Any]] = [] @@ -779,9 +822,79 @@ def cleanup_streaming_request(request_id: str, seq_id: int) -> None: _seq_id_to_request_id.pop(seq_id, None) _stream_loops.pop(request_id, None) _request_start_times.pop(request_id, None) + # If the stream ended early (client disconnected) the seq may still be + # generating in the engine core -> tell it to stop so it doesn't run to + # max_tokens and pile up. No-op if the seq already finished. + try: + engine.core_mgr.abort_request(seq_id) + except Exception: + pass engine.io_processor.requests.pop(seq_id, None) + + +class _ClientDisconnected(Exception): + """Raised when a non-streaming client hangs up mid-generation.""" + + def __init__(self, request_id: str): + super().__init__(request_id) + self.request_id = request_id + + +async def _run_nonstream_with_disconnect(agen, raw_request, request_id): + """Drive a non-stream ``generate_async*`` async-generator while actively + watching for client disconnect. + + Starlette does NOT cancel a *non-streaming* request handler when the client + goes away (unlike StreamingResponse, which is cancelled on http.disconnect). + Without this, an abandoned non-stream request keeps ``await``-ing the engine + until it hits ``max_tokens`` -- burning GPU on output nobody will read AND + leaking the seq in ``io_processor.requests`` (its finally never fires). + + We run the generator in a task and poll ``raw_request.is_disconnected()`` + concurrently. On disconnect we cancel the task; the cancellation propagates + into ``generate_async``'s ``await token_queue.get()`` so its ``finally`` + runs -> ``abort_request`` + ``io_processor.requests.pop``. Returns the last + yielded output, or raises ``_ClientDisconnected``. + """ + final_output = None + + async def _collect(): + nonlocal final_output + async for output in agen: + final_output = output + return final_output + + task = asyncio.ensure_future(_collect()) + try: + while True: + done, _ = await asyncio.wait({task}, timeout=0.5) + if task in done: + return task.result() + disconnected = False + if raw_request is not None: + try: + disconnected = await raw_request.is_disconnected() + except Exception: + disconnected = False + if disconnected: + logger.info( + f"Client disconnected (non-stream), aborting request " + f"{request_id}" + ) + raise _ClientDisconnected(request_id) + finally: + if not task.done(): + task.cancel() + try: + await task + except BaseException: + # CancelledError (expected) or any error from generator + # teardown; generate_async's finally already aborted + popped. + pass + + async def setup_streaming_request_fanout( prompt_or_tokens: str | List[int], sampling_params: SamplingParams, @@ -908,7 +1021,7 @@ async def general_error_handler(request: Request, exc: Exception): @app.post("/v1/chat/completions") -async def chat_completions(request: ChatCompletionRequest): +async def chat_completions(request: ChatCompletionRequest, raw_request: Request): """Handle chat completion requests (OpenAI-compatible).""" global engine, tokenizer, model_name @@ -1046,14 +1159,16 @@ async def chat_completions(request: ChatCompletionRequest): request_id, model_name, outputs, tools=request.tools ) else: - final_output = None - async for output in generate_async( - prompt, - sampling_params, + final_output = await _run_nonstream_with_disconnect( + generate_async( + prompt, + sampling_params, + request_id, + kv_transfer_params=request.kv_transfer_params, + ), + raw_request, request_id, - kv_transfer_params=request.kv_transfer_params, - ): - final_output = output + ) if final_output is None: raise RuntimeError("No output generated") resp = build_chat_response( @@ -1066,6 +1181,9 @@ async def chat_completions(request: ChatCompletionRequest): _log_request_event("response", request_id, resp.model_dump()) return resp + except _ClientDisconnected: + # Client hung up; seq already aborted + popped. Nothing to return. + return JSONResponse(status_code=499, content={"detail": "client disconnected"}) except ValueError as e: logger.error(f"Validation error in chat_completions: {e}") raise HTTPException(status_code=400, detail=str(e)) @@ -1075,7 +1193,7 @@ async def chat_completions(request: ChatCompletionRequest): @app.post("/v1/completions") -async def completions(request: CompletionRequest): +async def completions(request: CompletionRequest, raw_request: Request): """Handle text completion requests (OpenAI-compatible).""" global engine, tokenizer, model_name @@ -1149,15 +1267,17 @@ async def completions(request: CompletionRequest): raise RuntimeError("No output generated") resp = build_completion_response_multi(request_id, model_name, outputs) else: - final_output = None - async for output in generate_async( - request.prompt, - sampling_params, + final_output = await _run_nonstream_with_disconnect( + generate_async( + request.prompt, + sampling_params, + request_id, + kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=request.data_parallel_rank, + ), + raw_request, request_id, - kv_transfer_params=request.kv_transfer_params, - data_parallel_rank=request.data_parallel_rank, - ): - final_output = output + ) if final_output is None: raise RuntimeError("No output generated") @@ -1166,6 +1286,9 @@ async def completions(request: CompletionRequest): _log_request_event("response", request_id, resp.model_dump()) return resp + except _ClientDisconnected: + # Client hung up; seq already aborted + popped. Nothing to return. + return JSONResponse(status_code=499, content={"detail": "client disconnected"}) except ValueError as e: logger.error(f"Validation error in completions: {e}") raise HTTPException(status_code=400, detail=str(e)) diff --git a/atom/model_engine/engine_core_mgr.py b/atom/model_engine/engine_core_mgr.py index b254a792bd..eddebc82ce 100644 --- a/atom/model_engine/engine_core_mgr.py +++ b/atom/model_engine/engine_core_mgr.py @@ -460,6 +460,18 @@ def send_utility_command(self, cmd: str, dp_rank: int = None): copy=False, ) + def abort_request(self, req_id): + """Tell the engine core(s) to drop a request (client disconnected). + + Broadcast to every DP rank (only the one holding ``req_id`` acts). The + scheduler finishes the seq at its next step via the normal stop path, + freeing its KV blocks. Fire-and-forget; safe if the seq already finished. + """ + try: + self.broadcast_utility_command("abort_request", req_id=req_id) + except Exception as e: + logger.warning(f"{self.label}: abort_request({req_id}) failed: {e}") + def broadcast_utility_command(self, cmd: str, **kwargs): payload = {"cmd": cmd, **kwargs} # Serialize once and reuse for all ranks (optimization: avoid repeated pickle.dumps) diff --git a/atom/model_engine/engine_utility.py b/atom/model_engine/engine_utility.py index 49fb9219d5..c1aed18b2e 100644 --- a/atom/model_engine/engine_utility.py +++ b/atom/model_engine/engine_utility.py @@ -41,6 +41,7 @@ class EngineUtilityHandler: "stop_profile": "_handle_stop_profile", "get_mtp_stats": "_handle_get_mtp_stats", "get_mtp_statistics": "_handle_get_mtp_statistics", + "abort_request": "_handle_abort_request", } def __init__( @@ -210,6 +211,19 @@ def _handle_clear_kv_cache(self, args: dict): ("UTILITY_RESPONSE", {"cmd": "clear_kv_cache", "result": result}) ) + def _handle_abort_request(self, args: dict): + """Mark a sequence aborted (client disconnected) so the scheduler finishes + it at the next step via the normal stop path (frees KV, drops it).""" + req_id = args.get("req_id") if isinstance(args, dict) else None + if req_id is None or self.scheduler is None: + return + found = False + for seq in list(self.scheduler.running) + list(self.scheduler.waiting): + if seq.id == req_id: + seq.aborted = True + found = True + logger.info(f"{self.label}: abort_request req_id={req_id} found={found}") + def _handle_configure_hidden_states(self, args: dict): """Configure hidden states extraction on all model runners (TorchSpec).""" aux_layer_ids = args.get("aux_layer_ids", []) diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 295a76f3e9..166ae79f2d 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -1430,6 +1430,11 @@ def postprocess( num_tokens = seq.num_tokens - self.mtp_k - num_rejected leave_reason = None + # Client disconnected -> finish now via the normal stop path (frees + # KV blocks, emits a finished RequestOutput). A natural stop below + # may still overwrite the reason; either way the seq terminates. + if getattr(seq, "aborted", False): + leave_reason = "aborted" # MTP edge case: `rejection_sampler` does NOT inspect EOS — it # only compares draft vs target_argmax for acceptance. So when # the verified token is EOS the kernel still emits 1+ accepted diff --git a/atom/model_engine/sequence.py b/atom/model_engine/sequence.py index e925823acf..ea73665df1 100644 --- a/atom/model_engine/sequence.py +++ b/atom/model_engine/sequence.py @@ -86,6 +86,9 @@ def __init__( # out-of-window SWA blocks can be freed while compressed blocks persist). # Empty / unused for non-SWA models. self.swa_block_table = [] + # Set True when the client disconnected; the scheduler finishes the seq + # at the next step via the normal stop path (frees KV, emits finished). + self.aborted = False # Per-request cache slot index (filled by BlockManager.allocate()). # -1 = unallocated. The slot indexes into the per-req cache tensors # owned by ModelRunner (e.g. mamba_k_cache for GDN). From af0ae4552d54ef6595856391916e5213ee09f0b0 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Fri, 10 Jul 2026 11:34:55 -0500 Subject: [PATCH 2/6] [Frontend] Detect client disconnect via ASGI event, not polling Replace the 0.5s is_disconnected() polling loop in _run_nonstream_with_disconnect with vLLM-style event-driven cancellation: race the generator-collector task against a task that awaits the ASGI http.disconnect event (request.receive()), FIRST_COMPLETED wins. Detection is now immediate (0ms vs up to 500ms) and costs nothing while the client stays connected (no periodic wakeups). The abort path is unchanged: cancelling the collector still propagates into generate_async's finally -> abort_request + io_processor.requests.pop, freeing leaked KV. request.receive() is safe here because FastAPI parses the request body into a pydantic model before the handler runs, so there is no unread body to race against. Verified on DeepSeek-V4-Pro tp8: curl --max-time drop -> immediate "Client disconnected ... aborting request" + abort_request found=True; normal non-stream and streaming requests unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: yihonglie --- atom/entrypoints/openai/api_server.py | 92 +++++++++++++++++---------- 1 file changed, 57 insertions(+), 35 deletions(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index e9aec50035..f6cc8ab093 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -842,9 +842,22 @@ def __init__(self, request_id: str): self.request_id = request_id +async def _listen_for_disconnect(request) -> None: + """Block until the client sends an ``http.disconnect`` ASGI event. + + Unlike polling ``request.is_disconnected()`` on a timer, this awaits the + disconnect event directly, so detection is immediate and costs nothing while + the client stays connected. + """ + while True: + message = await request.receive() + if message["type"] == "http.disconnect": + break + + async def _run_nonstream_with_disconnect(agen, raw_request, request_id): """Drive a non-stream ``generate_async*`` async-generator while actively - watching for client disconnect. + watching for client disconnect (vLLM ``with_cancellation`` style). Starlette does NOT cancel a *non-streaming* request handler when the client goes away (unlike StreamingResponse, which is cancelled on http.disconnect). @@ -852,47 +865,56 @@ async def _run_nonstream_with_disconnect(agen, raw_request, request_id): until it hits ``max_tokens`` -- burning GPU on output nobody will read AND leaking the seq in ``io_processor.requests`` (its finally never fires). - We run the generator in a task and poll ``raw_request.is_disconnected()`` - concurrently. On disconnect we cancel the task; the cancellation propagates - into ``generate_async``'s ``await token_queue.get()`` so its ``finally`` - runs -> ``abort_request`` + ``io_processor.requests.pop``. Returns the last - yielded output, or raises ``_ClientDisconnected``. + We race two tasks: one collects the generator, one listens for the ASGI + ``http.disconnect`` event. Whichever finishes first wins; the loser is + cancelled. On disconnect the collector's cancellation propagates into + ``generate_async``'s ``await token_queue.get()`` so its ``finally`` runs -> + ``abort_request`` + ``io_processor.requests.pop``. Returns the last yielded + output, or raises ``_ClientDisconnected``. + + ``request.receive()`` is safe here because FastAPI has already parsed the + request body into a pydantic model before this handler runs, so there is no + unread body for ``receive()`` to race against. """ - final_output = None async def _collect(): - nonlocal final_output + final_output = None async for output in agen: final_output = output return final_output - task = asyncio.ensure_future(_collect()) - try: - while True: - done, _ = await asyncio.wait({task}, timeout=0.5) - if task in done: - return task.result() - disconnected = False - if raw_request is not None: - try: - disconnected = await raw_request.is_disconnected() - except Exception: - disconnected = False - if disconnected: - logger.info( - f"Client disconnected (non-stream), aborting request " - f"{request_id}" - ) - raise _ClientDisconnected(request_id) - finally: - if not task.done(): - task.cancel() - try: - await task - except BaseException: - # CancelledError (expected) or any error from generator - # teardown; generate_async's finally already aborted + popped. - pass + handler_task = asyncio.ensure_future(_collect()) + + # No ASGI request object (e.g. internal call) -> just await the generator. + if raw_request is None: + return await handler_task + + disconnect_task = asyncio.ensure_future(_listen_for_disconnect(raw_request)) + + done, pending = await asyncio.wait( + [handler_task, disconnect_task], + return_when=asyncio.FIRST_COMPLETED, + ) + + # Cancel the loser and let its cancellation settle (drives generate_async's + # finally -> abort_request when the collector is the loser). + for task in pending: + task.cancel() + for task in pending: + try: + await task + except BaseException: + # CancelledError (expected) or any teardown error; the generator's + # own finally already aborted + popped the seq. + pass + + if handler_task in done: + return handler_task.result() + + logger.info( + f"Client disconnected (non-stream), aborting request {request_id}" + ) + raise _ClientDisconnected(request_id) async def setup_streaming_request_fanout( From 2556e6f8353406c0c375596395f6ead39e7ed6f0 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Fri, 10 Jul 2026 18:36:07 -0500 Subject: [PATCH 3/6] [Frontend] Extend disconnect-abort to fan-out/multimodal; narrow cancel except Address review of the client-disconnect abort: - Factor the disconnect race into `_race_disconnect(coro, ...)` and make `_run_nonstream_with_disconnect` a thin wrapper that collects the async-generator. This lets the fan-out path (whose `generate_async_fanout` is a coroutine returning a list, not an async generator) reuse the same cancellation machinery. - Wrap the previously-unguarded non-stream branches so an abandoned request is aborted instead of running to max_tokens: multimodal (chat), n>1 fan-out (chat, both multimodal and text), and n>1 fan-out (completions). generate_async_fanout's try/finally aborts every sibling on cancel, so a disconnect frees all n sibling seqs. - Narrow the post-cancel teardown handler from `except BaseException` to `except asyncio.CancelledError` (expected) + `except Exception` (logged), letting KeyboardInterrupt/SystemExit propagate. Verified on DeepSeek-V4-Flash tp4: plain and n=2 fan-out non-stream drops both log "Client disconnected ... aborting" with abort_request found=True for every sibling (2/2 for n=2); normal plain, n=2 fan-out, and streaming unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: yihonglie --- atom/entrypoints/openai/api_server.py | 120 ++++++++++++++++---------- 1 file changed, 75 insertions(+), 45 deletions(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index f6cc8ab093..6c844fc4c3 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -855,37 +855,30 @@ async def _listen_for_disconnect(request) -> None: break -async def _run_nonstream_with_disconnect(agen, raw_request, request_id): - """Drive a non-stream ``generate_async*`` async-generator while actively - watching for client disconnect (vLLM ``with_cancellation`` style). +async def _race_disconnect(coro, raw_request, request_id): + """Race an awaitable against client disconnect (vLLM ``with_cancellation`` + style). Starlette does NOT cancel a *non-streaming* request handler when the client goes away (unlike StreamingResponse, which is cancelled on http.disconnect). Without this, an abandoned non-stream request keeps ``await``-ing the engine until it hits ``max_tokens`` -- burning GPU on output nobody will read AND - leaking the seq in ``io_processor.requests`` (its finally never fires). + leaking the seq(s) in ``io_processor.requests`` (their finally never fires). - We race two tasks: one collects the generator, one listens for the ASGI - ``http.disconnect`` event. Whichever finishes first wins; the loser is - cancelled. On disconnect the collector's cancellation propagates into - ``generate_async``'s ``await token_queue.get()`` so its ``finally`` runs -> - ``abort_request`` + ``io_processor.requests.pop``. Returns the last yielded - output, or raises ``_ClientDisconnected``. + We run ``coro`` (which produces the final result) as a task alongside a task + that awaits the ASGI ``http.disconnect`` event. Whichever finishes first + wins; the loser is cancelled. On disconnect, the coro's cancellation + propagates into its ``await`` points so its own ``try/finally`` runs -> + ``abort_request`` + ``io_processor.requests.pop`` (for fan-out, this aborts + every sibling). We then raise ``_ClientDisconnected``. ``request.receive()`` is safe here because FastAPI has already parsed the request body into a pydantic model before this handler runs, so there is no unread body for ``receive()`` to race against. """ + handler_task = asyncio.ensure_future(coro) - async def _collect(): - final_output = None - async for output in agen: - final_output = output - return final_output - - handler_task = asyncio.ensure_future(_collect()) - - # No ASGI request object (e.g. internal call) -> just await the generator. + # No ASGI request object (e.g. internal call) -> just await the coro. if raw_request is None: return await handler_task @@ -896,17 +889,22 @@ async def _collect(): return_when=asyncio.FIRST_COMPLETED, ) - # Cancel the loser and let its cancellation settle (drives generate_async's - # finally -> abort_request when the collector is the loser). + # Cancel the loser and let its cancellation settle (drives the coro's own + # finally -> abort_request when the handler is the loser). Only swallow the + # expected CancelledError; log anything else, and let BaseException + # (KeyboardInterrupt/SystemExit) propagate. for task in pending: task.cancel() for task in pending: try: await task - except BaseException: - # CancelledError (expected) or any teardown error; the generator's - # own finally already aborted + popped the seq. + except asyncio.CancelledError: pass + except Exception: + logger.warning( + f"Error tearing down cancelled task for request {request_id}", + exc_info=True, + ) if handler_task in done: return handler_task.result() @@ -917,6 +915,24 @@ async def _collect(): raise _ClientDisconnected(request_id) +async def _run_nonstream_with_disconnect(agen, raw_request, request_id): + """Drive a non-stream ``generate_async*`` async-*generator* while watching + for client disconnect. + + Thin wrapper over :func:`_race_disconnect` that collects the generator's + last yielded output. Use :func:`_race_disconnect` directly for the fan-out + path, whose ``generate_async_fanout`` is a coroutine returning a list. + """ + + async def _collect(): + final_output = None + async for output in agen: + final_output = output + return final_output + + return await _race_disconnect(_collect(), raw_request, request_id) + + async def setup_streaming_request_fanout( prompt_or_tokens: str | List[int], sampling_params: SamplingParams, @@ -1138,12 +1154,16 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request) # Non-streaming if is_multimodal and effective_n > 1: - outputs = await generate_async_fanout( - token_ids, - sampling_params, + outputs = await _race_disconnect( + generate_async_fanout( + token_ids, + sampling_params, + request_id, + multimodal_data=multimodal_data, + kv_transfer_params=request.kv_transfer_params, + ), + raw_request, request_id, - multimodal_data=multimodal_data, - kv_transfer_params=request.kv_transfer_params, ) if not outputs: raise RuntimeError("No output generated") @@ -1151,14 +1171,16 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request) request_id, model_name, outputs, tools=request.tools ) elif is_multimodal: - final_output = None - async for output in generate_async_multimodal( - token_ids, - multimodal_data, - sampling_params, + final_output = await _run_nonstream_with_disconnect( + generate_async_multimodal( + token_ids, + multimodal_data, + sampling_params, + request_id, + ), + raw_request, request_id, - ): - final_output = output + ) if final_output is None: raise RuntimeError("No output generated") resp = build_chat_response( @@ -1169,11 +1191,15 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request) tools=request.tools, ) elif effective_n > 1: - outputs = await generate_async_fanout( - prompt, - sampling_params, + outputs = await _race_disconnect( + generate_async_fanout( + prompt, + sampling_params, + request_id, + kv_transfer_params=request.kv_transfer_params, + ), + raw_request, request_id, - kv_transfer_params=request.kv_transfer_params, ) if not outputs: raise RuntimeError("No output generated") @@ -1278,12 +1304,16 @@ async def completions(request: CompletionRequest, raw_request: Request): # Non-streaming if effective_n > 1: - outputs = await generate_async_fanout( - request.prompt, - sampling_params, + outputs = await _race_disconnect( + generate_async_fanout( + request.prompt, + sampling_params, + request_id, + kv_transfer_params=request.kv_transfer_params, + data_parallel_rank=request.data_parallel_rank, + ), + raw_request, request_id, - kv_transfer_params=request.kv_transfer_params, - data_parallel_rank=request.data_parallel_rank, ) if not outputs: raise RuntimeError("No output generated") From e9ebbd63f2d4da2c4236c90e17a3281b5eb8d042 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Sun, 12 Jul 2026 22:22:57 -0500 Subject: [PATCH 4/6] style: apply black formatting to api_server.py Fixes the "Check Code Style with Black" CI check on the disconnect-abort changes (extra blank lines, single-line logger call). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: yihonglie --- atom/entrypoints/openai/api_server.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 6c844fc4c3..3a4c8940bd 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -832,8 +832,6 @@ def cleanup_streaming_request(request_id: str, seq_id: int) -> None: engine.io_processor.requests.pop(seq_id, None) - - class _ClientDisconnected(Exception): """Raised when a non-streaming client hangs up mid-generation.""" @@ -909,9 +907,7 @@ async def _race_disconnect(coro, raw_request, request_id): if handler_task in done: return handler_task.result() - logger.info( - f"Client disconnected (non-stream), aborting request {request_id}" - ) + logger.info(f"Client disconnected (non-stream), aborting request {request_id}") raise _ClientDisconnected(request_id) From 505d2b8e716e48b2db0805d3a38c78a5bd7a00e0 Mon Sep 17 00:00:00 2001 From: yihonglie Date: Fri, 10 Jul 2026 10:19:44 -0500 Subject: [PATCH 5/6] [Frontend] DeepSeek-V4 native OpenAI/Anthropic/Responses API + DSML tool parser Serve DeepSeek-V4 (Pro/Flash) so Claude Code (Anthropic /v1/messages) and Codex CLI (OpenAI Responses /v1/responses) talk to ATOM directly, with DSML tool-call parsing shared across /v1/chat/completions, /v1/messages and /v1/responses. - tool_parser: DeepSeek-V4 DSML tool-call format (<|DSML|invoke ...>) with marker-less / self-closing / direct-JSON recovery, schema-driven type coercion and key aliases; streaming + non-streaming (alongside Qwen/GLM/MiniMax). - serving_chat: reasoning-filter + tool-call streaming for /v1/chat/completions. - serving_responses (new) + /v1/responses: OpenAI Responses translation (input/tools<->chat, SSE emitter, reasoning/function_call items). Codex compat: inject a mandatory DSML tool-format instruction and normalize shell tool name/param aliases (exec/shell_exec/... -> registered exec tool; command/script -> cmd) so Codex tool calls execute instead of erroring. - /v1/messages: pass the tool schema to the parser so tool args are typed correctly (fixes "Invalid tool parameters"); streaming UTF-8 incremental detokenization (vLLM-style sliding window) so multi-byte chars (CJK, box-drawing) aren't split into U+FFFD. - reasoning / chat_encoders: DeepSeek-V4 message encoding + reasoning tags. Builds on the abort-on-disconnect path from the previous commit. Co-Authored-By: Claude Opus 4.8 (1M context) --- atom/entrypoints/openai/api_server.py | 325 ++++++++++- atom/entrypoints/openai/chat_encoders.py | 53 +- atom/entrypoints/openai/reasoning.py | 7 + atom/entrypoints/openai/serving_chat.py | 429 +++++++------- atom/entrypoints/openai/serving_responses.py | 572 ++++++++++++++++++ atom/entrypoints/openai/tool_parser.py | 585 ++++++++++++++++++- 6 files changed, 1744 insertions(+), 227 deletions(-) create mode 100644 atom/entrypoints/openai/serving_responses.py diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 3a4c8940bd..a37bdf14fa 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -68,6 +68,18 @@ stream_completion_response, stream_completion_response_fanout, ) +from .serving_responses import ( + ResponsesStreamEmitter, + inject_tool_format_instruction, + shell_arg_key, + build_responses_object, + remap_tool_name, + extract_cwd, + responses_input_to_messages, + responses_tools_to_openai, + tool_name_lookup, + translate_client_tool, +) # Configure logging logger = logging.getLogger("atom") @@ -324,6 +336,24 @@ def _prepare_multimodal_inputs( return inputs["input_ids"][0].tolist(), multimodal_data +# ── Batched stream dispatch ────────────────────────────────────────────── +# Per-seq `call_soon_threadsafe` floods the API event loop at high batch size +# (one call per token). Instead the callback only buffers the raw chunk; the +# mgr flushes a whole step with a single `tokenizer.batch_decode` (one +# GIL-released call instead of one decode per seq) plus one scheduled call per +# loop (see `flush_stream_batch`). +import threading as _threading # noqa: E402 + +_stream_batch_tls = _threading.local() + +# Per-request incremental detokenization state (vLLM-style sliding window). +# Decoding each step's new tokens in isolation splits multi-byte UTF-8 chars +# (byte-BPE tokenizers like DeepSeek-V4 split one CJK char across several +# byte-tokens) into U+FFFD. Keep accumulated tokens + prefix/read offsets so +# we only emit fully-formed characters. +_stream_detok_state: Dict[str, dict] = {} + + def _send_stream_chunk_direct( request_output: RequestOutput, request_id: str, @@ -346,7 +376,70 @@ def _send_stream_chunk_direct( } if getattr(request_output, "kv_transfer_params_output", None): chunk_data["kv_transfer_params"] = request_output.kv_transfer_params_output - loop.call_soon_threadsafe(stream_queue.put_nowait, chunk_data) + + chunk_data["request_id"] = request_id + buf = getattr(_stream_batch_tls, "buf", None) + if buf is None: + buf = _stream_batch_tls.buf = [] + buf.append((loop, stream_queue, chunk_data)) + + +def _drain_batch_into_queues(items: list) -> None: + """Runs ON the event loop: push each chunk into its per-request queue. + One scheduled call handles a whole step's worth of chunks.""" + for _loop, q, chunk in items: + q.put_nowait(chunk) + + +def flush_stream_batch() -> None: + """Flush a step's buffered chunks: one ``batch_decode`` for the whole step, + then one call_soon_threadsafe per loop (normally one — all requests on a + rank share the API loop).""" + global tokenizer + + buf = getattr(_stream_batch_tls, "buf", None) + if not buf: + return + _stream_batch_tls.buf = [] + # Decode the whole step in a single call. batch_decode is element-wise + # identical to per-seq decode but acquires/releases the GIL once instead of + # once per seq, cutting GIL ping-pong against the other rank output threads + # and the API event loop at high batch size. + # Incremental per-request detokenization: correct UTF-8 at token + # boundaries (see _stream_detok_state). Emits only fully-formed chars; + # a trailing partial multi-byte char is held until the next step. + for (_loop, _q, chunk) in buf: + rid = chunk.get("request_id") + st = _stream_detok_state.get(rid) + if st is None: + st = _stream_detok_state[rid] = { + "tokens": [], "prefix_offset": 0, "read_offset": 0, + } + toks = st["tokens"] + toks.extend(chunk["token_ids"]) + prefix_text = tokenizer.decode( + toks[st["prefix_offset"]:st["read_offset"]], skip_special_tokens=True + ) + new_text = tokenizer.decode( + toks[st["prefix_offset"]:], skip_special_tokens=True + ) + if len(new_text) > len(prefix_text) and not new_text.endswith("\ufffd"): + chunk["text"] = new_text[len(prefix_text):] + st["prefix_offset"] = st["read_offset"] + st["read_offset"] = len(toks) + elif chunk["finished"]: + chunk["text"] = new_text[len(prefix_text):] + else: + chunk["text"] = "" + if chunk["finished"]: + _stream_detok_state.pop(rid, None) + # Group by loop (normally a single loop). dict preserves insertion order + # so per-request chunk ordering within the step is maintained. + by_loop: Dict[AbstractEventLoop, list] = {} + for loop, q, chunk in buf: + by_loop.setdefault(loop, []).append((loop, q, chunk)) + for loop, items in by_loop.items(): + loop.call_soon_threadsafe(_drain_batch_into_queues, items) def _send_stream_chunk_tagged( @@ -1426,6 +1519,7 @@ async def generate_anthropic_stream(): if prompt.rstrip().endswith(""): reasoning_filter.state = 1 tool_parser = ToolCallStreamParser() + tool_parser.tools = anthropic_to_openai_tools(request.tools) block_index = 0 started_text = False started_thinking = False @@ -1590,15 +1684,17 @@ async def generate_anthropic_stream(): from .reasoning import separate_reasoning from .tool_parser import parse_tool_calls - final_output = None - async for output in generate_async(prompt, sampling_params, request_id): - final_output = output + final_output = await _run_nonstream_with_disconnect( + generate_async(prompt, sampling_params, request_id), + raw_request, + request_id, + ) if final_output is None: raise RuntimeError("No output generated") raw_text = final_output["text"] reasoning_content, content_with_tools = separate_reasoning(raw_text) - content_text, tool_calls = parse_tool_calls(content_with_tools) + content_text, tool_calls = parse_tool_calls(content_with_tools, anthropic_to_openai_tools(request.tools)) output_tokens = len(tokenizer.encode(raw_text)) cache_read_input_tokens = final_output.get("num_cached_tokens", 0) if not getattr(request, "thinking", None): @@ -1615,6 +1711,9 @@ async def generate_anthropic_stream(): cache_read_input_tokens=cache_read_input_tokens, ) + except _ClientDisconnected: + # Client hung up; seq already aborted + popped. Nothing to return. + return JSONResponse(status_code=499, content={"detail": "client disconnected"}) except Exception as e: logger.error(f"Error in anthropic_messages: {e}", exc_info=True) return JSONResponse( @@ -1626,6 +1725,222 @@ async def generate_anthropic_stream(): ) +@app.post("/v1/responses") +async def responses_endpoint(raw_request: Request): + """Handle OpenAI **Responses API** requests (`/v1/responses`). + + Native support so OpenAI Codex CLI (>= 0.14x, which speaks only the Responses + API) can talk to ATOM directly — no external responses->chat proxy needed. + Reuses ATOM's proven streaming path (setup_streaming_request + ReasoningFilter + + ToolCallStreamParser), the same one /v1/messages (claude-local) uses, so it + streams correctly for reasoning models. Stateless (full input sent each turn, + as Codex does); reasoning items are dropped from visible output. + """ + global engine, tokenizer, model_name + + try: + body = await raw_request.json() + model = body.get("model") or model_name + + from .protocol import ChatMessage + from .reasoning import ReasoningFilter, separate_reasoning + from .tool_parser import ToolCallStreamParser, parse_tool_calls + + openai_tools = responses_tools_to_openai(body.get("tools")) + valid_names, shell_tool = tool_name_lookup(openai_tools) + shell_param = shell_arg_key(openai_tools, shell_tool) + req_cwd = extract_cwd(body) # Codex — used to fix hallucinated paths + openai_messages = responses_input_to_messages( + body.get("instructions"), body.get("input") + ) + if openai_tools: + openai_messages = inject_tool_format_instruction(openai_messages) + messages = [ChatMessage(**m) for m in openai_messages] + + merged_kwargs = dict(default_chat_template_kwargs) + prompt = apply_chat_template( + tokenizer, + custom_message_encoder, + [msg.to_template_dict() for msg in messages], + tools=openai_tools or None, + **merged_kwargs, + ) + + max_out = int(body.get("max_output_tokens") or 32768) + sampling_params = _build_sampling_params( + temperature=body.get("temperature") if body.get("temperature") is not None else 1.0, + max_tokens=max_out, + stop_strings=None, + ignore_eos=False, + top_k=-1, + top_p=body.get("top_p") if body.get("top_p") is not None else 1.0, + ) + + request_id = "resp_" + uuid.uuid4().hex[:24] + input_tokens = len(tokenizer.encode(prompt)) + + # Resolve max context to bound the prompt (same probes as anthropic). + max_ctx = None + for _path in ( + lambda: engine.config.max_model_len, + lambda: engine.model_config.max_model_len, + lambda: engine.scheduler.max_model_len, + lambda: getattr(engine, "max_model_len"), + ): + try: + _v = _path() + if _v: + max_ctx = int(_v) + break + except Exception: + continue + if not max_ctx: + max_ctx = 30720 + headroom = min(max_out, max(1024, max_ctx // 8)) + max_input = max_ctx - headroom + if input_tokens > max_input: + logger.warning( + f"[responses] prompt too long ({input_tokens} > {max_input}), truncating" + ) + token_ids = tokenizer.encode(prompt)[:max_input] + prompt = tokenizer.decode(token_ids, skip_special_tokens=False) + input_tokens = max_input + + if body.get("stream"): + seq_id, stream_queue, _num_prompt_tokens = await setup_streaming_request( + prompt, sampling_params, request_id + ) + + async def generate_responses_stream(): + emitter = ResponsesStreamEmitter(request_id, model) + reasoning_filter = ReasoningFilter() + if prompt.rstrip().endswith(""): + reasoning_filter.state = 1 + tool_parser = ToolCallStreamParser() + tool_parser.tools = openai_tools or None # enables schema-based + # type coercion + key-alias (command->cmd) in _parse_dsml + output_tokens = 0 + + # Buffer each tool call (name + full args) so read/grep/ls/find + # can be translated to exec_command before emitting (name AND + # args change). See translate_client_tool. + _pending = {"tc": None} + + def _flush_pending(): + tc = _pending["tc"] + if tc is None: + return [] + _pending["tc"] = None + name, args = translate_client_tool( + tc["name"], tc["args"], valid_names, shell_tool, req_cwd, shell_param + ) + out = emitter.tool_start(tc["id"], name) + if args: + out += emitter.tool_args(args) + out += emitter.tool_end() + return out + + def handle(etype, edata): + # Map ToolCallStreamParser events -> Responses SSE strings, + # buffering tool calls for client-tool translation. + if etype == "content": + out = _flush_pending() + return out + emitter.text_delta(edata) + if etype == "tool_call_start": + out = _flush_pending() + fn = edata.get("function", {}) + _pending["tc"] = { + "id": edata.get("id", ""), + "name": fn.get("name", ""), + "args": "", + } + return out + if etype == "tool_call_args": + if _pending["tc"] is not None: + _pending["tc"]["args"] += ( + edata.get("function", {}).get("arguments", "") or "" + ) + return [] + if etype == "tool_call_end": + return _flush_pending() + return [] + + try: + for s in emitter.created(): + yield s + while True: + chunk_data = await stream_queue.get() + new_text = chunk_data["text"] + output_tokens += len(chunk_data.get("token_ids", [])) + finished = chunk_data.get("finished", False) + + segments = reasoning_filter.process(new_text) + if finished: + segments.extend(reasoning_filter.flush()) + for field, text in segments: + if not text or field == "reasoning_content": + continue # drop reasoning from visible output + for etype, edata in tool_parser.process(text): + for s in handle(etype, edata): + yield s + + if finished: + for etype, edata in tool_parser.flush(): + for s in handle(etype, edata): + yield s + for s in _flush_pending(): # emit any unclosed tool call + yield s + for s in emitter.finish(input_tokens, output_tokens): + yield s + yield "data: [DONE]\n\n" + break + finally: + cleanup_streaming_request(request_id, seq_id) + + return StreamingResponse( + generate_responses_stream(), + media_type="text/event-stream", + headers={"x-request-id": request_id}, + ) + + # Non-streaming response + final_output = await _run_nonstream_with_disconnect( + generate_async(prompt, sampling_params, request_id), + raw_request, + request_id, + ) + if final_output is None: + raise RuntimeError("No output generated") + + raw_text = final_output["text"] + _reasoning, content_with_tools = separate_reasoning(raw_text) + content_text, tool_calls = parse_tool_calls(content_with_tools, openai_tools or None) + output_tokens = len(tokenizer.encode(raw_text)) + + return JSONResponse( + content=build_responses_object( + resp_id=request_id, + model=model, + content_text=content_text, + tool_calls=tool_calls, + input_tokens=input_tokens, + output_tokens=output_tokens, + valid=valid_names, + shell_tool=shell_tool, + cwd=req_cwd, + ) + ) + + except _ClientDisconnected: + return JSONResponse(status_code=499, content={"detail": "client disconnected"}) + except Exception as e: + logger.error(f"Error in responses_endpoint: {e}", exc_info=True) + return JSONResponse( + status_code=500, + content={"error": {"type": "api_error", "message": str(e)}}, + ) + + @app.get("/v1/models") async def list_models(): """List available models.""" diff --git a/atom/entrypoints/openai/chat_encoders.py b/atom/entrypoints/openai/chat_encoders.py index d3fb0462ce..863d045275 100644 --- a/atom/entrypoints/openai/chat_encoders.py +++ b/atom/entrypoints/openai/chat_encoders.py @@ -85,6 +85,46 @@ def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoder]: return _load_encoder_from_dir(_resolve_model_path(model_path)) +def _content_str(c: Any) -> str: + if isinstance(c, list): + return "\n".join( + b.get("text", "") for b in c if isinstance(b, dict) and b.get("type") == "text" + ) + return c or "" + + +def _normalize_for_v4(messages: List[dict], tools: Optional[List[dict]]) -> List[dict]: + """Prepare messages for DeepSeek-V4's ``encode_messages``. + + Two things: + 1. **Hoist system messages to the front.** Clients (notably Claude Code) send + a trailing ``system``-role message (its "skills" list) AFTER the user turn. + ``encode_messages`` only appends the ``<|Assistant|>`` generation marker + after a *user*/developer message, so a trailing system message leaves the + prompt ending mid-system-text and the model just *continues* it instead of + answering. Merging all system content into one leading system message keeps + the final turn a user turn, so the assistant marker is emitted. + 2. **Attach tools** to that leading system message (``encode_messages`` reads + tool schemas from a system message's ``tools`` field). + Does not mutate the input. + """ + sys_parts, others = [], [] + for m in messages: + (sys_parts if m.get("role") == "system" else others).append(dict(m)) + + if not sys_parts and not tools: + return [dict(m) for m in messages] + + merged = "\n\n".join(s for s in (_content_str(m.get("content")) for m in sys_parts) if s) + sys_msg: dict = {"role": "system", "content": merged} + for m in sys_parts: # preserve any pre-attached tools + if m.get("tools"): + sys_msg["tools"] = m["tools"] + if tools: + sys_msg["tools"] = tools + return [sys_msg] + others + + def apply_chat_template( tokenizer: Any, custom_encoder: Optional[MessageEncoder], @@ -97,18 +137,15 @@ def apply_chat_template( Dispatches to ``custom_encoder`` if one was discovered for this model, otherwise to ``tokenizer.apply_chat_template``. Jinja-only kwargs - (``tokenize``, ``add_generation_prompt``) are stripped on the custom - path; ``tools`` are forwarded only on the Jinja path (custom encoders - don't currently have a tools API — caller is warned and tools are - dropped). + (``tokenize``, ``add_generation_prompt``) are stripped on the custom path. + ``tools`` are supported on both paths: custom encoders (e.g. DeepSeek-V4's + ``encode_messages``) read tool schemas from a system message's ``tools`` + field, so we attach them there before encoding. """ if custom_encoder is not None: for k in ("tokenize", "add_generation_prompt"): kwargs.pop(k, None) - if tools: - logger.warning( - "tools= is not supported with the custom message encoder; ignoring." - ) + messages = _normalize_for_v4(messages, tools) return custom_encoder(messages, **kwargs) kwargs["tokenize"] = False diff --git a/atom/entrypoints/openai/reasoning.py b/atom/entrypoints/openai/reasoning.py index 6fb8a8e004..f7cd8f0b41 100644 --- a/atom/entrypoints/openai/reasoning.py +++ b/atom/entrypoints/openai/reasoning.py @@ -23,6 +23,9 @@ def separate_reasoning(text: str) -> Tuple[Optional[str], str]: Tuple of (reasoning_content, content). reasoning_content is None if no thinking block was found. """ + # MiniMax M3 emits ... instead of ...; + # normalize so the shared logic below handles both. + text = text.replace("", "").replace("", "") # Check for closed thinking block: ... match = re.match(r"(.*?)\s*(.*)", text, flags=re.DOTALL) if match: @@ -75,6 +78,10 @@ def process(self, text: str) -> list: List of (field_name, text) tuples where field_name is "reasoning_content" or "content". """ + # MiniMax M3 uses /; normalize to the tags + # the state machine below keys on. These are single special tokens, so + # each arrives whole in one chunk — a plain replace is safe. + text = text.replace("", "").replace("", "") results = [] if self.state == 0: diff --git a/atom/entrypoints/openai/serving_chat.py b/atom/entrypoints/openai/serving_chat.py index 7c707e9683..9c21e8534f 100644 --- a/atom/entrypoints/openai/serving_chat.py +++ b/atom/entrypoints/openai/serving_chat.py @@ -61,56 +61,78 @@ async def stream_chat_response( cleanup_fn, tools=None, ) -> AsyncGenerator[str, None]: - """Generate streaming chat completion response with reasoning and tool calls. - - Yields SSE chunks with: - - reasoning_content deltas during thinking phase - - content deltas for the answer - - tool_calls deltas when model invokes tools - - ``num_prompt_tokens`` is the engine-computed prompt length (``Sequence. - num_prompt_tokens``); reusing it avoids re-tokenizing the prompt on the - event loop at stream start. - """ - num_tokens_input = num_prompt_tokens - num_tokens_output = 0 - num_cached_tokens = 0 - reasoning_filter = ReasoningFilter() - tool_parser = ToolCallStreamParser(tools=tools) - has_tool_calls = False - - # Send initial role chunk - yield create_chat_chunk(request_id, model, delta={"role": "assistant"}) - - kv_transfer_params_value = None - - while True: - chunk_data = await stream_queue.get() - new_text = chunk_data["text"] - num_tokens_output += len(chunk_data.get("token_ids", [])) - _ct = chunk_data.get("num_cached_tokens", 0) - if _ct: - num_cached_tokens = _ct - - if "kv_transfer_params" in chunk_data: - kv_transfer_params_value = chunk_data["kv_transfer_params"] - - # Phase 1: Process through reasoning filter - segments = reasoning_filter.process(new_text) - if chunk_data.get("finished", False): - segments.extend(reasoning_filter.flush()) - - # Phase 2: For content segments, check for tool calls - for field, text in segments: - if field == "reasoning_content": - if text: - yield create_chat_chunk( - request_id, model, delta={"reasoning_content": text} - ) - elif field == "content": - # Run through tool parser - events = tool_parser.process(text) - for event_type, data in events: + try: + """Generate streaming chat completion response with reasoning and tool calls. + + Yields SSE chunks with: + - reasoning_content deltas during thinking phase + - content deltas for the answer + - tool_calls deltas when model invokes tools + + ``num_prompt_tokens`` is the engine-computed prompt length (``Sequence. + num_prompt_tokens``); reusing it avoids re-tokenizing the prompt on the + event loop at stream start. + """ + num_tokens_input = num_prompt_tokens + num_tokens_output = 0 + num_cached_tokens = 0 + reasoning_filter = ReasoningFilter() + tool_parser = ToolCallStreamParser(tools=tools) + has_tool_calls = False + + # Send initial role chunk + yield create_chat_chunk(request_id, model, delta={"role": "assistant"}) + + kv_transfer_params_value = None + + while True: + chunk_data = await stream_queue.get() + new_text = chunk_data["text"] + num_tokens_output += len(chunk_data.get("token_ids", [])) + _ct = chunk_data.get("num_cached_tokens", 0) + if _ct: + num_cached_tokens = _ct + + if "kv_transfer_params" in chunk_data: + kv_transfer_params_value = chunk_data["kv_transfer_params"] + + # Phase 1: Process through reasoning filter + segments = reasoning_filter.process(new_text) + if chunk_data.get("finished", False): + segments.extend(reasoning_filter.flush()) + + # Phase 2: For content segments, check for tool calls + for field, text in segments: + if field == "reasoning_content": + if text: + yield create_chat_chunk( + request_id, model, delta={"reasoning_content": text} + ) + elif field == "content": + # Run through tool parser + events = tool_parser.process(text) + for event_type, data in events: + if event_type == "content": + yield create_chat_chunk( + request_id, model, delta={"content": data} + ) + elif event_type == "tool_call_start": + has_tool_calls = True + yield create_chat_chunk( + request_id, + model, + delta={"tool_calls": [data]}, + ) + elif event_type == "tool_call_args": + yield create_chat_chunk( + request_id, + model, + delta={"tool_calls": [data]}, + ) + + if chunk_data.get("finished", False): + # Flush tool parser + for event_type, data in tool_parser.flush(): if event_type == "content": yield create_chat_chunk( request_id, model, delta={"content": data} @@ -118,60 +140,44 @@ async def stream_chat_response( elif event_type == "tool_call_start": has_tool_calls = True yield create_chat_chunk( - request_id, - model, - delta={"tool_calls": [data]}, + request_id, model, delta={"tool_calls": [data]} ) elif event_type == "tool_call_args": yield create_chat_chunk( - request_id, - model, - delta={"tool_calls": [data]}, + request_id, model, delta={"tool_calls": [data]} ) - - if chunk_data.get("finished", False): - # Flush tool parser - for event_type, data in tool_parser.flush(): - if event_type == "content": - yield create_chat_chunk(request_id, model, delta={"content": data}) - elif event_type == "tool_call_start": - has_tool_calls = True - yield create_chat_chunk( - request_id, model, delta={"tool_calls": [data]} - ) - elif event_type == "tool_call_args": - yield create_chat_chunk( - request_id, model, delta={"tool_calls": [data]} - ) - break - - cleanup_fn(request_id, seq_id) - - # Final chunks - finish_reason = "tool_calls" if has_tool_calls else "stop" - usage = { - "prompt_tokens": num_tokens_input, - "completion_tokens": num_tokens_output, - "total_tokens": num_tokens_input + num_tokens_output, - "prompt_tokens_details": {"cached_tokens": num_cached_tokens}, - } - usage_chunk = { - "id": request_id, - "object": CHAT_COMPLETION_CHUNK_OBJECT, - "created": int(time.time()), - "model": model, - "usage": usage, - } - if kv_transfer_params_value is not None: - usage_chunk["kv_transfer_params"] = kv_transfer_params_value - # Coalesce finish + usage + [DONE] into one send: at a wave boundary many - # requests finalize at once, so collapsing 3 socket writes/req to 1 cuts - # the syscalls that saturate the API event loop. - yield ( - create_chat_chunk(request_id, model, finish_reason=finish_reason) - + f"data: {json.dumps(usage_chunk)}\n\n" - + STREAM_DONE_MESSAGE - ) + break + + cleanup_fn(request_id, seq_id) + + # Final chunks + finish_reason = "tool_calls" if has_tool_calls else "stop" + usage = { + "prompt_tokens": num_tokens_input, + "completion_tokens": num_tokens_output, + "total_tokens": num_tokens_input + num_tokens_output, + "prompt_tokens_details": {"cached_tokens": num_cached_tokens}, + } + usage_chunk = { + "id": request_id, + "object": CHAT_COMPLETION_CHUNK_OBJECT, + "created": int(time.time()), + "model": model, + "choices": [], + "usage": usage, + } + if kv_transfer_params_value is not None: + usage_chunk["kv_transfer_params"] = kv_transfer_params_value + # Coalesce finish + usage + [DONE] into one send: at a wave boundary many + # requests finalize at once, so collapsing 3 socket writes/req to 1 cuts + # the syscalls that saturate the API event loop. + yield ( + create_chat_chunk(request_id, model, finish_reason=finish_reason) + + f"data: {json.dumps(usage_chunk)}\n\n" + + STREAM_DONE_MESSAGE + ) + finally: + cleanup_fn(request_id, seq_id) def _build_chat_choice( @@ -298,59 +304,84 @@ async def stream_chat_response_fanout( cleanup_fn, tools=None, ) -> AsyncGenerator[str, None]: - """Streaming variant that multiplexes ``len(seq_ids)`` fan-out siblings - into a single SSE stream, tagging every chunk with ``choices[0].index``. - - The shared queue receives ``(sibling_index, chunk_data)`` tuples from - the engine callbacks registered in :func:`setup_streaming_request_fanout`. - Reasoning + tool-call state is kept independently per sibling. + try: + """Streaming variant that multiplexes ``len(seq_ids)`` fan-out siblings + into a single SSE stream, tagging every chunk with ``choices[0].index``. + + The shared queue receives ``(sibling_index, chunk_data)`` tuples from + the engine callbacks registered in :func:`setup_streaming_request_fanout`. + Reasoning + tool-call state is kept independently per sibling. + + ``num_prompt_tokens`` is the engine-computed prompt length shared by all + siblings (they tokenize the same prompt once); reusing it avoids + re-tokenizing on the event loop at stream start. + """ + n = len(seq_ids) + num_tokens_input = num_prompt_tokens + num_tokens_output = [0] * n + num_cached_tokens = 0 + reasoning_filters = [ReasoningFilter() for _ in range(n)] + tool_parsers = [ToolCallStreamParser(tools=tools) for _ in range(n)] + has_tool_calls = [False] * n + finished = [False] * n + kv_transfer_params_value = None + + for i in range(n): + yield create_chat_chunk( + request_id, model, delta={"role": "assistant"}, index=i + ) - ``num_prompt_tokens`` is the engine-computed prompt length shared by all - siblings (they tokenize the same prompt once); reusing it avoids - re-tokenizing on the event loop at stream start. - """ - n = len(seq_ids) - num_tokens_input = num_prompt_tokens - num_tokens_output = [0] * n - reasoning_filters = [ReasoningFilter() for _ in range(n)] - tool_parsers = [ToolCallStreamParser(tools=tools) for _ in range(n)] - has_tool_calls = [False] * n - finished = [False] * n - kv_transfer_params_value = None - num_cached_tokens = 0 - - for i in range(n): - yield create_chat_chunk(request_id, model, delta={"role": "assistant"}, index=i) - - while not all(finished): - idx, chunk_data = await shared_queue.get() - if finished[idx]: - # Defensive: should not happen, engine emits finished once per seq. - continue - new_text = chunk_data["text"] - num_tokens_output[idx] += len(chunk_data.get("token_ids", [])) - _ct = chunk_data.get("num_cached_tokens", 0) - if _ct: - num_cached_tokens = _ct - - if "kv_transfer_params" in chunk_data: - kv_transfer_params_value = chunk_data["kv_transfer_params"] - - segments = reasoning_filters[idx].process(new_text) - if chunk_data.get("finished", False): - segments.extend(reasoning_filters[idx].flush()) - - for field, text in segments: - if field == "reasoning_content": - if text: - yield create_chat_chunk( - request_id, - model, - delta={"reasoning_content": text}, - index=idx, - ) - elif field == "content": - for event_type, data in tool_parsers[idx].process(text): + while not all(finished): + idx, chunk_data = await shared_queue.get() + if finished[idx]: + # Defensive: should not happen, engine emits finished once per seq. + continue + new_text = chunk_data["text"] + num_tokens_output[idx] += len(chunk_data.get("token_ids", [])) + _ct = chunk_data.get("num_cached_tokens", 0) + if _ct: + num_cached_tokens = _ct + + if "kv_transfer_params" in chunk_data: + kv_transfer_params_value = chunk_data["kv_transfer_params"] + + segments = reasoning_filters[idx].process(new_text) + if chunk_data.get("finished", False): + segments.extend(reasoning_filters[idx].flush()) + + for field, text in segments: + if field == "reasoning_content": + if text: + yield create_chat_chunk( + request_id, + model, + delta={"reasoning_content": text}, + index=idx, + ) + elif field == "content": + for event_type, data in tool_parsers[idx].process(text): + if event_type == "content": + yield create_chat_chunk( + request_id, model, delta={"content": data}, index=idx + ) + elif event_type == "tool_call_start": + has_tool_calls[idx] = True + yield create_chat_chunk( + request_id, + model, + delta={"tool_calls": [data]}, + index=idx, + ) + elif event_type == "tool_call_args": + yield create_chat_chunk( + request_id, + model, + delta={"tool_calls": [data]}, + index=idx, + ) + + if chunk_data.get("finished", False): + for event_type, data in tool_parsers[idx].flush(): if event_type == "content": yield create_chat_chunk( request_id, model, delta={"content": data}, index=idx @@ -370,61 +401,43 @@ async def stream_chat_response_fanout( delta={"tool_calls": [data]}, index=idx, ) - - if chunk_data.get("finished", False): - for event_type, data in tool_parsers[idx].flush(): - if event_type == "content": - yield create_chat_chunk( - request_id, model, delta={"content": data}, index=idx - ) - elif event_type == "tool_call_start": - has_tool_calls[idx] = True - yield create_chat_chunk( - request_id, - model, - delta={"tool_calls": [data]}, - index=idx, - ) - elif event_type == "tool_call_args": - yield create_chat_chunk( - request_id, - model, - delta={"tool_calls": [data]}, - index=idx, - ) - finished[idx] = True - - # Clean up all sibling seq_id entries then the shared request state. - for sid in seq_ids: - cleanup_fn(request_id, sid) - - usage = { - "prompt_tokens": num_tokens_input, - "completion_tokens": sum(num_tokens_output), - "total_tokens": num_tokens_input + sum(num_tokens_output), - "num_choices": n, - "prompt_tokens_details": {"cached_tokens": num_cached_tokens}, - } - usage_chunk = { - "id": request_id, - "object": CHAT_COMPLETION_CHUNK_OBJECT, - "created": int(time.time()), - "model": model, - "usage": usage, - } - if kv_transfer_params_value is not None: - usage_chunk["kv_transfer_params"] = kv_transfer_params_value - # Coalesce the per-sibling finish chunks + usage + [DONE] into one send. - yield ( - "".join( - create_chat_chunk( - request_id, - model, - finish_reason="tool_calls" if has_tool_calls[i] else "stop", - index=i, + finished[idx] = True + + # Clean up all sibling seq_id entries then the shared request state. + for sid in seq_ids: + cleanup_fn(request_id, sid) + + usage = { + "prompt_tokens": num_tokens_input, + "completion_tokens": sum(num_tokens_output), + "total_tokens": num_tokens_input + sum(num_tokens_output), + "num_choices": n, + "prompt_tokens_details": {"cached_tokens": num_cached_tokens}, + } + usage_chunk = { + "id": request_id, + "object": CHAT_COMPLETION_CHUNK_OBJECT, + "created": int(time.time()), + "model": model, + "choices": [], + "usage": usage, + } + if kv_transfer_params_value is not None: + usage_chunk["kv_transfer_params"] = kv_transfer_params_value + # Coalesce the per-sibling finish chunks + usage + [DONE] into one send. + yield ( + "".join( + create_chat_chunk( + request_id, + model, + finish_reason="tool_calls" if has_tool_calls[i] else "stop", + index=i, + ) + for i in range(n) ) - for i in range(n) + + f"data: {json.dumps(usage_chunk)}\n\n" + + STREAM_DONE_MESSAGE ) - + f"data: {json.dumps(usage_chunk)}\n\n" - + STREAM_DONE_MESSAGE - ) + finally: + for _sid in seq_ids: + cleanup_fn(request_id, _sid) diff --git a/atom/entrypoints/openai/serving_responses.py b/atom/entrypoints/openai/serving_responses.py new file mode 100644 index 0000000000..bc14091ad0 --- /dev/null +++ b/atom/entrypoints/openai/serving_responses.py @@ -0,0 +1,572 @@ +"""OpenAI **Responses API** (`/v1/responses`) support for the ATOM server. + +OpenAI Codex CLI (>= 0.14x) dropped `wire_api = "chat"` and only speaks the +Responses API (streaming SSE). This module provides the translation between +Responses request/response shapes and ATOM's internal chat/engine machinery, +plus a streaming SSE event emitter. The `/v1/responses` route handler in +``api_server.py`` reuses ATOM's proven streaming path (``setup_streaming_request`` ++ ``ReasoningFilter`` + ``ToolCallStreamParser``) so it streams correctly for +reasoning models — the same path claude-local uses via ``/v1/messages``. + +This makes the external ``codex_responses_proxy.py`` unnecessary: Codex can point +straight at ATOM's ``:9700/v1``. +""" +import itertools +import json +import time +from typing import Any, Dict, List, Optional, Tuple + +_ids = itertools.count(1) + + +def _rid(prefix: str) -> str: + return f"{prefix}_{int(time.time() * 1000)}{next(_ids):04d}" + + +# --------------------------------------------------------------- request xlate +def _text_of(content: Any) -> str: + """Flatten Responses content (str | list of parts) to a plain string.""" + if isinstance(content, str): + return content + if isinstance(content, list): + out = [] + for p in content: + if isinstance(p, dict): + if "text" in p and isinstance(p["text"], str): + out.append(p["text"]) + elif p.get("type") in ("input_text", "output_text", "text"): + out.append(p.get("text", "")) + elif isinstance(p, str): + out.append(p) + return "".join(out) + return "" + + +def responses_input_to_messages( + instructions: Any, inp: Any +) -> List[Dict[str, Any]]: + """Translate Responses ``instructions`` + ``input`` into OpenAI chat messages. + + ``input`` may be a plain string or a list of items: ``message`` / + ``function_call`` (assistant tool call) / ``function_call_output`` (tool + result). ``reasoning`` items are dropped. + """ + messages: List[Dict[str, Any]] = [] + if instructions: + messages.append({"role": "system", "content": _text_of(instructions)}) + + if isinstance(inp, str): + messages.append({"role": "user", "content": inp}) + elif isinstance(inp, list): + for item in inp: + if not isinstance(item, dict): + continue + t = item.get("type", "message") + if t == "message": + role = item.get("role", "user") + messages.append( + {"role": role, "content": _text_of(item.get("content", ""))} + ) + elif t == "function_call": + messages.append( + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": item.get("call_id") + or item.get("id") + or _rid("call"), + "type": "function", + "function": { + "name": item.get("name", ""), + "arguments": item.get("arguments", "") or "", + }, + } + ], + } + ) + elif t == "function_call_output": + out = item.get("output", "") + messages.append( + { + "role": "tool", + "tool_call_id": item.get("call_id") or item.get("id") or "", + "content": out + if isinstance(out, str) + else json.dumps(out), + } + ) + elif t == "reasoning": + continue + return messages + + +_DSML_TOOL_INSTRUCTION = ( + "\n\n# Tool-call format (MANDATORY — overrides any other format instruction)\n" + "When you call a tool, output ONLY a DSML tool-call block and NOTHING else " + "in that message — no markdown, no ```json, no ///" + " tags, no prose. Use EXACTLY this syntax (the \uff5c characters " + "are U+FF5C fullwidth vertical bars, not ASCII '|'):\n" + "<\uff5cDSML\uff5ctool_calls>\n" + "<\uff5cDSML\uff5cinvoke name=\"TOOL_NAME\">\n" + "<\uff5cDSML\uff5cparameter name=\"PARAM_NAME\" string=\"true\">VALUE" + "\n" + "\n" + "\n" + "Use the exact tool and parameter names from the tools provided to you. " + "For a shell/exec tool, put the whole shell command string in its " + "command/cmd parameter. Emit one <\uff5cDSML\uff5cinvoke> per tool call." +) + + +def inject_tool_format_instruction(messages): + """Append the mandatory DSML tool-call format to the system message so the + model emits parseable DSML instead of ad-hoc /```json text. Codex + (/v1/responses) path only. Idempotent per request.""" + for m in messages: + if m.get("role") == "system": + base = m.get("content") or "" + if "\uff5cDSML\uff5ctool_calls" not in base: + m["content"] = _text_of(base) + _DSML_TOOL_INSTRUCTION + return messages + return [{"role": "system", "content": _DSML_TOOL_INSTRUCTION.strip()}] + list(messages) + + +def responses_tools_to_openai(tools: Any) -> List[Dict[str, Any]]: + """Translate Responses function tools into OpenAI chat tool defs. + + Responses puts ``name``/``description``/``parameters`` at the top level of a + ``{"type": "function", ...}`` tool (chat nests them under ``function``). + Non-function tool types (``web_search``, ``namespace``, ...) are dropped — + ATOM only executes model-emitted function/DSML tool calls; the client + (Codex) owns actual tool execution. + """ + if not tools: + return [] + ct: List[Dict[str, Any]] = [] + for tl in tools: + if not isinstance(tl, dict): + continue + if tl.get("type") == "function": + fn = tl.get("function", tl) + ct.append( + { + "type": "function", + "function": { + "name": fn.get("name"), + "description": fn.get("description", ""), + "parameters": fn.get("parameters", {}) or {}, + }, + } + ) + return ct + + +# ------------------------------------------------- tool-name normalization +# Non-codex-tuned models (e.g. DeepSeek-V4-Pro) don't reliably emit the EXACT +# tool name the client registered for shell exec. Codex 0.142.x names it +# ``exec_command`` (args ``{"cmd": "..."}``), but the model habitually calls it +# ``exec`` / ``exec_run`` / ``shell`` / ``bash``, which Codex's tool router +# rejects ("unsupported call: exec"). The ARGUMENTS are correct; only the name +# is wrong. So remap known shell-exec aliases onto whatever shell tool the +# client actually registered this turn. Pure rename; arguments untouched. +_SHELL_ALIASES = { + "exec", "exec_run", "exec_command", "execute_command", "shell", "bash", + "sh", "run", "run_command", "run_shell", "execute", "command", + "container.exec", "local_shell", "shell_command", "shell_exec", + "run_bash", "execute_shell", "bash_command", "run_terminal_cmd", + "terminal", "console", "run_shell_command", "runshell", +} +# Substrings that mark an unknown tool name as a shell/exec call (Codex's +# non-shell tools contain none of these). +_SHELL_NAME_TOKENS = ("shell", "exec", "bash", "cmd", "command", "termin", "console") +_SHELL_TOOL_PREFERENCE = ( + "exec_command", "shell", "local_shell", "bash", "container.exec", +) + + +def tool_name_lookup(openai_tools: List[Dict[str, Any]]) -> Tuple[set, Optional[str]]: + """Return (valid tool names, preferred shell tool name) for remapping.""" + valid = { + (t.get("function") or {}).get("name") + for t in (openai_tools or []) + if isinstance(t, dict) + } + valid.discard(None) + shell_tool = next((n for n in _SHELL_TOOL_PREFERENCE if n in valid), None) + return valid, shell_tool + + +def remap_tool_name(name: str, valid: set, shell_tool: Optional[str]) -> str: + """Fix a model tool_call name that doesn't match a registered tool. + + Only remaps shell-exec aliases to the registered shell tool; other + mismatches pass through unchanged (the client will surface them).""" + if name in valid: + return name + if not shell_tool: + return name + n = (name or "").lower().replace("-", "_").replace(".", "_") + if n in _SHELL_ALIASES: + return shell_tool + # Fuzzy: unknown tool whose name signals shell/exec -> the shell tool. + if any(k in n for k in _SHELL_NAME_TOKENS): + return shell_tool + return name + + +# ------------------------------------------------ Claude-tool -> shell adapter +# DeepSeek-V4-Pro is trained on Claude-Code's toolset, so under Codex it keeps +# calling read/grep/ls/find (which Codex doesn't have) instead of exec_command, +# and flails ("unsupported call: read"). When Codex registered an exec/shell +# tool, translate these read-only Claude tools into an equivalent shell command +# so the model's intent goes through. Only fires for names NOT in the registered +# set; exec_command's real calls are untouched. Codex-path only (never applied +# to /v1/messages, where read/grep ARE native tools). +def _q(s: Any) -> str: + import shlex + return shlex.quote(str(s)) + + +def _resolve_dir(path: str, cwd: Optional[str]) -> str: + """Pick a directory that actually exists: keep relative paths and paths under + cwd; otherwise fall back to cwd (the model often invents absolute prefixes + like /sglang or /sgl-workspace that don't exist here).""" + if not path or path == ".": + return cwd or "." + if not path.startswith("/"): + return path # relative — shell runs in cwd, fine + if cwd and path.startswith(cwd): + return path + return cwd or path + + +def _read_cmd(fp: str, cwd: Optional[str], start: int, end: int) -> str: + """Resilient file read: try the literal path, then cwd-prefixed, then locate + by basename under cwd — so a hallucinated absolute prefix still finds the file.""" + sed = f"sed -n {start},{end}p" + if not cwd: + return f"{sed} {_q(fp)}" + # f = literal; if missing, cwd + '/' + (f without leading /); if still + # missing, first match of `find cwd -name basename`. + return ( + f'f={_q(fp)}; [ -f "$f" ] || f={_q(cwd)}"/${{f#/}}"; ' + f'[ -f "$f" ] || f=$(find {_q(cwd)} -type f -name "$(basename {_q(fp)})" ' + f'2>/dev/null | head -1); {sed} "$f"' + ) + + +def _claude_tool_to_shell( + name: str, a: Dict[str, Any], cwd: Optional[str] = None +) -> Optional[str]: + n = (name or "").lower() + fp = (a.get("file_path") or a.get("path") or a.get("filePath") + or a.get("filename") or a.get("target_file") or a.get("file")) + pattern = a.get("pattern") or a.get("query") or a.get("regex") + path = (a.get("path") or a.get("directory") or a.get("target_directory") + or a.get("dir") or ".") + if n in ("read", "cat", "view", "view_file", "open", "read_file", + "readfile", "openfile"): + if not fp: + return None + off, lim = a.get("offset"), a.get("limit") + if off or lim: + start = int(off or 0) + 1 + end = start + int(lim or 200) - 1 + else: + start, end = 1, 400 + return _read_cmd(str(fp), cwd, start, end) + if n in ("grep", "search", "search_file", "ripgrep", "rg", "grep_search", + "codebase_search"): + if not pattern: + return None + return f"grep -rn -- {_q(pattern)} {_q(_resolve_dir(path, cwd))}" + if n in ("ls", "list", "list_dir", "list_directory", "listdir"): + return f"ls -la {_q(_resolve_dir(path, cwd))}" + if n in ("find", "glob", "glob_file_search", "file_search"): + base = _resolve_dir(path, cwd) + if pattern: + return f"find {_q(base)} -name {_q(pattern)}" + return f"find {_q(base)} -maxdepth 3" + return None + + +_SHELL_ARG_ALIASES = ( + "cmd", "command", "commandline", "command_line", "script", "bash", "sh", + "shell", "shell_command", "code", "input", "run", "cmd_string", +) + + +def shell_arg_key(openai_tools, shell_tool): + """Required (or first) param name of the registered shell tool, e.g. Codex's + exec_command -> "cmd". Used to normalize model arg keys.""" + if not shell_tool: + return None + for t in openai_tools or []: + fn = t.get("function", t) + if (fn.get("name") or t.get("name")) != shell_tool: + continue + params = fn.get("parameters") or {} + props = params.get("properties") or {} + for r in (params.get("required") or []): + if props.get(r, {}).get("type") in (None, "string"): + return r + if props: + return next(iter(props)) + return "cmd" + + +def _is_shell_name(name, shell_tool): + return name == shell_tool or name in _SHELL_ALIASES + + +def _normalize_shell_args(args_json, req): + """If the shell tool's required param `req` is absent but a known alias is + present, rename it. Keeps exec_command calls valid when the model uses + `command`/`script`/... instead of `cmd`.""" + if not req: + return args_json + try: + a = json.loads(args_json) if args_json else {} + except Exception: + return args_json + if not isinstance(a, dict) or req in a: + return args_json + for alias in _SHELL_ARG_ALIASES: + if alias != req and isinstance(a.get(alias), str): + a[req] = a.pop(alias) + return json.dumps(a) + return args_json + + +def translate_client_tool( + name: str, args_json: str, valid: set, shell_tool: Optional[str], + cwd: Optional[str] = None, shell_param: Optional[str] = None, +): + """Return (name, args_json) with Claude read-only tools rewritten to the + registered shell tool. exec_command's own calls and any already-valid tool + pass through untouched; unknown non-shell tools fall back to name-remap. + ``cwd`` (from the request's ) makes hallucinated absolute paths resolve.""" + if name in valid: + if shell_param and _is_shell_name(name, shell_tool): + args_json = _normalize_shell_args(args_json, shell_param) + return name, args_json + exec_tool = "exec_command" if "exec_command" in valid else shell_tool + if exec_tool: + try: + a = json.loads(args_json) if args_json else {} + except Exception: + a = {} + if isinstance(a, dict): + cmd = _claude_tool_to_shell(name, a, cwd) + if cmd is not None: + return exec_tool, json.dumps({(shell_param or "cmd"): cmd}) + _remapped = remap_tool_name(name, valid, shell_tool) + if shell_param and _is_shell_name(_remapped, shell_tool): + args_json = _normalize_shell_args(args_json, shell_param) + return _remapped, args_json + + +_CWD_RE = None + + +def extract_cwd(body: Dict[str, Any]) -> Optional[str]: + """Pull the working directory from Codex's ... environment_context + (sent in instructions/input), so path fix-ups target the real directory.""" + import re + global _CWD_RE + if _CWD_RE is None: + _CWD_RE = re.compile(r"\s*([^<\s]+)\s*") + blob = _text_of(body.get("instructions")) + inp = body.get("input") + if isinstance(inp, str): + blob += "\n" + inp + elif isinstance(inp, list): + for it in inp: + if isinstance(it, dict): + blob += "\n" + _text_of(it.get("content", "")) + m = _CWD_RE.search(blob or "") + return m.group(1) if m else None + + +# --------------------------------------------------------------- SSE emitter +class ResponsesStreamEmitter: + """Builds the ordered Responses SSE event stream from incremental text / + tool-call events. Each method returns a list of SSE strings to yield. + + Output-item lifecycle (Codex expects these exact event types): + message: output_item.added -> content_part.added -> output_text.delta* + -> output_text.done -> content_part.done -> output_item.done + function_call: output_item.added -> function_call_arguments.delta* + -> function_call_arguments.done -> output_item.done + end: response.completed + """ + + def __init__(self, resp_id: str, model: str): + self.resp_id = resp_id + self.model = model + self._seq = itertools.count(0) + self.out_index = 0 + self.final_output: List[Dict[str, Any]] = [] + self._open: Optional[Dict[str, Any]] = None # current open output item + + def _ev(self, ev: str, extra: Dict[str, Any]) -> str: + d = {"type": ev, "sequence_number": next(self._seq)} + d.update(extra) + return f"event: {ev}\ndata: {json.dumps(d)}\n\n" + + def _base(self, status: str, output: List[Dict[str, Any]]) -> Dict[str, Any]: + return { + "id": self.resp_id, + "object": "response", + "status": status, + "model": self.model, + "output": output, + } + + def created(self) -> List[str]: + return [ + self._ev("response.created", {"response": self._base("in_progress", [])}), + self._ev( + "response.in_progress", {"response": self._base("in_progress", [])} + ), + ] + + def _close_open(self) -> List[str]: + if self._open is None: + return [] + o = self._open + self._open = None + if o["kind"] == "message": + mid, cur, txt = o["id"], o["index"], o["text"] + item = { + "id": mid, "type": "message", "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": txt}], + } + self.final_output.append(item) + return [ + self._ev("response.output_text.done", { + "item_id": mid, "output_index": cur, + "content_index": 0, "text": txt}), + self._ev("response.content_part.done", { + "item_id": mid, "output_index": cur, "content_index": 0, + "part": {"type": "output_text", "text": txt}}), + self._ev("response.output_item.done", { + "output_index": cur, "item": item}), + ] + # function_call + fid, cur = o["id"], o["index"] + item = { + "id": fid, "type": "function_call", "status": "completed", + "call_id": o["call_id"], "name": o["name"], "arguments": o["args"], + } + self.final_output.append(item) + return [ + self._ev("response.function_call_arguments.done", { + "item_id": fid, "output_index": cur, "arguments": o["args"]}), + self._ev("response.output_item.done", { + "output_index": cur, "item": item}), + ] + + def text_delta(self, delta: str) -> List[str]: + out: List[str] = [] + if self._open and self._open["kind"] != "message": + out += self._close_open() + if not self._open: + mid, cur = _rid("msg"), self.out_index + self.out_index += 1 + self._open = {"kind": "message", "id": mid, "index": cur, "text": ""} + out.append(self._ev("response.output_item.added", { + "output_index": cur, + "item": {"id": mid, "type": "message", "role": "assistant", + "status": "in_progress", "content": []}})) + out.append(self._ev("response.content_part.added", { + "item_id": mid, "output_index": cur, "content_index": 0, + "part": {"type": "output_text", "text": ""}})) + self._open["text"] += delta + out.append(self._ev("response.output_text.delta", { + "item_id": self._open["id"], "output_index": self._open["index"], + "content_index": 0, "delta": delta})) + return out + + def tool_start(self, call_id: str, name: str) -> List[str]: + out = self._close_open() + fid, cur = _rid("fc"), self.out_index + self.out_index += 1 + self._open = { + "kind": "fc", "id": fid, "index": cur, + "call_id": call_id or _rid("call"), "name": name, "args": "", + } + out.append(self._ev("response.output_item.added", { + "output_index": cur, + "item": {"id": fid, "type": "function_call", "status": "in_progress", + "call_id": self._open["call_id"], "name": name, + "arguments": ""}})) + return out + + def tool_args(self, delta: str) -> List[str]: + if not self._open or self._open["kind"] != "fc" or not delta: + return [] + self._open["args"] += delta + return [self._ev("response.function_call_arguments.delta", { + "item_id": self._open["id"], "output_index": self._open["index"], + "delta": delta})] + + def tool_end(self) -> List[str]: + return self._close_open() + + def finish(self, input_tokens: int, output_tokens: int) -> List[str]: + out = self._close_open() + completed = self._base("completed", self.final_output) + completed["usage"] = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + } + out.append(self._ev("response.completed", {"response": completed})) + return out + + +# ------------------------------------------------------- non-stream response +def build_responses_object( + resp_id: str, + model: str, + content_text: str, + tool_calls: List[Any], + input_tokens: int, + output_tokens: int, + valid: set, + shell_tool: Optional[str], + cwd: Optional[str] = None, +) -> Dict[str, Any]: + """Build a full (non-streaming) Responses object from parsed output. + + ``tool_calls`` are ATOM ``ToolCall`` objects (``.id``, ``.function`` dict).""" + output: List[Dict[str, Any]] = [] + if content_text: + output.append({ + "id": _rid("msg"), "type": "message", "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": content_text}], + }) + for tc in tool_calls or []: + fn = getattr(tc, "function", None) or {} + name, args = translate_client_tool( + fn.get("name", ""), fn.get("arguments", "") or "", valid, shell_tool, cwd + ) + output.append({ + "id": _rid("fc"), "type": "function_call", "status": "completed", + "call_id": getattr(tc, "id", None) or _rid("call"), + "name": name, "arguments": args, + }) + return { + "id": resp_id, "object": "response", "status": "completed", + "model": model, "output": output, + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + } diff --git a/atom/entrypoints/openai/tool_parser.py b/atom/entrypoints/openai/tool_parser.py index 549277e8d9..8f0b992b0e 100644 --- a/atom/entrypoints/openai/tool_parser.py +++ b/atom/entrypoints/openai/tool_parser.py @@ -183,21 +183,382 @@ def _parse_qwen_xml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCal return content.strip(), tool_calls +# --------------------------------------------------------------------------- +# DeepSeek-V4 DSML tool-call format +# --------------------------------------------------------------------------- +# +# <|DSML|tool_calls> +# <|DSML|invoke name="NAME"> +# <|DSML|parameter name="PNAME" string="true|false">VALUE +# ... +# +# +# +# string="true" -> value is a raw string; string="false" -> value is JSON. +# DeepSeek-V4-Flash occasionally malforms this (singular ``tool_call``, a missing +# ``invoke`` wrapper, or params without ``string=``); the parser recovers those +# best-effort: it infers a dropped tool name from the parameter signature vs the +# request's ``tools`` and infers a missing value type from the schema / JSON. + +_DSML = "|DSML|" +# The model often DROPS the ``|DSML|`` marker and emits bare +# ````/````/```` tags, so the marker +# is matched OPTIONALLY everywhere. +_OPT = r"(?:" + re.escape(_DSML) + r")?" # optional |DSML| prefix +_DSML_PARAM_RE = re.compile( + r"<" + _OPT + r'parameter\s+name="(.*?)"(?:\s+string="(true|false)")?\s*>' + r"(.*?)", + re.DOTALL, +) +# Long-form `...` OR self-closing `` +# (the zero-arg shape; group(2) is None for self-closing). Matches SGLang's V4 +# detector, which accepts both. +_DSML_INVOKE_RE = re.compile( + r"<" + _OPT + r'invoke\s+name="(.*?)"\s*(?:/>|>(.*?))", + re.DOTALL, +) +# Region-start markers, both marked and marker-less variants. +_DSML_STARTS = ( + "<" + _DSML + "tool_call", # marked (covers tool_call / tool_calls) + "<" + _DSML + "invoke", # marked invoke + "", # marker-less section open +) + + +def _dsml_start(text: str) -> int: + """Index of the earliest DSML tool-call marker (marked or marker-less), or -1.""" + positions = [i for i in (text.find(m) for m in _DSML_STARTS) if i != -1] + return min(positions) if positions else -1 + + +def _is_dsml(text: str) -> bool: + return _dsml_start(text) != -1 + + +def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: + """Strip spurious ``{"arguments": {...}}`` / ``{"input": {...}}`` envelopes. + + Non-tuned models (DeepSeek-V4-Pro) frequently wrap the real args in an extra + ``arguments``/``input`` object — sometimes nested 2-3 deep, or stringified — + so a call meant as ``{"cmd": "ls"}`` arrives as ``{"arguments": {"cmd": + "ls"}}`` and the client (Codex) rejects it ("missing field cmd"). Recursively + unwrap while the sole key is a wrapper that is NOT itself a declared param of + the tool. Mirrors vLLM's ``_unwrap_wrapper_args`` (deepseek_v4.py).""" + for _ in range(4): # bounded against pathological nesting + if not (isinstance(args, dict) and len(args) == 1): + break + (k, v), = args.items() + if k not in ("arguments", "input"): + break + if allowed and k in allowed: + break # this tool really has a param named arguments/input + if isinstance(v, str): + try: + v = json.loads(v) + except Exception: + break + if not isinstance(v, dict): + break + args = v + return args + + +# Canonical param key -> emitted synonyms the model uses interchangeably. Only +# applied when the canonical key is in the tool's declared schema and the model +# used a synonym instead (e.g. Codex's exec_command wants `cmd`, but the model +# habitually emits `command` -> "missing field cmd" retry loop). Scoped to the +# schema so it never renames a key a tool legitimately declares. +_KEY_ALIASES: Dict[str, Tuple[str, ...]] = { + "cmd": ("command",), +} + + +def _apply_key_aliases(args: Any, allowed: set) -> Any: + if not (isinstance(args, dict) and allowed): + return args + for canon, syns in _KEY_ALIASES.items(): + if canon in allowed and canon not in args: + for s in syns: + if s in args: + args[canon] = args.pop(s) + break + return args + + +def _dsml_coerce(value: str, string_attr: Optional[str], ptype: Any) -> Any: + if string_attr == "true": + return value + if string_attr == "false": + try: + return json.loads(value) + except Exception: + return value + # attr absent -> use declared schema type if known, else infer via JSON. + if ptype is not None: + return _coerce_param_value(value, ptype) + v = value.strip() + try: + return json.loads(v) + except Exception: + return v + + +def _infer_dsml_name(arg_names: set, param_types: Dict[str, Dict[str, Any]]) -> Optional[str]: + """Pick the request tool whose parameter set best matches ``arg_names``.""" + best, best_score = None, -1e9 + for name, props in param_types.items(): + p = set(props) + if not p: + continue + score = len(p & arg_names) - 0.1 * len(p ^ arg_names) + if score > best_score: + best_score, best = score, name + return best + + +def _parse_dsml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: + """Parse DeepSeek-V4 DSML tool calls; return (leading_content, tool_calls).""" + param_types = _build_param_types(tools) + start = _dsml_start(text) + if start == -1: + return text.strip(), [] + content = text[:start] + region = text[start:] + + calls: List[Tuple[str, Dict[str, Any]]] = [] + invokes = list(_DSML_INVOKE_RE.finditer(region)) + if invokes: + for m in invokes: + name = m.group(1) + body = m.group(2) or "" # None for self-closing + types = param_types.get(name, {}) + args: Dict[str, Any] = { + pm.group(1): _dsml_coerce(pm.group(3), pm.group(2), types.get(pm.group(1))) + for pm in _DSML_PARAM_RE.finditer(body) + } + # Direct-JSON parameter body (DSML "Format 2", also accepted by + # vLLM/SGLang): ` { "k": "v" } ` with no + # tags. Falls through here with empty args; recover them. + if not args: + stripped = body.strip() + if stripped.startswith("{"): + try: + parsed = json.loads(stripped) + if isinstance(parsed, dict): + args = parsed + except Exception: + pass + args = _unwrap_wrapper_args(args, set(types)) + args = _apply_key_aliases(args, set(types)) + calls.append((name, args)) + else: + # malformed: no complete invoke wrapper -> collect params, infer tool name + raw = {pm.group(1): (pm.group(3), pm.group(2)) for pm in _DSML_PARAM_RE.finditer(region)} + if raw: + name = _infer_dsml_name(set(raw), param_types) or "unknown" + types = param_types.get(name, {}) + args = {k: _dsml_coerce(v, s, types.get(k)) for k, (v, s) in raw.items()} + args = _unwrap_wrapper_args(args, set(types)) + args = _apply_key_aliases(args, set(types)) + calls.append((name, args)) + + tool_calls = [ + ToolCall( + id=_unique_tool_call_id(), + type="function", + function={"name": name, "arguments": json.dumps(args, ensure_ascii=False)}, + ) + for name, args in calls + ] + if _DSML in content: # scrub any stray marker fragment + content = content.split("<" + _DSML, 1)[0] + return content.strip(), tool_calls + + +# --------------------------------------------------------------------------- +# GLM-4.5 / 4.6 / 5.x tool-call format +# --------------------------------------------------------------------------- +# +# NAME +# K1V1 +# K2V2 +# ... +# +# The function name follows the opening tag directly (no ``(.*?)|(.*)$", re.DOTALL +) +_GLM_ARG_RE = re.compile( + r"(.*?)\s*" + r"(.*?)(?:|(?=)|(?=)|$)", + re.DOTALL, +) + + +def _is_glm(text: str) -> bool: + """Detect the GLM ``...`` format (never Qwen/DSML).""" + if _QWEN_TOOL_PREFIX in text: # ' Qwen, not GLM + return False + return "" in text or "" in text + + +def _glm_coerce(value: str, ptype: Any) -> Any: + """Decode one GLM ````: schema type wins, else JSON, else raw.""" + v = value.strip("\n") + if ptype is not None: + return _coerce_param_value(v, ptype) + s = v.strip() + try: + return json.loads(s) + except Exception: + return v + + +def _parse_glm(text: str, tools: Optional[list] = None) -> Tuple[str, List[ToolCall]]: + """Parse GLM tool calls; return (leading_content, tool_calls).""" + param_types = _build_param_types(tools) + start = text.find("") + if start == -1: + return text.strip(), [] + content = text[:start] + tool_calls: List[ToolCall] = [] + for m in _GLM_TOOLCALL_RE.finditer(text): + body = m.group(1) if m.group(1) is not None else m.group(2) + if not body: + continue + ak = body.find("") + name = (body if ak == -1 else body[:ak]).strip() + if not name: + continue + types = param_types.get(name, {}) + args: Dict[str, Any] = {} + for pm in _GLM_ARG_RE.finditer(body): + k = pm.group(1).strip() + if k: + args[k] = _glm_coerce(pm.group(2), types.get(k)) + tool_calls.append( + ToolCall( + id=_unique_tool_call_id(), + type="function", + function={ + "name": name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + ) + ) + return content.strip(), tool_calls + + +# --------------------------------------------------------------------------- +# MiniMax-M3 tool-call format +# --------------------------------------------------------------------------- +# +# Every tag is prefixed by the ns_token ``]<]minimax[>[``: +# +# ]<]minimax[>[ +# ]<]minimax[>[ +# ]<]minimax[>[value]<]minimax[>[ +# ... +# ]<]minimax[>[ +# ]<]minimax[>[ +# +# Unlike DSML, parameters are named by the TAG itself (``Paris``), +# not a ``name="..."`` attribute. Strip the ns_token first, then parse +# / pairs. Values: schema type wins, else JSON, else raw string. + +_MINIMAX_NS = "]<]minimax[>[" +_MINIMAX_INVOKE_RE = re.compile( + r'(.*?)|(.*)$', + re.DOTALL, +) +_MINIMAX_PARAM_RE = re.compile(r"<([\w-]+)>(.*?)", re.DOTALL) + + +def _is_minimax(text: str) -> bool: + """Detect the MiniMax-M3 ns_token tool-call format.""" + return _MINIMAX_NS in text + + +def _minimax_coerce(value: str, ptype: Any) -> Any: + v = value.strip("\n") + if ptype is not None: + return _coerce_param_value(v, ptype) + s = v.strip() + try: + return json.loads(s) + except Exception: + return v + + +def _parse_minimax(text: str, tools: Optional[list] = None) -> Tuple[str, List[ToolCall]]: + """Parse MiniMax-M3 tool calls; return (leading_content, tool_calls).""" + param_types = _build_param_types(tools) + clean = text.replace(_MINIMAX_NS, "") + tc = clean.find("") + content = clean[:tc] if tc > 0 else ("" if tc == 0 else clean) + tool_calls: List[ToolCall] = [] + for m in _MINIMAX_INVOKE_RE.finditer(clean): + name = m.group(1) if m.group(1) is not None else m.group(3) + body = m.group(2) if m.group(2) is not None else (m.group(4) or "") + if not name: + continue + name = name.strip() + types = param_types.get(name, {}) + args: Dict[str, Any] = {} + for pm in _MINIMAX_PARAM_RE.finditer(body): + k = pm.group(1).strip() + if k: + args[k] = _minimax_coerce(pm.group(2), types.get(k)) + tool_calls.append( + ToolCall( + id=_unique_tool_call_id(), + type="function", + function={ + "name": name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + ) + ) + for mk in ("", ""): + content = content.replace(mk, "") + return content.strip(), tool_calls + + def parse_tool_calls( text: str, tools: Optional[list] = None ) -> Tuple[str, List[ToolCall]]: """Parse tool calls from model output text. Args: - text: Raw model output that may contain tool calls (Kimi token format - or Qwen3 XML format). - tools: Optional request tool definitions; used to type-coerce Qwen XML - parameter values to their declared JSON-Schema types. + text: Raw model output that may contain tool calls (DeepSeek-V4 DSML, + Kimi token format, or Qwen3 XML format). + tools: Optional request tool definitions; used to type-coerce parameter + values to their declared JSON-Schema types. Returns: Tuple of (content_text, list_of_tool_calls). ``content_text`` has the tool-call sections removed. """ + # MiniMax-M3 ns_token format (checked before DSML: both use , + # but MiniMax names params by tag and prefixes every tag with ]<]minimax[>[) + if _is_minimax(text): + return _parse_minimax(text, tools) + + # DeepSeek-V4 DSML format + if _is_dsml(text): + return _parse_dsml(text, tools) + + # GLM / format (checked before Qwen: both use + # , but GLM never emits the Qwen ' list: """Process a text chunk and return list of (event_type, data) tuples.""" if self.fmt is None: self.buf += text - if _QWEN_TOOL_PREFIX in self.buf or "" in self.buf: + if _MINIMAX_NS in self.buf: + self.fmt = "minimax" + elif _is_dsml(self.buf): + self.fmt = "dsml" + elif "" in self.buf: + self.fmt = "glm" + elif _QWEN_TOOL_PREFIX in self.buf: self.fmt = "qwen" + elif "" in self.buf: + # '' seen but neither '' (GLM) yet. A no-arg GLM call is complete once the + # closing tag arrives; otherwise wait for the sub-marker. + if "" in self.buf: + self.fmt = "glm" + else: + return [] elif "<|tool_calls_section_begin|>" in self.buf: self.fmt = "kimi" elif "<" not in self.buf and len(self.buf) > 8: @@ -297,10 +672,202 @@ def process(self, text: str) -> list: # Format decided: replay the accumulated buffer through the handler. text, self.buf = self.buf, "" + if self.fmt == "minimax": + return self._process_minimax(text) + if self.fmt == "dsml": + return self._process_dsml(text) + if self.fmt == "glm": + return self._process_glm(text) if self.fmt == "qwen": return self._process_qwen(text) return self._process_kimi(text) + # -- MiniMax-M3 ns_token ------------------------------------------------ + def _process_minimax(self, text: str) -> list: + results: list = [] + self.buf += text + if self.state == 0: + markers = [ + i + for i in (self.buf.find(_MINIMAX_NS), self.buf.find("")) + if i != -1 + ] + if markers: + m = min(markers) + before = self.buf[:m] + if before: + results.append(("content", before)) + self.buf = self.buf[m:] + self.state = 1 + else: + cut = self.buf.rfind("<") + cut = max(cut, self.buf.rfind("]")) # ns_token starts with ']' + if cut == -1: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + elif cut > 0: + results.append(("content", self.buf[:cut])) + self.buf = self.buf[cut:] + return results + + def _flush_minimax(self) -> list: + results: list = [] + if self.state == 0: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + return results + _content, tool_calls = _parse_minimax(self.buf, self.tools) + self.buf = "" + for tc in tool_calls: + results.append( + ( + "tool_call_start", + { + "index": self.current_index, + "id": tc.id, + "type": "function", + "function": {"name": tc.function["name"], "arguments": ""}, + }, + ) + ) + results.append( + ( + "tool_call_args", + { + "index": self.current_index, + "function": {"arguments": tc.function["arguments"]}, + }, + ) + ) + self.current_index += 1 + self._emitted_calls += 1 + if self._emitted_calls > 0: + results.append(("tool_call_end", None)) + return results + + # -- DeepSeek-V4 DSML --------------------------------------------------- + def _process_dsml(self, text: str) -> list: + results: list = [] + self.buf += text + if self.state == 0: + m = _dsml_start(self.buf) + if m != -1: + before = self.buf[:m] + if before: + results.append(("content", before)) + self.buf = self.buf[m:] + self.state = 1 + else: + # Emit content but hold back a possible partial '<...' marker tail. + cut = self.buf.rfind("<") + if cut == -1: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + elif cut > 0: + results.append(("content", self.buf[:cut])) + self.buf = self.buf[cut:] + return results + + def _flush_dsml(self) -> list: + results: list = [] + if self.state == 0: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + return results + _content, tool_calls = _parse_dsml(self.buf, self.tools) + self.buf = "" + for tc in tool_calls: + results.append( + ( + "tool_call_start", + { + "index": self.current_index, + "id": tc.id, + "type": "function", + "function": {"name": tc.function["name"], "arguments": ""}, + }, + ) + ) + results.append( + ( + "tool_call_args", + { + "index": self.current_index, + "function": {"arguments": tc.function["arguments"]}, + }, + ) + ) + self.current_index += 1 + self._emitted_calls += 1 + if self._emitted_calls > 0: + results.append(("tool_call_end", None)) + return results + + # -- Qwen3 XML ---------------------------------------------------------- + # -- GLM / ----------------------------------------- + def _process_glm(self, text: str) -> list: + results: list = [] + self.buf += text + if self.state == 0: + m = self.buf.find("") + if m != -1: + before = self.buf[:m] + if before: + results.append(("content", before)) + self.buf = self.buf[m:] + self.state = 1 + else: + # Emit content but hold back a possible partial '<...' marker tail. + cut = self.buf.rfind("<") + if cut == -1: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + elif cut > 0: + results.append(("content", self.buf[:cut])) + self.buf = self.buf[cut:] + return results + + def _flush_glm(self) -> list: + results: list = [] + if self.state == 0: + if self.buf: + results.append(("content", self.buf)) + self.buf = "" + return results + _content, tool_calls = _parse_glm(self.buf, self.tools) + self.buf = "" + for tc in tool_calls: + results.append( + ( + "tool_call_start", + { + "index": self.current_index, + "id": tc.id, + "type": "function", + "function": {"name": tc.function["name"], "arguments": ""}, + }, + ) + ) + results.append( + ( + "tool_call_args", + { + "index": self.current_index, + "function": {"arguments": tc.function["arguments"]}, + }, + ) + ) + self.current_index += 1 + self._emitted_calls += 1 + if self._emitted_calls > 0: + results.append(("tool_call_end", None)) + return results + # -- Qwen3 XML ---------------------------------------------------------- def _process_qwen(self, text: str) -> list: results: list = [] @@ -449,6 +1016,12 @@ def _process_buffer(self) -> list: def flush(self) -> list: """Flush remaining buffer content.""" + if self.fmt == "minimax": + return self._flush_minimax() + if self.fmt == "dsml": + return self._flush_dsml() + if self.fmt == "glm": + return self._flush_glm() if self.fmt == "qwen": return self._flush_qwen() results = [] From 999168e97ecea9f7c69b87942ae536e2802674cc Mon Sep 17 00:00:00 2001 From: yihonglie Date: Sun, 12 Jul 2026 22:24:17 -0500 Subject: [PATCH 6/6] style: apply black formatting Fixes "Check Code Style with Black" CI on the DSV4 API + DSML tool parser changes (serving_responses.py, tool_parser.py, api_server.py). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: yihonglie --- atom/entrypoints/openai/api_server.py | 33 ++- atom/entrypoints/openai/chat_encoders.py | 10 +- atom/entrypoints/openai/serving_responses.py | 296 ++++++++++++++----- atom/entrypoints/openai/tool_parser.py | 29 +- 4 files changed, 270 insertions(+), 98 deletions(-) diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index a37bdf14fa..a276726c6a 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -408,27 +408,29 @@ def flush_stream_batch() -> None: # Incremental per-request detokenization: correct UTF-8 at token # boundaries (see _stream_detok_state). Emits only fully-formed chars; # a trailing partial multi-byte char is held until the next step. - for (_loop, _q, chunk) in buf: + for _loop, _q, chunk in buf: rid = chunk.get("request_id") st = _stream_detok_state.get(rid) if st is None: st = _stream_detok_state[rid] = { - "tokens": [], "prefix_offset": 0, "read_offset": 0, + "tokens": [], + "prefix_offset": 0, + "read_offset": 0, } toks = st["tokens"] toks.extend(chunk["token_ids"]) prefix_text = tokenizer.decode( - toks[st["prefix_offset"]:st["read_offset"]], skip_special_tokens=True + toks[st["prefix_offset"] : st["read_offset"]], skip_special_tokens=True ) new_text = tokenizer.decode( - toks[st["prefix_offset"]:], skip_special_tokens=True + toks[st["prefix_offset"] :], skip_special_tokens=True ) if len(new_text) > len(prefix_text) and not new_text.endswith("\ufffd"): - chunk["text"] = new_text[len(prefix_text):] + chunk["text"] = new_text[len(prefix_text) :] st["prefix_offset"] = st["read_offset"] st["read_offset"] = len(toks) elif chunk["finished"]: - chunk["text"] = new_text[len(prefix_text):] + chunk["text"] = new_text[len(prefix_text) :] else: chunk["text"] = "" if chunk["finished"]: @@ -1694,7 +1696,9 @@ async def generate_anthropic_stream(): raw_text = final_output["text"] reasoning_content, content_with_tools = separate_reasoning(raw_text) - content_text, tool_calls = parse_tool_calls(content_with_tools, anthropic_to_openai_tools(request.tools)) + content_text, tool_calls = parse_tool_calls( + content_with_tools, anthropic_to_openai_tools(request.tools) + ) output_tokens = len(tokenizer.encode(raw_text)) cache_read_input_tokens = final_output.get("num_cached_tokens", 0) if not getattr(request, "thinking", None): @@ -1768,7 +1772,9 @@ async def responses_endpoint(raw_request: Request): max_out = int(body.get("max_output_tokens") or 32768) sampling_params = _build_sampling_params( - temperature=body.get("temperature") if body.get("temperature") is not None else 1.0, + temperature=( + body.get("temperature") if body.get("temperature") is not None else 1.0 + ), max_tokens=max_out, stop_strings=None, ignore_eos=False, @@ -1832,7 +1838,12 @@ def _flush_pending(): return [] _pending["tc"] = None name, args = translate_client_tool( - tc["name"], tc["args"], valid_names, shell_tool, req_cwd, shell_param + tc["name"], + tc["args"], + valid_names, + shell_tool, + req_cwd, + shell_param, ) out = emitter.tool_start(tc["id"], name) if args: @@ -1914,7 +1925,9 @@ def handle(etype, edata): raw_text = final_output["text"] _reasoning, content_with_tools = separate_reasoning(raw_text) - content_text, tool_calls = parse_tool_calls(content_with_tools, openai_tools or None) + content_text, tool_calls = parse_tool_calls( + content_with_tools, openai_tools or None + ) output_tokens = len(tokenizer.encode(raw_text)) return JSONResponse( diff --git a/atom/entrypoints/openai/chat_encoders.py b/atom/entrypoints/openai/chat_encoders.py index 863d045275..9021ab85d3 100644 --- a/atom/entrypoints/openai/chat_encoders.py +++ b/atom/entrypoints/openai/chat_encoders.py @@ -88,7 +88,9 @@ def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoder]: def _content_str(c: Any) -> str: if isinstance(c, list): return "\n".join( - b.get("text", "") for b in c if isinstance(b, dict) and b.get("type") == "text" + b.get("text", "") + for b in c + if isinstance(b, dict) and b.get("type") == "text" ) return c or "" @@ -115,9 +117,11 @@ def _normalize_for_v4(messages: List[dict], tools: Optional[List[dict]]) -> List if not sys_parts and not tools: return [dict(m) for m in messages] - merged = "\n\n".join(s for s in (_content_str(m.get("content")) for m in sys_parts) if s) + merged = "\n\n".join( + s for s in (_content_str(m.get("content")) for m in sys_parts) if s + ) sys_msg: dict = {"role": "system", "content": merged} - for m in sys_parts: # preserve any pre-attached tools + for m in sys_parts: # preserve any pre-attached tools if m.get("tools"): sys_msg["tools"] = m["tools"] if tools: diff --git a/atom/entrypoints/openai/serving_responses.py b/atom/entrypoints/openai/serving_responses.py index bc14091ad0..26d22af821 100644 --- a/atom/entrypoints/openai/serving_responses.py +++ b/atom/entrypoints/openai/serving_responses.py @@ -11,6 +11,7 @@ This makes the external ``codex_responses_proxy.py`` unnecessary: Codex can point straight at ATOM's ``:9700/v1``. """ + import itertools import json import time @@ -42,9 +43,7 @@ def _text_of(content: Any) -> str: return "" -def responses_input_to_messages( - instructions: Any, inp: Any -) -> List[Dict[str, Any]]: +def responses_input_to_messages(instructions: Any, inp: Any) -> List[Dict[str, Any]]: """Translate Responses ``instructions`` + ``input`` into OpenAI chat messages. ``input`` may be a plain string or a list of items: ``message`` / @@ -92,9 +91,7 @@ def responses_input_to_messages( { "role": "tool", "tool_call_id": item.get("call_id") or item.get("id") or "", - "content": out - if isinstance(out, str) - else json.dumps(out), + "content": out if isinstance(out, str) else json.dumps(out), } ) elif t == "reasoning": @@ -109,8 +106,8 @@ def responses_input_to_messages( " tags, no prose. Use EXACTLY this syntax (the \uff5c characters " "are U+FF5C fullwidth vertical bars, not ASCII '|'):\n" "<\uff5cDSML\uff5ctool_calls>\n" - "<\uff5cDSML\uff5cinvoke name=\"TOOL_NAME\">\n" - "<\uff5cDSML\uff5cparameter name=\"PARAM_NAME\" string=\"true\">VALUE" + '<\uff5cDSML\uff5cinvoke name="TOOL_NAME">\n' + '<\uff5cDSML\uff5cparameter name="PARAM_NAME" string="true">VALUE' "\n" "\n" "\n" @@ -130,7 +127,9 @@ def inject_tool_format_instruction(messages): if "\uff5cDSML\uff5ctool_calls" not in base: m["content"] = _text_of(base) + _DSML_TOOL_INSTRUCTION return messages - return [{"role": "system", "content": _DSML_TOOL_INSTRUCTION.strip()}] + list(messages) + return [{"role": "system", "content": _DSML_TOOL_INSTRUCTION.strip()}] + list( + messages + ) def responses_tools_to_openai(tools: Any) -> List[Dict[str, Any]]: @@ -172,17 +171,40 @@ def responses_tools_to_openai(tools: Any) -> List[Dict[str, Any]]: # is wrong. So remap known shell-exec aliases onto whatever shell tool the # client actually registered this turn. Pure rename; arguments untouched. _SHELL_ALIASES = { - "exec", "exec_run", "exec_command", "execute_command", "shell", "bash", - "sh", "run", "run_command", "run_shell", "execute", "command", - "container.exec", "local_shell", "shell_command", "shell_exec", - "run_bash", "execute_shell", "bash_command", "run_terminal_cmd", - "terminal", "console", "run_shell_command", "runshell", + "exec", + "exec_run", + "exec_command", + "execute_command", + "shell", + "bash", + "sh", + "run", + "run_command", + "run_shell", + "execute", + "command", + "container.exec", + "local_shell", + "shell_command", + "shell_exec", + "run_bash", + "execute_shell", + "bash_command", + "run_terminal_cmd", + "terminal", + "console", + "run_shell_command", + "runshell", } # Substrings that mark an unknown tool name as a shell/exec call (Codex's # non-shell tools contain none of these). _SHELL_NAME_TOKENS = ("shell", "exec", "bash", "cmd", "command", "termin", "console") _SHELL_TOOL_PREFERENCE = ( - "exec_command", "shell", "local_shell", "bash", "container.exec", + "exec_command", + "shell", + "local_shell", + "bash", + "container.exec", ) @@ -226,6 +248,7 @@ def remap_tool_name(name: str, valid: set, shell_tool: Optional[str]) -> str: # to /v1/messages, where read/grep ARE native tools). def _q(s: Any) -> str: import shlex + return shlex.quote(str(s)) @@ -261,13 +284,32 @@ def _claude_tool_to_shell( name: str, a: Dict[str, Any], cwd: Optional[str] = None ) -> Optional[str]: n = (name or "").lower() - fp = (a.get("file_path") or a.get("path") or a.get("filePath") - or a.get("filename") or a.get("target_file") or a.get("file")) + fp = ( + a.get("file_path") + or a.get("path") + or a.get("filePath") + or a.get("filename") + or a.get("target_file") + or a.get("file") + ) pattern = a.get("pattern") or a.get("query") or a.get("regex") - path = (a.get("path") or a.get("directory") or a.get("target_directory") - or a.get("dir") or ".") - if n in ("read", "cat", "view", "view_file", "open", "read_file", - "readfile", "openfile"): + path = ( + a.get("path") + or a.get("directory") + or a.get("target_directory") + or a.get("dir") + or "." + ) + if n in ( + "read", + "cat", + "view", + "view_file", + "open", + "read_file", + "readfile", + "openfile", + ): if not fp: return None off, lim = a.get("offset"), a.get("limit") @@ -277,8 +319,15 @@ def _claude_tool_to_shell( else: start, end = 1, 400 return _read_cmd(str(fp), cwd, start, end) - if n in ("grep", "search", "search_file", "ripgrep", "rg", "grep_search", - "codebase_search"): + if n in ( + "grep", + "search", + "search_file", + "ripgrep", + "rg", + "grep_search", + "codebase_search", + ): if not pattern: return None return f"grep -rn -- {_q(pattern)} {_q(_resolve_dir(path, cwd))}" @@ -293,8 +342,19 @@ def _claude_tool_to_shell( _SHELL_ARG_ALIASES = ( - "cmd", "command", "commandline", "command_line", "script", "bash", "sh", - "shell", "shell_command", "code", "input", "run", "cmd_string", + "cmd", + "command", + "commandline", + "command_line", + "script", + "bash", + "sh", + "shell", + "shell_command", + "code", + "input", + "run", + "cmd_string", ) @@ -309,7 +369,7 @@ def shell_arg_key(openai_tools, shell_tool): continue params = fn.get("parameters") or {} props = params.get("properties") or {} - for r in (params.get("required") or []): + for r in params.get("required") or []: if props.get(r, {}).get("type") in (None, "string"): return r if props: @@ -341,8 +401,12 @@ def _normalize_shell_args(args_json, req): def translate_client_tool( - name: str, args_json: str, valid: set, shell_tool: Optional[str], - cwd: Optional[str] = None, shell_param: Optional[str] = None, + name: str, + args_json: str, + valid: set, + shell_tool: Optional[str], + cwd: Optional[str] = None, + shell_param: Optional[str] = None, ): """Return (name, args_json) with Claude read-only tools rewritten to the registered shell tool. exec_command's own calls and any already-valid tool @@ -375,6 +439,7 @@ def extract_cwd(body: Dict[str, Any]) -> Optional[str]: """Pull the working directory from Codex's ... environment_context (sent in instructions/input), so path fix-ups target the real directory.""" import re + global _CWD_RE if _CWD_RE is None: _CWD_RE = re.compile(r"\s*([^<\s]+)\s*") @@ -441,33 +506,53 @@ def _close_open(self) -> List[str]: if o["kind"] == "message": mid, cur, txt = o["id"], o["index"], o["text"] item = { - "id": mid, "type": "message", "role": "assistant", + "id": mid, + "type": "message", + "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": txt}], } self.final_output.append(item) return [ - self._ev("response.output_text.done", { - "item_id": mid, "output_index": cur, - "content_index": 0, "text": txt}), - self._ev("response.content_part.done", { - "item_id": mid, "output_index": cur, "content_index": 0, - "part": {"type": "output_text", "text": txt}}), - self._ev("response.output_item.done", { - "output_index": cur, "item": item}), + self._ev( + "response.output_text.done", + { + "item_id": mid, + "output_index": cur, + "content_index": 0, + "text": txt, + }, + ), + self._ev( + "response.content_part.done", + { + "item_id": mid, + "output_index": cur, + "content_index": 0, + "part": {"type": "output_text", "text": txt}, + }, + ), + self._ev( + "response.output_item.done", {"output_index": cur, "item": item} + ), ] # function_call fid, cur = o["id"], o["index"] item = { - "id": fid, "type": "function_call", "status": "completed", - "call_id": o["call_id"], "name": o["name"], "arguments": o["args"], + "id": fid, + "type": "function_call", + "status": "completed", + "call_id": o["call_id"], + "name": o["name"], + "arguments": o["args"], } self.final_output.append(item) return [ - self._ev("response.function_call_arguments.done", { - "item_id": fid, "output_index": cur, "arguments": o["args"]}), - self._ev("response.output_item.done", { - "output_index": cur, "item": item}), + self._ev( + "response.function_call_arguments.done", + {"item_id": fid, "output_index": cur, "arguments": o["args"]}, + ), + self._ev("response.output_item.done", {"output_index": cur, "item": item}), ] def text_delta(self, delta: str) -> List[str]: @@ -478,17 +563,44 @@ def text_delta(self, delta: str) -> List[str]: mid, cur = _rid("msg"), self.out_index self.out_index += 1 self._open = {"kind": "message", "id": mid, "index": cur, "text": ""} - out.append(self._ev("response.output_item.added", { - "output_index": cur, - "item": {"id": mid, "type": "message", "role": "assistant", - "status": "in_progress", "content": []}})) - out.append(self._ev("response.content_part.added", { - "item_id": mid, "output_index": cur, "content_index": 0, - "part": {"type": "output_text", "text": ""}})) + out.append( + self._ev( + "response.output_item.added", + { + "output_index": cur, + "item": { + "id": mid, + "type": "message", + "role": "assistant", + "status": "in_progress", + "content": [], + }, + }, + ) + ) + out.append( + self._ev( + "response.content_part.added", + { + "item_id": mid, + "output_index": cur, + "content_index": 0, + "part": {"type": "output_text", "text": ""}, + }, + ) + ) self._open["text"] += delta - out.append(self._ev("response.output_text.delta", { - "item_id": self._open["id"], "output_index": self._open["index"], - "content_index": 0, "delta": delta})) + out.append( + self._ev( + "response.output_text.delta", + { + "item_id": self._open["id"], + "output_index": self._open["index"], + "content_index": 0, + "delta": delta, + }, + ) + ) return out def tool_start(self, call_id: str, name: str) -> List[str]: @@ -496,23 +608,45 @@ def tool_start(self, call_id: str, name: str) -> List[str]: fid, cur = _rid("fc"), self.out_index self.out_index += 1 self._open = { - "kind": "fc", "id": fid, "index": cur, - "call_id": call_id or _rid("call"), "name": name, "args": "", + "kind": "fc", + "id": fid, + "index": cur, + "call_id": call_id or _rid("call"), + "name": name, + "args": "", } - out.append(self._ev("response.output_item.added", { - "output_index": cur, - "item": {"id": fid, "type": "function_call", "status": "in_progress", - "call_id": self._open["call_id"], "name": name, - "arguments": ""}})) + out.append( + self._ev( + "response.output_item.added", + { + "output_index": cur, + "item": { + "id": fid, + "type": "function_call", + "status": "in_progress", + "call_id": self._open["call_id"], + "name": name, + "arguments": "", + }, + }, + ) + ) return out def tool_args(self, delta: str) -> List[str]: if not self._open or self._open["kind"] != "fc" or not delta: return [] self._open["args"] += delta - return [self._ev("response.function_call_arguments.delta", { - "item_id": self._open["id"], "output_index": self._open["index"], - "delta": delta})] + return [ + self._ev( + "response.function_call_arguments.delta", + { + "item_id": self._open["id"], + "output_index": self._open["index"], + "delta": delta, + }, + ) + ] def tool_end(self) -> List[str]: return self._close_open() @@ -546,24 +680,36 @@ def build_responses_object( ``tool_calls`` are ATOM ``ToolCall`` objects (``.id``, ``.function`` dict).""" output: List[Dict[str, Any]] = [] if content_text: - output.append({ - "id": _rid("msg"), "type": "message", "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": content_text}], - }) + output.append( + { + "id": _rid("msg"), + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": content_text}], + } + ) for tc in tool_calls or []: fn = getattr(tc, "function", None) or {} name, args = translate_client_tool( fn.get("name", ""), fn.get("arguments", "") or "", valid, shell_tool, cwd ) - output.append({ - "id": _rid("fc"), "type": "function_call", "status": "completed", - "call_id": getattr(tc, "id", None) or _rid("call"), - "name": name, "arguments": args, - }) + output.append( + { + "id": _rid("fc"), + "type": "function_call", + "status": "completed", + "call_id": getattr(tc, "id", None) or _rid("call"), + "name": name, + "arguments": args, + } + ) return { - "id": resp_id, "object": "response", "status": "completed", - "model": model, "output": output, + "id": resp_id, + "object": "response", + "status": "completed", + "model": model, + "output": output, "usage": { "input_tokens": input_tokens, "output_tokens": output_tokens, diff --git a/atom/entrypoints/openai/tool_parser.py b/atom/entrypoints/openai/tool_parser.py index 8f0b992b0e..21a98f5df4 100644 --- a/atom/entrypoints/openai/tool_parser.py +++ b/atom/entrypoints/openai/tool_parser.py @@ -204,7 +204,7 @@ def _parse_qwen_xml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCal # The model often DROPS the ``|DSML|`` marker and emits bare # ````/````/```` tags, so the marker # is matched OPTIONALLY everywhere. -_OPT = r"(?:" + re.escape(_DSML) + r")?" # optional |DSML| prefix +_OPT = r"(?:" + re.escape(_DSML) + r")?" # optional |DSML| prefix _DSML_PARAM_RE = re.compile( r"<" + _OPT + r'parameter\s+name="(.*?)"(?:\s+string="(true|false)")?\s*>' r"(.*?)", @@ -219,10 +219,10 @@ def _parse_qwen_xml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCal ) # Region-start markers, both marked and marker-less variants. _DSML_STARTS = ( - "<" + _DSML + "tool_call", # marked (covers tool_call / tool_calls) - "<" + _DSML + "invoke", # marked invoke - "", # marker-less section open + "<" + _DSML + "tool_call", # marked (covers tool_call / tool_calls) + "<" + _DSML + "invoke", # marked invoke + "", # marker-less section open ) @@ -248,7 +248,7 @@ def _unwrap_wrapper_args(args: Any, allowed: set) -> Any: for _ in range(4): # bounded against pathological nesting if not (isinstance(args, dict) and len(args) == 1): break - (k, v), = args.items() + ((k, v),) = args.items() if k not in ("arguments", "input"): break if allowed and k in allowed: @@ -304,7 +304,9 @@ def _dsml_coerce(value: str, string_attr: Optional[str], ptype: Any) -> Any: return v -def _infer_dsml_name(arg_names: set, param_types: Dict[str, Dict[str, Any]]) -> Optional[str]: +def _infer_dsml_name( + arg_names: set, param_types: Dict[str, Dict[str, Any]] +) -> Optional[str]: """Pick the request tool whose parameter set best matches ``arg_names``.""" best, best_score = None, -1e9 for name, props in param_types.items(): @@ -334,7 +336,9 @@ def _parse_dsml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: body = m.group(2) or "" # None for self-closing types = param_types.get(name, {}) args: Dict[str, Any] = { - pm.group(1): _dsml_coerce(pm.group(3), pm.group(2), types.get(pm.group(1))) + pm.group(1): _dsml_coerce( + pm.group(3), pm.group(2), types.get(pm.group(1)) + ) for pm in _DSML_PARAM_RE.finditer(body) } # Direct-JSON parameter body (DSML "Format 2", also accepted by @@ -354,7 +358,10 @@ def _parse_dsml(text: str, tools: Optional[list]) -> Tuple[str, List[ToolCall]]: calls.append((name, args)) else: # malformed: no complete invoke wrapper -> collect params, infer tool name - raw = {pm.group(1): (pm.group(3), pm.group(2)) for pm in _DSML_PARAM_RE.finditer(region)} + raw = { + pm.group(1): (pm.group(3), pm.group(2)) + for pm in _DSML_PARAM_RE.finditer(region) + } if raw: name = _infer_dsml_name(set(raw), param_types) or "unknown" types = param_types.get(name, {}) @@ -496,7 +503,9 @@ def _minimax_coerce(value: str, ptype: Any) -> Any: return v -def _parse_minimax(text: str, tools: Optional[list] = None) -> Tuple[str, List[ToolCall]]: +def _parse_minimax( + text: str, tools: Optional[list] = None +) -> Tuple[str, List[ToolCall]]: """Parse MiniMax-M3 tool calls; return (leading_content, tool_calls).""" param_types = _build_param_types(tools) clean = text.replace(_MINIMAX_NS, "")