Adding discriminators with hooks

import { isPlainObject } from "../../lib/is-plain-object.js";
import { TypeA$inboundSchema } from "../../models/components/typea.js";
import { TypeB$inboundSchema } from "../../models/components/typeb.js";
import { ResponseType } from "../../models/components/responsetype.js";
import { TypeC$inboundSchema } from "../../models/components/typec.js";
import type { AfterSuccessContext, AfterSuccessHook } from "../types.js";

export default class InjectDiscriminator implements AfterSuccessHook {
  /**
   * A hook that is called after the SDK receives a response. The hook can
   * introduce instrumentation code such as logging, tracing and metrics or
   * modify the response before it is handled or throw an error to stop the
   * response from being handled.
   */
  async afterSuccess(
    hookCtx: AfterSuccessContext,
    response: Response,
  ): Promise<Response> {
    if (hookCtx.operationID === "createResource") {
      return _injectResourceDiscriminator(response);
    }

    return response;
  }
}

export async function _injectResourceDiscriminator(
  response: Response,
): Promise<Response> {
  // Clone the response to leave the original body untouched
  const responseClone = response.clone();

  const data: { _type?: ResponseType } = await responseClone.json();

  if (!isPlainObject(data)) {
    return response;
  }

  // Use status code and response data to determine the type
  if (response.status === 200) {
    const isTypeA = TypeA$inboundSchema.safeParse({
      ...data,
      _type: ResponseType.TypeA,
    }).success;

    if (isTypeA) {
      const newData = {
        ...data,
        _type: ResponseType.TypeA,
      };
      return new Response(JSON.stringify(newData), response);
    }
  }

  if (response.status === 201) {
    const isTypeB = TypeB$inboundSchema.safeParse({
      ...data,
      _type: ResponseType.TypeB,
    }).success;

    if (isTypeB) {
      const newData = {
        ...data,
        _type: ResponseType.TypeB,
      };
      return new Response(JSON.stringify(newData), response);
    }
  }

  if (response.status === 202) {
    const isTypeC = TypeC$inboundSchema.safeParse({
      ...data,
      _type: ResponseType.TypeC,
    }).success;

    if (isTypeC) {
      const newData = {
        ...data,
        _type: ResponseType.TypeC,
      };
      return new Response(JSON.stringify(newData), response);
    }
  }

  // Default case - if we couldn't determine the type, return the original response
  return response;
}

# Example implementation for Python
from typing import Optional, Union
import httpx
from ..types import AfterSuccessContext, AfterSuccessHook
from ..models.components.typea import TypeA, TypeA_inbound_schema
from ..models.components.typeb import TypeB, TypeB_inbound_schema
from ..models.components.typec import TypeC, TypeC_inbound_schema
from ..models.components.responsetype import ResponseType

class InjectDiscriminator(AfterSuccessHook):
    async def after_success(
        self, hook_ctx: AfterSuccessContext, response: httpx.Response
    ) -> Union[httpx.Response, Exception]:
        if hook_ctx.operation_id == "create_resource":
            return await self._inject_resource_discriminator(response)
        
        return response
    
    async def _inject_resource_discriminator(
        self, response: httpx.Response
    ) -> httpx.Response:
        # Clone the response by reading the content
        content = response.content
        data = response.json()
        
        if not isinstance(data, dict):
            return response
        
        # Use status code and response data to determine the type
        if response.status_code == 200:
            # Try to validate as TypeA
            data_with_type = {**data, "_type": ResponseType.TYPE_A}
            if TypeA_inbound_schema.is_valid(data_with_type):
                # Create a new response with the discriminator
                new_content = json.dumps(data_with_type).encode()
                return httpx.Response(
                    status_code=response.status_code,
                    headers=response.headers,
                    content=new_content
                )
        
        # Similar logic for other status codes and types
        # ...
        
        # Default case - return the original response
        return response

// Example implementation for Go
package hooks

import (
    "bytes"
    "encoding/json"
    "io/ioutil"
    "net/http"
    
    "github.com/your-org/your-sdk/models/components"
)

type InjectDiscriminator struct{}

var _ AfterSuccessHook = (*InjectDiscriminator)(nil)

func (h *InjectDiscriminator) AfterSuccess(hookCtx AfterSuccessContext, res *http.Response) (*http.Response, error) {
    if hookCtx.OperationID == "createResource" {
        return injectResourceDiscriminator(res)
    }

    return res, nil
}

func injectResourceDiscriminator(res *http.Response) (*http.Response, error) {
    // Read the response body
    body, err := ioutil.ReadAll(res.Body)
    if err != nil {
        return res, err
    }
    
    // Close the original body
    res.Body.Close()
    
    // Parse the JSON
    var data map[string]interface{}
    if err := json.Unmarshal(body, &data); err != nil {
        // If we can't parse the JSON, restore the original response
        res.Body = ioutil.NopCloser(bytes.NewBuffer(body))
        return res, nil
    }
    
    // Determine the type based on status code and content
    var discriminator string
    
    switch res.StatusCode {
    case 200:
        // Check if it matches TypeA
        typeA := components.TypeA{}
        dataWithType := make(map[string]interface{})
        for k, v := range data {
            dataWithType[k] = v
        }
        dataWithType["_type"] = components.ResponseTypeTypeA
        
        typeABytes, _ := json.Marshal(dataWithType)
        if err := json.Unmarshal(typeABytes, &typeA); err == nil {
            // It's a valid TypeA, add the discriminator
            data["_type"] = components.ResponseTypeTypeA
            discriminator = components.ResponseTypeTypeA
        }
    // Similar logic for other status codes
    // ...
    }
    
    // If we determined a type, create a new response with the discriminator
    if discriminator != "" {
        newBody, _ := json.Marshal(data)
        newRes := *res
        newRes.Body = ioutil.NopCloser(bytes.NewBuffer(newBody))
        return &newRes, nil
    }
    
    // Default case - restore the original response
    res.Body = ioutil.NopCloser(bytes.NewBuffer(body))
    return res, nil
}

// Example implementation for Java
package com.example.hooks;

import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.net.http.HttpResponse;
import java.util.Map;
import java.util.Optional;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.example.models.components.TypeA;
import com.example.models.components.TypeB;
import com.example.models.components.TypeC;
import com.example.models.components.ResponseType;
import com.example.utils.Hook.AfterSuccessContext;
import com.example.utils.Hook.AfterSuccessHook;

public class InjectDiscriminator implements AfterSuccessHook {
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    @Override
    public HttpResponse<InputStream> afterSuccess(AfterSuccessContext context, HttpResponse<InputStream> response) throws Exception {
        if (context.getOperationId().equals("createResource")) {
            return injectResourceDiscriminator(response);
        }
        
        return response;
    }
    
    private HttpResponse<InputStream> injectResourceDiscriminator(HttpResponse<InputStream> response) throws Exception {
        // Read the response body
        byte[] body = response.body().readAllBytes();
        
        // Parse the JSON
        Map<String, Object> data = objectMapper.readValue(body, Map.class);
        
        // Determine the type based on status code and content
        String discriminator = null;
        
        if (response.statusCode() == 200) {
            // Try to validate as TypeA
            Map<String, Object> dataWithType = new HashMap<>(data);
            dataWithType.put("_type", ResponseType.TYPE_A);
            
            try {
                objectMapper.convertValue(dataWithType, TypeA.class);
                // It's a valid TypeA, add the discriminator
                data.put("_type", ResponseType.TYPE_A);
                discriminator = ResponseType.TYPE_A;
            } catch (Exception e) {
                // Not a valid TypeA
            }
        }
        
        // Similar logic for other status codes
        // ...
        
        // If we determined a type, create a new response with the discriminator
        if (discriminator != null) {
            byte[] newBody = objectMapper.writeValueAsBytes(data);
            return HttpResponse.BodyHandlers.ofInputStream()
                .apply(response.statusCode(), response.headers(), new ByteArrayInputStream(newBody));
        }
        
        // Default case - return the original response
        return HttpResponse.BodyHandlers.ofInputStream()
            .apply(response.statusCode(), response.headers(), new ByteArrayInputStream(body));
    }
}

// Example implementation for C#
using System.Net.Http;
using System.Text;
using System.Text.Json;
using System.Threading.Tasks;
using YourNamespace.Models.Components;
using YourNamespace.Utils;

namespace YourNamespace.Hooks
{
    public class InjectDiscriminator : IAfterSuccessHook
    {
        public async Task<HttpResponseMessage> AfterSuccessAsync(AfterSuccessContext hookCtx, HttpResponseMessage response)
        {
            if (hookCtx.OperationId == "createResource")
            {
                return await InjectResourceDiscriminatorAsync(response);
            }
            
            return response;
        }
        
        private async Task<HttpResponseMessage> InjectResourceDiscriminatorAsync(HttpResponseMessage response)
        {
            // Read the response content
            string content = await response.Content.ReadAsStringAsync();
            
            // Parse the JSON
            JsonDocument doc = JsonDocument.Parse(content);
            JsonElement root = doc.RootElement;
            
            // Convert to a mutable object
            var data = JsonSerializer.Deserialize<Dictionary<string, JsonElement>>(content);
            
            // Determine the type based on status code and content
            string discriminator = null;
            
            if (response.StatusCode == System.Net.HttpStatusCode.OK)
            {
                // Try to validate as TypeA
                var dataWithType = new Dictionary<string, JsonElement>(data);
                dataWithType["_type"] = JsonSerializer.SerializeToElement(ResponseType.TypeA);
                
                try
                {
                    var typeA = JsonSerializer.Deserialize<TypeA>(JsonSerializer.Serialize(dataWithType));
                    // It's a valid TypeA, add the discriminator
                    data["_type"] = JsonSerializer.SerializeToElement(ResponseType.TypeA);
                    discriminator = ResponseType.TypeA;
                }
                catch
                {
                    // Not a valid TypeA
                }
            }
            
            // Similar logic for other status codes
            // ...
            
            // If we determined a type, create a new response with the discriminator
            if (discriminator != null)
            {
                var newContent = new StringContent(
                    JsonSerializer.Serialize(data),
                    Encoding.UTF8,
                    "application/json");
                
                var newResponse = new HttpResponseMessage(response.StatusCode)
                {
                    Content = newContent
                };
                
                foreach (var header in response.Headers)
                {
                    newResponse.Headers.TryAddWithoutValidation(header.Key, header.Value);
                }
                
                return newResponse;
            }
            
            // Default case - return the original response
            return response;
        }
    }
}

# Example implementation for Ruby
# typed: true
# frozen_string_literal: true

require_relative './types'
require 'sorbet-runtime'
require 'json'

module YourSDK
  module SDKHooks
    class InjectDiscriminator
      extend T::Sig
      include AbstractAfterSuccessHook
      
      sig do
        override.params(
          hook_ctx: AfterSuccessHookContext,
          response: Faraday::Response
        ).returns(Faraday::Response)
      end
      def after_success(hook_ctx:, response:)
        if hook_ctx.operation_id == 'create_resource'
          return inject_resource_discriminator(response)
        end
        
        response
      end
      
      private
      
      sig do
        params(
          response: Faraday::Response
        ).returns(Faraday::Response)
      end
      def inject_resource_discriminator(response)
        # Parse the JSON
        data = JSON.parse(response.body)
        
        # Only process objects
        return response unless data.is_a?(Hash)
        
        # Determine the type based on status code and content
        discriminator = nil
        
        if response.status == 200
          # Try to validate as TypeA
          data_with_type = data.dup
          data_with_type['_type'] = ResponseType::TYPE_A
          
          begin
            # Validate using schema
            TypeA.from_json(data_with_type.to_json)
            # It's a valid TypeA, add the discriminator
            data['_type'] = ResponseType::TYPE_A
            discriminator = ResponseType::TYPE_A
          rescue
            # Not a valid TypeA
          end
        end
        
        # Similar logic for other status codes
        # ...
        
        # If we determined a type, create a new response with the discriminator
        if discriminator
          # Create a new response with the modified body
          new_response = response.dup
          new_response.env.body = data.to_json
          return new_response
        end
        
        # Default case - return the original response
        response
      end
    end
  end
end