diff --git a/internal/cli/gateway_commands.go b/internal/cli/gateway_commands.go index 13101328..09ba9731 100644 --- a/internal/cli/gateway_commands.go +++ b/internal/cli/gateway_commands.go @@ -124,10 +124,14 @@ func defaultGatewayCommandRunner(ctx context.Context, options gatewayCommandOpti signalContext, stop := signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM) defer stop() + relay := gateway.NewStreamRelay(gateway.StreamRelayOptions{ + Logger: logger, + }) ipcServer, err := newGatewayServer(gateway.ServerOptions{ ListenAddress: options.ListenAddress, Logger: logger, + Relay: relay, }) if err != nil { return err @@ -135,12 +139,14 @@ func defaultGatewayCommandRunner(ctx context.Context, options gatewayCommandOpti networkServer, err := newGatewayNetwork(gateway.NetworkServerOptions{ ListenAddress: options.HTTPAddress, Logger: logger, + Relay: relay, }) if err != nil { _ = ipcServer.Close(context.Background()) return err } defer func() { + relay.Stop() _ = networkServer.Close(context.Background()) _ = ipcServer.Close(context.Background()) }() diff --git a/internal/gateway/bootstrap.go b/internal/gateway/bootstrap.go index 703d4b25..081c8d2f 100644 --- a/internal/gateway/bootstrap.go +++ b/internal/gateway/bootstrap.go @@ -10,26 +10,27 @@ import ( "neo-code/internal/gateway/protocol" ) -type requestFrameHandler func(frame MessageFrame) MessageFrame +type requestFrameHandler func(ctx context.Context, frame MessageFrame) MessageFrame var wakeOpenURLHandler = handlers.NewWakeOpenURLHandler() var requestFrameHandlers = map[FrameAction]requestFrameHandler{ FrameActionPing: handlePingFrame, + FrameActionBindStream: handleBindStreamFrame, FrameActionWakeOpenURL: handleWakeOpenURLFrame, } // dispatchRequestFrame 统一分发 request 帧到对应动作处理器。 -func dispatchRequestFrame(_ context.Context, frame MessageFrame, _ RuntimePort) MessageFrame { +func dispatchRequestFrame(ctx context.Context, frame MessageFrame, _ RuntimePort) MessageFrame { handler, ok := requestFrameHandlers[frame.Action] if !ok { return errorFrame(frame, NewFrameError(ErrorCodeUnsupportedAction, "action is not implemented in gateway step 2")) } - return handler(frame) + return handler(ctx, frame) } // handlePingFrame 处理 ping 探活请求并返回 pong 响应。 -func handlePingFrame(frame MessageFrame) MessageFrame { +func handlePingFrame(_ context.Context, frame MessageFrame) MessageFrame { return MessageFrame{ Type: FrameTypeAck, Action: FrameActionPing, @@ -40,8 +41,43 @@ func handlePingFrame(frame MessageFrame) MessageFrame { } } +// handleBindStreamFrame 处理 gateway.bindStream 请求并登记连接到会话路由表。 +func handleBindStreamFrame(ctx context.Context, frame MessageFrame) MessageFrame { + params, err := decodeBindStreamParams(frame.Payload) + if err != nil { + return errorFrame(frame, err) + } + + relay, relayExists := StreamRelayFromContext(ctx) + connectionID, connectionExists := ConnectionIDFromContext(ctx) + if !relayExists || !connectionExists { + return errorFrame(frame, NewFrameError(ErrorCodeInternalError, "stream relay context is unavailable")) + } + + if bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: params.SessionID, + RunID: params.RunID, + Channel: params.Channel, + Explicit: true, + }); bindErr != nil { + return errorFrame(frame, bindErr) + } + + return MessageFrame{ + Type: FrameTypeAck, + Action: FrameActionBindStream, + RequestID: frame.RequestID, + SessionID: params.SessionID, + RunID: params.RunID, + Payload: map[string]any{ + "message": "stream binding updated", + "channel": params.Channel, + }, + } +} + // handleWakeOpenURLFrame 解析并处理 wake.openUrl 请求。 -func handleWakeOpenURLFrame(frame MessageFrame) MessageFrame { +func handleWakeOpenURLFrame(_ context.Context, frame MessageFrame) MessageFrame { intent, err := decodeWakeIntent(frame.Payload) if err != nil { return errorFrame(frame, NewFrameError(ErrorCodeInvalidFrame, "invalid wake payload")) @@ -51,16 +87,92 @@ func handleWakeOpenURLFrame(frame MessageFrame) MessageFrame { if wakeErr != nil { return errorFrame(frame, toFrameError(wakeErr)) } + sessionID := intent.SessionID + if strings.TrimSpace(sessionID) == "" { + sessionID = strings.TrimSpace(frame.SessionID) + } return MessageFrame{ Type: FrameTypeAck, Action: FrameActionWakeOpenURL, RequestID: frame.RequestID, - SessionID: intent.SessionID, + SessionID: sessionID, Payload: result, } } +type bindStreamParams struct { + SessionID string + RunID string + Channel StreamChannel +} + +// decodeBindStreamParams 将 payload 解析为 bind_stream 所需参数。 +func decodeBindStreamParams(payload any) (bindStreamParams, *FrameError) { + switch typed := payload.(type) { + case protocol.BindStreamParams: + return normalizeBindStreamParams(typed) + case *protocol.BindStreamParams: + if typed == nil { + return bindStreamParams{}, NewFrameError(ErrorCodeInvalidFrame, "invalid bind_stream payload") + } + return normalizeBindStreamParams(*typed) + case map[string]any: + return normalizeBindStreamParams(protocol.BindStreamParams{ + SessionID: readStringValue(typed, "session_id"), + RunID: readStringValue(typed, "run_id"), + Channel: readStringValue(typed, "channel"), + }) + default: + raw, marshalErr := json.Marshal(payload) + if marshalErr != nil { + return bindStreamParams{}, NewFrameError(ErrorCodeInvalidFrame, "invalid bind_stream payload") + } + var decoded protocol.BindStreamParams + if unmarshalErr := json.Unmarshal(raw, &decoded); unmarshalErr != nil { + return bindStreamParams{}, NewFrameError(ErrorCodeInvalidFrame, "invalid bind_stream payload") + } + return normalizeBindStreamParams(decoded) + } +} + +// normalizeBindStreamParams 对 bind_stream 参数执行归一化与有效性校验。 +func normalizeBindStreamParams(params protocol.BindStreamParams) (bindStreamParams, *FrameError) { + sessionID := strings.TrimSpace(params.SessionID) + if sessionID == "" { + return bindStreamParams{}, NewMissingRequiredFieldError("payload.session_id") + } + + runID := strings.TrimSpace(params.RunID) + channel := strings.ToLower(strings.TrimSpace(params.Channel)) + if channel == "" { + channel = string(StreamChannelAll) + } + parsedChannel, validChannel := ParseStreamChannel(channel) + if !validChannel { + return bindStreamParams{}, NewFrameError(ErrorCodeInvalidAction, "invalid bind_stream channel") + } + + return bindStreamParams{ + SessionID: sessionID, + RunID: runID, + Channel: parsedChannel, + }, nil +} + +// readStringValue 从 map 负载中读取并归一化字符串字段。 +func readStringValue(payload map[string]any, key string) string { + rawValue, exists := payload[key] + if !exists { + return "" + } + stringValue, ok := rawValue.(string) + if !ok { + return "" + } + return strings.TrimSpace(stringValue) +} + // decodeWakeIntent 将任意 payload 解码为 WakeIntent 结构。 func decodeWakeIntent(payload any) (protocol.WakeIntent, error) { if payload == nil { diff --git a/internal/gateway/bootstrap_test.go b/internal/gateway/bootstrap_test.go index d182c01d..745b325d 100644 --- a/internal/gateway/bootstrap_test.go +++ b/internal/gateway/bootstrap_test.go @@ -95,6 +95,108 @@ func TestDispatchRequestFrameUnsupportedAction(t *testing.T) { } } +func TestDispatchRequestFrameBindStream(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := NewConnectionID() + connectionCtx := WithConnectionID(ctx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { + _ = message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + defer relay.dropConnection(connectionID) + + response := dispatchRequestFrame(connectionCtx, MessageFrame{ + Type: FrameTypeRequest, + Action: FrameActionBindStream, + RequestID: "bind-1", + Payload: protocol.BindStreamParams{ + SessionID: "session-1", + RunID: "run-1", + Channel: "ipc", + }, + }, nil) + if response.Type != FrameTypeAck { + t.Fatalf("response type = %q, want %q", response.Type, FrameTypeAck) + } + if response.Action != FrameActionBindStream { + t.Fatalf("response action = %q, want %q", response.Action, FrameActionBindStream) + } + if response.SessionID != "session-1" { + t.Fatalf("session_id = %q, want %q", response.SessionID, "session-1") + } +} + +func TestHandleBindStreamFrameErrors(t *testing.T) { + t.Run("missing relay context", func(t *testing.T) { + response := handleBindStreamFrame(context.Background(), MessageFrame{ + Type: FrameTypeRequest, + Action: FrameActionBindStream, + Payload: protocol.BindStreamParams{ + SessionID: "session-1", + }, + }) + if response.Type != FrameTypeError { + t.Fatalf("response type = %q, want %q", response.Type, FrameTypeError) + } + if response.Error == nil || response.Error.Code != ErrorCodeInternalError.String() { + t.Fatalf("response error = %#v, want %q", response.Error, ErrorCodeInternalError.String()) + } + }) + + t.Run("channel mismatch", func(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := NewConnectionID() + connectionCtx := WithConnectionID(ctx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelWS, + Context: connectionCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { + _ = message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + defer relay.dropConnection(connectionID) + + response := handleBindStreamFrame(connectionCtx, MessageFrame{ + Type: FrameTypeRequest, + Action: FrameActionBindStream, + Payload: protocol.BindStreamParams{ + SessionID: "session-1", + Channel: "ipc", + }, + }) + if response.Type != FrameTypeError { + t.Fatalf("response type = %q, want %q", response.Type, FrameTypeError) + } + if response.Error == nil || response.Error.Code != ErrorCodeInvalidAction.String() { + t.Fatalf("response error = %#v, want %q", response.Error, ErrorCodeInvalidAction.String()) + } + }) +} + func TestDecodeWakeIntentAdditionalBranches(t *testing.T) { t.Run("nil payload", func(t *testing.T) { _, err := decodeWakeIntent(nil) diff --git a/internal/gateway/connection_context.go b/internal/gateway/connection_context.go new file mode 100644 index 00000000..26deae2d --- /dev/null +++ b/internal/gateway/connection_context.go @@ -0,0 +1,100 @@ +package gateway + +import ( + "context" + "fmt" + "strings" + "sync/atomic" + "time" +) + +// StreamChannel 表示连接所属的流式通道类型。 +type StreamChannel string + +const ( + // StreamChannelAll 表示绑定对当前连接所属通道不过滤。 + StreamChannelAll StreamChannel = "all" + // StreamChannelIPC 表示绑定仅用于本地 IPC 连接。 + StreamChannelIPC StreamChannel = "ipc" + // StreamChannelWS 表示绑定仅用于 WebSocket 连接。 + StreamChannelWS StreamChannel = "ws" + // StreamChannelSSE 表示绑定仅用于 SSE 连接。 + StreamChannelSSE StreamChannel = "sse" +) + +// ConnectionID 表示网关侧分配给物理连接的全局唯一标识。 +type ConnectionID string + +type connectionIDContextKey struct{} +type streamRelayContextKey struct{} + +var ( + connectionSequence uint64 + connectionStartEpoch = time.Now().Unix() +) + +// NewConnectionID 生成全局唯一 ConnectionID，用于连接绑定和路由兜底。 +func NewConnectionID() ConnectionID { + sequence := atomic.AddUint64(&connectionSequence, 1) + return ConnectionID(fmt.Sprintf("cid_%d_%d", connectionStartEpoch, sequence)) +} + +// WithConnectionID 将 ConnectionID 注入上下文，供后续路由和提取逻辑读取。 +func WithConnectionID(ctx context.Context, connectionID ConnectionID) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, connectionIDContextKey{}, NormalizeConnectionID(connectionID)) +} + +// ConnectionIDFromContext 从上下文读取 ConnectionID。 +func ConnectionIDFromContext(ctx context.Context) (ConnectionID, bool) { + if ctx == nil { + return "", false + } + value, ok := ctx.Value(connectionIDContextKey{}).(ConnectionID) + if !ok { + return "", false + } + value = NormalizeConnectionID(value) + if value == "" { + return "", false + } + return value, true +} + +// WithStreamRelay 将流式中继实例注入上下文，供请求处理阶段读取。 +func WithStreamRelay(ctx context.Context, relay *StreamRelay) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, streamRelayContextKey{}, relay) +} + +// StreamRelayFromContext 从上下文中读取流式中继实例。 +func StreamRelayFromContext(ctx context.Context) (*StreamRelay, bool) { + if ctx == nil { + return nil, false + } + relay, ok := ctx.Value(streamRelayContextKey{}).(*StreamRelay) + if !ok || relay == nil { + return nil, false + } + return relay, true +} + +// ParseStreamChannel 解析并校验连接通道参数。 +func ParseStreamChannel(raw string) (StreamChannel, bool) { + normalized := StreamChannel(strings.ToLower(strings.TrimSpace(raw))) + switch normalized { + case StreamChannelAll, StreamChannelIPC, StreamChannelWS, StreamChannelSSE: + return normalized, true + default: + return "", false + } +} + +// NormalizeConnectionID 将连接标识归一化为空白裁剪后的稳定值。 +func NormalizeConnectionID(connectionID ConnectionID) ConnectionID { + return ConnectionID(strings.TrimSpace(string(connectionID))) +} diff --git a/internal/gateway/connection_context_test.go b/internal/gateway/connection_context_test.go new file mode 100644 index 00000000..70985f5d --- /dev/null +++ b/internal/gateway/connection_context_test.go @@ -0,0 +1,46 @@ +package gateway + +import ( + "context" + "strings" + "testing" +) + +func TestConnectionContextRoundTrip(t *testing.T) { + connectionID := NewConnectionID() + if !strings.HasPrefix(string(connectionID), "cid_") { + t.Fatalf("connection id = %q, want prefix %q", connectionID, "cid_") + } + + ctx := WithConnectionID(context.Background(), connectionID) + resolvedID, exists := ConnectionIDFromContext(ctx) + if !exists { + t.Fatal("connection id should exist in context") + } + if resolvedID != connectionID { + t.Fatalf("resolved connection id = %q, want %q", resolvedID, connectionID) + } +} + +func TestStreamRelayContextRoundTrip(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + ctx := WithStreamRelay(context.Background(), relay) + + resolvedRelay, exists := StreamRelayFromContext(ctx) + if !exists { + t.Fatal("stream relay should exist in context") + } + if resolvedRelay != relay { + t.Fatal("resolved relay should match original relay") + } +} + +func TestParseStreamChannel(t *testing.T) { + if channel, ok := ParseStreamChannel("ws"); !ok || channel != StreamChannelWS { + t.Fatalf("parse channel = %q ok=%v, want %q true", channel, ok, StreamChannelWS) + } + + if _, ok := ParseStreamChannel("tcp"); ok { + t.Fatal("invalid channel should be rejected") + } +} diff --git a/internal/gateway/coverage_boost_test.go b/internal/gateway/coverage_boost_test.go new file mode 100644 index 00000000..6e8a20e1 --- /dev/null +++ b/internal/gateway/coverage_boost_test.go @@ -0,0 +1,660 @@ +package gateway + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "neo-code/internal/gateway/protocol" +) + +func TestDecodeBindStreamParamsAndReadStringValueBranches(t *testing.T) { + t.Run("nil bind stream pointer", func(t *testing.T) { + _, frameErr := decodeBindStreamParams((*protocol.BindStreamParams)(nil)) + if frameErr == nil || frameErr.Code != ErrorCodeInvalidFrame.String() { + t.Fatalf("frameErr = %#v, want invalid_frame", frameErr) + } + }) + + t.Run("map payload and trim", func(t *testing.T) { + params, frameErr := decodeBindStreamParams(map[string]any{ + "session_id": " s-1 ", + "run_id": 123, + "channel": " ws ", + }) + if frameErr != nil { + t.Fatalf("decode bind stream params: %v", frameErr) + } + if params.SessionID != "s-1" { + t.Fatalf("session_id = %q, want %q", params.SessionID, "s-1") + } + if params.RunID != "" { + t.Fatalf("run_id = %q, want empty", params.RunID) + } + if params.Channel != StreamChannelWS { + t.Fatalf("channel = %q, want %q", params.Channel, StreamChannelWS) + } + }) + + t.Run("marshal error", func(t *testing.T) { + type badPayload struct { + Bad chan int `json:"bad"` + } + _, frameErr := decodeBindStreamParams(badPayload{Bad: make(chan int)}) + if frameErr == nil || frameErr.Code != ErrorCodeInvalidFrame.String() { + t.Fatalf("frameErr = %#v, want invalid_frame", frameErr) + } + }) + + t.Run("invalid bind stream channel", func(t *testing.T) { + _, frameErr := decodeBindStreamParams(protocol.BindStreamParams{SessionID: "s-1", Channel: "tcp"}) + if frameErr == nil || frameErr.Code != ErrorCodeInvalidAction.String() { + t.Fatalf("frameErr = %#v, want invalid_action", frameErr) + } + }) + + t.Run("read string helper", func(t *testing.T) { + payload := map[string]any{ + "str": " value ", + "num": 1, + } + if got := readStringValue(payload, "missing"); got != "" { + t.Fatalf("missing value = %q, want empty", got) + } + if got := readStringValue(payload, "num"); got != "" { + t.Fatalf("non-string value = %q, want empty", got) + } + if got := readStringValue(payload, "str"); got != "value" { + t.Fatalf("string value = %q, want %q", got, "value") + } + }) +} + +func TestConnectionContextAdditionalBranches(t *testing.T) { + ctx := WithConnectionID(nil, ConnectionID(" ")) + if _, exists := ConnectionIDFromContext(ctx); exists { + t.Fatal("blank normalized connection id should not exist") + } + + ctx = context.WithValue(context.Background(), connectionIDContextKey{}, "cid-raw") + if _, exists := ConnectionIDFromContext(ctx); exists { + t.Fatal("non-ConnectionID type should not exist") + } + + relayCtx := WithStreamRelay(nil, nil) + if _, exists := StreamRelayFromContext(relayCtx); exists { + t.Fatal("nil relay should not be resolved") + } + + if channel, ok := ParseStreamChannel(" ALL "); !ok || channel != StreamChannelAll { + t.Fatalf("channel = %q ok=%v, want all true", channel, ok) + } +} + +func TestStreamRelayRegisterConnectionValidationBranches(t *testing.T) { + var nilRelay *StreamRelay + if err := nilRelay.RegisterConnection(ConnectionRegistration{}); err == nil { + t.Fatal("expected nil relay register error") + } + + relay := NewStreamRelay(StreamRelayOptions{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + cases := []ConnectionRegistration{ + {Context: ctx, Cancel: cancel, Write: func(message RelayMessage) error { return nil }, Close: func() {}}, + {ConnectionID: "cid", Cancel: cancel, Write: func(message RelayMessage) error { return nil }, Close: func() {}}, + {ConnectionID: "cid", Context: ctx, Write: func(message RelayMessage) error { return nil }, Close: func() {}}, + {ConnectionID: "cid", Context: ctx, Cancel: cancel, Close: func() {}}, + {ConnectionID: "cid", Context: ctx, Cancel: cancel, Write: func(message RelayMessage) error { return nil }}, + {ConnectionID: "cid", Channel: StreamChannelAll, Context: ctx, Cancel: cancel, Write: func(message RelayMessage) error { return nil }, Close: func() {}}, + } + for i, registration := range cases { + if err := relay.RegisterConnection(registration); err == nil { + t.Fatalf("case %d: expected register error", i) + } + } + + registration := ConnectionRegistration{ + ConnectionID: " cid-dup ", + Channel: StreamChannelIPC, + Context: ctx, + Cancel: cancel, + Write: func(message RelayMessage) error { return nil }, + Close: func() {}, + } + if err := relay.RegisterConnection(registration); err != nil { + t.Fatalf("register connection: %v", err) + } + t.Cleanup(func() { relay.dropConnection("cid-dup") }) + if err := relay.RegisterConnection(registration); err == nil { + t.Fatal("expected duplicate registration error") + } +} + +func TestStreamRelayBindingAndRefreshBranches(t *testing.T) { + var nilRelay *StreamRelay + if frameErr := nilRelay.BindConnection("cid", StreamBinding{SessionID: "s"}); frameErr == nil { + t.Fatal("expected nil relay bind error") + } + if nilRelay.ResolveFallbackSessionID("cid") != "" { + t.Fatal("nil relay fallback should be empty") + } + if nilRelay.RefreshConnectionBindings("cid") { + t.Fatal("nil relay refresh should be false") + } + + relay := NewStreamRelay(StreamRelayOptions{BindingTTL: 20 * time.Millisecond, MaxBindingsPerConnection: 1}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := ConnectionID("cid-binding") + if frameErr := relay.BindConnection("", StreamBinding{SessionID: "s"}); frameErr == nil || frameErr.Code != ErrorCodeMissingRequiredField.String() { + t.Fatalf("blank connection id error = %#v", frameErr) + } + if frameErr := relay.BindConnection(connectionID, StreamBinding{}); frameErr == nil || frameErr.Code != ErrorCodeMissingRequiredField.String() { + t.Fatalf("missing session error = %#v", frameErr) + } + if frameErr := relay.BindConnection(connectionID, StreamBinding{SessionID: "s", Channel: "tcp"}); frameErr == nil || frameErr.Code != ErrorCodeInvalidAction.String() { + t.Fatalf("invalid channel error = %#v", frameErr) + } + if frameErr := relay.BindConnection(connectionID, StreamBinding{SessionID: "s"}); frameErr == nil || frameErr.Code != ErrorCodeInvalidAction.String() { + t.Fatalf("unregistered connection error = %#v", frameErr) + } + + connectionCtx := WithConnectionID(ctx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelWS, + Context: connectionCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { return nil }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + t.Cleanup(func() { relay.dropConnection(connectionID) }) + + if frameErr := relay.BindConnection(connectionID, StreamBinding{SessionID: "s-1", Channel: StreamChannelIPC}); frameErr == nil || frameErr.Code != ErrorCodeInvalidAction.String() { + t.Fatalf("channel mismatch error = %#v", frameErr) + } + if frameErr := relay.BindConnection(connectionID, StreamBinding{SessionID: "s-1", RunID: "r-1", Channel: StreamChannelAll}); frameErr != nil { + t.Fatalf("first bind error: %v", frameErr) + } + if frameErr := relay.BindConnection(connectionID, StreamBinding{SessionID: "s-1", RunID: "r-1", Channel: StreamChannelAll}); frameErr != nil { + t.Fatalf("upsert bind should pass: %v", frameErr) + } + if frameErr := relay.BindConnection(connectionID, StreamBinding{SessionID: "s-1", RunID: "r-2", Channel: StreamChannelAll}); frameErr == nil || frameErr.Code != ErrorCodeInvalidAction.String() { + t.Fatalf("max bindings error = %#v", frameErr) + } + + relay.mu.Lock() + bindings := relay.connectionBindings[connectionID] + bindings[bindingKey{sessionID: "s-2", runID: "r-expired"}] = &bindingState{sessionID: "s-2", runID: "r-expired", expireAt: time.Now().Add(-time.Second)} + bindings[bindingKey{sessionID: "s-3", runID: "r-nil"}] = nil + relay.addConnectionToSessionIndexLocked("s-2", connectionID) + relay.addConnectionToSessionRunIndexLocked("s-2", "r-expired", connectionID) + relay.mu.Unlock() + + if !relay.RefreshConnectionBindings(connectionID) { + t.Fatal("expected refresh to succeed with active bindings") + } + if fallback := relay.ResolveFallbackSessionID(connectionID); fallback == "" { + t.Fatal("fallback session should not be empty") + } + + relay.cleanupExpiredBindings() + relay.mu.RLock() + _, hasExpired := relay.connectionBindings[connectionID][bindingKey{sessionID: "s-2", runID: "r-expired"}] + relay.mu.RUnlock() + if hasExpired { + t.Fatal("expired binding should be removed after cleanup") + } +} + +func TestStreamRelayAutoBindAndExtractSessionBranches(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := ConnectionID("cid-auto-bind") + connectionCtx := WithConnectionID(ctx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { return nil }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + t.Cleanup(func() { relay.dropConnection(connectionID) }) + + relay.AutoBindFromFrame(connectionID, MessageFrame{SessionID: "session-direct", RunID: "run-direct"}) + relay.AutoBindFromFrame(connectionID, MessageFrame{Payload: protocol.BindStreamParams{SessionID: "session-bind"}}) + relay.AutoBindFromFrame(connectionID, MessageFrame{Payload: &protocol.WakeIntent{SessionID: "session-wake"}}) + relay.AutoBindFromFrame(connectionID, MessageFrame{Payload: map[string]any{"session_id": " session-map "}}) + relay.AutoBindFromFrame(connectionID, MessageFrame{Payload: map[string]any{"session_id": 1}}) + relay.AutoBindFromFrame(connectionID, MessageFrame{}) + + relay.mu.RLock() + bindingCount := len(relay.connectionBindings[connectionID]) + relay.mu.RUnlock() + if bindingCount < 4 { + t.Fatalf("binding count = %d, want >= 4", bindingCount) + } + + if got := extractSessionIDFromPayload(protocol.WakeIntent{SessionID: "s1"}); got != "s1" { + t.Fatalf("session from wake intent = %q, want s1", got) + } + if got := extractSessionIDFromPayload((*protocol.WakeIntent)(nil)); got != "" { + t.Fatalf("session from nil wake ptr = %q, want empty", got) + } + if got := extractSessionIDFromPayload(protocol.BindStreamParams{SessionID: "s2"}); got != "s2" { + t.Fatalf("session from bind params = %q, want s2", got) + } + if got := extractSessionIDFromPayload((*protocol.BindStreamParams)(nil)); got != "" { + t.Fatalf("session from nil bind ptr = %q, want empty", got) + } + if got := extractSessionIDFromPayload(map[string]any{"session_id": " s3 "}); got != "s3" { + t.Fatalf("session from map = %q, want s3", got) + } +} + +func TestStreamRelayRuntimeAndWriterBranches(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{QueueSize: 1, BindingTTL: 20 * time.Millisecond, CleanupInterval: 5 * time.Millisecond}) + baseCtx := context.Background() + + if relay.SendJSONRPCPayload("", map[string]string{"x": "y"}) { + t.Fatal("blank connection send should fail") + } + + closedCount := int32(0) + writeErrConnID := ConnectionID("cid-write-err") + writeErrCtx, writeErrCancel := context.WithCancel(baseCtx) + t.Cleanup(writeErrCancel) + writeErrCtx = WithConnectionID(writeErrCtx, writeErrConnID) + writeErrCtx = WithStreamRelay(writeErrCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: writeErrConnID, + Channel: StreamChannelWS, + Context: writeErrCtx, + Cancel: writeErrCancel, + Write: func(message RelayMessage) error { + return io.ErrClosedPipe + }, + Close: func() { + atomic.AddInt32(&closedCount, 1) + }, + }); err != nil { + t.Fatalf("register write error connection: %v", err) + } + if !relay.SendJSONRPCPayload(writeErrConnID, map[string]string{"trigger": "drop"}) { + t.Fatal("send payload should enqueue") + } + time.Sleep(60 * time.Millisecond) + if atomic.LoadInt32(&closedCount) == 0 { + t.Fatal("connection close should be called after write failure") + } + + sseMessageCh := make(chan RelayMessage, 4) + sseConnID := ConnectionID("cid-sse") + sseCtx, sseRootCancel := context.WithCancel(baseCtx) + t.Cleanup(sseRootCancel) + sseCtx = WithConnectionID(sseCtx, sseConnID) + sseCtx = WithStreamRelay(sseCtx, relay) + sseCancelCtx, sseCancel := context.WithCancel(sseCtx) + t.Cleanup(sseCancel) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: sseConnID, + Channel: StreamChannelSSE, + Context: sseCancelCtx, + Cancel: sseCancel, + Write: func(message RelayMessage) error { + sseMessageCh <- message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register sse connection: %v", err) + } + t.Cleanup(func() { relay.dropConnection(sseConnID) }) + + if bindErr := relay.BindConnection(sseConnID, StreamBinding{SessionID: "sse-session", Channel: StreamChannelSSE, Explicit: true}); bindErr != nil { + t.Fatalf("bind sse connection: %v", bindErr) + } + + relay.PublishRuntimeEvent(RuntimeEvent{Type: RuntimeEventTypeRunProgress, SessionID: "sse-session", Payload: map[string]string{"chunk": "ok"}}) + select { + case message := <-sseMessageCh: + if message.Kind != relayMessageKindSSE { + t.Fatalf("message kind = %q, want %q", message.Kind, relayMessageKindSSE) + } + if message.Event != protocol.MethodGatewayEvent { + t.Fatalf("event = %q, want %q", message.Event, protocol.MethodGatewayEvent) + } + case <-time.After(time.Second): + t.Fatal("expected sse relay event") + } + + var nilRelay *StreamRelay + nilRelay.PublishRuntimeEvent(RuntimeEvent{}) + nilRelay.AutoBindFromFrame("cid", MessageFrame{SessionID: "s"}) + if channel, ok := nilRelay.connectionChannel("cid"); ok || channel != "" { + t.Fatal("nil relay channel should not exist") + } + + if channel, ok := relay.connectionChannel("missing"); ok || channel != "" { + t.Fatal("missing connection should not exist") + } +} + +func TestStreamRelayMatchingAndLifecycleBranches(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{BindingTTL: 20 * time.Millisecond, CleanupInterval: 5 * time.Millisecond}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connID := ConnectionID("cid-match") + connCtx := WithConnectionID(ctx, connID) + connCtx = WithStreamRelay(connCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connID, + Channel: StreamChannelWS, + Context: connCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { return nil }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + t.Cleanup(func() { relay.dropConnection(connID) }) + + if bindErr := relay.BindConnection(connID, StreamBinding{SessionID: "session-x", RunID: "run-x", Channel: StreamChannelWS}); bindErr != nil { + t.Fatalf("bind connection: %v", bindErr) + } + + now := time.Now() + relay.mu.RLock() + if !relay.connectionMatchesEventLocked(connID, StreamChannelWS, "session-x", "run-x", now) { + t.Fatal("expected exact run to match") + } + if relay.connectionMatchesEventLocked(connID, StreamChannelWS, "session-x", "", now) { + t.Fatal("event with empty run_id should not match run-scoped binding") + } + relay.mu.RUnlock() + + relay.mu.Lock() + relay.connectionBindings[connID][bindingKey{sessionID: "session-x", runID: "run-x"}].channel = StreamChannelSSE + relay.mu.Unlock() + relay.mu.RLock() + if relay.connectionMatchesEventLocked(connID, StreamChannelWS, "session-x", "run-x", now) { + t.Fatal("channel mismatch should not match") + } + relay.mu.RUnlock() + + if matches := relay.matchConnectionsForEvent("", "run-x"); len(matches) != 0 { + t.Fatalf("empty session should not match, got %d", len(matches)) + } + if key := buildSessionRunKey(" s ", " r "); key != "s\x00r" { + t.Fatalf("session run key = %q, want %q", key, "s\\x00r") + } + + var nilRelay *StreamRelay + nilRelay.Start(nil, nil) + nilRelay.Stop() + + stubPort := &runtimePortEventStub{events: make(chan RuntimeEvent)} + relay.Start(ctx, stubPort) + cancel() + waitForStreamRelayState(t, relay, false) + + relay.runRuntimeEventLoop(context.Background(), nil, 0) + relay.runRuntimeEventLoop(context.Background(), &runtimePortEventStub{events: nil}, 0) +} + +func TestRPCDispatchAdditionalBranches(t *testing.T) { + if !requiresSession(FrameActionRun) { + t.Fatal("run should require session") + } + if requiresSession(FrameActionPing) { + t.Fatal("ping should not require session") + } + + ctx := context.Background() + frame := MessageFrame{Type: FrameTypeRequest, Action: FrameActionPing} + applyAutomaticBinding(ctx, frame) + + relay := NewStreamRelay(StreamRelayOptions{}) + connCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + connectionID := ConnectionID("cid-rpc") + connCtx = WithConnectionID(connCtx, connectionID) + connCtx = WithStreamRelay(connCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { return nil }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + t.Cleanup(func() { relay.dropConnection(connectionID) }) + + applyAutomaticBinding(connCtx, MessageFrame{Action: FrameActionBindStream, SessionID: "session-skip"}) + applyAutomaticBinding(connCtx, MessageFrame{Action: FrameActionRun, Payload: map[string]any{"session_id": "session-auto"}}) + if fallback := relay.ResolveFallbackSessionID(connectionID); fallback != "session-auto" { + t.Fatalf("fallback session = %q, want %q", fallback, "session-auto") + } +} + +func TestNetworkServerHelperBranches(t *testing.T) { + request := httptest.NewRequest("GET", "/sse?method=gateway.ping&id=%20custom-id%20", nil) + trigger := buildSSETriggerRequest(request) + if trigger.Method != protocol.MethodGatewayPing { + t.Fatalf("trigger method = %q, want %q", trigger.Method, protocol.MethodGatewayPing) + } + if string(trigger.ID) != `"custom-id"` { + t.Fatalf("trigger id = %s, want %s", trigger.ID, `"custom-id"`) + } + + raw := []byte(`{"jsonrpc":"2.0","id":"x","method":"gateway.ping","params":{}}`) + if _, rpcErr := decodeJSONRPCRequestFromBytes(raw); rpcErr != nil { + t.Fatalf("decode jsonrpc bytes: %v", rpcErr) + } + + if _, rpcErr := decodeJSONRPCRequestFromBytes([]byte(`{"jsonrpc":"2.0"} {"extra":1}`)); rpcErr == nil { + t.Fatal("expected trailing json parse error") + } + + recorder := httptest.NewRecorder() + writeJSONRPCHTTPResponse(recorder, protocol.NewJSONRPCErrorResponse(json.RawMessage(`"id-1"`), protocol.NewJSONRPCError( + protocol.JSONRPCCodeInternalError, + "boom", + protocol.GatewayCodeInternalError, + ))) + if contentType := recorder.Header().Get("Content-Type"); contentType != "application/json" { + t.Fatalf("content type = %q, want application/json", contentType) + } + if !bytes.Contains(recorder.Body.Bytes(), []byte(`"jsonrpc":"2.0"`)) { + t.Fatalf("response body = %s, want jsonrpc payload", recorder.Body.String()) + } +} + +func TestStreamRelayInternalBranchCoverage(t *testing.T) { + t.Run("resolve fallback skips expired and nil", func(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + connectionID := ConnectionID("cid-fallback") + now := time.Now() + + relay.mu.Lock() + relay.connectionBindings[connectionID] = map[bindingKey]*bindingState{ + {sessionID: "s-expired", runID: "r"}: { + sessionID: "s-expired", + expireAt: now.Add(-time.Second), + lastSeen: now.Add(time.Second), + }, + {sessionID: "s-old", runID: "r"}: { + sessionID: "s-old", + expireAt: now.Add(time.Second), + lastSeen: now.Add(-2 * time.Second), + }, + {sessionID: "s-new", runID: "r"}: { + sessionID: "s-new", + expireAt: now.Add(time.Second), + lastSeen: now, + }, + {sessionID: "s-nil", runID: "r"}: nil, + } + relay.mu.Unlock() + + if got := relay.ResolveFallbackSessionID(connectionID); got != "s-new" { + t.Fatalf("fallback session = %q, want %q", got, "s-new") + } + }) + + t.Run("run connection writer exits on closed queue", func(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connection := &relayConnection{ + id: "cid-queue-close", + channel: StreamChannelIPC, + ctx: ctx, + cancel: cancel, + writeFn: func(message RelayMessage) error { return nil }, + closeFn: func() {}, + queue: make(chan RelayMessage, 1), + } + relay.mu.Lock() + relay.connections[connection.id] = connection + relay.mu.Unlock() + + done := make(chan struct{}) + go func() { + defer close(done) + relay.runConnectionWriter(connection) + }() + close(connection.queue) + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("writer did not exit after queue closed") + } + }) + + t.Run("unregister connection without close callback", func(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := ConnectionID("cid-unregister") + relay.mu.Lock() + relay.connections[connectionID] = &relayConnection{ + id: connectionID, + channel: StreamChannelIPC, + ctx: ctx, + cancel: cancel, + writeFn: func(message RelayMessage) error { return nil }, + closeFn: func() {}, + queue: make(chan RelayMessage, 1), + } + relay.connectionBindings[connectionID] = map[bindingKey]*bindingState{ + {sessionID: "s", runID: "r"}: { + sessionID: "s", + runID: "r", + expireAt: time.Now().Add(time.Minute), + }, + } + relay.addConnectionToSessionIndexLocked("s", connectionID) + relay.addConnectionToSessionRunIndexLocked("s", "r", connectionID) + relay.mu.Unlock() + + if connection := relay.unregisterConnection(connectionID, false); connection == nil { + t.Fatal("expected unregister to return connection") + } + relay.mu.RLock() + _, hasSession := relay.sessionIndex["s"] + _, hasRun := relay.sessionRunIndex[buildSessionRunKey("s", "r")] + relay.mu.RUnlock() + if hasSession || hasRun { + t.Fatal("indexes should be cleaned after unregister") + } + + if connection := relay.unregisterConnection(connectionID, true); connection != nil { + t.Fatal("second unregister should return nil") + } + }) + + t.Run("connection match branch matrix", func(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + connectionID := ConnectionID("cid-match-matrix") + channelOnlyConnectionID := ConnectionID("cid-match-channel-only") + now := time.Now() + + relay.mu.Lock() + relay.connectionBindings[connectionID] = map[bindingKey]*bindingState{ + {sessionID: "s", runID: ""}: { + sessionID: "s", + runID: "", + channel: StreamChannelAll, + expireAt: now.Add(time.Minute), + }, + {sessionID: "other", runID: "r"}: { + sessionID: "other", + runID: "r", + channel: StreamChannelWS, + expireAt: now.Add(time.Minute), + }, + {sessionID: "s", runID: "expired"}: { + sessionID: "s", + runID: "expired", + channel: StreamChannelWS, + expireAt: now.Add(-time.Minute), + }, + {sessionID: "s", runID: "channel-mismatch"}: { + sessionID: "s", + runID: "channel-mismatch", + channel: StreamChannelSSE, + expireAt: now.Add(time.Minute), + }, + } + relay.connectionBindings[channelOnlyConnectionID] = map[bindingKey]*bindingState{ + {sessionID: "s", runID: "channel-mismatch"}: { + sessionID: "s", + runID: "channel-mismatch", + channel: StreamChannelSSE, + expireAt: now.Add(time.Minute), + }, + } + relay.mu.Unlock() + + relay.mu.RLock() + if !relay.connectionMatchesEventLocked(connectionID, StreamChannelWS, "s", "", now) { + t.Fatal("session-only binding should match event with empty run_id") + } + if !relay.connectionMatchesEventLocked(connectionID, StreamChannelWS, "s", "run-1", now) { + t.Fatal("session-only binding should match event with concrete run_id") + } + if relay.connectionMatchesEventLocked(channelOnlyConnectionID, StreamChannelWS, "s", "channel-mismatch", now) { + t.Fatal("channel-mismatched binding should not match") + } + relay.mu.RUnlock() + }) + + t.Run("send event notification skips missing connection", func(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + relay.sendEventNotification("missing", protocol.NewJSONRPCNotification(protocol.MethodGatewayEvent, map[string]string{"x": "y"})) + }) +} diff --git a/internal/gateway/network_server.go b/internal/gateway/network_server.go index da3cdf6b..891fadd8 100644 --- a/internal/gateway/network_server.go +++ b/internal/gateway/network_server.go @@ -54,6 +54,7 @@ type NetworkServerOptions struct { HeartbeatInterval time.Duration MaxRequestBytes int64 MaxStreamConnections int + Relay *StreamRelay listenFn func(network, address string) (net.Listener, error) } @@ -68,6 +69,7 @@ type NetworkServer struct { maxRequestBytes int64 maxStreamConnections int listenFn func(network, address string) (net.Listener, error) + relay *StreamRelay mu sync.Mutex server *http.Server @@ -124,6 +126,13 @@ func NewNetworkServer(options NetworkServerOptions) (*NetworkServer, error) { maxStreamConnections = DefaultNetworkMaxStreamConnections } + relay := options.Relay + if relay == nil { + relay = NewStreamRelay(StreamRelayOptions{ + Logger: logger, + }) + } + return &NetworkServer{ listenAddress: listenAddress, logger: logger, @@ -134,6 +143,7 @@ func NewNetworkServer(options NetworkServerOptions) (*NetworkServer, error) { maxRequestBytes: maxRequestBytes, maxStreamConnections: maxStreamConnections, listenFn: listenFn, + relay: relay, wsConns: make(map[*websocket.Conn]context.CancelFunc), sseCancels: make(map[int]context.CancelFunc), }, nil @@ -190,6 +200,10 @@ func (s *NetworkServer) ListenAddress() string { // Serve 启动网络访问面服务，并注册 HTTP/WebSocket/SSE 三类入口。 func (s *NetworkServer) Serve(ctx context.Context, runtimePort RuntimePort) error { + if s.relay == nil { + s.relay = NewStreamRelay(StreamRelayOptions{Logger: s.logger}) + } + listener, err := s.listenFn("tcp", s.listenAddress) if err != nil { return fmt.Errorf("gateway network listen failed: %w", err) @@ -212,6 +226,7 @@ func (s *NetworkServer) Serve(ctx context.Context, runtimePort RuntimePort) erro s.listenAddress = listener.Addr().String() s.mu.Unlock() + s.relay.Start(ctx, runtimePort) s.logger.Printf("network listening on %s", listener.Addr().String()) go func() { @@ -233,10 +248,15 @@ func (s *NetworkServer) Close(ctx context.Context) error { s.mu.Lock() httpServer := s.server listener := s.listener + relay := s.relay s.server = nil s.listener = nil s.mu.Unlock() + if relay != nil { + relay.Stop() + } + if httpServer == nil && listener == nil { return nil } @@ -344,6 +364,13 @@ func (s *NetworkServer) handleWebSocket(conn *websocket.Conn, runtimePort Runtim parentContext = request.Context() } connectionContext, cancelConnection := context.WithCancel(parentContext) + relay := s.relay + if relay == nil { + relay = NewStreamRelay(StreamRelayOptions{Logger: s.logger}) + } + connectionID := NewConnectionID() + connectionContext = WithConnectionID(connectionContext, connectionID) + connectionContext = WithStreamRelay(connectionContext, relay) if !s.registerWSConnection(conn, cancelConnection) { cancelConnection() @@ -352,9 +379,43 @@ func (s *NetworkServer) handleWebSocket(conn *websocket.Conn, runtimePort Runtim _ = conn.Close() return } + + registerErr := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelWS, + Context: connectionContext, + Cancel: cancelConnection, + Write: func(message RelayMessage) error { + if message.Kind != relayMessageKindJSON { + return fmt.Errorf("websocket connection only supports json payload") + } + if s.writeTimeout > 0 { + if err := conn.SetWriteDeadline(time.Now().Add(s.writeTimeout)); err != nil { + return err + } + } + rawPayload, err := json.Marshal(message.Payload) + if err != nil { + return err + } + return websocket.Message.Send(conn, string(rawPayload)) + }, + Close: func() { + _ = conn.Close() + }, + }) + if registerErr != nil { + cancelConnection() + s.unregisterWSConnection(conn) + _ = conn.Close() + s.logger.Printf("register websocket connection failed: %v", registerErr) + return + } + defer func() { cancelConnection() s.unregisterWSConnection(conn) + relay.dropConnection(connectionID) _ = conn.Close() }() @@ -363,10 +424,9 @@ func (s *NetworkServer) handleWebSocket(conn *websocket.Conn, runtimePort Runtim conn.MaxPayloadBytes = maxPayloadBytes } - var writeMu sync.Mutex stopHeartbeat := make(chan struct{}) defer close(stopHeartbeat) - go s.runWSHeartbeatLoop(conn, &writeMu, stopHeartbeat) + go s.runWSHeartbeatLoop(relay, connectionID, stopHeartbeat) for { // 注意：此处不再强制上行读超时，避免单向推送场景下误杀健康连接。 @@ -388,18 +448,15 @@ func (s *NetworkServer) handleWebSocket(conn *websocket.Conn, runtimePort Runtim rpcResponse = dispatchRPCRequestFn(connectionContext, rpcRequest, runtimePort) } - if err := s.writeWebSocketMessage(conn, &writeMu, rpcResponse); err != nil { + if !relay.SendJSONRPCResponse(connectionID, rpcResponse) { cancelConnection() - if !isConnectionClosedError(err) { - s.logger.Printf("websocket write failed: %v", err) - } return } } } // runWSHeartbeatLoop 周期性推送 WebSocket 心跳帧，保证长连接可观测与保活。 -func (s *NetworkServer) runWSHeartbeatLoop(conn *websocket.Conn, writeMu *sync.Mutex, stop <-chan struct{}) { +func (s *NetworkServer) runWSHeartbeatLoop(relay *StreamRelay, connectionID ConnectionID, stop <-chan struct{}) { ticker := time.NewTicker(s.heartbeatInterval) defer ticker.Stop() @@ -408,33 +465,16 @@ func (s *NetworkServer) runWSHeartbeatLoop(conn *websocket.Conn, writeMu *sync.M case <-stop: return case <-ticker.C: - if err := s.writeWebSocketMessage(conn, writeMu, map[string]any{ + if !relay.SendJSONRPCPayload(connectionID, map[string]any{ "type": "heartbeat", "timestamp": time.Now().UTC().Format(time.RFC3339Nano), - }); err != nil { + }) { return } } } } -// writeWebSocketMessage 将任意结构序列化后写入 WS 连接，并施加写超时约束。 -func (s *NetworkServer) writeWebSocketMessage(conn *websocket.Conn, writeMu *sync.Mutex, payload any) error { - if s.writeTimeout > 0 { - if err := conn.SetWriteDeadline(time.Now().Add(s.writeTimeout)); err != nil { - return err - } - } - rawPayload, err := json.Marshal(payload) - if err != nil { - return err - } - - writeMu.Lock() - defer writeMu.Unlock() - return websocket.Message.Send(conn, string(rawPayload)) -} - // handleSSERequest 处理 SSE 入口请求，先返回一次结果事件，再持续发送心跳事件。 func (s *NetworkServer) handleSSERequest(writer http.ResponseWriter, request *http.Request, runtimePort RuntimePort) { if request.Method != http.MethodGet { @@ -455,11 +495,62 @@ func (s *NetworkServer) handleSSERequest(writer http.ResponseWriter, request *ht http.Error(writer, "stream connection limit exceeded", http.StatusServiceUnavailable) return } + sseMessageCh := make(chan RelayMessage, DefaultStreamQueueSize) + + relay := s.relay + if relay == nil { + relay = NewStreamRelay(StreamRelayOptions{Logger: s.logger}) + } + streamConnectionID := NewConnectionID() + streamCtx = WithConnectionID(streamCtx, streamConnectionID) + streamCtx = WithStreamRelay(streamCtx, relay) + + registerErr := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: streamConnectionID, + Channel: StreamChannelSSE, + Context: streamCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { + if message.Kind != relayMessageKindSSE { + return fmt.Errorf("sse connection only supports sse events") + } + select { + case <-streamCtx.Done(): + return context.Canceled + case sseMessageCh <- message: + return nil + } + }, + Close: func() {}, + }) + if registerErr != nil { + cancel() + s.unregisterSSEConnection(connectionID) + http.Error(writer, "failed to register stream connection", http.StatusInternalServerError) + return + } + defer func() { cancel() s.unregisterSSEConnection(connectionID) + relay.dropConnection(streamConnectionID) }() + queryValues := request.URL.Query() + sessionID := strings.TrimSpace(queryValues.Get("session_id")) + if sessionID != "" { + runID := strings.TrimSpace(queryValues.Get("run_id")) + if bindErr := relay.BindConnection(streamConnectionID, StreamBinding{ + SessionID: sessionID, + RunID: runID, + Channel: StreamChannelSSE, + Explicit: true, + }); bindErr != nil { + http.Error(writer, bindErr.Message, http.StatusBadRequest) + return + } + } + writer.Header().Set("Content-Type", "text/event-stream") writer.Header().Set("Cache-Control", "no-cache") writer.Header().Set("Connection", "keep-alive") @@ -467,7 +558,7 @@ func (s *NetworkServer) handleSSERequest(writer http.ResponseWriter, request *ht rpcRequest := buildSSETriggerRequest(request) rpcResponse := dispatchRPCRequestFn(streamCtx, rpcRequest, runtimePort) - if err := s.writeSSEEvent(writer, flusher, "result", rpcResponse); err != nil { + if !relay.SendSSEEvent(streamConnectionID, "result", rpcResponse) { return } @@ -479,9 +570,16 @@ func (s *NetworkServer) handleSSERequest(writer http.ResponseWriter, request *ht case <-streamCtx.Done(): return case <-ticker.C: - if err := s.writeSSEEvent(writer, flusher, "heartbeat", map[string]string{ + if !relay.SendSSEEvent(streamConnectionID, "heartbeat", map[string]string{ "timestamp": time.Now().UTC().Format(time.RFC3339Nano), - }); err != nil { + }) { + return + } + case message := <-sseMessageCh: + if strings.TrimSpace(message.Event) == "" { + return + } + if err := s.writeSSEEvent(writer, flusher, message.Event, message.Payload); err != nil { return } } diff --git a/internal/gateway/network_server_additional_test.go b/internal/gateway/network_server_additional_test.go index c1bf9f8d..dc6c60d1 100644 --- a/internal/gateway/network_server_additional_test.go +++ b/internal/gateway/network_server_additional_test.go @@ -126,6 +126,12 @@ func TestNetworkServerServeErrorBranches(t *testing.T) { if serveErr := server.Serve(context.Background(), nil); serveErr == nil || !strings.Contains(serveErr.Error(), "listen failed") { t.Fatalf("expected listen error, got %v", serveErr) } + server.relay.mu.RLock() + started := server.relay.cleanupStarted || server.relay.eventPumpStarted + server.relay.mu.RUnlock() + if started { + t.Fatal("relay loops should not start when listen failed") + } }) t.Run("already serving", func(t *testing.T) { @@ -171,6 +177,25 @@ func TestNetworkServerServeErrorBranches(t *testing.T) { }) } +func TestNetworkServerCloseStopsRelayWithoutActiveServer(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{ + CleanupInterval: 5 * time.Millisecond, + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + relay.Start(ctx, nil) + waitForStreamRelayState(t, relay, true) + + server := &NetworkServer{ + relay: relay, + } + if err := server.Close(context.Background()); err != nil { + t.Fatalf("close server: %v", err) + } + waitForStreamRelayState(t, relay, false) +} + func TestNetworkServerIsClosedState(t *testing.T) { server := &NetworkServer{} if !server.isClosed() { diff --git a/internal/gateway/network_server_test.go b/internal/gateway/network_server_test.go index db7ae2b6..6cde321f 100644 --- a/internal/gateway/network_server_test.go +++ b/internal/gateway/network_server_test.go @@ -574,6 +574,81 @@ func TestNetworkServerCloseInterruptsStreams(t *testing.T) { } } +func TestNetworkServerStreamsReceiveGatewayEventNotification(t *testing.T) { + eventCh := make(chan RuntimeEvent, 2) + runtimePort := &runtimePortEventStub{events: eventCh} + + server := newTestNetworkServer(t, NetworkServerOptions{}) + testContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := make(chan error, 1) + go func() { + serveDone <- server.Serve(testContext, runtimePort) + }() + t.Cleanup(func() { + _ = server.Close(context.Background()) + select { + case <-serveDone: + case <-time.After(2 * time.Second): + t.Fatal("network serve goroutine did not exit") + } + }) + + listenAddress := waitForNetworkAddress(t, server) + wsConn, err := websocket.Dial("ws://"+listenAddress+"/ws", "", "http://localhost:3000") + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + t.Cleanup(func() { _ = wsConn.Close() }) + + if err := websocket.Message.Send(wsConn, `{"jsonrpc":"2.0","id":"bind-ws-1","method":"gateway.bindStream","params":{"session_id":"session-relay","run_id":"run-relay","channel":"ws"}}`); err != nil { + t.Fatalf("send bindStream request: %v", err) + } + bindAck := receiveWSAckFrame(t, wsConn) + if bindAck.Action != FrameActionBindStream { + t.Fatalf("bind action = %q, want %q", bindAck.Action, FrameActionBindStream) + } + + sseRequest, err := http.NewRequest( + http.MethodGet, + "http://"+listenAddress+"/sse?method=gateway.ping&id=sse-relay-1&session_id=session-relay&run_id=run-relay", + nil, + ) + if err != nil { + t.Fatalf("new sse request: %v", err) + } + sseRequest.Header.Set("Origin", "http://localhost:3000") + sseResponse, err := http.DefaultClient.Do(sseRequest) + if err != nil { + t.Fatalf("open sse stream: %v", err) + } + t.Cleanup(func() { _ = sseResponse.Body.Close() }) + if sseResponse.StatusCode != http.StatusOK { + t.Fatalf("sse status = %d, want %d", sseResponse.StatusCode, http.StatusOK) + } + _ = readSSEResultFrame(t, sseResponse.Body) + + eventCh <- RuntimeEvent{ + Type: RuntimeEventTypeRunProgress, + SessionID: "session-relay", + RunID: "run-relay", + Payload: map[string]string{ + "chunk": "hello", + }, + } + + wsEvent := receiveWSGatewayEventNotification(t, wsConn) + if wsEvent.SessionID != "session-relay" || wsEvent.RunID != "run-relay" { + t.Fatalf("ws event frame mismatch: %#v", wsEvent) + } + + sseEvent := readSSEGatewayEventFrame(t, sseResponse.Body) + if sseEvent.SessionID != "session-relay" || sseEvent.RunID != "run-relay" { + t.Fatalf("sse event frame mismatch: %#v", sseEvent) + } +} + // newTestNetworkServer 创建默认测试网络服务实例，统一收敛测试参数。 func newTestNetworkServer(t *testing.T, overrides NetworkServerOptions) *NetworkServer { t.Helper() @@ -694,6 +769,87 @@ func readSSEResultFrame(t *testing.T, body io.Reader) MessageFrame { } } +// receiveWSGatewayEventNotification 读取 WS 消息直到拿到 gateway.event 通知并返回其内层事件帧。 +func receiveWSGatewayEventNotification(t *testing.T, wsConn *websocket.Conn) MessageFrame { + t.Helper() + _ = wsConn.SetReadDeadline(time.Now().Add(2 * time.Second)) + for attempt := 0; attempt < 20; attempt++ { + var rawResponse string + if err := websocket.Message.Receive(wsConn, &rawResponse); err != nil { + t.Fatalf("receive websocket message: %v", err) + } + + var envelope struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + Params json.RawMessage `json:"params"` + } + if err := json.Unmarshal([]byte(rawResponse), &envelope); err != nil { + continue + } + if envelope.Method != protocol.MethodGatewayEvent { + continue + } + var eventFrame MessageFrame + if err := json.Unmarshal(envelope.Params, &eventFrame); err != nil { + t.Fatalf("decode gateway.event params: %v", err) + } + return eventFrame + } + t.Fatal("did not receive websocket gateway.event notification") + return MessageFrame{} +} + +// readSSEGatewayEventFrame 读取 SSE 流直到捕获 gateway.event 事件并解析其内层事件帧。 +func readSSEGatewayEventFrame(t *testing.T, body io.Reader) MessageFrame { + t.Helper() + reader := bufio.NewReader(body) + timeout := time.After(3 * time.Second) + currentEvent := "" + + for { + select { + case <-timeout: + t.Fatal("timed out waiting for sse gateway.event") + default: + line, err := reader.ReadString('\n') + if err != nil { + t.Fatalf("read sse line: %v", err) + } + trimmed := strings.TrimSpace(line) + if trimmed == "" { + continue + } + if strings.HasPrefix(trimmed, "event:") { + currentEvent = strings.TrimSpace(strings.TrimPrefix(trimmed, "event:")) + continue + } + if currentEvent != protocol.MethodGatewayEvent || !strings.HasPrefix(trimmed, "data:") { + continue + } + + rawData := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + var notification struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + Params json.RawMessage `json:"params"` + } + if err := json.Unmarshal([]byte(rawData), ¬ification); err != nil { + t.Fatalf("decode sse gateway.event notification: %v", err) + } + if notification.Method != protocol.MethodGatewayEvent { + continue + } + + var eventFrame MessageFrame + if err := json.Unmarshal(notification.Params, &eventFrame); err != nil { + t.Fatalf("decode sse gateway.event params: %v", err) + } + return eventFrame + } + } +} + type noFlushResponseWriter struct { header http.Header status int diff --git a/internal/gateway/protocol/jsonrpc.go b/internal/gateway/protocol/jsonrpc.go index eb6ad305..45207324 100644 --- a/internal/gateway/protocol/jsonrpc.go +++ b/internal/gateway/protocol/jsonrpc.go @@ -14,6 +14,10 @@ const ( const ( // MethodGatewayPing 表示网关探活方法。 MethodGatewayPing = "gateway.ping" + // MethodGatewayBindStream 表示客户端向网关声明流式订阅绑定的方法。 + MethodGatewayBindStream = "gateway.bindStream" + // MethodGatewayEvent 表示网关向客户端推送运行时事件的通知方法。 + MethodGatewayEvent = "gateway.event" // MethodWakeOpenURL 表示 URL Scheme 唤醒方法。 MethodWakeOpenURL = "wake.openUrl" ) @@ -64,6 +68,13 @@ type JSONRPCResponse struct { Error *JSONRPCError `json:"error,omitempty"` } +// JSONRPCNotification 表示控制面向客户端主动推送的 JSON-RPC 通知。 +type JSONRPCNotification struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + Params any `json:"params,omitempty"` +} + // JSONRPCError 表示 JSON-RPC 错误负载。 type JSONRPCError struct { Code int `json:"code"` @@ -82,10 +93,18 @@ type NormalizedRequest struct { RequestID string Action string SessionID string + RunID string Workdir string Payload any } +// BindStreamParams 表示 gateway.bindStream 的标准化参数载荷。 +type BindStreamParams struct { + SessionID string `json:"session_id"` + RunID string `json:"run_id,omitempty"` + Channel string `json:"channel,omitempty"` +} + // NormalizeJSONRPCRequest 将 JSON-RPC 请求归一化为内部请求模型，并做方法级参数解析。 func NormalizeJSONRPCRequest(request JSONRPCRequest) (NormalizedRequest, *JSONRPCError) { normalized := NormalizedRequest{} @@ -118,6 +137,16 @@ func NormalizeJSONRPCRequest(request JSONRPCRequest) (NormalizedRequest, *JSONRP case MethodGatewayPing: normalized.Action = "ping" return normalized, nil + case MethodGatewayBindStream: + params, parseErr := decodeBindStreamParams(request.Params) + if parseErr != nil { + return normalized, parseErr + } + normalized.Action = "bind_stream" + normalized.SessionID = params.SessionID + normalized.RunID = params.RunID + normalized.Payload = params + return normalized, nil case MethodWakeOpenURL: intent, parseErr := decodeWakeIntentParams(request.Params) if parseErr != nil { @@ -164,6 +193,15 @@ func NewJSONRPCErrorResponse(id json.RawMessage, rpcError *JSONRPCError) JSONRPC } } +// NewJSONRPCNotification 创建 JSON-RPC 通知负载，供网关向客户端推送事件使用。 +func NewJSONRPCNotification(method string, params any) JSONRPCNotification { + return JSONRPCNotification{ + JSONRPC: JSONRPCVersion, + Method: strings.TrimSpace(method), + Params: params, + } +} + // NewJSONRPCError 创建带 gateway_code 的 JSON-RPC 错误对象。 func NewJSONRPCError(code int, message, gatewayCode string) *JSONRPCError { errorPayload := &JSONRPCError{ @@ -280,6 +318,54 @@ func decodeWakeIntentParams(raw json.RawMessage) (WakeIntent, *JSONRPCError) { return intent, nil } +// decodeBindStreamParams 对 gateway.bindStream 的 params 执行反序列化与最小参数校验。 +func decodeBindStreamParams(raw json.RawMessage) (BindStreamParams, *JSONRPCError) { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + return BindStreamParams{}, NewJSONRPCError( + JSONRPCCodeInvalidParams, + "missing required field: params", + GatewayCodeMissingRequiredField, + ) + } + + var params BindStreamParams + if err := json.Unmarshal(trimmed, ¶ms); err != nil { + return BindStreamParams{}, NewJSONRPCError( + JSONRPCCodeInvalidParams, + "invalid params for gateway.bindStream", + GatewayCodeInvalidFrame, + ) + } + + params.SessionID = strings.TrimSpace(params.SessionID) + params.RunID = strings.TrimSpace(params.RunID) + params.Channel = strings.ToLower(strings.TrimSpace(params.Channel)) + if params.Channel == "" { + params.Channel = "all" + } + + if params.SessionID == "" { + return BindStreamParams{}, NewJSONRPCError( + JSONRPCCodeInvalidParams, + "missing required field: params.session_id", + GatewayCodeMissingRequiredField, + ) + } + + switch params.Channel { + case "all", "ipc", "ws", "sse": + default: + return BindStreamParams{}, NewJSONRPCError( + JSONRPCCodeInvalidParams, + "invalid field: params.channel", + GatewayCodeInvalidAction, + ) + } + + return params, nil +} + // cloneJSONRawMessage 复制 RawMessage，避免共享底层切片导致的并发风险。 func cloneJSONRawMessage(raw json.RawMessage) json.RawMessage { if len(raw) == 0 { diff --git a/internal/gateway/protocol/jsonrpc_test.go b/internal/gateway/protocol/jsonrpc_test.go index e0cccc1f..7a7f5fad 100644 --- a/internal/gateway/protocol/jsonrpc_test.go +++ b/internal/gateway/protocol/jsonrpc_test.go @@ -71,6 +71,38 @@ func TestNormalizeJSONRPCRequestWakeOpenURL(t *testing.T) { } } +func TestNormalizeJSONRPCRequestBindStream(t *testing.T) { + normalized, rpcErr := NormalizeJSONRPCRequest(JSONRPCRequest{ + JSONRPC: JSONRPCVersion, + ID: json.RawMessage(`"bind-1"`), + Method: MethodGatewayBindStream, + Params: json.RawMessage(`{ + "session_id":"session-1", + "run_id":"run-1", + "channel":"ws" + }`), + }) + if rpcErr != nil { + t.Fatalf("normalize bindStream request: %v", rpcErr) + } + if normalized.Action != "bind_stream" { + t.Fatalf("action = %q, want %q", normalized.Action, "bind_stream") + } + if normalized.SessionID != "session-1" { + t.Fatalf("session_id = %q, want %q", normalized.SessionID, "session-1") + } + if normalized.RunID != "run-1" { + t.Fatalf("run_id = %q, want %q", normalized.RunID, "run-1") + } + params, ok := normalized.Payload.(BindStreamParams) + if !ok { + t.Fatalf("payload type = %T, want BindStreamParams", normalized.Payload) + } + if params.Channel != "ws" { + t.Fatalf("channel = %q, want %q", params.Channel, "ws") + } +} + func TestNormalizeJSONRPCRequestErrors(t *testing.T) { testCases := []struct { name string @@ -167,6 +199,38 @@ func TestNormalizeJSONRPCRequestErrors(t *testing.T) { wantCode: JSONRPCCodeInvalidParams, wantGatewayCode: GatewayCodeInvalidFrame, }, + { + name: "bindStream missing params", + request: JSONRPCRequest{ + JSONRPC: JSONRPCVersion, + ID: json.RawMessage(`"x"`), + Method: MethodGatewayBindStream, + }, + wantCode: JSONRPCCodeInvalidParams, + wantGatewayCode: GatewayCodeMissingRequiredField, + }, + { + name: "bindStream missing session", + request: JSONRPCRequest{ + JSONRPC: JSONRPCVersion, + ID: json.RawMessage(`"x"`), + Method: MethodGatewayBindStream, + Params: json.RawMessage(`{"channel":"all"}`), + }, + wantCode: JSONRPCCodeInvalidParams, + wantGatewayCode: GatewayCodeMissingRequiredField, + }, + { + name: "bindStream invalid channel", + request: JSONRPCRequest{ + JSONRPC: JSONRPCVersion, + ID: json.RawMessage(`"x"`), + Method: MethodGatewayBindStream, + Params: json.RawMessage(`{"session_id":"s-1","channel":"tcp"}`), + }, + wantCode: JSONRPCCodeInvalidParams, + wantGatewayCode: GatewayCodeInvalidAction, + }, } for _, tc := range testCases { @@ -254,6 +318,14 @@ func TestJSONRPCHelpers(t *testing.T) { if MapGatewayCodeToJSONRPCCode("unknown") != JSONRPCCodeInternalError { t.Fatal("unknown code should map to internal_error") } + + notification := NewJSONRPCNotification(MethodGatewayEvent, map[string]any{"message": "ok"}) + if notification.JSONRPC != JSONRPCVersion { + t.Fatalf("notification jsonrpc = %q, want %q", notification.JSONRPC, JSONRPCVersion) + } + if notification.Method != MethodGatewayEvent { + t.Fatalf("notification method = %q, want %q", notification.Method, MethodGatewayEvent) + } } func TestNewJSONRPCErrorResponseWithNilIDEncodesNull(t *testing.T) { diff --git a/internal/gateway/rpc_dispatch.go b/internal/gateway/rpc_dispatch.go index 612e49b7..20ad1f00 100644 --- a/internal/gateway/rpc_dispatch.go +++ b/internal/gateway/rpc_dispatch.go @@ -2,6 +2,7 @@ package gateway import ( "context" + "strings" "neo-code/internal/gateway/protocol" ) @@ -18,10 +19,24 @@ func dispatchRPCRequest(ctx context.Context, request protocol.JSONRPCRequest, ru Action: FrameAction(normalized.Action), RequestID: normalized.RequestID, SessionID: normalized.SessionID, + RunID: normalized.RunID, Workdir: normalized.Workdir, Payload: normalized.Payload, } + frame = hydrateFrameSessionFromConnection(ctx, frame) + if requiresSession(frame.Action) && strings.TrimSpace(frame.SessionID) == "" { + return protocol.NewJSONRPCErrorResponse( + normalized.ID, + protocol.NewJSONRPCError( + protocol.JSONRPCCodeInvalidParams, + "missing required field: session_id", + protocol.GatewayCodeMissingRequiredField, + ), + ) + } + applyAutomaticBinding(ctx, frame) + responseFrame := dispatchFrame(ctx, frame, runtimePort) if responseFrame.Type != FrameTypeError { rpcResponse, encodeErr := protocol.NewJSONRPCResultResponse(normalized.ID, responseFrame) @@ -57,3 +72,55 @@ func dispatchFrame(ctx context.Context, frame MessageFrame, runtimePort RuntimeP return dispatchRequestFrame(ctx, frame, runtimePort) } + +// hydrateFrameSessionFromConnection 根据统一优先级为请求帧补齐 session_id：显式字段 > payload 参数 > 连接绑定兜底。 +func hydrateFrameSessionFromConnection(ctx context.Context, frame MessageFrame) MessageFrame { + if strings.TrimSpace(frame.SessionID) != "" { + return frame + } + + payloadSessionID := strings.TrimSpace(extractSessionIDFromPayload(frame.Payload)) + if payloadSessionID != "" { + frame.SessionID = payloadSessionID + return frame + } + + relay, relayExists := StreamRelayFromContext(ctx) + connectionID, connectionExists := ConnectionIDFromContext(ctx) + if !relayExists || !connectionExists { + return frame + } + + frame.SessionID = strings.TrimSpace(relay.ResolveFallbackSessionID(connectionID)) + return frame +} + +// requiresSession 判断指定动作在分发阶段是否必须携带 session_id。 +func requiresSession(action FrameAction) bool { + switch action { + case FrameActionBindStream, FrameActionRun, FrameActionCompact, FrameActionLoadSession, FrameActionResolvePermission: + return true + default: + return false + } +} + +// applyAutomaticBinding 在请求分发前执行自动续绑与 ping 续期逻辑。 +func applyAutomaticBinding(ctx context.Context, frame MessageFrame) { + relay, relayExists := StreamRelayFromContext(ctx) + connectionID, connectionExists := ConnectionIDFromContext(ctx) + if !relayExists || !connectionExists { + return + } + + if frame.Action == FrameActionPing { + relay.RefreshConnectionBindings(connectionID) + return + } + + if frame.Action == FrameActionBindStream { + return + } + + relay.AutoBindFromFrame(connectionID, frame) +} diff --git a/internal/gateway/rpc_dispatch_test.go b/internal/gateway/rpc_dispatch_test.go index 1b296e36..8b8d64da 100644 --- a/internal/gateway/rpc_dispatch_test.go +++ b/internal/gateway/rpc_dispatch_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "testing" + "time" "neo-code/internal/gateway/protocol" ) @@ -11,7 +12,7 @@ import ( func TestDispatchRPCRequestResultEncodeError(t *testing.T) { originalHandlers := requestFrameHandlers requestFrameHandlers = map[FrameAction]requestFrameHandler{ - FrameActionPing: func(frame MessageFrame) MessageFrame { + FrameActionPing: func(_ context.Context, frame MessageFrame) MessageFrame { return MessageFrame{ Type: FrameTypeAck, Action: FrameActionPing, @@ -43,6 +44,90 @@ func TestDispatchRPCRequestResultEncodeError(t *testing.T) { } } +func TestHydrateFrameSessionFromConnectionFallback(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + baseContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := NewConnectionID() + connectionContext := WithConnectionID(baseContext, connectionID) + connectionContext = WithStreamRelay(connectionContext, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionContext, + Cancel: cancel, + Write: func(message RelayMessage) error { + _ = message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + defer relay.dropConnection(connectionID) + + if bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: "session-fallback", + Channel: StreamChannelAll, + Explicit: true, + }); bindErr != nil { + t.Fatalf("bind connection: %v", bindErr) + } + + hydrated := hydrateFrameSessionFromConnection(connectionContext, MessageFrame{ + Type: FrameTypeRequest, + Action: FrameActionPing, + }) + if hydrated.SessionID != "session-fallback" { + t.Fatalf("session_id = %q, want %q", hydrated.SessionID, "session-fallback") + } +} + +func TestApplyAutomaticBindingPingRefreshesTTL(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{ + BindingTTL: 20 * time.Millisecond, + }) + baseContext, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := NewConnectionID() + connectionContext := WithConnectionID(baseContext, connectionID) + connectionContext = WithStreamRelay(connectionContext, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionContext, + Cancel: cancel, + Write: func(message RelayMessage) error { + _ = message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + defer relay.dropConnection(connectionID) + + if bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: "session-ping", + Channel: StreamChannelAll, + Explicit: true, + }); bindErr != nil { + t.Fatalf("bind connection: %v", bindErr) + } + + time.Sleep(10 * time.Millisecond) + applyAutomaticBinding(connectionContext, MessageFrame{ + Type: FrameTypeRequest, + Action: FrameActionPing, + }) + time.Sleep(15 * time.Millisecond) + if !relay.RefreshConnectionBindings(connectionID) { + t.Fatal("expected ping to refresh existing bindings") + } +} + func TestDispatchFrameValidationBranches(t *testing.T) { response := dispatchFrame(context.Background(), MessageFrame{ Type: FrameType("invalid"), diff --git a/internal/gateway/server.go b/internal/gateway/server.go index d6a6360a..d507130c 100644 --- a/internal/gateway/server.go +++ b/internal/gateway/server.go @@ -44,6 +44,7 @@ type ServerOptions struct { MaxConnections int ReadTimeout time.Duration WriteTimeout time.Duration + Relay *StreamRelay listenFn func(address string) (net.Listener, error) } @@ -55,6 +56,7 @@ type Server struct { maxConnections int readTimeout time.Duration writeTimeout time.Duration + relay *StreamRelay mu sync.Mutex listener net.Listener @@ -102,6 +104,13 @@ func NewServer(options ServerOptions) (*Server, error) { writeTimeout = DefaultWriteTimeout } + relay := options.Relay + if relay == nil { + relay = NewStreamRelay(StreamRelayOptions{ + Logger: logger, + }) + } + return &Server{ listenAddress: listenAddress, logger: logger, @@ -109,6 +118,7 @@ func NewServer(options ServerOptions) (*Server, error) { maxConnections: maxConnections, readTimeout: readTimeout, writeTimeout: writeTimeout, + relay: relay, conns: make(map[net.Conn]struct{}), }, nil } @@ -135,6 +145,10 @@ func (s *Server) Serve(ctx context.Context, runtimePort RuntimePort) error { s.mu.Unlock() s.logger.Printf("listening on %s", s.listenAddress) + if s.relay == nil { + s.relay = NewStreamRelay(StreamRelayOptions{Logger: s.logger}) + } + s.relay.Start(ctx, runtimePort) go func() { <-ctx.Done() @@ -176,6 +190,10 @@ func (s *Server) Close(ctx context.Context) error { s.listener = nil s.mu.Unlock() + if s.relay != nil { + s.relay.Stop() + } + var closeErr error if listener != nil { closeErr = listener.Close() @@ -248,11 +266,47 @@ func (s *Server) handleConnection(ctx context.Context, conn net.Conn, runtimePor }() reader := bufio.NewReader(conn) + + connectionContext, cancelConnection := context.WithCancel(ctx) + defer cancelConnection() + + relay := s.relay + if relay == nil { + relay = NewStreamRelay(StreamRelayOptions{Logger: s.logger}) + } + + connectionID := NewConnectionID() + connectionContext = WithConnectionID(connectionContext, connectionID) + connectionContext = WithStreamRelay(connectionContext, relay) + encoder := json.NewEncoder(conn) + registerErr := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionContext, + Cancel: cancelConnection, + Write: func(message RelayMessage) error { + if message.Kind != relayMessageKindJSON { + return fmt.Errorf("ipc connection only supports json messages") + } + if err := s.applyWriteDeadline(conn); err != nil { + return err + } + return encoder.Encode(message.Payload) + }, + Close: func() { + _ = conn.Close() + }, + }) + if registerErr != nil { + s.logger.Printf("register ipc connection failed: %v", registerErr) + return + } + defer relay.dropConnection(connectionID) for { select { - case <-ctx.Done(): + case <-connectionContext.Done(): return default: } @@ -276,7 +330,7 @@ func (s *Server) handleConnection(ctx context.Context, conn net.Conn, runtimePor } if errors.Is(err, errFrameTooLarge) { s.logger.Printf("decode frame failed: %v", err) - _ = s.writeRPCResponse(conn, encoder, protocol.NewJSONRPCErrorResponse( + _ = relay.SendJSONRPCResponseSync(connectionID, protocol.NewJSONRPCErrorResponse( nil, protocol.NewJSONRPCError( protocol.JSONRPCCodeInvalidRequest, @@ -288,7 +342,7 @@ func (s *Server) handleConnection(ctx context.Context, conn net.Conn, runtimePor } s.logger.Printf("decode frame failed: %v", err) - _ = s.writeRPCResponse(conn, encoder, protocol.NewJSONRPCErrorResponse( + _ = relay.SendJSONRPCResponseSync(connectionID, protocol.NewJSONRPCErrorResponse( nil, protocol.NewJSONRPCError( protocol.JSONRPCCodeParseError, @@ -299,8 +353,8 @@ func (s *Server) handleConnection(ctx context.Context, conn net.Conn, runtimePor return } - rpcResponse := s.dispatchRPCRequest(ctx, rpcRequest, runtimePort) - if !s.writeRPCResponse(conn, encoder, rpcResponse) { + rpcResponse := s.dispatchRPCRequest(connectionContext, rpcRequest, runtimePort) + if !relay.SendJSONRPCResponse(connectionID, rpcResponse) { return } } @@ -322,19 +376,6 @@ func (s *Server) applyWriteDeadline(conn net.Conn) error { return conn.SetWriteDeadline(time.Now().Add(s.writeTimeout)) } -// writeRPCResponse 统一处理 JSON-RPC 响应写回及写超时设置，失败时返回 false 供上层快速终止连接循环。 -func (s *Server) writeRPCResponse(conn net.Conn, encoder *json.Encoder, response protocol.JSONRPCResponse) bool { - if err := s.applyWriteDeadline(conn); err != nil { - s.logger.Printf("set write deadline failed: %v", err) - return false - } - if err := encoder.Encode(response); err != nil { - s.logger.Printf("write frame failed: %v", err) - return false - } - return true -} - // isTimeoutError 判断错误是否为网络超时，用于区分慢连接超时与协议错误。 func isTimeoutError(err error) bool { var netErr net.Error diff --git a/internal/gateway/server_additional_test.go b/internal/gateway/server_additional_test.go index 52466de9..5c5e88e5 100644 --- a/internal/gateway/server_additional_test.go +++ b/internal/gateway/server_additional_test.go @@ -238,6 +238,26 @@ func TestCloseReturnsContextErrorWhenWaitCanceled(t *testing.T) { server.wg.Done() } +func TestCloseStopsRelayBackgroundLoops(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{ + CleanupInterval: 5 * time.Millisecond, + }) + relay.Start(context.Background(), nil) + waitForStreamRelayState(t, relay, true) + + server := &Server{ + relay: relay, + conns: make(map[net.Conn]struct{}), + } + + closeCtx, cancelClose := context.WithTimeout(context.Background(), time.Second) + defer cancelClose() + if err := server.Close(closeCtx); err != nil { + t.Fatalf("close server: %v", err) + } + waitForStreamRelayState(t, relay, false) +} + func TestDecodeRPCRequestTrailingJSON(t *testing.T) { reader := bufio.NewReader(strings.NewReader(`{"jsonrpc":"2.0","id":"x","method":"gateway.ping"} {"extra":1}` + "\n")) _, err := decodeRPCRequest(reader) @@ -338,7 +358,7 @@ func TestDispatchRPCRequestConvertsFrameErrorWithoutPayload(t *testing.T) { server := &Server{} originalHandlers := requestFrameHandlers requestFrameHandlers = map[FrameAction]requestFrameHandler{ - FrameActionPing: func(frame MessageFrame) MessageFrame { + FrameActionPing: func(_ context.Context, frame MessageFrame) MessageFrame { return MessageFrame{ Type: FrameTypeError, Action: frame.Action, diff --git a/internal/gateway/server_test.go b/internal/gateway/server_test.go index 69c3b154..dd9261e0 100644 --- a/internal/gateway/server_test.go +++ b/internal/gateway/server_test.go @@ -199,6 +199,119 @@ func TestServerHandleConnectionRejectsOversizedFrame(t *testing.T) { } } +func TestServerHandleConnectionRelaysRuntimeEventAfterBindStream(t *testing.T) { + t.Parallel() + + eventCh := make(chan RuntimeEvent, 1) + relay := NewStreamRelay(StreamRelayOptions{ + Logger: log.New(io.Discard, "", 0), + }) + server := &Server{ + logger: log.New(io.Discard, "", 0), + relay: relay, + } + relay.Start(context.Background(), &runtimePortEventStub{events: eventCh}) + + serverConn, clientConn := net.Pipe() + done := make(chan struct{}) + go func() { + defer close(done) + server.handleConnection(context.Background(), serverConn, nil) + }() + + encoder := json.NewEncoder(clientConn) + decoder := json.NewDecoder(clientConn) + if err := encoder.Encode(protocol.JSONRPCRequest{ + JSONRPC: protocol.JSONRPCVersion, + ID: json.RawMessage(`"bind-1"`), + Method: protocol.MethodGatewayBindStream, + Params: json.RawMessage(`{ + "session_id":"session-1", + "run_id":"run-1", + "channel":"ipc" + }`), + }); err != nil { + t.Fatalf("encode bind request: %v", err) + } + + var bindResponse protocol.JSONRPCResponse + if err := decoder.Decode(&bindResponse); err != nil { + t.Fatalf("decode bind response: %v", err) + } + if bindResponse.Error != nil { + t.Fatalf("unexpected bind rpc error: %+v", bindResponse.Error) + } + + eventCh <- RuntimeEvent{ + Type: RuntimeEventTypeRunProgress, + SessionID: "session-1", + RunID: "run-1", + Payload: map[string]string{ + "chunk": "hello", + }, + } + + type notificationPayload struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + Params json.RawMessage `json:"params"` + } + var notification notificationPayload + if err := decoder.Decode(¬ification); err != nil { + t.Fatalf("decode notification: %v", err) + } + if notification.Method != protocol.MethodGatewayEvent { + t.Fatalf("notification method = %q, want %q", notification.Method, protocol.MethodGatewayEvent) + } + + var eventFrame MessageFrame + if err := json.Unmarshal(notification.Params, &eventFrame); err != nil { + t.Fatalf("decode notification params frame: %v", err) + } + if eventFrame.SessionID != "session-1" || eventFrame.RunID != "run-1" { + t.Fatalf("event frame session/run mismatch: %#v", eventFrame) + } + + _ = clientConn.Close() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("handleConnection did not exit") + } +} + +type runtimePortEventStub struct { + events <-chan RuntimeEvent +} + +func (s *runtimePortEventStub) Run(_ context.Context, _ RunInput) error { + return nil +} + +func (s *runtimePortEventStub) Compact(_ context.Context, _ CompactInput) (CompactResult, error) { + return CompactResult{}, nil +} + +func (s *runtimePortEventStub) ResolvePermission(_ context.Context, _ PermissionResolutionInput) error { + return nil +} + +func (s *runtimePortEventStub) CancelActiveRun() bool { + return false +} + +func (s *runtimePortEventStub) Events() <-chan RuntimeEvent { + return s.events +} + +func (s *runtimePortEventStub) ListSessions(_ context.Context) ([]SessionSummary, error) { + return nil, nil +} + +func (s *runtimePortEventStub) LoadSession(_ context.Context, _ string) (Session, error) { + return Session{}, nil +} + func decodeJSONRPCResultFrame(response protocol.JSONRPCResponse) (MessageFrame, error) { if response.Result == nil { return MessageFrame{}, errors.New("rpc result is nil") diff --git a/internal/gateway/stream_relay.go b/internal/gateway/stream_relay.go new file mode 100644 index 00000000..33588b67 --- /dev/null +++ b/internal/gateway/stream_relay.go @@ -0,0 +1,881 @@ +package gateway + +import ( + "context" + "fmt" + "log" + "os" + "strings" + "sync" + "time" + + "neo-code/internal/gateway/protocol" +) + +const ( + // DefaultStreamBindingTTL 定义连接绑定关系的默认生存时长（滑动续期）。 + DefaultStreamBindingTTL = 15 * time.Minute + // DefaultStreamCleanupInterval 定义路由表过期清理扫描间隔。 + DefaultStreamCleanupInterval = 30 * time.Second + // DefaultStreamQueueSize 定义每个连接的默认发送队列容量。 + DefaultStreamQueueSize = 64 + // DefaultStreamMaxBindingsPerConnection 定义单连接可维护的最大绑定数，防止路由表被无限放大。 + DefaultStreamMaxBindingsPerConnection = 128 +) + +type relayMessageKind string + +const ( + relayMessageKindJSON relayMessageKind = "json" + relayMessageKindSSE relayMessageKind = "sse" +) + +// RelayMessage 表示发往具体连接的统一出站消息。 +type RelayMessage struct { + Kind relayMessageKind + Event string + Payload any + Enqueued time.Time +} + +// ConnectionRegistration 描述连接注册到中继器所需的写入与关闭钩子。 +type ConnectionRegistration struct { + ConnectionID ConnectionID + Channel StreamChannel + Context context.Context + Cancel context.CancelFunc + Write func(message RelayMessage) error + Close func() +} + +// StreamBinding 描述连接绑定到会话路由表的一条订阅关系。 +type StreamBinding struct { + SessionID string + RunID string + Channel StreamChannel + Explicit bool +} + +// StreamRelayOptions 描述会话路由与流式中继的可选配置。 +type StreamRelayOptions struct { + Logger *log.Logger + BindingTTL time.Duration + CleanupInterval time.Duration + QueueSize int + // MaxBindingsPerConnection 控制单连接可建立的会话绑定上限。 + MaxBindingsPerConnection int +} + +type relayConnection struct { + id ConnectionID + channel StreamChannel + ctx context.Context + cancel context.CancelFunc + writeFn func(message RelayMessage) error + closeFn func() + queue chan RelayMessage + writeMu sync.Mutex +} + +type bindingKey struct { + sessionID string + runID string +} + +type bindingState struct { + sessionID string + runID string + channel StreamChannel + explicit bool + expireAt time.Time + lastSeen time.Time +} + +// StreamRelay 维护连接-会话-运行态映射，并负责运行事件的精确中继。 +type StreamRelay struct { + logger *log.Logger + bindingTTL time.Duration + cleanupInterval time.Duration + queueSize int + maxBindings int + + mu sync.RWMutex + connections map[ConnectionID]*relayConnection + connectionBindings map[ConnectionID]map[bindingKey]*bindingState + sessionIndex map[string]map[ConnectionID]struct{} + sessionRunIndex map[string]map[ConnectionID]struct{} + cleanupStarted bool + eventPumpStarted bool + cleanupLoopCancel context.CancelFunc + runtimeEventLoopCancel context.CancelFunc + cleanupLoopGeneration uint64 + eventLoopGeneration uint64 +} + +// NewStreamRelay 创建会话路由与流式中继实例。 +func NewStreamRelay(options StreamRelayOptions) *StreamRelay { + logger := options.Logger + if logger == nil { + logger = log.New(os.Stderr, "gateway-relay: ", log.LstdFlags) + } + + bindingTTL := options.BindingTTL + if bindingTTL <= 0 { + bindingTTL = DefaultStreamBindingTTL + } + + cleanupInterval := options.CleanupInterval + if cleanupInterval <= 0 { + cleanupInterval = DefaultStreamCleanupInterval + } + + queueSize := options.QueueSize + if queueSize <= 0 { + queueSize = DefaultStreamQueueSize + } + + maxBindings := options.MaxBindingsPerConnection + if maxBindings <= 0 { + maxBindings = DefaultStreamMaxBindingsPerConnection + } + + return &StreamRelay{ + logger: logger, + bindingTTL: bindingTTL, + cleanupInterval: cleanupInterval, + queueSize: queueSize, + maxBindings: maxBindings, + connections: make(map[ConnectionID]*relayConnection), + connectionBindings: make(map[ConnectionID]map[bindingKey]*bindingState), + sessionIndex: make(map[string]map[ConnectionID]struct{}), + sessionRunIndex: make(map[string]map[ConnectionID]struct{}), + } +} + +// Start 启动中继器后台任务（过期清理与运行事件消费），多次调用可重入。 +func (r *StreamRelay) Start(ctx context.Context, runtimePort RuntimePort) { + if r == nil { + return + } + if ctx == nil { + ctx = context.Background() + } + + r.mu.Lock() + if !r.cleanupStarted { + cleanupCtx, cleanupCancel := context.WithCancel(ctx) + r.cleanupLoopCancel = cleanupCancel + r.cleanupStarted = true + r.cleanupLoopGeneration++ + cleanupGeneration := r.cleanupLoopGeneration + go r.runCleanupLoop(cleanupCtx, cleanupGeneration) + } + + if runtimePort != nil && !r.eventPumpStarted { + eventCtx, eventCancel := context.WithCancel(ctx) + r.runtimeEventLoopCancel = eventCancel + r.eventPumpStarted = true + r.eventLoopGeneration++ + eventGeneration := r.eventLoopGeneration + go r.runRuntimeEventLoop(eventCtx, runtimePort, eventGeneration) + } + r.mu.Unlock() +} + +// Stop 停止后台任务并主动断开全部登记连接，保证网关退出时状态可回收。 +func (r *StreamRelay) Stop() { + if r == nil { + return + } + + r.mu.Lock() + cleanupCancel := r.cleanupLoopCancel + r.cleanupLoopCancel = nil + r.cleanupStarted = false + r.cleanupLoopGeneration++ + + eventCancel := r.runtimeEventLoopCancel + r.runtimeEventLoopCancel = nil + r.eventPumpStarted = false + r.eventLoopGeneration++ + + connectionIDs := make([]ConnectionID, 0, len(r.connections)) + for connectionID := range r.connections { + connectionIDs = append(connectionIDs, connectionID) + } + r.mu.Unlock() + + if cleanupCancel != nil { + cleanupCancel() + } + if eventCancel != nil { + eventCancel() + } + + for _, connectionID := range connectionIDs { + r.dropConnection(connectionID) + } +} + +// RegisterConnection 注册一个物理连接，并为其启动单写协程和有界队列。 +func (r *StreamRelay) RegisterConnection(registration ConnectionRegistration) error { + if r == nil { + return fmt.Errorf("stream relay is nil") + } + + connectionID := NormalizeConnectionID(registration.ConnectionID) + if connectionID == "" { + return fmt.Errorf("connection_id is required") + } + if registration.Context == nil { + return fmt.Errorf("connection context is required") + } + if registration.Cancel == nil { + return fmt.Errorf("connection cancel function is required") + } + if registration.Write == nil { + return fmt.Errorf("connection write function is required") + } + if registration.Close == nil { + return fmt.Errorf("connection close function is required") + } + if _, ok := ParseStreamChannel(string(registration.Channel)); !ok || registration.Channel == StreamChannelAll { + return fmt.Errorf("invalid connection channel %q", registration.Channel) + } + + r.mu.Lock() + if _, exists := r.connections[connectionID]; exists { + r.mu.Unlock() + return fmt.Errorf("connection %s already registered", connectionID) + } + + connection := &relayConnection{ + id: connectionID, + channel: registration.Channel, + ctx: registration.Context, + cancel: registration.Cancel, + writeFn: registration.Write, + closeFn: registration.Close, + queue: make(chan RelayMessage, r.queueSize), + } + r.connections[connectionID] = connection + r.mu.Unlock() + + go r.runConnectionWriter(connection) + return nil +} + +// SendJSONRPCResponse 将 JSON-RPC 响应写入连接发送队列。 +func (r *StreamRelay) SendJSONRPCResponse(connectionID ConnectionID, response protocol.JSONRPCResponse) bool { + return r.enqueueMessage(connectionID, RelayMessage{ + Kind: relayMessageKindJSON, + Payload: response, + Enqueued: time.Now(), + }) +} + +// SendJSONRPCResponseSync 通过连接统一串行写路径同步发送 JSON-RPC 响应，适用于协议错误等即时反馈场景。 +func (r *StreamRelay) SendJSONRPCResponseSync(connectionID ConnectionID, response protocol.JSONRPCResponse) bool { + if r == nil { + return false + } + + normalizedConnectionID := NormalizeConnectionID(connectionID) + if normalizedConnectionID == "" { + return false + } + + r.mu.RLock() + connection := r.connections[normalizedConnectionID] + r.mu.RUnlock() + if connection == nil { + return false + } + + writeErr := r.writeConnectionMessage(connection, RelayMessage{ + Kind: relayMessageKindJSON, + Payload: response, + Enqueued: time.Now(), + }) + if writeErr == nil { + return true + } + + r.logger.Printf("connection %s sync write failed: %v", normalizedConnectionID, writeErr) + r.dropConnection(normalizedConnectionID) + return false +} + +// SendJSONRPCPayload 将任意 JSON 载荷写入连接发送队列，适用于 IPC/WS 心跳与响应。 +func (r *StreamRelay) SendJSONRPCPayload(connectionID ConnectionID, payload any) bool { + return r.enqueueMessage(connectionID, RelayMessage{ + Kind: relayMessageKindJSON, + Payload: payload, + Enqueued: time.Now(), + }) +} + +// SendSSEEvent 将结构化负载作为指定事件写入 SSE 连接发送队列。 +func (r *StreamRelay) SendSSEEvent(connectionID ConnectionID, eventName string, payload any) bool { + return r.enqueueMessage(connectionID, RelayMessage{ + Kind: relayMessageKindSSE, + Event: strings.TrimSpace(eventName), + Payload: payload, + Enqueued: time.Now(), + }) +} + +// BindConnection 将连接绑定到指定会话与运行态，并刷新绑定 TTL。 +func (r *StreamRelay) BindConnection(connectionID ConnectionID, binding StreamBinding) *FrameError { + if r == nil { + return NewFrameError(ErrorCodeInternalError, "stream relay is nil") + } + + normalizedConnectionID := NormalizeConnectionID(connectionID) + if normalizedConnectionID == "" { + return NewMissingRequiredFieldError("connection_id") + } + + sessionID := strings.TrimSpace(binding.SessionID) + if sessionID == "" { + return NewMissingRequiredFieldError("session_id") + } + runID := strings.TrimSpace(binding.RunID) + + channel := binding.Channel + if channel == "" { + channel = StreamChannelAll + } + if _, ok := ParseStreamChannel(string(channel)); !ok { + return NewFrameError(ErrorCodeInvalidAction, "invalid bind channel") + } + + now := time.Now() + + r.mu.Lock() + connection := r.connections[normalizedConnectionID] + if connection == nil { + r.mu.Unlock() + return NewFrameError(ErrorCodeInvalidAction, "connection is not registered") + } + if channel != StreamChannelAll && channel != connection.channel { + r.mu.Unlock() + return NewFrameError(ErrorCodeInvalidAction, "bind channel does not match connection channel") + } + + key := bindingKey{sessionID: sessionID, runID: runID} + connectionBindingMap := r.connectionBindings[normalizedConnectionID] + if connectionBindingMap == nil { + connectionBindingMap = make(map[bindingKey]*bindingState) + r.connectionBindings[normalizedConnectionID] = connectionBindingMap + } + for existingKey, existingState := range connectionBindingMap { + if existingState == nil || existingState.expireAt.Before(now) { + delete(connectionBindingMap, existingKey) + if existingState != nil { + r.removeConnectionFromIndexesLocked(normalizedConnectionID, existingState.sessionID, existingState.runID) + } + } + } + if _, exists := connectionBindingMap[key]; !exists && len(connectionBindingMap) >= r.maxBindings { + r.mu.Unlock() + return NewFrameError(ErrorCodeInvalidAction, "too many stream bindings for connection") + } + connectionBindingMap[key] = &bindingState{ + sessionID: sessionID, + runID: runID, + channel: channel, + explicit: binding.Explicit, + expireAt: now.Add(r.bindingTTL), + lastSeen: now, + } + r.addConnectionToSessionIndexLocked(sessionID, normalizedConnectionID) + if runID != "" { + r.addConnectionToSessionRunIndexLocked(sessionID, runID, normalizedConnectionID) + } + r.mu.Unlock() + return nil +} + +// ResolveFallbackSessionID 返回连接当前可用绑定中的会话兜底值（取最近续期的绑定）。 +func (r *StreamRelay) ResolveFallbackSessionID(connectionID ConnectionID) string { + if r == nil { + return "" + } + + normalizedConnectionID := NormalizeConnectionID(connectionID) + if normalizedConnectionID == "" { + return "" + } + + now := time.Now() + + r.mu.RLock() + connectionBindingMap := r.connectionBindings[normalizedConnectionID] + var ( + latestSessionID string + latestSeen time.Time + ) + for _, state := range connectionBindingMap { + if state == nil || state.expireAt.Before(now) { + continue + } + if state.lastSeen.After(latestSeen) { + latestSeen = state.lastSeen + latestSessionID = state.sessionID + } + } + r.mu.RUnlock() + return strings.TrimSpace(latestSessionID) +} + +// RefreshConnectionBindings 刷新连接下全部绑定 TTL，供 ping 保活和活跃续期使用。 +func (r *StreamRelay) RefreshConnectionBindings(connectionID ConnectionID) bool { + if r == nil { + return false + } + + normalizedConnectionID := NormalizeConnectionID(connectionID) + if normalizedConnectionID == "" { + return false + } + + now := time.Now() + refreshed := false + + r.mu.Lock() + connectionBindingMap := r.connectionBindings[normalizedConnectionID] + for key, state := range connectionBindingMap { + if state == nil { + delete(connectionBindingMap, key) + continue + } + if state.expireAt.Before(now) { + delete(connectionBindingMap, key) + r.removeConnectionFromIndexesLocked(normalizedConnectionID, state.sessionID, state.runID) + continue + } + state.expireAt = now.Add(r.bindingTTL) + refreshed = true + } + r.mu.Unlock() + return refreshed +} + +// AutoBindFromFrame 根据请求帧中的 session/run 信息执行自动续绑。 +func (r *StreamRelay) AutoBindFromFrame(connectionID ConnectionID, frame MessageFrame) { + if r == nil { + return + } + + sessionID := strings.TrimSpace(frame.SessionID) + if sessionID == "" { + sessionID = strings.TrimSpace(extractSessionIDFromPayload(frame.Payload)) + } + if sessionID == "" { + return + } + + _ = r.BindConnection(connectionID, StreamBinding{ + SessionID: sessionID, + RunID: strings.TrimSpace(frame.RunID), + Channel: StreamChannelAll, + Explicit: false, + }) +} + +// PublishRuntimeEvent 将运行时事件转换为标准通知，并按会话路由表精准投递。 +func (r *StreamRelay) PublishRuntimeEvent(event RuntimeEvent) { + if r == nil { + return + } + + sessionID := strings.TrimSpace(event.SessionID) + if sessionID == "" { + return + } + runID := strings.TrimSpace(event.RunID) + + eventFrame := MessageFrame{ + Type: FrameTypeEvent, + Action: FrameActionRun, + SessionID: sessionID, + RunID: runID, + Payload: map[string]any{ + "event_type": event.Type, + "payload": event.Payload, + }, + } + notification := protocol.NewJSONRPCNotification(protocol.MethodGatewayEvent, eventFrame) + + for _, connectionID := range r.matchConnectionsForEvent(sessionID, runID) { + r.sendEventNotification(connectionID, notification) + } +} + +// runConnectionWriter 串行消费连接发送队列并执行底层写入，保证单连接单写协程。 +func (r *StreamRelay) runConnectionWriter(connection *relayConnection) { + if connection == nil { + return + } + + defer r.unregisterConnection(connection.id, false) + + for { + select { + case <-connection.ctx.Done(): + return + case message, ok := <-connection.queue: + if !ok { + return + } + if err := r.writeConnectionMessage(connection, message); err != nil { + r.logger.Printf("connection %s write failed: %v", connection.id, err) + r.dropConnection(connection.id) + return + } + } + } +} + +// writeConnectionMessage 复用每连接单写互斥，保证队列写与同步写不会并发交错。 +func (r *StreamRelay) writeConnectionMessage(connection *relayConnection, message RelayMessage) error { + connection.writeMu.Lock() + defer connection.writeMu.Unlock() + return connection.writeFn(message) +} + +// enqueueMessage 将消息尝试放入连接队列；队列满会触发慢连接剔除。 +func (r *StreamRelay) enqueueMessage(connectionID ConnectionID, message RelayMessage) bool { + if r == nil { + return false + } + + normalizedConnectionID := NormalizeConnectionID(connectionID) + if normalizedConnectionID == "" { + return false + } + + r.mu.RLock() + connection := r.connections[normalizedConnectionID] + r.mu.RUnlock() + if connection == nil { + return false + } + + select { + case connection.queue <- message: + return true + default: + r.logger.Printf("connection %s queue is full, dropping slow connection", normalizedConnectionID) + r.dropConnection(normalizedConnectionID) + return false + } +} + +// sendEventNotification 按连接通道选择合适封装发送 gateway.event 通知。 +func (r *StreamRelay) sendEventNotification(connectionID ConnectionID, notification protocol.JSONRPCNotification) { + channel, exists := r.connectionChannel(connectionID) + if !exists { + return + } + + if channel == StreamChannelSSE { + _ = r.SendSSEEvent(connectionID, protocol.MethodGatewayEvent, notification) + return + } + _ = r.SendJSONRPCPayload(connectionID, notification) +} + +// connectionChannel 返回连接所属通道类型。 +func (r *StreamRelay) connectionChannel(connectionID ConnectionID) (StreamChannel, bool) { + if r == nil { + return "", false + } + + normalizedConnectionID := NormalizeConnectionID(connectionID) + if normalizedConnectionID == "" { + return "", false + } + + r.mu.RLock() + connection := r.connections[normalizedConnectionID] + r.mu.RUnlock() + if connection == nil { + return "", false + } + return connection.channel, true +} + +// dropConnection 主动剔除连接并触发取消、关闭与路由清理。 +func (r *StreamRelay) dropConnection(connectionID ConnectionID) { + connection := r.unregisterConnection(connectionID, true) + if connection == nil { + return + } +} + +// unregisterConnection 从中继器注销连接并回收所有路由索引。 +func (r *StreamRelay) unregisterConnection(connectionID ConnectionID, shouldClose bool) *relayConnection { + if r == nil { + return nil + } + + normalizedConnectionID := NormalizeConnectionID(connectionID) + if normalizedConnectionID == "" { + return nil + } + + r.mu.Lock() + connection := r.connections[normalizedConnectionID] + if connection == nil { + r.mu.Unlock() + return nil + } + + delete(r.connections, normalizedConnectionID) + + connectionBindingMap := r.connectionBindings[normalizedConnectionID] + delete(r.connectionBindings, normalizedConnectionID) + for _, state := range connectionBindingMap { + if state == nil { + continue + } + r.removeConnectionFromIndexesLocked(normalizedConnectionID, state.sessionID, state.runID) + } + r.mu.Unlock() + + if shouldClose { + connection.cancel() + connection.closeFn() + } + return connection +} + +// runCleanupLoop 周期性扫描并清理过期绑定，避免路由表长期膨胀。 +func (r *StreamRelay) runCleanupLoop(ctx context.Context, generation uint64) { + ticker := time.NewTicker(r.cleanupInterval) + defer ticker.Stop() + defer func() { + r.mu.Lock() + if r.cleanupLoopGeneration == generation { + r.cleanupLoopCancel = nil + r.cleanupStarted = false + } + r.mu.Unlock() + }() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + r.cleanupExpiredBindings() + } + } +} + +// runRuntimeEventLoop 订阅运行时事件流并触发中继投递。 +func (r *StreamRelay) runRuntimeEventLoop(ctx context.Context, runtimePort RuntimePort, generation uint64) { + defer func() { + r.mu.Lock() + if r.eventLoopGeneration == generation { + r.runtimeEventLoopCancel = nil + r.eventPumpStarted = false + } + r.mu.Unlock() + }() + + if runtimePort == nil { + return + } + eventChannel := runtimePort.Events() + if eventChannel == nil { + return + } + + for { + select { + case <-ctx.Done(): + return + case event, ok := <-eventChannel: + if !ok { + return + } + r.PublishRuntimeEvent(event) + } + } +} + +// cleanupExpiredBindings 移除所有到期绑定并同步更新反向索引。 +func (r *StreamRelay) cleanupExpiredBindings() { + if r == nil { + return + } + + now := time.Now() + r.mu.Lock() + for connectionID, connectionBindingMap := range r.connectionBindings { + for key, state := range connectionBindingMap { + if state == nil || state.expireAt.Before(now) { + delete(connectionBindingMap, key) + if state != nil { + r.removeConnectionFromIndexesLocked(connectionID, state.sessionID, state.runID) + } + } + } + if len(connectionBindingMap) == 0 { + delete(r.connectionBindings, connectionID) + } + } + r.mu.Unlock() +} + +// matchConnectionsForEvent 返回满足 session/run/channel 条件的目标连接列表。 +func (r *StreamRelay) matchConnectionsForEvent(sessionID, runID string) []ConnectionID { + if r == nil { + return nil + } + + normalizedSessionID := strings.TrimSpace(sessionID) + if normalizedSessionID == "" { + return nil + } + normalizedRunID := strings.TrimSpace(runID) + now := time.Now() + + r.mu.RLock() + candidateSet := make(map[ConnectionID]struct{}) + for connectionID := range r.sessionIndex[normalizedSessionID] { + candidateSet[connectionID] = struct{}{} + } + if normalizedRunID != "" { + for connectionID := range r.sessionRunIndex[buildSessionRunKey(normalizedSessionID, normalizedRunID)] { + candidateSet[connectionID] = struct{}{} + } + } + + matched := make([]ConnectionID, 0, len(candidateSet)) + for connectionID := range candidateSet { + connection := r.connections[connectionID] + if connection == nil { + continue + } + if r.connectionMatchesEventLocked(connectionID, connection.channel, normalizedSessionID, normalizedRunID, now) { + matched = append(matched, connectionID) + } + } + r.mu.RUnlock() + return matched +} + +// connectionMatchesEventLocked 判断连接在当前时刻是否命中目标事件的路由条件。 +func (r *StreamRelay) connectionMatchesEventLocked( + connectionID ConnectionID, + connectionChannel StreamChannel, + sessionID string, + runID string, + now time.Time, +) bool { + connectionBindingMap := r.connectionBindings[connectionID] + for _, state := range connectionBindingMap { + if state == nil { + continue + } + if state.expireAt.Before(now) { + continue + } + if state.sessionID != sessionID { + continue + } + if state.channel != StreamChannelAll && state.channel != connectionChannel { + continue + } + if runID == "" { + if state.runID == "" { + return true + } + continue + } + if state.runID == "" || state.runID == runID { + return true + } + } + return false +} + +// addConnectionToSessionIndexLocked 将连接加入 session 维度索引。 +func (r *StreamRelay) addConnectionToSessionIndexLocked(sessionID string, connectionID ConnectionID) { + sessionSet := r.sessionIndex[sessionID] + if sessionSet == nil { + sessionSet = make(map[ConnectionID]struct{}) + r.sessionIndex[sessionID] = sessionSet + } + sessionSet[connectionID] = struct{}{} +} + +// addConnectionToSessionRunIndexLocked 将连接加入 session+run 维度索引。 +func (r *StreamRelay) addConnectionToSessionRunIndexLocked(sessionID, runID string, connectionID ConnectionID) { + runSet := r.sessionRunIndex[buildSessionRunKey(sessionID, runID)] + if runSet == nil { + runSet = make(map[ConnectionID]struct{}) + r.sessionRunIndex[buildSessionRunKey(sessionID, runID)] = runSet + } + runSet[connectionID] = struct{}{} +} + +// removeConnectionFromIndexesLocked 将连接从 session 与 session+run 索引中移除。 +func (r *StreamRelay) removeConnectionFromIndexesLocked(connectionID ConnectionID, sessionID, runID string) { + normalizedSessionID := strings.TrimSpace(sessionID) + if normalizedSessionID != "" { + if sessionSet := r.sessionIndex[normalizedSessionID]; sessionSet != nil { + delete(sessionSet, connectionID) + if len(sessionSet) == 0 { + delete(r.sessionIndex, normalizedSessionID) + } + } + } + + normalizedRunID := strings.TrimSpace(runID) + if normalizedSessionID != "" && normalizedRunID != "" { + runKey := buildSessionRunKey(normalizedSessionID, normalizedRunID) + if runSet := r.sessionRunIndex[runKey]; runSet != nil { + delete(runSet, connectionID) + if len(runSet) == 0 { + delete(r.sessionRunIndex, runKey) + } + } + } +} + +// buildSessionRunKey 构建 session+run 组合索引键。 +func buildSessionRunKey(sessionID, runID string) string { + return strings.TrimSpace(sessionID) + "\x00" + strings.TrimSpace(runID) +} + +// extractSessionIDFromPayload 尝试从不同 payload 结构中提取 session_id 字段。 +func extractSessionIDFromPayload(payload any) string { + switch typed := payload.(type) { + case protocol.WakeIntent: + return strings.TrimSpace(typed.SessionID) + case *protocol.WakeIntent: + if typed == nil { + return "" + } + return strings.TrimSpace(typed.SessionID) + case protocol.BindStreamParams: + return strings.TrimSpace(typed.SessionID) + case *protocol.BindStreamParams: + if typed == nil { + return "" + } + return strings.TrimSpace(typed.SessionID) + case map[string]any: + if rawSessionID, exists := typed["session_id"]; exists { + if sessionID, ok := rawSessionID.(string); ok { + return strings.TrimSpace(sessionID) + } + } + } + return "" +} diff --git a/internal/gateway/stream_relay_test.go b/internal/gateway/stream_relay_test.go new file mode 100644 index 00000000..2e980bcf --- /dev/null +++ b/internal/gateway/stream_relay_test.go @@ -0,0 +1,487 @@ +package gateway + +import ( + "context" + "encoding/json" + "testing" + "time" + + "neo-code/internal/gateway/protocol" +) + +func TestStreamRelayBindAndFallbackSession(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := NewConnectionID() + connectionCtx := WithConnectionID(ctx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { + _ = message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + defer relay.dropConnection(connectionID) + + if bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: "session-1", + RunID: "run-1", + Channel: StreamChannelAll, + Explicit: true, + }); bindErr != nil { + t.Fatalf("bind connection: %v", bindErr) + } + + fallbackSessionID := relay.ResolveFallbackSessionID(connectionID) + if fallbackSessionID != "session-1" { + t.Fatalf("fallback session id = %q, want %q", fallbackSessionID, "session-1") + } +} + +func TestStreamRelayPublishRuntimeEventNoCrossSession(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + messageChA := make(chan RelayMessage, 2) + messageChB := make(chan RelayMessage, 2) + registerConnectionForRelayTest(t, relay, ctx, StreamChannelWS, "session-a", "run-a", messageChA) + registerConnectionForRelayTest(t, relay, ctx, StreamChannelWS, "session-b", "run-b", messageChB) + + relay.PublishRuntimeEvent(RuntimeEvent{ + Type: RuntimeEventTypeRunProgress, + SessionID: "session-a", + RunID: "run-a", + Payload: map[string]string{ + "chunk": "hello", + }, + }) + + select { + case message := <-messageChA: + if message.Kind != relayMessageKindJSON { + t.Fatalf("message kind = %q, want %q", message.Kind, relayMessageKindJSON) + } + notification, ok := message.Payload.(protocol.JSONRPCNotification) + if !ok { + t.Fatalf("payload type = %T, want protocol.JSONRPCNotification", message.Payload) + } + if notification.Method != protocol.MethodGatewayEvent { + t.Fatalf("method = %q, want %q", notification.Method, protocol.MethodGatewayEvent) + } + case <-time.After(time.Second): + t.Fatal("expected session-a to receive runtime event") + } + + select { + case <-messageChB: + t.Fatal("session-b should not receive session-a event") + default: + } +} + +func TestStreamRelayQueueOverflowDropsSlowConnection(t *testing.T) { + blockWrite := make(chan struct{}) + relay := NewStreamRelay(StreamRelayOptions{ + QueueSize: 1, + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := NewConnectionID() + connectionCtx := WithConnectionID(ctx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { + _ = message + <-blockWrite + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + defer close(blockWrite) + + response := protocol.NewJSONRPCErrorResponse( + json.RawMessage(`"queue-1"`), + protocol.NewJSONRPCError(protocol.JSONRPCCodeInternalError, "boom", protocol.GatewayCodeInternalError), + ) + if !relay.SendJSONRPCResponse(connectionID, response) { + t.Fatal("first enqueue should succeed") + } + + dropped := false + for attempt := 0; attempt < 8; attempt++ { + if !relay.SendJSONRPCResponse(connectionID, response) { + dropped = true + break + } + } + if !dropped { + t.Fatal("expected slow connection to be dropped when queue overflows") + } +} + +func TestStreamRelayCleanupExpiredBindings(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{ + BindingTTL: 20 * time.Millisecond, + CleanupInterval: 5 * time.Millisecond, + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + relay.Start(ctx, nil) + + connectionID := NewConnectionID() + connectionCtx := WithConnectionID(ctx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { + _ = message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + defer relay.dropConnection(connectionID) + + if bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: "session-expire", + Channel: StreamChannelAll, + Explicit: true, + }); bindErr != nil { + t.Fatalf("bind connection: %v", bindErr) + } + + time.Sleep(60 * time.Millisecond) + if sessionID := relay.ResolveFallbackSessionID(connectionID); sessionID != "" { + t.Fatalf("expired fallback session id = %q, want empty", sessionID) + } +} + +func TestStreamRelayRestartCleanupLoopAfterContextDone(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{ + CleanupInterval: 5 * time.Millisecond, + }) + + firstCtx, firstCancel := context.WithCancel(context.Background()) + relay.Start(firstCtx, nil) + firstCancel() + + waitForStreamRelayState(t, relay, false) + + secondCtx, secondCancel := context.WithCancel(context.Background()) + relay.Start(secondCtx, nil) + waitForStreamRelayState(t, relay, true) + + secondCancel() + waitForStreamRelayState(t, relay, false) +} + +func TestStreamRelayBindConnectionLimit(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{ + MaxBindingsPerConnection: 1, + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := NewConnectionID() + connectionCtx := WithConnectionID(ctx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { + _ = message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + defer relay.dropConnection(connectionID) + + if bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: "session-1", + RunID: "run-1", + Channel: StreamChannelAll, + Explicit: true, + }); bindErr != nil { + t.Fatalf("first bind: %v", bindErr) + } + + bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: "session-1", + RunID: "run-2", + Channel: StreamChannelAll, + Explicit: true, + }) + if bindErr == nil || bindErr.Code != ErrorCodeInvalidAction.String() { + t.Fatalf("second bind error = %#v, want invalid_action", bindErr) + } +} + +func TestStreamRelayBindConnectionPrunesExpiredBindingsBeforeLimitCheck(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{ + MaxBindingsPerConnection: 1, + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := NewConnectionID() + connectionCtx := WithConnectionID(ctx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { + _ = message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + defer relay.dropConnection(connectionID) + + if bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: "session-old", + RunID: "run-1", + Channel: StreamChannelAll, + Explicit: true, + }); bindErr != nil { + t.Fatalf("first bind: %v", bindErr) + } + + relay.mu.Lock() + connectionBindings := relay.connectionBindings[connectionID] + state := connectionBindings[bindingKey{sessionID: "session-old", runID: "run-1"}] + if state == nil { + relay.mu.Unlock() + t.Fatal("expected old binding state") + } + state.expireAt = time.Now().Add(-time.Second) + relay.mu.Unlock() + + if bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: "session-new", + RunID: "run-2", + Channel: StreamChannelAll, + Explicit: true, + }); bindErr != nil { + t.Fatalf("second bind after pruning expired binding: %v", bindErr) + } + + relay.mu.RLock() + defer relay.mu.RUnlock() + if _, exists := relay.sessionIndex["session-old"]; exists { + t.Fatalf("expired session index should be removed: %#v", relay.sessionIndex["session-old"]) + } + if _, exists := relay.sessionIndex["session-new"][connectionID]; !exists { + t.Fatal("new session index should include current connection") + } +} + +func TestStreamRelayRefreshConnectionBindingsPreservesLastSeen(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := NewConnectionID() + connectionCtx := WithConnectionID(ctx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { + _ = message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + defer relay.dropConnection(connectionID) + + if bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: "session-a", + RunID: "run-a", + Channel: StreamChannelAll, + Explicit: true, + }); bindErr != nil { + t.Fatalf("bind session-a: %v", bindErr) + } + if bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: "session-b", + RunID: "run-b", + Channel: StreamChannelAll, + Explicit: true, + }); bindErr != nil { + t.Fatalf("bind session-b: %v", bindErr) + } + + expectedLastSeenA := time.Now().Add(-2 * time.Minute) + expectedLastSeenB := time.Now().Add(-1 * time.Minute) + + relay.mu.Lock() + for key, state := range relay.connectionBindings[connectionID] { + if state == nil { + continue + } + state.expireAt = time.Now().Add(2 * time.Minute) + switch key.sessionID { + case "session-a": + state.lastSeen = expectedLastSeenA + case "session-b": + state.lastSeen = expectedLastSeenB + } + } + relay.mu.Unlock() + + if refreshed := relay.RefreshConnectionBindings(connectionID); !refreshed { + t.Fatal("expected refresh to succeed") + } + + relay.mu.RLock() + lastSeenA := relay.connectionBindings[connectionID][bindingKey{sessionID: "session-a", runID: "run-a"}].lastSeen + lastSeenB := relay.connectionBindings[connectionID][bindingKey{sessionID: "session-b", runID: "run-b"}].lastSeen + relay.mu.RUnlock() + + if !lastSeenA.Equal(expectedLastSeenA) { + t.Fatalf("session-a lastSeen changed: got %v, want %v", lastSeenA, expectedLastSeenA) + } + if !lastSeenB.Equal(expectedLastSeenB) { + t.Fatalf("session-b lastSeen changed: got %v, want %v", lastSeenB, expectedLastSeenB) + } + + if fallbackSessionID := relay.ResolveFallbackSessionID(connectionID); fallbackSessionID != "session-b" { + t.Fatalf("fallback session id = %q, want %q", fallbackSessionID, "session-b") + } +} + +func TestStreamRelaySendJSONRPCResponseSync(t *testing.T) { + relay := NewStreamRelay(StreamRelayOptions{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + connectionID := NewConnectionID() + connectionCtx := WithConnectionID(ctx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + messageCh := make(chan RelayMessage, 1) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: StreamChannelIPC, + Context: connectionCtx, + Cancel: cancel, + Write: func(message RelayMessage) error { + messageCh <- message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + defer relay.dropConnection(connectionID) + + response, rpcErr := protocol.NewJSONRPCResultResponse(json.RawMessage(`"sync-1"`), MessageFrame{ + Type: FrameTypeAck, + Action: FrameActionPing, + RequestID: "sync-1", + }) + if rpcErr != nil { + t.Fatalf("build rpc response: %v", rpcErr) + } + if ok := relay.SendJSONRPCResponseSync(connectionID, response); !ok { + t.Fatal("expected sync response send to succeed") + } + + select { + case message := <-messageCh: + if message.Kind != relayMessageKindJSON { + t.Fatalf("message kind = %q, want %q", message.Kind, relayMessageKindJSON) + } + case <-time.After(time.Second): + t.Fatal("did not receive sync response message") + } +} + +func waitForStreamRelayState(t *testing.T, relay *StreamRelay, wantStarted bool) { + t.Helper() + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + relay.mu.RLock() + started := relay.cleanupStarted + relay.mu.RUnlock() + if started == wantStarted { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("cleanupStarted state mismatch: want %v", wantStarted) +} + +func registerConnectionForRelayTest( + t *testing.T, + relay *StreamRelay, + ctx context.Context, + channel StreamChannel, + sessionID string, + runID string, + messageCh chan RelayMessage, +) { + t.Helper() + + connectionID := NewConnectionID() + connectionCtx, cancelConnection := context.WithCancel(ctx) + connectionCtx = WithConnectionID(connectionCtx, connectionID) + connectionCtx = WithStreamRelay(connectionCtx, relay) + if err := relay.RegisterConnection(ConnectionRegistration{ + ConnectionID: connectionID, + Channel: channel, + Context: connectionCtx, + Cancel: cancelConnection, + Write: func(message RelayMessage) error { + messageCh <- message + return nil + }, + Close: func() {}, + }); err != nil { + t.Fatalf("register connection: %v", err) + } + t.Cleanup(func() { + relay.dropConnection(connectionID) + }) + + if bindErr := relay.BindConnection(connectionID, StreamBinding{ + SessionID: sessionID, + RunID: runID, + Channel: StreamChannelAll, + Explicit: true, + }); bindErr != nil { + t.Fatalf("bind connection: %v", bindErr) + } +} diff --git a/internal/gateway/types.go b/internal/gateway/types.go index 4230ea15..c2e3e43a 100644 --- a/internal/gateway/types.go +++ b/internal/gateway/types.go @@ -20,6 +20,8 @@ type FrameAction string const ( // FrameActionPing 表示探活动作，用于验证网关可用性。 FrameActionPing FrameAction = "ping" + // FrameActionBindStream 表示声明流式事件订阅绑定。 + FrameActionBindStream FrameAction = "bind_stream" // FrameActionRun 表示发起一次运行。 FrameActionRun FrameAction = "run" // FrameActionCompact 表示触发一次手动压缩。 diff --git a/internal/gateway/validate.go b/internal/gateway/validate.go index 6d2b658d..4e47090d 100644 --- a/internal/gateway/validate.go +++ b/internal/gateway/validate.go @@ -32,6 +32,11 @@ func validateRequestFrame(frame MessageFrame) *FrameError { switch frame.Action { case FrameActionPing: return nil + case FrameActionBindStream: + if frame.Payload == nil { + return NewMissingRequiredFieldError("payload") + } + return nil case FrameActionWakeOpenURL: if frame.Payload == nil { return NewMissingRequiredFieldError("payload") @@ -176,6 +181,7 @@ func isValidFrameType(frameType FrameType) bool { func isValidFrameAction(action FrameAction) bool { switch action { case FrameActionPing, + FrameActionBindStream, FrameActionWakeOpenURL, FrameActionRun, FrameActionCompact, diff --git a/internal/gateway/validate_additional_test.go b/internal/gateway/validate_additional_test.go index e8611842..1086bc44 100644 --- a/internal/gateway/validate_additional_test.go +++ b/internal/gateway/validate_additional_test.go @@ -70,6 +70,28 @@ func TestValidateFrameCancelAndListSessions(t *testing.T) { if listErr != nil { t.Fatalf("list_sessions request should be valid, got %v", listErr) } + + bindErr := ValidateFrame(MessageFrame{ + Type: FrameTypeRequest, + Action: FrameActionBindStream, + }) + if bindErr == nil { + t.Fatal("bind_stream missing payload should be invalid") + } + if bindErr.Code != ErrorCodeMissingRequiredField.String() { + t.Fatalf("error code = %q, want %q", bindErr.Code, ErrorCodeMissingRequiredField.String()) + } + + bindValidErr := ValidateFrame(MessageFrame{ + Type: FrameTypeRequest, + Action: FrameActionBindStream, + Payload: map[string]string{ + "session_id": "sess-1", + }, + }) + if bindValidErr != nil { + t.Fatalf("bind_stream request should be valid, got %v", bindValidErr) + } } func TestValidateResolvePermissionInvalidPayloadType(t *testing.T) {