Skip to content

fix: populate missing tee_id, tee_endpoint, tee_payment_address in StreamChunk.from_sse_data#197

Open
amathxbt wants to merge 1 commit intoOpenGradient:mainfrom
amathxbt:fix/streamchunk-missing-tee-metadata-fields
Open

fix: populate missing tee_id, tee_endpoint, tee_payment_address in StreamChunk.from_sse_data#197
amathxbt wants to merge 1 commit intoOpenGradient:mainfrom
amathxbt:fix/streamchunk-missing-tee-metadata-fields

Conversation

@amathxbt
Copy link
Contributor

Bug

StreamChunk declares three TEE metadata fields for identifying the enclave that served a streaming request:

tee_id: Optional[str] = None
tee_endpoint: Optional[str] = None  
tee_payment_address: Optional[str] = None

However, StreamChunk.from_sse_data() never reads these from the SSE data dictionary:

# BEFORE (broken) — tee_id, tee_endpoint, tee_payment_address silently dropped
return cls(choices=choices, model=..., usage=usage, is_final=is_final)

This means all three fields are always None in streaming responses, even when the TEE server sends them — making it impossible to:

  • Identify which TEE served the request
  • Verify the registered endpoint on-chain
  • Look up the payment address

Fix

Expand the return cls(...) call to read the three missing fields from the SSE data dictionary:

# AFTER (fixed)
return cls(
    choices=choices,
    model=data.get("model", "unknown"),
    usage=usage,
    is_final=is_final,
    tee_signature=data.get("tee_signature"),
    tee_timestamp=data.get("tee_timestamp"),
    tee_id=data.get("tee_id"),              # was missing
    tee_endpoint=data.get("tee_endpoint"),  # was missing
    tee_payment_address=data.get("tee_payment_address"),  # was missing
)

Impact

  • Streaming TEE metadata is now accessible — callers can identify and verify the serving TEE from the final chunk
  • Non-streaming path (TextGenerationOutput) already populates these fields correctly — this aligns the two code paths
  • Backward compatible: fields still default to None if not present in the response

…k.from_sse_data

These three TEE metadata fields exist on the StreamChunk dataclass but
were never read from the SSE data dict in from_sse_data().
As a result, tee_id/tee_endpoint/tee_payment_address are always None
in streaming responses, making it impossible to identify or verify
which TEE served the request.
@amathxbt
Copy link
Contributor Author

@adambalogh

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant