mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-08 01:16:47 +08:00
feat(interceptor, jshandler): enhance request/response handling with original request support
- Added `opts.OriginalRequest` handling to `applyResponseInterceptors` for improved context passing. - Introduced new test `TestApplyJSBeforeRequestUsesReturnedCtxBody` to validate JavaScript interceptor behavior. - Updated JavaScript-based handler to safely rewrite sensitive content and headers in requests. - Refined interceptor logic to ensure consistent state retention across request processing.
This commit is contained in:
@@ -4,9 +4,45 @@ import (
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestApplyJSBeforeRequestUsesReturnedCtxBody(t *testing.T) {
|
||||
scriptPath := filepath.Join(t.TempDir(), "before.js")
|
||||
script := `
|
||||
function on_before_request(ctx) {
|
||||
var req = JSON.parse(ctx.body);
|
||||
req.messages[0].content = req.messages[0].content.replace("sensitive_word", "safe_word");
|
||||
ctx.body = JSON.stringify(req);
|
||||
ctx.headers["X-Plugin"] = "updated";
|
||||
return ctx;
|
||||
}
|
||||
`
|
||||
if errWrite := os.WriteFile(scriptPath, []byte(script), 0600); errWrite != nil {
|
||||
t.Fatalf("os.WriteFile() error = %v", errWrite)
|
||||
}
|
||||
|
||||
plugin := &jsHandlerPlugin{cfg: defaultJSHandlerConfig()}
|
||||
headers := http.Header{"X-Plugin": []string{"original"}}
|
||||
processed, _, errApply := plugin.applyJSBeforeRequest(
|
||||
scriptPath,
|
||||
[]byte(`{"messages":[{"role":"user","content":"contains sensitive_word"}]}`),
|
||||
"gpt-test",
|
||||
"openai",
|
||||
headers,
|
||||
)
|
||||
if errApply != nil {
|
||||
t.Fatalf("applyJSBeforeRequest() error = %v", errApply)
|
||||
}
|
||||
if body := string(processed); !strings.Contains(body, "safe_word") || strings.Contains(body, "sensitive_word") {
|
||||
t.Fatalf("processed body = %q, want sensitive word rewritten", body)
|
||||
}
|
||||
if got := headers.Get("X-Plugin"); got != "updated" {
|
||||
t.Fatalf("header X-Plugin = %q, want updated", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyJSAfterResponseUsesFrozenNativeHistoryChunks(t *testing.T) {
|
||||
scriptPath := filepath.Join(t.TempDir(), "stream.js")
|
||||
script := `
|
||||
|
||||
@@ -667,7 +667,7 @@ func (h *BaseAPIHandler) executeWithAuthManager(ctx context.Context, handlerType
|
||||
}
|
||||
rawResponseHeaders := cloneHeader(resp.Headers)
|
||||
responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg))
|
||||
body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, modelName, opts, rawResponseHeaders, responseHeaders, rawJSON, req.Payload, resp.Payload, http.StatusOK)
|
||||
body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, modelName, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK)
|
||||
return body, responseHeaders, nil
|
||||
}
|
||||
|
||||
@@ -718,7 +718,7 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle
|
||||
}
|
||||
rawResponseHeaders := cloneHeader(resp.Headers)
|
||||
responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg))
|
||||
body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, modelName, opts, rawResponseHeaders, responseHeaders, rawJSON, req.Payload, resp.Payload, http.StatusOK)
|
||||
body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, modelName, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK)
|
||||
return body, responseHeaders, nil
|
||||
}
|
||||
|
||||
@@ -819,7 +819,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl
|
||||
RequestedModel: modelName,
|
||||
RequestHeaders: cloneHeader(opts.Headers),
|
||||
ResponseHeaders: cloneHeader(rawStreamHeaders),
|
||||
OriginalRequest: cloneBytes(rawJSON),
|
||||
OriginalRequest: cloneBytes(opts.OriginalRequest),
|
||||
RequestBody: cloneBytes(req.Payload),
|
||||
ChunkIndex: pluginapi.StreamChunkHeaderInitIndex,
|
||||
Metadata: opts.Metadata,
|
||||
@@ -973,7 +973,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl
|
||||
RequestedModel: modelName,
|
||||
RequestHeaders: cloneHeader(opts.Headers),
|
||||
ResponseHeaders: cloneHeader(rawStreamHeaders),
|
||||
OriginalRequest: cloneBytes(rawJSON),
|
||||
OriginalRequest: cloneBytes(opts.OriginalRequest),
|
||||
RequestBody: cloneBytes(req.Payload),
|
||||
Body: payload,
|
||||
HistoryChunks: cloneByteSlices(historyChunks),
|
||||
@@ -1304,6 +1304,7 @@ func (h *BaseAPIHandler) applyRequestInterceptors(ctx context.Context, handlerTy
|
||||
opts.Headers = finalInterceptorHeaders(opts.Headers, resp.Headers)
|
||||
if len(resp.Body) > 0 {
|
||||
req.Payload = cloneBytes(resp.Body)
|
||||
opts.OriginalRequest = cloneBytes(resp.Body)
|
||||
}
|
||||
return req, opts
|
||||
}
|
||||
|
||||
@@ -214,6 +214,9 @@ func TestHandlerRequestInterceptorRewritesExecutorRequest(t *testing.T) {
|
||||
if string(gotReq.Payload) != wantPayload {
|
||||
t.Fatalf("executor payload = %q, want %q", gotReq.Payload, wantPayload)
|
||||
}
|
||||
if string(gotOpts.OriginalRequest) != wantPayload {
|
||||
t.Fatalf("executor original request = %q, want %q", gotOpts.OriginalRequest, wantPayload)
|
||||
}
|
||||
if gotOpts.Headers.Get("X-Original") != "plugin" || gotOpts.Headers.Get("X-Plugin") != "1" {
|
||||
t.Fatalf("executor headers = %#v, want plugin rewrite", gotOpts.Headers)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user