fix: Improve cancellation for OpenAI response streams

[pskiran1]

What does the PR do?

This PR ensures that Triton inference is cancelled when an OpenAI chat or completion stream terminates before completion. It centralizes best-effort cancellation and error reporting, while adding backend-observed tests for tool-call truncation and abandoned streams.

• Adds shared early-termination cleanup for chat and completion streams.
• Routes tool-call parser truncation through the common cancellation helper.
• Preserves normal stream completion without unnecessary cancellation.
• Enhances the test model with delayed output and cancellation-state tracking.
• Adds coverage for truncation, cancellation failures, and abandoned chat and completion streams.

Checklist

• PR title reflects the change and is of format <commit_type>: <Title>
• Changes are described in the pull request.
• Related issues are referenced.
• Populated github labels field
• Added test plan and verified test passes.
• Verified that the PR passes existing CI.
• Verified copyright is correct on all changed files.
• Added succinct git squash message before merging ref.
• All template sections are filled out.
• Optional: Additional screenshots for behavior/output changes with before/after.

Commit Type:

Check the conventional commit type
box here and add the label to the github PR.

• build
• ci
• docs
• feat
• fix
• perf
• refactor
• revert
• style
• test

Related PRs:

Where should the reviewer start?

Test plan:

• CI Pipeline ID: 63699899

Caveats:

Background

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• closes GitHub issue: #xxx

[greptile-apps]

Greptile Summary

The PR centralizes best-effort cancellation of Triton inference when OpenAI chat or completion streams terminate early.

• Wraps chat and completion streams in shared early-termination cleanup.
• Routes tool-call truncation through the common cancellation helper.
• Adds backend-observed coverage for abandoned streams, truncation, and cancellation errors.
• Updates the engine protocol to reflect its asynchronous interface.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

No blocking failure remains.

Important Files Changed

Filename	Overview
python/openai/openai_frontend/engine/engine.py	Updates the engine protocol to declare asynchronous chat, completion, and embedding methods and asynchronous streaming return types.
python/openai/openai_frontend/engine/triton_engine.py	Adds shared cancellation and cleanup around prematurely terminated chat and completion response streams.
python/openai/tests/test_tool_calling.py	Adds coverage for cancellation after truncation, explicit stream abandonment, normal completion, and cancellation errors.
python/openai/tests/tool_parser_models/tool_stream_py/1/model.py	Extends the test model with delayed output and observable cancellation state.

Sequence Diagram

sequenceDiagram
    participant Client
    participant FastAPI
    participant Engine
    participant Triton
    Client->>FastAPI: Streaming chat/completion request
    FastAPI->>Engine: chat() / completion()
    Engine->>Triton: async_infer()
    Triton-->>Engine: Response iterator
    Engine-->>FastAPI: Wrapped async stream
    FastAPI-->>Client: SSE chunks
    alt Stream completes normally
        Engine-->>Client: Final chunk and [DONE]
    else Stream terminates early
        Engine->>Triton: cancel()
        Engine->>Engine: Close inner stream
    end

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into spolisetty/tri-..." | Re-trigger Greptile}

[yinggeh]

_stream_with_inference_cancellation returns AsynsGenerator, not matching the return type of completion function. How does it work?

[pskiran1]

completion() is awaited by the router. For a streaming request, the result of that await is the AsyncGenerator returned by _stream_with_inference_cancellation, which is then consumed by StreamingResponse.

AsyncGenerator[str, None] is a subtype of AsyncIterator[str], so it satisfies the declared return type of completion(). The shared LLMEngine interface has also been updated to declare these methods as asynchronous and use AsyncIterator[str] for streaming results.