116 lines
3.3 KiB
Go
116 lines
3.3 KiB
Go
package openai
|
|
|
|
import (
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
edgeservice "iop/apps/edge/internal/service"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
var anthropicResponseHeaderAllowlist = map[string]struct{}{
|
|
"Cache-Control": {},
|
|
"Content-Length": {},
|
|
"Content-Type": {},
|
|
"Request-Id": {},
|
|
"Retry-After": {},
|
|
"X-Request-Id": {},
|
|
"X-Robots-Tag": {},
|
|
}
|
|
|
|
func (s *Server) writeAnthropicNativeTunnelResponse(w http.ResponseWriter, r *http.Request, handle edgeservice.ProviderTunnelResult) {
|
|
frames := handle.Stream().Frames
|
|
if frames == nil {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel is unavailable")
|
|
return
|
|
}
|
|
flusher, _ := w.(http.Flusher)
|
|
timer := time.NewTimer(handle.WaitTimeout())
|
|
defer timer.Stop()
|
|
wroteHeader := false
|
|
|
|
for {
|
|
select {
|
|
case <-r.Context().Done():
|
|
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
|
|
return
|
|
case <-timer.C:
|
|
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
|
|
if !wroteHeader {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider response timed out")
|
|
}
|
|
return
|
|
case frame, ok := <-frames:
|
|
if !ok {
|
|
if !wroteHeader {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel closed before a response")
|
|
}
|
|
return
|
|
}
|
|
switch frame.GetKind() {
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START:
|
|
if wroteHeader {
|
|
continue
|
|
}
|
|
copyAnthropicResponseHeaders(w.Header(), frame.GetHeaders())
|
|
status := int(frame.GetStatusCode())
|
|
if status == 0 {
|
|
status = http.StatusOK
|
|
}
|
|
w.WriteHeader(status)
|
|
wroteHeader = true
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY:
|
|
if len(frame.GetBody()) == 0 {
|
|
continue
|
|
}
|
|
if !wroteHeader {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
wroteHeader = true
|
|
}
|
|
if _, err := w.Write(frame.GetBody()); err != nil {
|
|
s.sendCancelRun(handle.Dispatch())
|
|
return
|
|
}
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
|
|
if !wroteHeader {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel failed")
|
|
} else if strings.Contains(strings.ToLower(w.Header().Get("Content-Type")), "text/event-stream") {
|
|
_ = writeAnthropicSSEEvent(w, "error", anthropicErrorResponse{
|
|
Type: "error", Error: errorBody{Type: "api_error", Message: "provider tunnel failed"},
|
|
})
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
}
|
|
return
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END:
|
|
if !wroteHeader {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel ended before a response")
|
|
}
|
|
return
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE:
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func copyAnthropicResponseHeaders(dst http.Header, headers map[string]string) {
|
|
for key, value := range headers {
|
|
canonical := http.CanonicalHeaderKey(key)
|
|
_, exact := anthropicResponseHeaderAllowlist[canonical]
|
|
lower := strings.ToLower(key)
|
|
if !exact && !strings.HasPrefix(lower, "anthropic-ratelimit-") && !strings.HasPrefix(lower, "ratelimit-") {
|
|
continue
|
|
}
|
|
dst.Set(canonical, value)
|
|
}
|
|
}
|