From 4b5f1eab25fca4b3815369a826e958e7c070a69e Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Thu, 27 Aug 2026 05:30:19 +0800 Subject: [PATCH] feat(plugin): support observing upstream websocket response events - Introduce `WebSocketResponseObserver` capability and bump plugin ABI schema version to 4. - Forward upstream WebSocket response frames from Codex and xAI executors to configured observers. - Wire `WebSocketResponseObserver` across API handlers and plugin host dispatchers. Closes: #5248 --- internal/pluginhost/adapters_interceptors.go | 55 ++++ internal/pluginhost/host.go | 1 + internal/pluginhost/rpc_client.go | 19 ++ internal/pluginhost/rpc_schema.go | 7 + .../pluginhost/websocket_observer_test.go | 247 ++++++++++++++++++ .../executor/codex_websockets_execute.go | 1 + .../codex_websockets_executor_test.go | 161 ++++++++++++ .../executor/codex_websockets_stream.go | 2 + .../helps/websocket_observer_helpers.go | 48 ++++ .../helps/websocket_observer_helpers_test.go | 73 ++++++ .../executor/xai_websockets_executor.go | 1 + sdk/api/handlers/handlers_execution.go | 19 +- sdk/api/handlers/handlers_interceptors.go | 63 +++++ .../handlers/handlers_interceptors_test.go | 61 ++++- sdk/api/handlers/handlers_stream.go | 1 + sdk/cliproxy/executor/types.go | 21 ++ sdk/pluginabi/types.go | 8 +- sdk/pluginabi/types_test.go | 10 +- sdk/pluginapi/types.go | 23 ++ sdk/pluginapi/types_test.go | 5 + 20 files changed, 810 insertions(+), 16 deletions(-) create mode 100644 internal/pluginhost/websocket_observer_test.go create mode 100644 internal/runtime/executor/helps/websocket_observer_helpers.go create mode 100644 internal/runtime/executor/helps/websocket_observer_helpers_test.go diff --git a/internal/pluginhost/adapters_interceptors.go b/internal/pluginhost/adapters_interceptors.go index 7239d4d4..ffc8c337 100644 --- a/internal/pluginhost/adapters_interceptors.go +++ b/internal/pluginhost/adapters_interceptors.go @@ -72,6 +72,20 @@ func (h *Host) callStreamChunkInterceptor(ctx context.Context, record capability return resp, true } +func (h *Host) callWebSocketResponseObserver(ctx context.Context, record capabilityRecord, observer pluginapi.WebSocketResponseObserver, event pluginapi.WebSocketResponseEvent) { + if h == nil || observer == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { + return + } + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(record.id, "WebSocketResponseObserver.ObserveWebSocketResponseEvent", recovered) + } + }() + if errObserve := observer.ObserveWebSocketResponseEvent(ctx, event); errObserve != nil { + log.Warnf("pluginhost: websocket response observer %s failed: %v", record.id, errObserve) + } +} + func (h *Host) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { return h.InterceptRequestBeforeAuthExcept(ctx, req, "") } @@ -131,6 +145,47 @@ func (h *Host) CompleteRequest(ctx context.Context, completion pluginapi.Request h.CompleteRequestExcept(ctx, completion, "") } +func (h *Host) HasWebSocketResponseObservers() bool { + if h == nil { + return false + } + for _, record := range h.activeRecords() { + if h.isPluginFused(record.id) { + continue + } + if record.plugin.Capabilities.WebSocketResponseObserver != nil { + return true + } + } + return false +} + +func (h *Host) ObserveWebSocketResponseEvent(ctx context.Context, event pluginapi.WebSocketResponseEvent) { + h.ObserveWebSocketResponseEventExcept(ctx, event, "") +} + +func (h *Host) ObserveWebSocketResponseEventExcept(ctx context.Context, event pluginapi.WebSocketResponseEvent, skipPluginID string) { + if h == nil { + return + } + if ctx == nil { + ctx = context.Background() + } else { + ctx = context.WithoutCancel(ctx) + } + skipPluginID = strings.TrimSpace(skipPluginID) + for _, record := range h.activeRecords() { + observer := record.plugin.Capabilities.WebSocketResponseObserver + if h.isPluginFused(record.id) || observer == nil || record.id == skipPluginID || !h.recordCurrent(record) { + continue + } + next := event + next.Payload = bytes.Clone(event.Payload) + next.Metadata = cloneInterceptorMetadata(event.Metadata) + h.callWebSocketResponseObserver(ctx, record, observer, next) + } +} + // CompleteRequestExcept notifies lifecycle plugins except the plugin that initiated a nested host execution. func (h *Host) CompleteRequestExcept(ctx context.Context, completion pluginapi.RequestCompletion, skipPluginID string) { if h == nil { diff --git a/internal/pluginhost/host.go b/internal/pluginhost/host.go index dd6cfd1a..52e56cfd 100644 --- a/internal/pluginhost/host.go +++ b/internal/pluginhost/host.go @@ -1052,6 +1052,7 @@ func validPlugin(plugin pluginapi.Plugin) bool { caps.ResponseAfterTranslator != nil || caps.ResponseInterceptor != nil || caps.StreamChunkInterceptor != nil || + caps.WebSocketResponseObserver != nil || caps.ThinkingApplier != nil || caps.UsagePlugin != nil || caps.CommandLinePlugin != nil || diff --git a/internal/pluginhost/rpc_client.go b/internal/pluginhost/rpc_client.go index 845a4c20..0806701c 100644 --- a/internal/pluginhost/rpc_client.go +++ b/internal/pluginhost/rpc_client.go @@ -133,6 +133,9 @@ func registerRPCPlugin(ctx context.Context, host *Host, id string, client plugin if resp.Capabilities.StreamChunkInterceptor { plugin.Capabilities.StreamChunkInterceptor = adapter } + if resp.Capabilities.WebSocketResponseObserver { + plugin.Capabilities.WebSocketResponseObserver = adapter + } if resp.Capabilities.ThinkingApplier { plugin.Capabilities.ThinkingApplier = rpcThinkingApplier{rpcPluginAdapter: adapter} } @@ -211,6 +214,9 @@ func sanitizePluginRequest(request any) any { case pluginapi.StreamChunkInterceptRequest: req.Metadata = sanitizePluginMetadata(req.Metadata) return req + case pluginapi.WebSocketResponseEvent: + req.Metadata = sanitizePluginMetadata(req.Metadata) + return req case rpcRequestInterceptRequest: req.Metadata = sanitizePluginMetadata(req.Metadata) return req @@ -226,6 +232,9 @@ func sanitizePluginRequest(request any) any { case rpcStreamChunkInterceptRequest: req.Metadata = sanitizePluginMetadata(req.Metadata) return req + case rpcWebSocketResponseEvent: + req.Metadata = sanitizePluginMetadata(req.Metadata) + return req case pluginapi.ExecutorHTTPRequest: req.HTTPClient = nil return req @@ -530,6 +539,16 @@ func (a *rpcPluginAdapter) InterceptStreamChunk(ctx context.Context, req plugina }) } +func (a *rpcPluginAdapter) ObserveWebSocketResponseEvent(ctx context.Context, event pluginapi.WebSocketResponseEvent) error { + callbackID, closeCallback := a.openHostCallbackContext(ctx) + defer closeCallback() + _, errCall := callPlugin[rpcEmptyResponse](ctx, a.client, pluginabi.MethodWebSocketResponseEvent, rpcWebSocketResponseEvent{ + WebSocketResponseEvent: event, + HostCallbackID: callbackID, + }) + return errCall +} + func (a rpcThinkingApplier) ApplyThinking(ctx context.Context, req pluginapi.ThinkingApplyRequest) (pluginapi.PayloadResponse, error) { callbackID, closeCallback := a.openHostCallbackContext(ctx) defer closeCallback() diff --git a/internal/pluginhost/rpc_schema.go b/internal/pluginhost/rpc_schema.go index 306d9166..2a3ca3fd 100644 --- a/internal/pluginhost/rpc_schema.go +++ b/internal/pluginhost/rpc_schema.go @@ -39,6 +39,7 @@ type rpcCapabilities struct { ResponseAfterTranslator bool `json:"response_after_translator"` ResponseInterceptor bool `json:"response_interceptor"` StreamChunkInterceptor bool `json:"response_stream_interceptor"` + WebSocketResponseObserver bool `json:"websocket_response_observer"` ThinkingApplier bool `json:"thinking_applier"` UsagePlugin bool `json:"usage_plugin"` CommandLinePlugin bool `json:"command_line_plugin"` @@ -110,6 +111,11 @@ type rpcStreamChunkInterceptRequest struct { HostCallbackID string `json:"host_callback_id,omitempty"` } +type rpcWebSocketResponseEvent struct { + pluginapi.WebSocketResponseEvent + HostCallbackID string `json:"host_callback_id,omitempty"` +} + type rpcThinkingApplyRequest struct { pluginapi.ThinkingApplyRequest HostCallbackID string `json:"host_callback_id,omitempty"` @@ -150,6 +156,7 @@ func rpcCapabilitiesFromPlugin(plugin pluginapi.Plugin) rpcCapabilities { ResponseAfterTranslator: caps.ResponseAfterTranslator != nil, ResponseInterceptor: caps.ResponseInterceptor != nil, StreamChunkInterceptor: caps.StreamChunkInterceptor != nil, + WebSocketResponseObserver: caps.WebSocketResponseObserver != nil, ThinkingApplier: caps.ThinkingApplier != nil, UsagePlugin: caps.UsagePlugin != nil, CommandLinePlugin: caps.CommandLinePlugin != nil, diff --git a/internal/pluginhost/websocket_observer_test.go b/internal/pluginhost/websocket_observer_test.go new file mode 100644 index 00000000..034dae4e --- /dev/null +++ b/internal/pluginhost/websocket_observer_test.go @@ -0,0 +1,247 @@ +package pluginhost + +import ( + "bytes" + "context" + "encoding/json" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +type testWebSocketObserverFunc func(context.Context, pluginapi.WebSocketResponseEvent) error + +func (f testWebSocketObserverFunc) ObserveWebSocketResponseEvent(ctx context.Context, event pluginapi.WebSocketResponseEvent) error { + return f(ctx, event) +} + +func TestObserveWebSocketResponseEventInvokesPlugin(t *testing.T) { + var got pluginapi.WebSocketResponseEvent + called := false + host := newHostWithRecords(capabilityRecord{ + id: "quota-tracker", + plugin: pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + WebSocketResponseObserver: testWebSocketObserverFunc(func(_ context.Context, event pluginapi.WebSocketResponseEvent) error { + got = event + called = true + return nil + }), + }, + }, + }) + + rawPayload := []byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42}}}`) + host.ObserveWebSocketResponseEvent(context.Background(), pluginapi.WebSocketResponseEvent{ + RequestID: "req-123", + SourceFormat: "openai", + Model: "gpt-5.3-codex", + RequestedModel: "gpt-5.3-codex", + Provider: "codex", + AuthID: "auth-abc", + AuthLabel: "test-auth", + AuthType: "oauth", + EventType: "codex.rate_limits", + Payload: rawPayload, + }) + + if !called { + t.Fatal("observer callback was not invoked") + } + if got.RequestID != "req-123" { + t.Fatalf("RequestID = %q, want req-123", got.RequestID) + } + if got.Provider != "codex" { + t.Fatalf("Provider = %q, want codex", got.Provider) + } + if got.AuthID != "auth-abc" || got.AuthLabel != "test-auth" { + t.Fatalf("Auth = (%q, %q), want (auth-abc, test-auth)", got.AuthID, got.AuthLabel) + } + if got.EventType != "codex.rate_limits" { + t.Fatalf("EventType = %q, want codex.rate_limits", got.EventType) + } + if !bytes.Equal(got.Payload, rawPayload) { + t.Fatalf("Payload = %s, want %s", got.Payload, rawPayload) + } +} + +func TestObserveWebSocketResponseEventClonesPayloadAndMetadata(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + originalPayload := []byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":1}}}`) + originalMetadata := map[string]any{"key": "value"} + called := false + + host := newHostWithRecords(capabilityRecord{ + id: "observer-clone", + plugin: pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + WebSocketResponseObserver: testWebSocketObserverFunc(func(_ context.Context, event pluginapi.WebSocketResponseEvent) error { + event.Payload[0] = 'X' + event.Metadata["key"] = "mutated" + called = true + return nil + }), + }, + }, + }) + + host.ObserveWebSocketResponseEvent(ctx, pluginapi.WebSocketResponseEvent{ + RequestID: "req-clone", + Payload: originalPayload, + Metadata: originalMetadata, + }) + + if !called { + t.Fatal("observer callback was not invoked") + } + if originalPayload[0] != '{' { + t.Fatalf("original payload was mutated: %s", originalPayload) + } + if originalMetadata["key"] != "value" { + t.Fatalf("original metadata was mutated: %#v", originalMetadata) + } +} + +func TestObserveWebSocketResponseEventFusesOnPanic(t *testing.T) { + host := newHostWithRecords(capabilityRecord{ + id: "panicking-observer", + plugin: pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + WebSocketResponseObserver: testWebSocketObserverFunc(func(context.Context, pluginapi.WebSocketResponseEvent) error { + panic("observer panic") + }), + }, + }, + }) + + if !host.HasWebSocketResponseObservers() { + t.Fatal("HasWebSocketResponseObservers() = false, want true") + } + + host.ObserveWebSocketResponseEvent(context.Background(), pluginapi.WebSocketResponseEvent{ + RequestID: "req-panic", + Payload: []byte(`{"type":"test"}`), + }) + + if !host.isPluginFused("panicking-observer") { + t.Fatal("isPluginFused(panicking-observer) = false, want true") + } + if host.HasWebSocketResponseObservers() { + t.Fatal("HasWebSocketResponseObservers() = true after fusing, want false") + } +} + +func TestObserveWebSocketResponseEventSkipPlugin(t *testing.T) { + calls := 0 + host := newHostWithRecords(capabilityRecord{ + id: "skipped-plugin", + plugin: pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + WebSocketResponseObserver: testWebSocketObserverFunc(func(context.Context, pluginapi.WebSocketResponseEvent) error { + calls++ + return nil + }), + }, + }, + }) + + host.ObserveWebSocketResponseEventExcept(context.Background(), pluginapi.WebSocketResponseEvent{ + RequestID: "req-skip", + }, "skipped-plugin") + + if calls != 0 { + t.Fatalf("observer calls = %d, want 0 when skipped", calls) + } +} + +func TestRegisterRPCPluginRegistersWebSocketResponseObserver(t *testing.T) { + lookup := newTestSymbolLookup(&testPlugin{ + registerResult: pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + WebSocketResponseObserver: testWebSocketObserverFunc(func(context.Context, pluginapi.WebSocketResponseEvent) error { + return nil + }), + }, + }, + }) + + registered, errRegister := registerRPCPlugin(context.Background(), nil, "rpc-observer", lookup, pluginabi.MethodPluginRegister, nil) + if errRegister != nil { + t.Fatalf("registerRPCPlugin() error = %v", errRegister) + } + if registered.Capabilities.WebSocketResponseObserver == nil { + t.Fatal("WebSocketResponseObserver = nil, want RPC adapter") + } +} + +type rpcObserverRecordingClient struct { + lastMethod string + lastRequest []byte +} + +func (c *rpcObserverRecordingClient) Call(_ context.Context, method string, request []byte) ([]byte, error) { + c.lastMethod = method + c.lastRequest = bytes.Clone(request) + return json.Marshal(pluginabi.Envelope{OK: true, Result: json.RawMessage(`{}`)}) +} + +func (c *rpcObserverRecordingClient) Shutdown() {} + +func TestObserveWebSocketResponseEventRPCSanitizesMetadata(t *testing.T) { + client := &rpcObserverRecordingClient{} + adapter := &rpcPluginAdapter{ + id: "rpc-observer-sanitize", + client: client, + } + + rawPayload := []byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":99}}}`) + unserializableMetadata := map[string]any{ + "safe_key": "safe_value", + "func_field": func() {}, + "chan_field": make(chan int), + } + + err := adapter.ObserveWebSocketResponseEvent(context.Background(), pluginapi.WebSocketResponseEvent{ + RequestID: "req-rpc-sanitize", + SourceFormat: "openai", + Model: "gpt-5.3-codex", + RequestedModel: "gpt-5.3-codex", + Provider: "codex", + AuthID: "auth-xyz", + AuthLabel: "test-auth-rpc", + AuthType: "oauth", + EventType: "codex.rate_limits", + Payload: rawPayload, + Metadata: unserializableMetadata, + }) + if err != nil { + t.Fatalf("ObserveWebSocketResponseEvent() error = %v", err) + } + + if client.lastMethod != pluginabi.MethodWebSocketResponseEvent { + t.Fatalf("lastMethod = %q, want %q", client.lastMethod, pluginabi.MethodWebSocketResponseEvent) + } + + var decoded rpcWebSocketResponseEvent + if errUnmarshal := json.Unmarshal(client.lastRequest, &decoded); errUnmarshal != nil { + t.Fatalf("unmarshal rpc request: %v", errUnmarshal) + } + + if decoded.RequestID != "req-rpc-sanitize" { + t.Fatalf("decoded RequestID = %q, want req-rpc-sanitize", decoded.RequestID) + } + if decoded.EventType != "codex.rate_limits" { + t.Fatalf("decoded EventType = %q, want codex.rate_limits", decoded.EventType) + } + if decoded.Metadata["safe_key"] != "safe_value" { + t.Fatalf("decoded safe_key = %v, want safe_value", decoded.Metadata["safe_key"]) + } + if _, exists := decoded.Metadata["func_field"]; exists { + t.Fatalf("func_field was not sanitized: %#v", decoded.Metadata) + } + if _, exists := decoded.Metadata["chan_field"]; exists { + t.Fatalf("chan_field was not sanitized: %#v", decoded.Metadata) + } +} diff --git a/internal/runtime/executor/codex_websockets_execute.go b/internal/runtime/executor/codex_websockets_execute.go index ef887272..a7369ee7 100644 --- a/internal/runtime/executor/codex_websockets_execute.go +++ b/internal/runtime/executor/codex_websockets_execute.go @@ -286,6 +286,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut reporter.MarkFirstResponseByte() payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) helps.AppendCodexAPIWebsocketResponse(ctx, e.cfg, payload) + helps.EmitWebSocketResponseEvent(ctx, opts, auth, e.Identifier(), req.Model, payload) payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2) if wsErr, ok := parseCodexWebsocketError(payload); ok { diff --git a/internal/runtime/executor/codex_websockets_executor_test.go b/internal/runtime/executor/codex_websockets_executor_test.go index 755bf1fb..d019950e 100644 --- a/internal/runtime/executor/codex_websockets_executor_test.go +++ b/internal/runtime/executor/codex_websockets_executor_test.go @@ -2145,3 +2145,164 @@ func TestCodexWebsocketLifecycleBindFailureReleasesSessionRequestLock(t *testing t.Fatal("lifecycle bind failure left the session request lock held") } } + +func TestCodexWebsocketsExecuteObservesWebSocketResponseEvents(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Fatalf("upgrade websocket: %v", err) + } + defer func() { _ = conn.Close() }() + + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Fatalf("read upstream websocket message: %v", errRead) + } + + rateLimitMsg := []byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42}}}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, rateLimitMsg); errWrite != nil { + t.Fatalf("write rate limit websocket message: %v", errWrite) + } + + completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { + t.Fatalf("write completed websocket message: %v", errWrite) + } + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + ID: "auth-codex-1", + Label: "codex-account", + Provider: "codex", + Attributes: map[string]string{ + "api_key": "sk-test", + "base_url": server.URL, + "auth_kind": "oauth", + }, + Metadata: map[string]any{ + "email": "user@example.com", + }, + } + req := cliproxyexecutor.Request{ + Model: "gpt-5.6-sol", + Payload: []byte(`{"model":"gpt-5.6-sol","input":[{"role":"user","content":"hello"}]}`), + } + + var observedEvents []cliproxyexecutor.WebSocketResponseEvent + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("codex"), + Metadata: map[string]any{ + "request_id": "test-req-codex", + }, + WebSocketResponseObserver: func(_ context.Context, ev cliproxyexecutor.WebSocketResponseEvent) { + observedEvents = append(observedEvents, ev) + }, + } + + if _, err := exec.Execute(context.Background(), auth, req, opts); err != nil { + t.Fatalf("Execute() error = %v", err) + } + + if len(observedEvents) < 2 { + t.Fatalf("observed %d events, want at least 2", len(observedEvents)) + } + + rateLimitEvent := observedEvents[0] + if rateLimitEvent.RequestID != "test-req-codex" { + t.Fatalf("RequestID = %q, want test-req-codex", rateLimitEvent.RequestID) + } + if rateLimitEvent.AuthID != "auth-codex-1" || rateLimitEvent.AuthLabel != "codex-account" { + t.Fatalf("Auth = (%q, %q), want (auth-codex-1, codex-account)", rateLimitEvent.AuthID, rateLimitEvent.AuthLabel) + } + if rateLimitEvent.EventType != "codex.rate_limits" { + t.Fatalf("EventType = %q, want codex.rate_limits", rateLimitEvent.EventType) + } + if !bytes.Contains(rateLimitEvent.Payload, []byte(`"used_percent":42`)) { + t.Fatalf("Payload = %s, want used_percent 42", rateLimitEvent.Payload) + } +} + +func TestCodexWebsocketsExecuteStreamObservesWebSocketResponseEvents(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Fatalf("upgrade websocket: %v", err) + } + defer func() { _ = conn.Close() }() + + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Fatalf("read upstream websocket message: %v", errRead) + } + + rateLimitMsg := []byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":75}}}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, rateLimitMsg); errWrite != nil { + t.Fatalf("write rate limit websocket message: %v", errWrite) + } + + completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { + t.Fatalf("write completed websocket message: %v", errWrite) + } + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + ID: "auth-stream-1", + Label: "codex-stream-account", + Provider: "codex", + Attributes: map[string]string{ + "api_key": "sk-test", + "base_url": server.URL, + "auth_kind": "oauth", + }, + } + req := cliproxyexecutor.Request{ + Model: "gpt-5.6-sol", + Payload: []byte(`{"model":"gpt-5.6-sol","input":[{"role":"user","content":"hello"}]}`), + } + + var observedEvents []cliproxyexecutor.WebSocketResponseEvent + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("codex"), + Metadata: map[string]any{ + "request_id": "test-req-stream", + }, + WebSocketResponseObserver: func(_ context.Context, ev cliproxyexecutor.WebSocketResponseEvent) { + observedEvents = append(observedEvents, ev) + }, + } + + streamResult, err := exec.ExecuteStream(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + + // Drain stream + for chunk := range streamResult.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error: %v", chunk.Err) + } + } + + if len(observedEvents) < 2 { + t.Fatalf("observed %d events, want at least 2", len(observedEvents)) + } + + rateLimitEvent := observedEvents[0] + if rateLimitEvent.RequestID != "test-req-stream" { + t.Fatalf("RequestID = %q, want test-req-stream", rateLimitEvent.RequestID) + } + if rateLimitEvent.AuthID != "auth-stream-1" || rateLimitEvent.AuthLabel != "codex-stream-account" { + t.Fatalf("Auth = (%q, %q), want (auth-stream-1, codex-stream-account)", rateLimitEvent.AuthID, rateLimitEvent.AuthLabel) + } + if rateLimitEvent.EventType != "codex.rate_limits" { + t.Fatalf("EventType = %q, want codex.rate_limits", rateLimitEvent.EventType) + } + if !bytes.Contains(rateLimitEvent.Payload, []byte(`"used_percent":75`)) { + t.Fatalf("Payload = %s, want used_percent 75", rateLimitEvent.Payload) + } +} diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index 5b859009..e1eb7627 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -328,6 +328,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr reporter.MarkFirstResponseByte() payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) helps.AppendCodexAPIWebsocketResponse(ctx, e.cfg, payload) + helps.EmitWebSocketResponseEvent(ctx, opts, auth, e.Identifier(), req.Model, payload) payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2) if wsErr, ok := parseCodexWebsocketError(payload); ok { @@ -539,6 +540,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr reporter.MarkFirstResponseByte() payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) helps.AppendCodexAPIWebsocketResponse(ctx, e.cfg, payload) + helps.EmitWebSocketResponseEvent(ctx, opts, auth, e.Identifier(), req.Model, payload) payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2) if wsErr, ok := parseCodexWebsocketError(payload); ok { diff --git a/internal/runtime/executor/helps/websocket_observer_helpers.go b/internal/runtime/executor/helps/websocket_observer_helpers.go new file mode 100644 index 00000000..cbb68578 --- /dev/null +++ b/internal/runtime/executor/helps/websocket_observer_helpers.go @@ -0,0 +1,48 @@ +package helps + +import ( + "context" + + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/tidwall/gjson" +) + +// EmitWebSocketResponseEvent delivers an upstream WebSocket response frame to any configured observer. +func EmitWebSocketResponseEvent(ctx context.Context, opts cliproxyexecutor.Options, auth *cliproxyauth.Auth, provider, model string, payload []byte) { + if opts.WebSocketResponseObserver == nil || len(payload) == 0 { + return + } + var authID, authLabel, authType string + if auth != nil { + authID = auth.ID + authLabel = auth.Label + authType, _ = auth.AccountInfo() + } + eventType := gjson.GetBytes(payload, "type").String() + var requestID string + var traceID string + if opts.Metadata != nil { + if reqID, ok := opts.Metadata["request_id"].(string); ok { + requestID = reqID + } + if trID, ok := opts.Metadata["trace_id"].(string); ok { + traceID = trID + } + } + requestedModel := PayloadRequestedModel(opts, model) + opts.WebSocketResponseObserver(ctx, cliproxyexecutor.WebSocketResponseEvent{ + RequestID: requestID, + TraceID: traceID, + SourceFormat: opts.SourceFormat.String(), + Model: model, + RequestedModel: requestedModel, + Provider: provider, + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + EventType: eventType, + Payload: payload, + Metadata: opts.Metadata, + }) +} diff --git a/internal/runtime/executor/helps/websocket_observer_helpers_test.go b/internal/runtime/executor/helps/websocket_observer_helpers_test.go new file mode 100644 index 00000000..7378e043 --- /dev/null +++ b/internal/runtime/executor/helps/websocket_observer_helpers_test.go @@ -0,0 +1,73 @@ +package helps + +import ( + "bytes" + "context" + "testing" + + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" +) + +func TestEmitWebSocketResponseEvent(t *testing.T) { + var received []cliproxyexecutor.WebSocketResponseEvent + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + Metadata: map[string]any{ + "request_id": "test-req-1", + "trace_id": "test-trace-1", + }, + WebSocketResponseObserver: func(_ context.Context, ev cliproxyexecutor.WebSocketResponseEvent) { + received = append(received, ev) + }, + } + + auth := &cliproxyauth.Auth{ + ID: "auth-123", + Label: "primary-auth", + Provider: "codex", + Metadata: map[string]any{ + "email": "alice@example.com", + }, + Attributes: map[string]string{ + "auth_kind": "oauth", + }, + } + + payload := []byte(`{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":42}}}`) + EmitWebSocketResponseEvent(context.Background(), opts, auth, "codex", "gpt-5.3-codex", payload) + + if len(received) != 1 { + t.Fatalf("received %d events, want 1", len(received)) + } + ev := received[0] + if ev.RequestID != "test-req-1" { + t.Fatalf("RequestID = %q, want test-req-1", ev.RequestID) + } + if ev.TraceID != "test-trace-1" { + t.Fatalf("TraceID = %q, want test-trace-1", ev.TraceID) + } + if ev.Provider != "codex" { + t.Fatalf("Provider = %q, want codex", ev.Provider) + } + if ev.AuthID != "auth-123" || ev.AuthLabel != "primary-auth" { + t.Fatalf("Auth = (%q, %q), want (auth-123, primary-auth)", ev.AuthID, ev.AuthLabel) + } + if ev.AuthType != "oauth" { + t.Fatalf("AuthType = %q, want oauth", ev.AuthType) + } + if ev.EventType != "codex.rate_limits" { + t.Fatalf("EventType = %q, want codex.rate_limits", ev.EventType) + } + if !bytes.Equal(ev.Payload, payload) { + t.Fatalf("Payload = %s, want %s", ev.Payload, payload) + } +} + +func TestEmitWebSocketResponseEventNilObserverNoOp(t *testing.T) { + opts := cliproxyexecutor.Options{} + payload := []byte(`{"type":"codex.rate_limits"}`) + // Should not panic + EmitWebSocketResponseEvent(context.Background(), opts, nil, "codex", "gpt-5.3-codex", payload) +} diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index fbe779d1..16b31e1e 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -746,6 +746,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } reporter.MarkFirstResponseByte() helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) + helps.EmitWebSocketResponseEvent(ctx, opts, auth, e.Identifier(), req.Model, payload) if wsErr, ok := parseXAIWebsocketError(payload); ok { terminateReason = "upstream_error" diff --git a/sdk/api/handlers/handlers_execution.go b/sdk/api/handlers/handlers_execution.go index e4a4c254..ac521ec9 100644 --- a/sdk/api/handlers/handlers_execution.go +++ b/sdk/api/handlers/handlers_execution.go @@ -81,6 +81,7 @@ func (h *BaseAPIHandler) executeWithAuthManagerFormats(ctx context.Context, entr Headers: modelExecutionHeaders(ctx, execOptions.Headers), Query: modelExecutionQuery(ctx, execOptions.Query), RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, lifecycle.requestID(), execOptions.SkipInterceptorPluginID), + WebSocketResponseObserver: h.webSocketResponseObserver(lifecycle.requestID(), execOptions.SkipInterceptorPluginID), } opts.Metadata = reqMeta var interceptErr *interfaces.ErrorMessage @@ -145,6 +146,7 @@ func (h *BaseAPIHandler) executeCountWithAuthManager(ctx context.Context, handle Headers: modelExecutionHeaders(ctx, execOptions.Headers), Query: modelExecutionQuery(ctx, execOptions.Query), RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, lifecycle.requestID(), execOptions.SkipInterceptorPluginID), + WebSocketResponseObserver: h.webSocketResponseObserver(lifecycle.requestID(), execOptions.SkipInterceptorPluginID), } opts.Metadata = reqMeta var interceptErr *interfaces.ErrorMessage @@ -264,14 +266,15 @@ func (h *BaseAPIHandler) pluginExecutorRequest(ctx context.Context, entryProtoco } req := coreexecutor.Request{Model: modelName, Payload: payload} opts := coreexecutor.Options{ - Stream: stream, - Alt: alt, - OriginalRequest: rawJSON, - SourceFormat: sdktranslator.FromString(entryProtocol), - ResponseFormat: sdktranslator.FromString(responseProtocol), - Headers: modelExecutionHeaders(ctx, execOptions.Headers), - Query: modelExecutionQuery(ctx, execOptions.Query), - Metadata: reqMeta, + Stream: stream, + Alt: alt, + OriginalRequest: rawJSON, + SourceFormat: sdktranslator.FromString(entryProtocol), + ResponseFormat: sdktranslator.FromString(responseProtocol), + Headers: modelExecutionHeaders(ctx, execOptions.Headers), + Query: modelExecutionQuery(ctx, execOptions.Query), + WebSocketResponseObserver: h.webSocketResponseObserver("", execOptions.SkipInterceptorPluginID), + Metadata: reqMeta, } return req, opts } diff --git a/sdk/api/handlers/handlers_interceptors.go b/sdk/api/handlers/handlers_interceptors.go index dfed4310..d1b5d442 100644 --- a/sdk/api/handlers/handlers_interceptors.go +++ b/sdk/api/handlers/handlers_interceptors.go @@ -63,6 +63,28 @@ type requestLifecycleSkipHost interface { CompleteRequestExcept(context.Context, pluginapi.RequestCompletion, string) } +type webSocketResponseObserverHost interface { + ObserveWebSocketResponseEvent(context.Context, pluginapi.WebSocketResponseEvent) +} + +type webSocketResponseObserverSkipHost interface { + ObserveWebSocketResponseEventExcept(context.Context, pluginapi.WebSocketResponseEvent, string) +} + +type webSocketResponseObserverDetector interface { + HasWebSocketResponseObservers() bool +} + +func webSocketResponseObserversEnabled(host PluginInterceptorHost) bool { + if host == nil { + return false + } + if detector, ok := host.(webSocketResponseObserverDetector); ok { + return detector.HasWebSocketResponseObservers() + } + return true +} + type requestLifecycleTracker struct { once sync.Once ctx context.Context @@ -463,6 +485,47 @@ func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCa } } +func (h *BaseAPIHandler) webSocketResponseObserver(requestID, skipPluginID string) coreexecutor.WebSocketResponseObserver { + host := h.interceptorHost() + if !webSocketResponseObserversEnabled(host) { + return nil + } + observerHost, ok := host.(webSocketResponseObserverHost) + if !ok { + return nil + } + skipHost, _ := host.(webSocketResponseObserverSkipHost) + return func(ctx context.Context, ev coreexecutor.WebSocketResponseEvent) { + traceID := ev.TraceID + if traceID == "" { + traceID = logging.GetRequestID(ctx) + } + reqID := ev.RequestID + if reqID == "" { + reqID = requestID + } + pluginEvent := pluginapi.WebSocketResponseEvent{ + RequestID: reqID, + TraceID: traceID, + SourceFormat: ev.SourceFormat, + Model: ev.Model, + RequestedModel: ev.RequestedModel, + Provider: ev.Provider, + AuthID: ev.AuthID, + AuthLabel: ev.AuthLabel, + AuthType: ev.AuthType, + EventType: ev.EventType, + Payload: ev.Payload, + Metadata: ev.Metadata, + } + if skipPluginID != "" && skipHost != nil { + skipHost.ObserveWebSocketResponseEventExcept(ctx, pluginEvent, skipPluginID) + return + } + observerHost.ObserveWebSocketResponseEvent(ctx, pluginEvent) + } +} + func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest, requestID, skipPluginID string) coreexecutor.RequestAfterAuthInterceptResponse { host := h.interceptorHost() if !requestInterceptorsEnabled(host) { diff --git a/sdk/api/handlers/handlers_interceptors_test.go b/sdk/api/handlers/handlers_interceptors_test.go index c328a620..42db2bcb 100644 --- a/sdk/api/handlers/handlers_interceptors_test.go +++ b/sdk/api/handlers/handlers_interceptors_test.go @@ -21,11 +21,12 @@ import ( ) type handlerInterceptorTestHost struct { - interceptRequestBeforeAuth func(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse - interceptRequestAfterAuth func(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse - interceptResponse func(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse - interceptStreamChunk func(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse - completeRequest func(context.Context, pluginapi.RequestCompletion) + interceptRequestBeforeAuth func(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse + interceptRequestAfterAuth func(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse + interceptResponse func(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse + interceptStreamChunk func(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse + observeWebSocketResponseEvent func(context.Context, pluginapi.WebSocketResponseEvent) + completeRequest func(context.Context, pluginapi.RequestCompletion) // includeStreamChunkRequestBodies simulates legacy schema_version < 3 plugins. includeStreamChunkRequestBodies bool } @@ -92,6 +93,12 @@ func (h *handlerInterceptorTestHost) CompleteRequest(ctx context.Context, comple } } +func (h *handlerInterceptorTestHost) ObserveWebSocketResponseEvent(ctx context.Context, event pluginapi.WebSocketResponseEvent) { + if h != nil && h.observeWebSocketResponseEvent != nil { + h.observeWebSocketResponseEvent(ctx, event) + } +} + // StreamChunkPayloadIncludesRequestBody implements streamChunkRequestBodyPolicy. // Default false simulates schema_version >= 3 (omit request bodies on payload chunks). func (h *handlerInterceptorTestHost) StreamChunkPayloadIncludesRequestBody() bool { @@ -1481,3 +1488,47 @@ func TestHandlerResponseInterceptorSeesRawHeadersWhenPassthroughDisabled(t *test t.Fatalf("headers leaked raw upstream header with passthrough disabled: %#v", headers) } } + +func TestHandlerWebSocketResponseObserverForwardsToPluginHost(t *testing.T) { + model := "handler-ws-observer-model" + var observed []pluginapi.WebSocketResponseEvent + executor := &interceptorCaptureExecutor{ + execute: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) { + if opts.WebSocketResponseObserver != nil { + opts.WebSocketResponseObserver(ctx, coreexecutor.WebSocketResponseEvent{ + SourceFormat: opts.SourceFormat.String(), + Model: req.Model, + Provider: "codex", + AuthID: "auth-test", + EventType: "codex.rate_limits", + Payload: []byte(`{"type":"codex.rate_limits"}`), + }) + } + return coreexecutor.Response{Payload: []byte(`{"id":"resp-1"}`)}, nil + }, + } + handler := newInterceptorHandler(t, model, executor, nil) + handler.SetPluginHost(&handlerInterceptorTestHost{ + observeWebSocketResponseEvent: func(ctx context.Context, event pluginapi.WebSocketResponseEvent) { + observed = append(observed, event) + }, + }) + + _, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", model, []byte(fmt.Sprintf(`{"model":%q}`, model)), "") + if errMsg != nil { + t.Fatalf("ExecuteWithAuthManager() error = %+v", errMsg) + } + + if len(observed) != 1 { + t.Fatalf("observed %d events, want 1", len(observed)) + } + if observed[0].EventType != "codex.rate_limits" { + t.Fatalf("EventType = %q, want codex.rate_limits", observed[0].EventType) + } + if observed[0].AuthID != "auth-test" { + t.Fatalf("AuthID = %q, want auth-test", observed[0].AuthID) + } + if observed[0].RequestID == "" { + t.Fatal("RequestID is empty, want populated request ID") + } +} diff --git a/sdk/api/handlers/handlers_stream.go b/sdk/api/handlers/handlers_stream.go index a6525fc3..f476e104 100644 --- a/sdk/api/handlers/handlers_stream.go +++ b/sdk/api/handlers/handlers_stream.go @@ -339,6 +339,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context Headers: modelExecutionHeaders(ctx, execOptions.Headers), Query: modelExecutionQuery(ctx, execOptions.Query), RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, lifecycle.requestID(), execOptions.SkipInterceptorPluginID), + WebSocketResponseObserver: h.webSocketResponseObserver(lifecycle.requestID(), execOptions.SkipInterceptorPluginID), } opts.Metadata = reqMeta var interceptErr *interfaces.ErrorMessage diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go index 9839f4ac..37755875 100644 --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -144,6 +144,25 @@ func (e *RequestTerminatedError) ResponseBody() []byte { return append([]byte(nil), e.Body...) } +// WebSocketResponseEvent describes an upstream WebSocket response event received during execution. +type WebSocketResponseEvent struct { + RequestID string + TraceID string + SourceFormat string + Model string + RequestedModel string + Provider string + AuthID string + AuthLabel string + AuthType string + EventType string + Payload []byte + Metadata map[string]any +} + +// WebSocketResponseObserver receives upstream WebSocket response events during execution. +type WebSocketResponseObserver func(context.Context, WebSocketResponseEvent) + // Options controls execution behavior for both streaming and non-streaming calls. type Options struct { // Stream toggles streaming mode. @@ -165,6 +184,8 @@ type Options struct { Metadata map[string]any // RequestAfterAuthInterceptor runs after credential selection and before executor translation. RequestAfterAuthInterceptor RequestAfterAuthInterceptor + // WebSocketResponseObserver receives upstream WebSocket response events during execution. + WebSocketResponseObserver WebSocketResponseObserver // ExecutionLifecycle owns Home-dispatched execution resources. Executors must not add it to request metadata. ExecutionLifecycle ExecutionLifecycle } diff --git a/sdk/pluginabi/types.go b/sdk/pluginabi/types.go index 5089b132..0df98f49 100644 --- a/sdk/pluginabi/types.go +++ b/sdk/pluginabi/types.go @@ -10,10 +10,14 @@ const ( // Version 3 omits OriginalRequest/RequestBody on payload stream chunks // (ChunkIndex >= 0); those fields remain on StreamChunkHeaderInitIndex only. // Plugins that still need per-chunk request bodies should keep schema_version < 3. - SchemaVersion uint32 = 3 + // Version 4 adds upstream WebSocket response event observation. + SchemaVersion uint32 = 4 // SchemaVersionStreamChunkOmitRequestBody is the first schema version that omits // request bodies on payload stream-chunk interceptor calls. SchemaVersionStreamChunkOmitRequestBody uint32 = 3 + // SchemaVersionWebSocketResponseObserver is the first schema version that supports + // upstream WebSocket response event observation. + SchemaVersionWebSocketResponseObserver uint32 = 4 ) const ( @@ -58,6 +62,8 @@ const ( MethodResponseInterceptAfter = "response.intercept_after" MethodResponseInterceptStreamChunk = "response.intercept_stream_chunk" + MethodWebSocketResponseEvent = "websocket.response_event" + MethodThinkingIdentifier = "thinking.identifier" MethodThinkingApply = "thinking.apply" diff --git a/sdk/pluginabi/types_test.go b/sdk/pluginabi/types_test.go index 45ac19f4..60489e72 100644 --- a/sdk/pluginabi/types_test.go +++ b/sdk/pluginabi/types_test.go @@ -27,8 +27,11 @@ func TestEnvelopeRoundTrip(t *testing.T) { } func TestMethodNamesAreStable(t *testing.T) { - if SchemaVersion != 3 { - t.Fatalf("SchemaVersion = %d, want 3", SchemaVersion) + if SchemaVersion != 4 { + t.Fatalf("SchemaVersion = %d, want 4", SchemaVersion) + } + if SchemaVersionWebSocketResponseObserver != 4 { + t.Fatalf("SchemaVersionWebSocketResponseObserver = %d, want 4", SchemaVersionWebSocketResponseObserver) } if SchemaVersionStreamChunkOmitRequestBody != 3 { t.Fatalf("SchemaVersionStreamChunkOmitRequestBody = %d, want 3", SchemaVersionStreamChunkOmitRequestBody) @@ -54,6 +57,9 @@ func TestMethodNamesAreStable(t *testing.T) { if MethodResponseInterceptStreamChunk != "response.intercept_stream_chunk" { t.Fatalf("MethodResponseInterceptStreamChunk = %q", MethodResponseInterceptStreamChunk) } + if MethodWebSocketResponseEvent != "websocket.response_event" { + t.Fatalf("MethodWebSocketResponseEvent = %q", MethodWebSocketResponseEvent) + } if MethodHostHTTPDo != "host.http.do" { t.Fatalf("MethodHostHTTPDo = %q", MethodHostHTTPDo) } diff --git a/sdk/pluginapi/types.go b/sdk/pluginapi/types.go index 6add5d69..adcf1ff1 100644 --- a/sdk/pluginapi/types.go +++ b/sdk/pluginapi/types.go @@ -112,6 +112,8 @@ type Capabilities struct { ResponseInterceptor ResponseInterceptor // StreamChunkInterceptor rewrites successful HTTP stream chunks before downstream delivery. StreamChunkInterceptor StreamChunkInterceptor + // WebSocketResponseObserver receives upstream WebSocket response events during execution. + WebSocketResponseObserver WebSocketResponseObserver // ThinkingApplier applies validated thinking configuration to provider payloads. ThinkingApplier ThinkingApplier // UsagePlugin receives completed usage records. @@ -949,6 +951,11 @@ type StreamChunkInterceptor interface { InterceptStreamChunk(context.Context, StreamChunkInterceptRequest) (StreamChunkInterceptResponse, error) } +// WebSocketResponseObserver observes upstream WebSocket response events received during execution. +type WebSocketResponseObserver interface { + ObserveWebSocketResponseEvent(context.Context, WebSocketResponseEvent) error +} + // StreamChunkHeaderInitIndex marks the header-only stream initialization interceptor call. const StreamChunkHeaderInitIndex = -1 @@ -1123,6 +1130,22 @@ type StreamChunkInterceptResponse struct { DropChunk bool } +// WebSocketResponseEvent describes an upstream WebSocket response event received during execution. +type WebSocketResponseEvent struct { + RequestID string + TraceID string + SourceFormat string + Model string + RequestedModel string + Provider string + AuthID string + AuthLabel string + AuthType string + EventType string + Payload []byte + Metadata map[string]any +} + // PayloadResponse returns a transformed raw payload. type PayloadResponse struct { // Body contains the transformed payload bytes. diff --git a/sdk/pluginapi/types_test.go b/sdk/pluginapi/types_test.go index 0cbd10ed..fe2b8cfa 100644 --- a/sdk/pluginapi/types_test.go +++ b/sdk/pluginapi/types_test.go @@ -27,6 +27,7 @@ var _ RequestInterceptor = (*compileTimePlugin)(nil) var _ RequestLifecyclePlugin = (*compileTimePlugin)(nil) var _ ResponseInterceptor = (*compileTimePlugin)(nil) var _ StreamChunkInterceptor = (*compileTimePlugin)(nil) +var _ WebSocketResponseObserver = (*compileTimePlugin)(nil) var _ ThinkingApplier = (*compileTimePlugin)(nil) var _ UsagePlugin = (*compileTimePlugin)(nil) var _ CommandLinePlugin = (*compileTimePlugin)(nil) @@ -529,6 +530,10 @@ func (compileTimePlugin) InterceptStreamChunk(context.Context, StreamChunkInterc return StreamChunkInterceptResponse{}, nil } +func (compileTimePlugin) ObserveWebSocketResponseEvent(context.Context, WebSocketResponseEvent) error { + return nil +} + func (compileTimePlugin) ApplyThinking(context.Context, ThinkingApplyRequest) (PayloadResponse, error) { return PayloadResponse{}, nil } -- 2.51.2