fix: missing components from output schema if it is a root model (#4733)

* fix: missing components from output schema if it is a root model

Signed-off-by: Frost Ming <[email protected]>

* fix: request hook to set trace id

Signed-off-by: Frost Ming <[email protected]>

* fix: use the same logging format for trace headers

Signed-off-by: Frost Ming <[email protected]>

---------

Signed-off-by: Frost Ming <[email protected]>

src/_bentoml_impl/server/app.py

@@ -220,7 +220,7 @@ def middlewares(self) -> list[Middleware]:
                 Middleware(CORSMiddleware, **self.access_control_options)
             )
 
-        def client_request_hook(span: Span | None, _scope: dict[str, t.Any]) -> None:
+        def server_request_hook(span: Span | None, _scope: dict[str, t.Any]) -> None:
             from bentoml._internal.context import trace_context
 
             if span is not None:
@@ -231,8 +231,8 @@ def client_request_hook(span: Span | None, _scope: dict[str, t.Any]) -> None:
                 OpenTelemetryMiddleware,
                 excluded_urls=BentoMLContainer.tracing_excluded_urls.get(),
                 default_span_details=None,
-                server_request_hook=None,
-                client_request_hook=client_request_hook,
+                server_request_hook=server_request_hook,
+                client_request_hook=None,
                 tracer_provider=BentoMLContainer.tracer_provider.get(),
             )
         )
@@ -267,16 +267,25 @@ def create_instance(self) -> None:
         self._service_instance = self.service()
         set_current_service(self._service_instance)
 
-    def _add_response_headers(self, resp: Response) -> None:
+    @inject
+    def _add_response_headers(
+        self,
+        resp: Response,
+        logging_format: dict[str, str] = Provide[BentoMLContainer.logging_formatting],
+    ) -> None:
         from bentoml._internal.context import trace_context
 
         if trace_context.request_id is not None:
-            resp.headers["X-BentoML-Request-ID"] = str(trace_context.request_id)
+            resp.headers["X-BentoML-Request-ID"] = format(
+                trace_context.request_id, logging_format["span_id"]
+            )
         if (
             BentoMLContainer.http.response.trace_id.get()
             and trace_context.trace_id is not None
         ):
-            resp.headers["X-BentoML-Trace-ID"] = str(trace_context.trace_id)
+            resp.headers["X-BentoML-Trace-ID"] = format(
+                trace_context.trace_id, logging_format["trace_id"]
+            )
 
     async def destroy_instance(self) -> None:
         from _bentoml_sdk.service.dependency import cleanup

src/_bentoml_sdk/io_models.py

@@ -80,18 +80,18 @@ class IOMixin:
     def openapi_components(cls, name: str) -> dict[str, Schema]:
         from .service.openapi import REF_TEMPLATE
 
-        if issubclass(cls, RootModel):
-            return {}
-        assert issubclass(cls, IOMixin) and issubclass(cls, BaseModel)
+        assert issubclass(cls, BaseModel)
         json_schema = cls.model_json_schema(ref_template=REF_TEMPLATE)
         defs = json_schema.pop("$defs", None)
-        main_name = (
-            f"{name}__{cls.__name__}"
-            if cls.__name__ in ("Input", "Output")
-            else cls.__name__
-        )
-        json_schema["title"] = main_name
-        components: dict[str, Schema] = {main_name: Schema(**json_schema)}
+        components: dict[str, Schema] = {}
+        if not issubclass(cls, RootModel):
+            main_name = (
+                f"{name}__{cls.__name__}"
+                if cls.__name__ in ("Input", "Output")
+                else cls.__name__
+            )
+            json_schema["title"] = main_name
+            components[main_name] = Schema(**json_schema)
         if defs is not None:
             # NOTE: This is a nested models, hence we will update the definitions
             components.update({k: Schema(**v) for k, v in defs.items()})