diff --git a/go.mod b/go.mod index b5e29061a..80159164c 100644 --- a/go.mod +++ b/go.mod @@ -34,6 +34,7 @@ require ( github.com/go-logr/logr v1.4.4 github.com/go-logr/zapr v1.3.0 github.com/go-playground/assert/v2 v2.2.0 + github.com/go-zeromq/zmq4 v0.17.0 github.com/gobwas/glob v1.0.0 github.com/golang-jwt/jwt/v4 v4.5.2 github.com/golang-jwt/jwt/v5 v5.3.1 @@ -70,6 +71,7 @@ require ( go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.71.0 go.opentelemetry.io/otel v1.46.0 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.46.0 + go.opentelemetry.io/otel/metric v1.46.0 go.opentelemetry.io/otel/sdk v1.46.0 go.opentelemetry.io/otel/trace v1.46.0 go.uber.org/dig v1.19.0 @@ -158,6 +160,7 @@ require ( github.com/go-playground/universal-translator v0.18.2 // indirect github.com/go-playground/validator/v10 v10.30.5 // indirect github.com/go-viper/mapstructure/v2 v2.5.0 // indirect + github.com/go-zeromq/goczmq/v4 v4.2.2 // indirect github.com/goccy/go-yaml v1.19.2 // indirect github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 // indirect github.com/golang/protobuf v1.5.4 // indirect @@ -228,7 +231,6 @@ require ( go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.71.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.46.0 // indirect - go.opentelemetry.io/otel/metric v1.46.0 // indirect go.opentelemetry.io/proto/otlp v1.11.0 // indirect go.yaml.in/yaml/v2 v2.4.4 // indirect go.yaml.in/yaml/v3 v3.0.5 // indirect diff --git a/go.sum b/go.sum index 7326b9492..ca4f653e8 100644 --- a/go.sum +++ b/go.sum @@ -234,6 +234,10 @@ github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1v github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8= github.com/go-viper/mapstructure/v2 v2.5.0 h1:vM5IJoUAy3d7zRSVtIwQgBj7BiWtMPfmPEgAXnvj1Ro= github.com/go-viper/mapstructure/v2 v2.5.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM= +github.com/go-zeromq/goczmq/v4 v4.2.2 h1:HAJN+i+3NW55ijMJJhk7oWxHKXgAuSBkoFfvr8bYj4U= +github.com/go-zeromq/goczmq/v4 v4.2.2/go.mod h1:Sm/lxrfxP/Oxqs0tnHD6WAhwkWrx+S+1MRrKzcxoaYE= +github.com/go-zeromq/zmq4 v0.17.0 h1:r12/XdqPeRbuaF4C3QZJeWCt7a5vpJbslDH1rTXF+Kc= +github.com/go-zeromq/zmq4 v0.17.0/go.mod h1:EQxjJD92qKnrsVMzAnx62giD6uJIPi1dMGZ781iCDtY= github.com/gobwas/glob v1.0.0 h1:p+FKbLEIsK1yZ39/OINwFvqNb5oyPY4H8xcy6uYu8dg= github.com/gobwas/glob v1.0.0/go.mod h1:oWCdo522i2P1n/hMXGNWs7yoV4wy/ciZuUIbvKj5rkc= github.com/goccy/go-yaml v1.19.2 h1:PmFC1S6h8ljIz6gMRBopkjP1TVT7xuwrButHID66PoM= diff --git a/pkg/agent/README.md b/pkg/agent/README.md index 767ed0c15..dee3610f4 100644 --- a/pkg/agent/README.md +++ b/pkg/agent/README.md @@ -2,8 +2,8 @@ The `pkg/agent` package contains the Go server used by the Runme web application. It provides runner and parser services, optional static web -application hosting, authentication, and Jupyter proxying. WebMCP is the -supported integration path for browser automation. +application hosting, authentication, and direct Jupyter kernel management. +WebMCP is the supported integration path for browser automation. Despite its historical package name, this is not an AI-only server. Runme Web depends on it for normal notebook and cell execution. @@ -18,8 +18,11 @@ The server provides the runtime services required by Runme Web: - The Runner service manages execution sessions and related runner operations. - The Parser service parses and serializes notebook content when `assistantServer.parserService` is enabled. -- The Jupyter proxy manages Jupyter servers and forwards kernel channel - WebSockets. +- The Jupyter kernel API starts client-specified kernels, manages their + lifecycle, and bridges browser channel WebSockets directly to kernel ZeroMQ + sockets. It is available only when both + `assistantServer.jupyterService` and `assistantServer.runnerService` are + enabled. Callers need `RunnerUserRole`; no separate Jupyter Server is used. - Optional OIDC authentication, authorization, and telemetry are provided by the same HTTP server. - A web application can optionally be served from the directory configured by @@ -44,6 +47,8 @@ logging: assistantServer: port: 8080 runnerService: true + # Opt in to client-selected Jupyter kernel commands. Requires runnerService. + jupyterService: true parserService: true # Optional: serve a web application from this directory. # staticAssets: /absolute/path/to/webapp/dist diff --git a/pkg/agent/config/config.go b/pkg/agent/config/config.go index 380d0d7f4..e9a8002d9 100644 --- a/pkg/agent/config/config.go +++ b/pkg/agent/config/config.go @@ -234,6 +234,9 @@ func (ac *AppConfig) GetConfigDir() string { // IsValid validates the configuration and returns any errors. func (c *Config) IsValid() []string { problems := make([]string, 0, 1) + if c.AssistantServer != nil && c.AssistantServer.JupyterService && !c.AssistantServer.RunnerService { + problems = append(problems, "assistantServer.jupyterService requires assistantServer.runnerService") + } return problems } @@ -472,6 +475,9 @@ type AssistantServerConfig struct { // RunnerService starts the Runme runner service if true otherwise it doesn't start the runner service. RunnerService bool `json:"runnerService" yaml:"runnerService"` + // JupyterService starts the direct Jupyter kernel service if true. It is disabled by default and requires RunnerService. + JupyterService bool `json:"jupyterService" yaml:"jupyterService"` + // RunnerReconnect is a flag to enable automatic reconnecting to the runner from the frontend. RunnerReconnect bool `json:"runnerReconnect,omitempty" yaml:"runnerReconnect,omitempty"` diff --git a/pkg/agent/config/config_test.go b/pkg/agent/config/config_test.go index 6c8b0867f..d46c02706 100644 --- a/pkg/agent/config/config_test.go +++ b/pkg/agent/config/config_test.go @@ -77,6 +77,25 @@ func Test_ConfigAutomaticEnvVars(t *testing.T) { } } +func Test_ConfigJupyterServiceRequiresRunnerService(t *testing.T) { + cfg := &Config{ + AssistantServer: &AssistantServerConfig{JupyterService: true}, + } + + problems := cfg.IsValid() + if len(problems) != 1 { + t.Fatalf("Expected one validation problem, got %v", problems) + } + if problems[0] != "assistantServer.jupyterService requires assistantServer.runnerService" { + t.Fatalf("Unexpected validation problem: %q", problems[0]) + } + + cfg.AssistantServer.RunnerService = true + if problems := cfg.IsValid(); len(problems) != 0 { + t.Fatalf("Expected runner-backed Jupyter service configuration to be valid, got %v", problems) + } +} + func Test_DeprecatedAgentConfigurationIsAccepted(t *testing.T) { v := viper.New() v.SetConfigType("yaml") diff --git a/pkg/agent/iam/oidc.go b/pkg/agent/iam/oidc.go index 247b8cda7..bbab799bd 100644 --- a/pkg/agent/iam/oidc.go +++ b/pkg/agent/iam/oidc.go @@ -384,10 +384,9 @@ func NewAuthMiddlewareForOIDC(oidc *OIDC) (func(http.Handler) http.Handler, erro // Prefer bearer token over session cookie. // Browser WebSocket clients cannot set custom Authorization headers, so // the Jupyter channels client sends `?authorization=Bearer ...` during - // the handshake. We use query auth for the Jupyter proxy because we want - // to preserve the kernel protocol payload as-is, instead of introducing - // a Runme-style in-band auth envelope in websocket messages. - // The Jupyter proxy strips this param before forwarding upstream. + // the handshake. We use query auth for the direct Jupyter bridge to + // preserve the kernel protocol payload as-is instead of introducing a + // Runme-style in-band auth envelope in WebSocket messages. bearerToken := r.Header.Get("Authorization") if !strings.HasPrefix(strings.ToLower(strings.TrimSpace(bearerToken)), "bearer ") { queryAuthorization := r.URL.Query().Get("authorization") diff --git a/pkg/agent/jupyter/bridge.go b/pkg/agent/jupyter/bridge.go new file mode 100644 index 000000000..40b72e2e7 --- /dev/null +++ b/pkg/agent/jupyter/bridge.go @@ -0,0 +1,688 @@ +package jupyter + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + "sync/atomic" + "time" +) + +var ( + ErrSlowJupyterClient = errors.New("jupyter WebSocket client is too slow") + ErrIOPubRateLimit = errors.New("jupyter IOPub rate limit exceeded") +) + +// ChannelsClient is the browser WebSocket boundary used by the direct bridge. +// Implementations must unblock Read and Write when ctx is canceled. +type ChannelsClient interface { + Read(ctx context.Context) ([]byte, error) + Write(ctx context.Context, payload []byte) error + Close(code int, reason string) error +} + +// BridgeConfig bounds one browser-to-kernel channel bridge. +type BridgeConfig struct { + Limits Limits + OutboundQueueMessages int + OutboundQueueBytes int64 + IOPubMessagesPerSec int + IOPubBytesPerSec int64 + RateWindow time.Duration + ReadinessTimeout time.Duration + ReadinessRetry time.Duration + HeartbeatInterval time.Duration + HeartbeatTimeout time.Duration + HeartbeatMisses int +} + +func (c BridgeConfig) normalized() BridgeConfig { + c.Limits = c.Limits.normalized() + if c.OutboundQueueMessages <= 0 { + c.OutboundQueueMessages = 256 + } + if c.OutboundQueueBytes <= 0 { + c.OutboundQueueBytes = 32 * 1024 * 1024 + } + if c.IOPubMessagesPerSec <= 0 { + c.IOPubMessagesPerSec = 1000 + } + if c.IOPubBytesPerSec <= 0 { + c.IOPubBytesPerSec = 1_000_000 + } + if c.RateWindow <= 0 { + c.RateWindow = 3 * time.Second + } + if c.ReadinessTimeout <= 0 { + c.ReadinessTimeout = 10 * time.Second + } + if c.ReadinessRetry <= 0 { + c.ReadinessRetry = 500 * time.Millisecond + } + if c.HeartbeatInterval <= 0 { + c.HeartbeatInterval = 10 * time.Second + } + if c.HeartbeatTimeout <= 0 { + c.HeartbeatTimeout = 3 * time.Second + } + if c.HeartbeatMisses <= 0 { + c.HeartbeatMisses = 3 + } + return c +} + +// KernelChannelsBridge connects one browser WebSocket to one immutable kernel +// generation using a dedicated set of ZeroMQ sockets. +type KernelChannelsBridge struct { + manager *KernelManager + factory SocketFactory + config BridgeConfig +} + +// NewKernelChannelsBridge creates a bounded direct channels bridge. +func NewKernelChannelsBridge(manager *KernelManager, factory SocketFactory, config BridgeConfig) (*KernelChannelsBridge, error) { + if manager == nil { + return nil, errors.New("jupyter kernel manager is required") + } + if factory == nil { + factory = manager.config.SocketFactory + } + if factory == nil { + return nil, errors.New("jupyter socket factory is required") + } + return &KernelChannelsBridge{manager: manager, factory: factory, config: config.normalized()}, nil +} + +// Bridge runs until the client disconnects, the kernel generation changes, a +// safety limit is reached, or ctx is canceled. It never stops the kernel. +func (b *KernelChannelsBridge) Bridge(ctx context.Context, kernelID string, client ChannelsClient) error { + if client == nil { + return errors.New("jupyter channels client is required") + } + addActiveBridge(1) + defer addActiveBridge(-1) + connection, release, err := b.manager.Connect(kernelID) + if err != nil { + return err + } + defer release() + + bridgeCtx, cancel := context.WithCancel(ctx) + defer cancel() + generationChanged := make(chan struct{}) + go func() { + select { + case <-connection.Context.Done(): + close(generationChanged) + cancel() + case <-bridgeCtx.Done(): + } + }() + + signer, err := NewSigner(connection.Info.Key, connection.Info.SignatureScheme) + if err != nil { + return err + } + sockets, err := b.openSocketSet(bridgeCtx, connection.Info) + if err != nil { + return fmt.Errorf("open Jupyter channel sockets: %w", err) + } + defer sockets.Close() + + readinessCtx, readinessCancel := context.WithTimeout(bridgeCtx, b.config.ReadinessTimeout) + err = b.readinessNudge(readinessCtx, sockets, connection.Info, signer) + readinessCancel() + if err != nil { + _ = client.Close(1013, "Jupyter kernel channels are not ready; retry") + return fmt.Errorf("jupyter channels readiness nudge: %w", err) + } + + queue := newOutboundQueue(b.config.OutboundQueueMessages, b.config.OutboundQueueBytes) + errCh := make(chan error, 8) + var wg sync.WaitGroup + start := func(run func() error) { + wg.Add(1) + go func() { + defer wg.Done() + if err := run(); err != nil && !errors.Is(err, context.Canceled) { + select { + case errCh <- err: + default: + } + cancel() + } + }() + } + + start(func() error { return b.writeClient(bridgeCtx, client, queue) }) + start(func() error { return b.readClient(bridgeCtx, client, sockets, signer) }) + limiter := newIOPubRateLimiter(b.config) + for channel, socket := range map[Channel]Socket{ + ChannelShell: sockets.shell, + ChannelControl: sockets.control, + ChannelStdin: sockets.stdin, + ChannelIOPub: sockets.iopub, + } { + channel, socket := channel, socket + start(func() error { + return b.readKernel(bridgeCtx, kernelID, connection.Generation, channel, socket, signer, queue, limiter) + }) + } + start(func() error { return b.monitorHeartbeat(bridgeCtx, sockets.heartbeat, connection.Info) }) + + var bridgeErr error + select { + case <-bridgeCtx.Done(): + select { + case <-generationChanged: + bridgeErr = errors.New("jupyter kernel connection generation changed") + default: + bridgeErr = bridgeCtx.Err() + } + case bridgeErr = <-errCh: + cancel() + } + sockets.Close() + queue.Close() + wg.Wait() + + closeCode, closeReason := 1011, "Jupyter channels bridge closed" + if bridgeErr == nil || errors.Is(bridgeErr, context.Canceled) { + closeCode, closeReason = 1000, "Jupyter channels client disconnected" + } else if errors.Is(bridgeErr, ErrBinaryBuffersUnsupported) { + closeCode, closeReason = 1003, "unsupported_binary_buffers" + } else if errors.Is(bridgeErr, ErrSlowJupyterClient) || errors.Is(bridgeErr, ErrIOPubRateLimit) { + closeCode, closeReason = 1008, bridgeErr.Error() + } else if isClosed(generationChanged) { + closeCode, closeReason = 1012, "Jupyter kernel restarted; reconnect" + } + closeMetricReason := "internal_error" + switch { + case bridgeErr == nil || errors.Is(bridgeErr, context.Canceled): + closeMetricReason = "client_disconnect" + case errors.Is(bridgeErr, ErrSlowJupyterClient): + closeMetricReason = "slow_client" + case errors.Is(bridgeErr, ErrIOPubRateLimit): + closeMetricReason = "rate_limit" + case errors.Is(bridgeErr, ErrBinaryBuffersUnsupported): + closeMetricReason = "unsupported_binary_buffers" + case isClosed(generationChanged): + closeMetricReason = "kernel_restart" + } + observeBridgeClose(closeMetricReason) + _ = client.Close(closeCode, closeReason) + return bridgeErr +} + +type bridgeSocketSet struct { + shell Socket + control Socket + stdin Socket + iopub Socket + heartbeat Socket + closeOnce sync.Once +} + +func (s *bridgeSocketSet) Close() { + s.closeOnce.Do(func() { + for _, socket := range []Socket{s.shell, s.control, s.stdin, s.iopub, s.heartbeat} { + if socket != nil { + _ = socket.Close() + } + } + }) +} + +func (b *KernelChannelsBridge) openSocketSet(ctx context.Context, connection ConnectionInfo) (*bridgeSocketSet, error) { + identity := []byte("runme-" + randomIdentifier()) + set := &bridgeSocketSet{} + var err error + if set.shell, err = b.openSocket(ctx, SocketTypeDealer, identity, connection, ChannelShell); err != nil { + set.Close() + return nil, err + } + if set.control, err = b.openSocket(ctx, SocketTypeDealer, identity, connection, ChannelControl); err != nil { + set.Close() + return nil, err + } + if set.stdin, err = b.openSocket(ctx, SocketTypeDealer, identity, connection, ChannelStdin); err != nil { + set.Close() + return nil, err + } + if set.iopub, err = b.openSocket(ctx, SocketTypeSub, nil, connection, ChannelIOPub); err != nil { + set.Close() + return nil, err + } + if set.heartbeat, err = b.openSocket(ctx, SocketTypeReq, nil, connection, ChannelHeartbeat); err != nil { + set.Close() + return nil, err + } + return set, nil +} + +func (b *KernelChannelsBridge) openSocket( + ctx context.Context, + socketType SocketType, + identity []byte, + connection ConnectionInfo, + channel Channel, +) (Socket, error) { + socket, err := b.factory.NewSocket(ctx, socketType, identity) + if err != nil { + return nil, err + } + endpoint, err := connection.Endpoint(channel) + if err != nil { + _ = socket.Close() + return nil, err + } + if err := socket.Dial(endpoint); err != nil { + _ = socket.Close() + return nil, err + } + return socket, nil +} + +func (b *KernelChannelsBridge) readinessNudge( + ctx context.Context, + sockets *bridgeSocketSet, + connection ConnectionInfo, + signer *Signer, +) error { + identity := []byte("runme-nudge-" + randomIdentifier()) + shell, err := b.openSocket(ctx, SocketTypeDealer, identity, connection, ChannelShell) + if err != nil { + return err + } + defer shell.Close() + control, err := b.openSocket(ctx, SocketTypeDealer, identity, connection, ChannelControl) + if err != nil { + return err + } + defer control.Close() + + infoReply := make(chan string, 2) + iopubMessage := make(chan struct{}, 1) + errorsCh := make(chan error, 3) + readInfo := func(socket Socket) { + for { + replyFrames, err := socket.ReceiveMultipart() + if err != nil { + if ctx.Err() == nil { + errorsCh <- err + } + return + } + reply, err := ParseMultipart(replyFrames, signer, b.config.Limits) + if err != nil { + errorsCh <- err + return + } + messageType, parentID, err := messageTypeAndParent(reply) + if err != nil { + errorsCh <- err + return + } + if messageType == "kernel_info_reply" { + select { + case infoReply <- parentID: + default: + } + return + } + } + } + go readInfo(shell) + go readInfo(control) + go func() { + replyFrames, err := sockets.iopub.ReceiveMultipart() + if err != nil { + if ctx.Err() == nil { + errorsCh <- err + } + return + } + if _, err := ParseMultipart(replyFrames, signer, b.config.Limits); err != nil { + errorsCh <- err + return + } + iopubMessage <- struct{}{} + }() + + ticker := time.NewTicker(b.config.ReadinessRetry) + defer ticker.Stop() + gotInfo, gotIOPub := false, false + sentRequestIDs := make(map[string]struct{}) + send := func() error { + request, requestID, err := newProtocolRequest("kernel_info_request", `{}`) + if err != nil { + return err + } + frames, err := MarshalMultipart(request, signer, b.config.Limits) + if err != nil { + return err + } + sentRequestIDs[requestID] = struct{}{} + if err := shell.SendMultipart(frames); err != nil { + return err + } + return control.SendMultipart(frames) + } + if err := send(); err != nil { + return err + } + for !gotInfo || !gotIOPub { + select { + case <-ctx.Done(): + return fmt.Errorf( + "readiness incomplete (kernel_info=%t iopub=%t): %w", + gotInfo, + gotIOPub, + ctx.Err(), + ) + case err := <-errorsCh: + return err + case parentID := <-infoReply: + if _, ok := sentRequestIDs[parentID]; ok { + gotInfo = true + } + case <-iopubMessage: + gotIOPub = true + case <-ticker.C: + if err := send(); err != nil { + return err + } + } + } + return nil +} + +func (b *KernelChannelsBridge) readClient( + ctx context.Context, + client ChannelsClient, + sockets *bridgeSocketSet, + signer *Signer, +) error { + for { + payload, err := client.Read(ctx) + if err != nil { + return err + } + channel, message, err := ParseChannelJSON(payload, b.config.Limits) + if err != nil { + observeProtocolError(err) + return err + } + frames, err := MarshalMultipart(message, signer, b.config.Limits) + if err != nil { + return err + } + var socket Socket + switch channel { + case ChannelShell: + socket = sockets.shell + case ChannelControl: + socket = sockets.control + case ChannelStdin: + socket = sockets.stdin + default: + return fmt.Errorf("unsupported browser Jupyter channel %q", channel) + } + if err := socket.SendMultipart(frames); err != nil { + return err + } + observeChannelMessage("browser_to_kernel", channel, len(payload)) + } +} + +func (b *KernelChannelsBridge) readKernel( + ctx context.Context, + kernelID string, + generation uint64, + channel Channel, + socket Socket, + signer *Signer, + queue *outboundQueue, + limiter *iopubRateLimiter, +) error { + for { + frames, err := socket.ReceiveMultipart() + if err != nil { + return err + } + message, err := ParseMultipart(frames, signer, b.config.Limits) + if err != nil { + observeProtocolError(err) + return err + } + payload, err := MarshalChannelJSON(channel, message, b.config.Limits) + if err != nil { + return err + } + if channel == ChannelIOPub { + if err := limiter.Allow(len(payload)); err != nil { + addCounter(jupyterRateLimits) + return err + } + b.observeKernelState(kernelID, generation, message) + } + if err := queue.Enqueue(ctx, payload); err != nil { + return err + } + observeChannelMessage("kernel_to_browser", channel, len(payload)) + } +} + +func (b *KernelChannelsBridge) observeKernelState(kernelID string, generation uint64, message Message) { + var header struct { + MessageType string `json:"msg_type"` + } + if json.Unmarshal(message.Header, &header) != nil || header.MessageType != "status" { + return + } + var content struct { + ExecutionState string `json:"execution_state"` + } + if json.Unmarshal(message.Content, &content) != nil { + return + } + switch content.ExecutionState { + case "busy": + b.manager.SetExecutionState(kernelID, generation, KernelStateBusy) + case "idle": + b.manager.SetExecutionState(kernelID, generation, KernelStateIdle) + } +} + +func (b *KernelChannelsBridge) writeClient(ctx context.Context, client ChannelsClient, queue *outboundQueue) error { + for { + payload, err := queue.Dequeue(ctx) + if err != nil { + return err + } + if err := client.Write(ctx, payload); err != nil { + return err + } + } +} + +func (b *KernelChannelsBridge) monitorHeartbeat(ctx context.Context, initial Socket, connection ConnectionInfo) error { + current := initial + misses := 0 + ticker := time.NewTicker(b.config.HeartbeatInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: + } + payload := []byte("runme-heartbeat-" + randomIdentifier()) + if err := current.SendMultipart([][]byte{payload}); err != nil { + misses++ + } else { + replyCh := make(chan [][]byte, 1) + errCh := make(chan error, 1) + go func(socket Socket) { + reply, err := socket.ReceiveMultipart() + if err != nil { + errCh <- err + return + } + replyCh <- reply + }(current) + timer := time.NewTimer(b.config.HeartbeatTimeout) + select { + case <-ctx.Done(): + timer.Stop() + return ctx.Err() + case reply := <-replyCh: + timer.Stop() + if len(reply) == 1 && string(reply[0]) == string(payload) { + misses = 0 + } else { + misses++ + } + case <-errCh: + timer.Stop() + misses++ + case <-timer.C: + misses++ + } + } + if misses < b.config.HeartbeatMisses { + if misses > 0 { + addCounter(jupyterHeartbeatMisses) + _ = current.Close() + replacement, err := b.openSocket(ctx, SocketTypeReq, nil, connection, ChannelHeartbeat) + if err != nil { + continue + } + current = replacement + defer current.Close() + } + continue + } + addCounter(jupyterHeartbeatMisses) + return fmt.Errorf("jupyter heartbeat missed %d times", misses) + } +} + +func messageTypeAndParent(message Message) (string, string, error) { + var header, parent struct { + MessageID string `json:"msg_id"` + MessageType string `json:"msg_type"` + } + if err := json.Unmarshal(message.Header, &header); err != nil { + return "", "", err + } + if err := json.Unmarshal(message.ParentHeader, &parent); err != nil { + return "", "", err + } + return header.MessageType, parent.MessageID, nil +} + +type outboundQueue struct { + messages chan []byte + bytes atomic.Int64 + maxBytes int64 + closed chan struct{} + once sync.Once +} + +func newOutboundQueue(maxMessages int, maxBytes int64) *outboundQueue { + return &outboundQueue{ + messages: make(chan []byte, maxMessages), + maxBytes: maxBytes, + closed: make(chan struct{}), + } +} + +func (q *outboundQueue) Enqueue(ctx context.Context, payload []byte) error { + size := int64(len(payload)) + if size > q.maxBytes || q.bytes.Add(size) > q.maxBytes { + q.bytes.Add(-size) + return ErrSlowJupyterClient + } + select { + case q.messages <- append([]byte(nil), payload...): + return nil + case <-ctx.Done(): + q.bytes.Add(-size) + return ctx.Err() + case <-q.closed: + q.bytes.Add(-size) + return context.Canceled + default: + q.bytes.Add(-size) + return ErrSlowJupyterClient + } +} + +func (q *outboundQueue) Dequeue(ctx context.Context) ([]byte, error) { + select { + case payload := <-q.messages: + q.bytes.Add(-int64(len(payload))) + return payload, nil + case <-ctx.Done(): + return nil, ctx.Err() + case <-q.closed: + return nil, context.Canceled + } +} + +func (q *outboundQueue) Close() { + q.once.Do(func() { close(q.closed) }) +} + +type iopubRateLimiter struct { + window time.Duration + messageLimit int + byteLimit int64 + started time.Time + messages int + bytes int64 +} + +func newIOPubRateLimiter(config BridgeConfig) *iopubRateLimiter { + messageLimit := int(float64(config.IOPubMessagesPerSec) * config.RateWindow.Seconds()) + if messageLimit < 1 { + messageLimit = 1 + } + byteLimit := int64(float64(config.IOPubBytesPerSec) * config.RateWindow.Seconds()) + if byteLimit < 1 { + byteLimit = 1 + } + return &iopubRateLimiter{ + window: config.RateWindow, + messageLimit: messageLimit, + byteLimit: byteLimit, + started: time.Now(), + } +} + +func (l *iopubRateLimiter) Allow(bytes int) error { + now := time.Now() + if now.Sub(l.started) >= l.window { + l.started = now + l.messages = 0 + l.bytes = 0 + } + l.messages++ + l.bytes += int64(bytes) + if l.messages > l.messageLimit || l.bytes > l.byteLimit { + return ErrIOPubRateLimit + } + return nil +} + +func isClosed(ch <-chan struct{}) bool { + select { + case <-ch: + return true + default: + return false + } +} diff --git a/pkg/agent/jupyter/bridge_integration_test.go b/pkg/agent/jupyter/bridge_integration_test.go new file mode 100644 index 000000000..2b47607ce --- /dev/null +++ b/pkg/agent/jupyter/bridge_integration_test.go @@ -0,0 +1,557 @@ +package jupyter + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "sync" + "testing" + "time" +) + +type memoryChannelsClient struct { + inbound chan []byte + outbound chan []byte + mu sync.Mutex + code int + reason string +} + +func newMemoryChannelsClient() *memoryChannelsClient { + return &memoryChannelsClient{ + inbound: make(chan []byte, 16), + outbound: make(chan []byte, 256), + } +} + +func (c *memoryChannelsClient) Read(ctx context.Context) ([]byte, error) { + select { + case payload := <-c.inbound: + return append([]byte(nil), payload...), nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (c *memoryChannelsClient) Write(ctx context.Context, payload []byte) error { + select { + case c.outbound <- append([]byte(nil), payload...): + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (c *memoryChannelsClient) Close(code int, reason string) error { + c.mu.Lock() + defer c.mu.Unlock() + c.code = code + c.reason = reason + return nil +} + +func TestKernelChannelsBridgeSharedStateAndTwoClients(t *testing.T) { + manager := newTestKernelManager(t) + defer manager.Close(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) + defer cancel() + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) + if err != nil { + t.Fatal(err) + } + bridge := newTestBridge(t, manager) + + clientA, cancelA, resultA := startTestBridge(ctx, bridge, model.ID) + clientB, cancelB, resultB := startTestBridge(ctx, bridge, model.ID) + defer stopTestBridge(t, cancelA, resultA) + defer stopTestBridge(t, cancelB, resultB) + + requestA := randomHex(t, 16) + requestB := randomHex(t, 16) + clientA.inbound <- browserExecuteRequest(t, requestA, "shared_bridge_value = 100\nshared_bridge_value + 1") + result := collectBrowserExecution(t, ctx, clientA, requestA) + if !result["result:101"] || !result["execute_reply"] || !result["status:idle"] { + t.Fatalf("client A execution result = %v", result) + } + + clientB.inbound <- browserExecuteRequest(t, requestB, "shared_bridge_value + 2") + result = collectBrowserExecution(t, ctx, clientB, requestB) + if !result["result:102"] || !result["execute_reply"] || !result["status:idle"] { + t.Fatalf("client B did not observe shared kernel state: %v", result) + } + + if got, err := manager.Get(model.ID); err != nil || got.ExecutionState != string(KernelStateIdle) || got.Connections != 2 { + t.Fatalf("kernel model while clients connected = %+v, error = %v", got, err) + } +} + +func TestKernelChannelsBridgeConcurrentExecutionsRouteByParent(t *testing.T) { + manager := newTestKernelManager(t) + defer manager.Close(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) + defer cancel() + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) + if err != nil { + t.Fatal(err) + } + bridge := newTestBridge(t, manager) + clientA, cancelA, resultA := startTestBridge(ctx, bridge, model.ID) + clientB, cancelB, resultB := startTestBridge(ctx, bridge, model.ID) + defer stopTestBridge(t, cancelA, resultA) + defer stopTestBridge(t, cancelB, resultB) + + requestA := randomHex(t, 16) + requestB := randomHex(t, 16) + clientA.inbound <- browserExecuteRequest(t, requestA, "'client-a-result'") + clientB.inbound <- browserExecuteRequest(t, requestB, "'client-b-result'") + + results := make(chan map[string]bool, 2) + go func() { results <- collectBrowserExecution(t, ctx, clientA, requestA) }() + go func() { results <- collectBrowserExecution(t, ctx, clientB, requestB) }() + first, second := <-results, <-results + if !(first["result:'client-a-result'"] || second["result:'client-a-result'"]) { + t.Fatalf("client A result missing from routed executions: %v / %v", first, second) + } + if !(first["result:'client-b-result'"] || second["result:'client-b-result'"]) { + t.Fatalf("client B result missing from routed executions: %v / %v", first, second) + } +} + +func TestKernelChannelsBridgeInterruptAndReuse(t *testing.T) { + manager := newTestKernelManager(t) + defer manager.Close(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) + defer cancel() + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) + if err != nil { + t.Fatal(err) + } + bridge := newTestBridge(t, manager) + client, cancelBridge, bridgeResult := startTestBridge(ctx, bridge, model.ID) + defer stopTestBridge(t, cancelBridge, bridgeResult) + + requestID := randomHex(t, 16) + client.inbound <- browserExecuteRequest(t, requestID, "import time\ntime.sleep(60)") + waitForBrowserStatus(t, ctx, client, requestID, "busy") + if err := manager.Interrupt(model.ID); err != nil { + t.Fatalf("Interrupt() error = %v", err) + } + result := collectBrowserExecution(t, ctx, client, requestID) + if !result["execute_reply"] || !result["status:idle"] { + t.Fatalf("interrupted execution did not settle: %v", result) + } + waitForKernelState(t, manager, model.ID, KernelStateIdle) + + reuseID := randomHex(t, 16) + client.inbound <- browserExecuteRequest(t, reuseID, "6 * 7") + result = collectBrowserExecution(t, ctx, client, reuseID) + if !result["result:42"] || !result["execute_reply"] || !result["status:idle"] { + t.Fatalf("kernel was not reusable after interrupt: %v", result) + } +} + +func TestKernelChannelsBridgeRestartClosesOldGeneration(t *testing.T) { + manager := newTestKernelManager(t) + defer manager.Close(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) + defer cancel() + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) + if err != nil { + t.Fatal(err) + } + bridge := newTestBridge(t, manager) + client, cancelBridge, bridgeResult := startTestBridge(ctx, bridge, model.ID) + + requestID := randomHex(t, 16) + client.inbound <- browserExecuteRequest(t, requestID, "restart_only_value = 9") + _ = collectBrowserExecution(t, ctx, client, requestID) + if _, err := manager.Restart(ctx, model.ID); err != nil { + t.Fatalf("Restart() error = %v", err) + } + + select { + case err := <-bridgeResult: + if err == nil || !strings.Contains(err.Error(), "generation changed") { + t.Fatalf("old bridge error = %v, want generation change", err) + } + case <-time.After(5 * time.Second): + t.Fatal("old bridge did not close after restart") + } + cancelBridge() + client.mu.Lock() + closeCode := client.code + client.mu.Unlock() + if closeCode != 1012 { + t.Fatalf("old bridge close code = %d, want 1012", closeCode) + } + + newClient, cancelNew, newResult := startTestBridge(ctx, bridge, model.ID) + defer stopTestBridge(t, cancelNew, newResult) + newRequestID := randomHex(t, 16) + newClient.inbound <- browserExecuteRequest(t, newRequestID, "restart_only_value") + result := collectBrowserExecution(t, ctx, newClient, newRequestID) + if !result["error:NameError"] || !result["execute_reply"] || !result["status:idle"] { + t.Fatalf("restart did not reset kernel state: %v", result) + } +} + +func TestKernelChannelsBridgeRichOutputsAndQuietExecution(t *testing.T) { + manager := newTestKernelManager(t) + defer manager.Close(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) + if err != nil { + t.Fatal(err) + } + bridge := newTestBridge(t, manager) + client, cancelBridge, bridgeResult := startTestBridge(ctx, bridge, model.ID) + defer stopTestBridge(t, cancelBridge, bridgeResult) + + quietID := randomHex(t, 16) + client.inbound <- browserExecuteRequest(t, quietID, "import time\ntime.sleep(1)\nprint('quiet-finished')") + quiet := collectBrowserMessages(t, ctx, client, quietID) + if !quiet["stdout:quiet-finished"] || !quiet["execute_reply:ok"] || !quiet["status:idle"] { + t.Fatalf("quiet execution = %v", quiet) + } + + richID := randomHex(t, 16) + client.inbound <- browserExecuteRequest(t, richID, strings.Join([]string{ + "import sys", + "from IPython.display import display", + "print('stdout-marker')", + "print('stderr-marker', file=sys.stderr)", + "display({'text/html': 'html-marker'}, raw=True)", + "display({'image/svg+xml': 'svg-marker'}, raw=True)", + "display({'image/png': 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAAB'}, raw=True)", + "raise ValueError('error-marker')", + }, "\n")) + rich := collectBrowserMessages(t, ctx, client, richID) + for _, key := range []string{ + "stdout:stdout-marker", "stderr:stderr-marker", "mime:text/html", "mime:image/svg+xml", + "mime:image/png", "error:ValueError", "execute_reply:error", "status:idle", + } { + if !rich[key] { + t.Fatalf("rich execution missing %q: %v", key, rich) + } + } +} + +func TestKernelChannelsBridgeStdinAndControlRouting(t *testing.T) { + manager := newTestKernelManager(t) + defer manager.Close(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) + defer cancel() + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) + if err != nil { + t.Fatal(err) + } + bridge := newTestBridge(t, manager) + client, cancelBridge, bridgeResult := startTestBridge(ctx, bridge, model.ID) + defer stopTestBridge(t, cancelBridge, bridgeResult) + + controlID := randomHex(t, 16) + client.inbound <- browserProtocolRequest(t, ChannelControl, controlID, "kernel_info_request", map[string]any{}) + for { + message := readBrowserMessage(t, ctx, client) + parent := message["parent_header"].(map[string]any) + if fmt.Sprint(parent["msg_id"]) == controlID && fmt.Sprint(message["msg_type"]) == "kernel_info_reply" { + break + } + } + + executeID := randomHex(t, 16) + request := browserExecuteRequest(t, executeID, "answer = input('prompt:')\nprint('input=' + answer)") + var execute map[string]any + if err := json.Unmarshal(request, &execute); err != nil { + t.Fatal(err) + } + execute["content"].(map[string]any)["allow_stdin"] = true + request, err = json.Marshal(execute) + if err != nil { + t.Fatal(err) + } + client.inbound <- request + seen := make(map[string]bool) + for !(seen["execute_reply:ok"] && seen["status:idle"]) { + message := readBrowserMessage(t, ctx, client) + parent := message["parent_header"].(map[string]any) + if fmt.Sprint(parent["msg_id"]) != executeID { + continue + } + messageType := fmt.Sprint(message["msg_type"]) + content := message["content"].(map[string]any) + switch messageType { + case "input_request": + header := message["header"].(map[string]any) + client.inbound <- browserProtocolRequest(t, ChannelStdin, randomHex(t, 16), "input_reply", map[string]any{"value": "forty-two"}, header) + case "stream": + if strings.Contains(fmt.Sprint(content["text"]), "input=forty-two") { + seen["stream"] = true + } + case "execute_reply": + seen["execute_reply:"+fmt.Sprint(content["status"])] = true + case "status": + seen["status:"+fmt.Sprint(content["execution_state"])] = true + } + } + if !seen["stream"] { + t.Fatalf("stdin reply was not routed: %v", seen) + } +} + +func TestKernelChannelsBridgeReconnectSoak(t *testing.T) { + manager := newTestKernelManager(t) + defer manager.Close(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second) + defer cancel() + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) + if err != nil { + t.Fatal(err) + } + bridge := newTestBridge(t, manager) + for i := 0; i < 10; i++ { + client, cancelBridge, bridgeResult := startTestBridge(ctx, bridge, model.ID) + requestID := randomHex(t, 16) + client.inbound <- browserExecuteRequest(t, requestID, fmt.Sprintf("%d + 1", i)) + result := collectBrowserExecution(t, ctx, client, requestID) + if !result[fmt.Sprintf("result:%d", i+1)] || !result["status:idle"] { + t.Fatalf("soak iteration %d = %v", i, result) + } + stopTestBridge(t, cancelBridge, bridgeResult) + } + if got, err := manager.Get(model.ID); err != nil || got.Connections != 0 || got.ExecutionState != string(KernelStateIdle) { + t.Fatalf("kernel after reconnect soak = %+v, error = %v", got, err) + } +} + +func newTestBridge(t *testing.T, manager *KernelManager) *KernelChannelsBridge { + t.Helper() + bridge, err := NewKernelChannelsBridge(manager, nil, BridgeConfig{ + ReadinessTimeout: 5 * time.Second, + ReadinessRetry: 100 * time.Millisecond, + HeartbeatInterval: 100 * time.Millisecond, + HeartbeatTimeout: 500 * time.Millisecond, + HeartbeatMisses: 3, + }) + if err != nil { + t.Fatal(err) + } + return bridge +} + +func startTestBridge( + parent context.Context, + bridge *KernelChannelsBridge, + kernelID string, +) (*memoryChannelsClient, context.CancelFunc, chan error) { + ctx, cancel := context.WithCancel(parent) + client := newMemoryChannelsClient() + result := make(chan error, 1) + go func() { result <- bridge.Bridge(ctx, kernelID, client) }() + return client, cancel, result +} + +func stopTestBridge(t *testing.T, cancel context.CancelFunc, result <-chan error) { + t.Helper() + cancel() + select { + case <-result: + case <-time.After(5 * time.Second): + t.Error("Jupyter bridge did not stop after cancellation") + } +} + +func browserExecuteRequest(t *testing.T, messageID, code string) []byte { + t.Helper() + payload, err := json.Marshal(map[string]any{ + "channel": "shell", + "header": map[string]any{ + "msg_id": messageID, + "username": "runme-test", + "session": randomHex(t, 16), + "date": time.Now().UTC().Format(time.RFC3339Nano), + "msg_type": "execute_request", + "version": "5.3", + }, + "parent_header": map[string]any{}, + "metadata": map[string]any{"runme_unknown": true}, + "content": map[string]any{ + "code": code, + "silent": false, + "store_history": true, + "user_expressions": map[string]any{}, + "allow_stdin": false, + "stop_on_error": true, + }, + }) + if err != nil { + t.Fatal(err) + } + return payload +} + +func browserProtocolRequest( + t *testing.T, + channel Channel, + messageID, messageType string, + content map[string]any, + parents ...map[string]any, +) []byte { + t.Helper() + parent := map[string]any{} + if len(parents) > 0 { + parent = parents[0] + } + payload, err := json.Marshal(map[string]any{ + "channel": channel, + "header": map[string]any{ + "msg_id": messageID, "username": "runme-test", "session": randomHex(t, 16), + "date": time.Now().UTC().Format(time.RFC3339Nano), "msg_type": messageType, "version": "5.3", + }, + "parent_header": parent, + "metadata": map[string]any{}, + "content": content, + }) + if err != nil { + t.Fatal(err) + } + return payload +} + +func collectBrowserMessages(t *testing.T, ctx context.Context, client *memoryChannelsClient, requestID string) map[string]bool { + t.Helper() + seen := make(map[string]bool) + for !(strings.HasPrefix(firstMapKey(seen, "execute_reply:"), "execute_reply:") && seen["status:idle"]) { + message := readBrowserMessage(t, ctx, client) + parent := message["parent_header"].(map[string]any) + if fmt.Sprint(parent["msg_id"]) != requestID { + continue + } + messageType := fmt.Sprint(message["msg_type"]) + content := message["content"].(map[string]any) + switch messageType { + case "status": + seen["status:"+fmt.Sprint(content["execution_state"])] = true + case "stream": + seen[fmt.Sprint(content["name"])+":"+strings.TrimSpace(fmt.Sprint(content["text"]))] = true + case "display_data", "execute_result": + if data, ok := content["data"].(map[string]any); ok { + for mime := range data { + seen["mime:"+mime] = true + } + } + case "error": + seen["error:"+fmt.Sprint(content["ename"])] = true + case "execute_reply": + seen["execute_reply:"+fmt.Sprint(content["status"])] = true + } + } + return seen +} + +func firstMapKey(values map[string]bool, prefix string) string { + for key := range values { + if strings.HasPrefix(key, prefix) { + return key + } + } + return "" +} + +func collectBrowserExecution( + t *testing.T, + ctx context.Context, + client *memoryChannelsClient, + requestID string, +) map[string]bool { + t.Helper() + seen := make(map[string]bool) + for !(seen["execute_reply"] && seen["status:idle"]) { + message := readBrowserMessage(t, ctx, client) + parent := message["parent_header"].(map[string]any) + if fmt.Sprint(parent["msg_id"]) != requestID { + continue + } + messageType := fmt.Sprint(message["msg_type"]) + content := message["content"].(map[string]any) + switch messageType { + case "status": + seen["status:"+fmt.Sprint(content["execution_state"])] = true + case "execute_result": + seen["execute_result"] = true + if data, ok := content["data"].(map[string]any); ok { + seen["result:"+fmt.Sprint(data["text/plain"])] = true + } + case "error": + seen["error:"+fmt.Sprint(content["ename"])] = true + default: + seen[messageType] = true + } + } + return seen +} + +func waitForBrowserStatus( + t *testing.T, + ctx context.Context, + client *memoryChannelsClient, + requestID, state string, +) { + t.Helper() + for { + message := readBrowserMessage(t, ctx, client) + parent := message["parent_header"].(map[string]any) + if fmt.Sprint(parent["msg_id"]) != requestID || fmt.Sprint(message["msg_type"]) != "status" { + continue + } + content := message["content"].(map[string]any) + if fmt.Sprint(content["execution_state"]) == state { + return + } + } +} + +func readBrowserMessage(t *testing.T, ctx context.Context, client *memoryChannelsClient) map[string]any { + t.Helper() + select { + case payload := <-client.outbound: + var message map[string]any + if err := json.Unmarshal(payload, &message); err != nil { + t.Fatalf("decode browser message: %v", err) + } + return message + case <-ctx.Done(): + t.Fatalf("wait for browser message: %v", ctx.Err()) + return nil + } +} + +func TestOutboundQueueAndIOPubRateLimits(t *testing.T) { + queue := newOutboundQueue(1, 4) + ctx := context.Background() + if err := queue.Enqueue(ctx, []byte("1234")); err != nil { + t.Fatal(err) + } + if err := queue.Enqueue(ctx, []byte("1")); !errors.Is(err, ErrSlowJupyterClient) { + t.Fatalf("Enqueue() error = %v, want slow client", err) + } + payload, err := queue.Dequeue(ctx) + if err != nil || string(payload) != "1234" { + t.Fatalf("Dequeue() = %q, %v", payload, err) + } + + limiter := newIOPubRateLimiter(BridgeConfig{ + IOPubMessagesPerSec: 1, + IOPubBytesPerSec: 4, + RateWindow: time.Second, + }.normalized()) + if err := limiter.Allow(4); err != nil { + t.Fatal(err) + } + if err := limiter.Allow(1); !errors.Is(err, ErrIOPubRateLimit) { + t.Fatalf("Allow() error = %v, want rate limit", err) + } +} diff --git a/pkg/agent/jupyter/channel_json.go b/pkg/agent/jupyter/channel_json.go new file mode 100644 index 000000000..c4bcef533 --- /dev/null +++ b/pkg/agent/jupyter/channel_json.go @@ -0,0 +1,119 @@ +package jupyter + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" +) + +// ErrBinaryBuffersUnsupported is returned explicitly instead of silently +// dropping Jupyter comm/widget buffers that the current browser client cannot +// consume over its legacy JSON WebSocket protocol. +var ErrBinaryBuffersUnsupported = errors.New("jupyter binary buffers require the binary WebSocket subprotocol") + +// ParseChannelJSON converts the existing browser-facing Jupyter JSON shape to +// a wire message while preserving all fields inside the four protocol objects. +func ParseChannelJSON(payload []byte, limits Limits) (Channel, Message, error) { + limits = limits.normalized() + if len(payload) == 0 { + return "", Message{}, errors.New("empty Jupyter WebSocket message") + } + if len(payload) > limits.MaxMessageBytes { + return "", Message{}, fmt.Errorf("jupyter WebSocket payload exceeds limit %d", limits.MaxMessageBytes) + } + var object map[string]json.RawMessage + if err := json.Unmarshal(payload, &object); err != nil { + return "", Message{}, fmt.Errorf("parse Jupyter WebSocket JSON: %w", err) + } + if object == nil { + return "", Message{}, errors.New("jupyter WebSocket message must be an object") + } + + var channel Channel + if err := json.Unmarshal(object["channel"], &channel); err != nil || channel == "" { + return "", Message{}, errors.New("jupyter WebSocket message requires a channel") + } + switch channel { + case ChannelShell, ChannelControl, ChannelStdin: + case ChannelIOPub: + return "", Message{}, errors.New("browser messages cannot target Jupyter IOPub") + default: + return "", Message{}, fmt.Errorf("unsupported browser Jupyter channel %q", channel) + } + + if rawBuffers, ok := object["buffers"]; ok && len(bytes.TrimSpace(rawBuffers)) > 0 && !bytes.Equal(bytes.TrimSpace(rawBuffers), []byte("null")) { + var buffers []json.RawMessage + if err := json.Unmarshal(rawBuffers, &buffers); err != nil { + return "", Message{}, errors.New("invalid Jupyter WebSocket buffers") + } + if len(buffers) > 0 { + return "", Message{}, ErrBinaryBuffersUnsupported + } + } + + message := Message{ + Header: bytes.Clone(object["header"]), + ParentHeader: bytes.Clone(object["parent_header"]), + Metadata: bytes.Clone(object["metadata"]), + Content: bytes.Clone(object["content"]), + } + for name, raw := range map[string][]byte{ + "header": message.Header, + "parent_header": message.ParentHeader, + "metadata": message.Metadata, + "content": message.Content, + } { + if err := validateJSONObject(raw); err != nil { + return "", Message{}, fmt.Errorf("invalid browser Jupyter %s: %w", name, err) + } + } + return channel, message, nil +} + +// MarshalChannelJSON converts an authenticated wire message to the legacy +// Jupyter Server JSON WebSocket shape already consumed by Runme Web. +func MarshalChannelJSON(channel Channel, message Message, limits Limits) ([]byte, error) { + if channel != ChannelShell && channel != ChannelControl && channel != ChannelStdin && channel != ChannelIOPub { + return nil, fmt.Errorf("unsupported kernel Jupyter channel %q", channel) + } + if len(message.Buffers) > 0 { + return nil, ErrBinaryBuffersUnsupported + } + for name, raw := range map[string][]byte{ + "header": message.Header, + "parent_header": message.ParentHeader, + "metadata": message.Metadata, + "content": message.Content, + } { + if err := validateJSONObject(raw); err != nil { + return nil, fmt.Errorf("invalid kernel Jupyter %s: %w", name, err) + } + } + var header struct { + MessageID string `json:"msg_id"` + MessageType string `json:"msg_type"` + } + if err := json.Unmarshal(message.Header, &header); err != nil { + return nil, fmt.Errorf("parse kernel Jupyter header: %w", err) + } + object := map[string]any{ + "channel": channel, + "header": json.RawMessage(message.Header), + "parent_header": json.RawMessage(message.ParentHeader), + "metadata": json.RawMessage(message.Metadata), + "content": json.RawMessage(message.Content), + "buffers": []any{}, + "msg_id": header.MessageID, + "msg_type": header.MessageType, + } + payload, err := json.Marshal(object) + if err != nil { + return nil, fmt.Errorf("marshal Jupyter WebSocket JSON: %w", err) + } + limits = limits.normalized() + if len(payload) > limits.MaxMessageBytes { + return nil, fmt.Errorf("jupyter WebSocket payload exceeds limit %d", limits.MaxMessageBytes) + } + return payload, nil +} diff --git a/pkg/agent/jupyter/channel_json_test.go b/pkg/agent/jupyter/channel_json_test.go new file mode 100644 index 000000000..9c221060d --- /dev/null +++ b/pkg/agent/jupyter/channel_json_test.go @@ -0,0 +1,87 @@ +package jupyter + +import ( + "encoding/json" + "errors" + "reflect" + "testing" +) + +func TestChannelJSONRoundTripPreservesProtocolObjects(t *testing.T) { + payload := []byte(`{ + "channel":"shell", + "header":{"msg_id":"m1","msg_type":"custom_request","extension":{"future":true}}, + "parent_header":{}, + "metadata":{"unknown":[1,2,3]}, + "content":{"custom":{"nested":"value"}}, + "buffers":[], + "future_top_level":"ignored by Jupyter Session.send" +}`) + channel, message, err := ParseChannelJSON(payload, Limits{}) + if err != nil { + t.Fatalf("ParseChannelJSON() error = %v", err) + } + if channel != ChannelShell { + t.Fatalf("channel = %q, want shell", channel) + } + + encoded, err := MarshalChannelJSON(channel, message, Limits{}) + if err != nil { + t.Fatalf("MarshalChannelJSON() error = %v", err) + } + var got map[string]any + if err := json.Unmarshal(encoded, &got); err != nil { + t.Fatal(err) + } + if got["msg_id"] != "m1" || got["msg_type"] != "custom_request" { + t.Fatalf("top-level compatibility fields = %v/%v", got["msg_id"], got["msg_type"]) + } + var original map[string]json.RawMessage + if err := json.Unmarshal(payload, &original); err != nil { + t.Fatal(err) + } + var roundTrip map[string]json.RawMessage + if err := json.Unmarshal(encoded, &roundTrip); err != nil { + t.Fatal(err) + } + for _, field := range []string{"header", "parent_header", "metadata", "content"} { + var wantValue, gotValue any + if err := json.Unmarshal(original[field], &wantValue); err != nil { + t.Fatal(err) + } + if err := json.Unmarshal(roundTrip[field], &gotValue); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(gotValue, wantValue) { + t.Fatalf("%s = %#v, want %#v", field, gotValue, wantValue) + } + } +} + +func TestChannelJSONRejectsUnsupportedTraffic(t *testing.T) { + base := `"header":{},"parent_header":{},"metadata":{},"content":{}` + tests := map[string]string{ + "iopub target": `{"channel":"iopub",` + base + `}`, + "unknown channel": `{"channel":"comm",` + base + `}`, + "binary buffer": `{"channel":"shell",` + base + `,"buffers":["AA=="]}`, + "missing header": `{"channel":"shell","parent_header":{},"metadata":{},"content":{}}`, + } + for name, payload := range tests { + t.Run(name, func(t *testing.T) { + if _, _, err := ParseChannelJSON([]byte(payload), Limits{}); err == nil { + t.Fatal("ParseChannelJSON() unexpectedly succeeded") + } + }) + } + + _, err := MarshalChannelJSON(ChannelIOPub, Message{ + Header: []byte(`{"msg_id":"m","msg_type":"comm_msg"}`), + ParentHeader: []byte(`{}`), + Metadata: []byte(`{}`), + Content: []byte(`{}`), + Buffers: [][]byte{{1, 2, 3}}, + }, Limits{}) + if !errors.Is(err, ErrBinaryBuffersUnsupported) { + t.Fatalf("MarshalChannelJSON() error = %v, want ErrBinaryBuffersUnsupported", err) + } +} diff --git a/pkg/agent/jupyter/connection.go b/pkg/agent/jupyter/connection.go new file mode 100644 index 000000000..6b927b3e2 --- /dev/null +++ b/pkg/agent/jupyter/connection.go @@ -0,0 +1,144 @@ +// Package jupyter implements the bounded Jupyter wire protocol used between +// Runme and locally managed kernels. +package jupyter + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "net" + "os" + "strconv" + "strings" +) + +const maxConnectionFileBytes = 64 * 1024 + +// Channel identifies one of the five sockets in a Jupyter connection file. +type Channel string + +const ( + ChannelShell Channel = "shell" + ChannelIOPub Channel = "iopub" + ChannelStdin Channel = "stdin" + ChannelControl Channel = "control" + ChannelHeartbeat Channel = "heartbeat" +) + +// ConnectionInfo is the validated subset of a Jupyter connection file needed +// by the direct kernel bridge. Direct connections are intentionally restricted +// to loopback TCP endpoints. +type ConnectionInfo struct { + Transport string `json:"transport"` + IP string `json:"ip"` + ShellPort int `json:"shell_port"` + IOPubPort int `json:"iopub_port"` + StdinPort int `json:"stdin_port"` + ControlPort int `json:"control_port"` + HeartbeatPort int `json:"hb_port"` + Key string `json:"key"` + SignatureScheme string `json:"signature_scheme"` +} + +// LoadConnectionFile reads and validates a bounded Jupyter connection file. +func LoadConnectionFile(path string) (ConnectionInfo, error) { + if strings.TrimSpace(path) == "" { + return ConnectionInfo{}, errors.New("connection file path is required") + } + + f, err := os.Open(path) + if err != nil { + return ConnectionInfo{}, fmt.Errorf("open connection file: %w", err) + } + defer f.Close() + + raw, err := io.ReadAll(io.LimitReader(f, maxConnectionFileBytes+1)) + if err != nil { + return ConnectionInfo{}, fmt.Errorf("read connection file: %w", err) + } + return ParseConnectionInfo(raw) +} + +// ParseConnectionInfo parses and validates a Jupyter connection file payload. +func ParseConnectionInfo(raw []byte) (ConnectionInfo, error) { + if len(raw) == 0 { + return ConnectionInfo{}, errors.New("connection file is empty") + } + if len(raw) > maxConnectionFileBytes { + return ConnectionInfo{}, fmt.Errorf("connection file exceeds %d bytes", maxConnectionFileBytes) + } + + var info ConnectionInfo + if err := json.Unmarshal(raw, &info); err != nil { + return ConnectionInfo{}, fmt.Errorf("parse connection file: %w", err) + } + if err := info.Validate(); err != nil { + return ConnectionInfo{}, err + } + return info, nil +} + +// Validate checks that the connection data is complete and safe for a local +// managed kernel. +func (c ConnectionInfo) Validate() error { + if c.Transport != "tcp" { + return fmt.Errorf("unsupported Jupyter transport %q", c.Transport) + } + + ip := net.ParseIP(strings.TrimSpace(c.IP)) + if ip == nil { + return fmt.Errorf("invalid Jupyter IP %q", c.IP) + } + if !ip.IsLoopback() { + return fmt.Errorf("jupyter IP %q is not loopback", c.IP) + } + if c.SignatureScheme != SignatureSchemeHMACSHA256 { + return fmt.Errorf("unsupported Jupyter signature scheme %q", c.SignatureScheme) + } + + ports := []struct { + name string + port int + }{ + {"shell_port", c.ShellPort}, + {"iopub_port", c.IOPubPort}, + {"stdin_port", c.StdinPort}, + {"control_port", c.ControlPort}, + {"hb_port", c.HeartbeatPort}, + } + seen := make(map[int]string, len(ports)) + for _, item := range ports { + if item.port < 1 || item.port > 65535 { + return fmt.Errorf("%s must be between 1 and 65535", item.name) + } + if previous, ok := seen[item.port]; ok { + return fmt.Errorf("%s duplicates %s", item.name, previous) + } + seen[item.port] = item.name + } + return nil +} + +// Endpoint returns the tcp:// endpoint for a connection-file channel. +func (c ConnectionInfo) Endpoint(channel Channel) (string, error) { + var port int + switch channel { + case ChannelShell: + port = c.ShellPort + case ChannelIOPub: + port = c.IOPubPort + case ChannelStdin: + port = c.StdinPort + case ChannelControl: + port = c.ControlPort + case ChannelHeartbeat: + port = c.HeartbeatPort + default: + return "", fmt.Errorf("unknown Jupyter channel %q", channel) + } + if err := c.Validate(); err != nil { + return "", err + } + return "tcp://" + net.JoinHostPort(c.IP, strconv.Itoa(port)), nil +} diff --git a/pkg/agent/jupyter/connection_test.go b/pkg/agent/jupyter/connection_test.go new file mode 100644 index 000000000..c90182adf --- /dev/null +++ b/pkg/agent/jupyter/connection_test.go @@ -0,0 +1,60 @@ +package jupyter + +import ( + "strings" + "testing" +) + +const validConnectionJSON = `{ + "transport": "tcp", + "ip": "127.0.0.1", + "shell_port": 51001, + "iopub_port": 51002, + "stdin_port": 51003, + "control_port": 51004, + "hb_port": 51005, + "key": "secret", + "signature_scheme": "hmac-sha256" +}` + +func TestParseConnectionInfo(t *testing.T) { + info, err := ParseConnectionInfo([]byte(validConnectionJSON)) + if err != nil { + t.Fatalf("ParseConnectionInfo() error = %v", err) + } + got, err := info.Endpoint(ChannelShell) + if err != nil { + t.Fatalf("Endpoint() error = %v", err) + } + if want := "tcp://127.0.0.1:51001"; got != want { + t.Fatalf("Endpoint() = %q, want %q", got, want) + } +} + +func TestParseConnectionInfoRejectsInvalidFiles(t *testing.T) { + tests := map[string]string{ + "malformed JSON": `{`, + "unsupported transport": strings.Replace(validConnectionJSON, `"tcp"`, `"ipc"`, 1), + "non-loopback IP": strings.Replace(validConnectionJSON, `"127.0.0.1"`, `"192.0.2.10"`, 1), + "invalid IP": strings.Replace(validConnectionJSON, `"127.0.0.1"`, `"localhost"`, 1), + "bad signature scheme": strings.Replace(validConnectionJSON, `"hmac-sha256"`, `"hmac-sha1"`, 1), + "invalid port": strings.Replace(validConnectionJSON, `51001`, `0`, 1), + "duplicate port": strings.Replace(validConnectionJSON, `51002`, `51001`, 1), + } + for name, raw := range tests { + t.Run(name, func(t *testing.T) { + if _, err := ParseConnectionInfo([]byte(raw)); err == nil { + t.Fatal("ParseConnectionInfo() unexpectedly succeeded") + } + }) + } +} + +func FuzzParseConnectionInfo(f *testing.F) { + f.Add([]byte(validConnectionJSON)) + f.Add([]byte(`{}`)) + f.Add([]byte{0xff, 0x00, '{'}) + f.Fuzz(func(t *testing.T, raw []byte) { + _, _ = ParseConnectionInfo(raw) + }) +} diff --git a/pkg/agent/jupyter/ipykernel_integration_test.go b/pkg/agent/jupyter/ipykernel_integration_test.go new file mode 100644 index 000000000..5b97e171c --- /dev/null +++ b/pkg/agent/jupyter/ipykernel_integration_test.go @@ -0,0 +1,405 @@ +package jupyter + +import ( + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "reflect" + "runtime" + "strings" + "testing" + "time" +) + +const integrationTestTimeout = 30 * time.Second + +type testHeader struct { + MessageID string `json:"msg_id"` + MessageType string `json:"msg_type"` +} + +type receivedMessage struct { + channel Channel + wire Message + header testHeader + parent testHeader +} + +func TestIPyKernelCompatibility(t *testing.T) { + python := findIPyKernelPython(t) + ctx, cancel := context.WithTimeout(context.Background(), integrationTestTimeout) + defer cancel() + + connectionPath := filepath.Join(t.TempDir(), "kernel.json") + command := exec.CommandContext(ctx, python, "-m", "ipykernel_launcher", "-f", connectionPath) + command.Stdout = nil + command.Stderr = nil + if err := command.Start(); err != nil { + t.Fatalf("start ipykernel: %v", err) + } + processDone := make(chan error, 1) + go func() { + processDone <- command.Wait() + }() + defer stopTestKernel(t, command, processDone) + + connection := waitForConnectionFile(t, ctx, connectionPath, processDone) + signer, err := NewSigner(connection.Key, connection.SignatureScheme) + if err != nil { + t.Fatal(err) + } + factory := NewGoZeroMQFactory(Limits{}) + identity := []byte("runme-test-" + randomHex(t, 8)) + + heartbeat := openTestSocket(t, ctx, factory, SocketTypeReq, nil, connection, ChannelHeartbeat) + control := openTestSocket(t, ctx, factory, SocketTypeDealer, identity, connection, ChannelControl) + shell := openTestSocket(t, ctx, factory, SocketTypeDealer, identity, connection, ChannelShell) + iopub := openTestSocket(t, ctx, factory, SocketTypeSub, nil, connection, ChannelIOPub) + stdin := openTestSocket(t, ctx, factory, SocketTypeDealer, identity, connection, ChannelStdin) + defer heartbeat.Close() + defer control.Close() + defer shell.Close() + defer iopub.Close() + defer stdin.Close() + + heartbeatPayload := []byte("runme-heartbeat") + if err := heartbeat.SendMultipart([][]byte{heartbeatPayload}); err != nil { + t.Fatalf("send heartbeat: %v", err) + } + heartbeatReply, err := heartbeat.ReceiveMultipart() + if err != nil { + t.Fatalf("receive heartbeat: %v", err) + } + if !reflect.DeepEqual(heartbeatReply, [][]byte{heartbeatPayload}) { + t.Fatalf("heartbeat reply = %q, want echo %q", heartbeatReply, heartbeatPayload) + } + + sessionID := randomHex(t, 16) + infoRequest := newTestRequest(t, "kernel_info_request", sessionID, `{}`) + infoRequestID := messageID(t, infoRequest) + sendTestMessage(t, control, signer, infoRequest) + infoReply := receiveTestMessage(t, control, signer, ChannelControl) + if infoReply.header.MessageType != "kernel_info_reply" { + t.Fatalf("control reply type = %q, want kernel_info_reply", infoReply.header.MessageType) + } + if infoReply.parent.MessageID != infoRequestID { + t.Fatalf("control parent msg_id = %q, want %q", infoReply.parent.MessageID, infoRequestID) + } + + messages := make(chan receivedMessage, 64) + errorsCh := make(chan error, 2) + go readTestMessages(ctx, shell, signer, ChannelShell, messages, errorsCh) + go readTestMessages(ctx, iopub, signer, ChannelIOPub, messages, errorsCh) + + first := executeTestCode(t, ctx, shell, signer, sessionID, + "shared_value = 41\nprint('runme-stream')\nshared_value + 1", messages, errorsCh) + for _, messageType := range []string{"status:busy", "execute_input", "stream", "execute_result", "result:42", "execute_reply", "status:idle"} { + if !first[messageType] { + t.Errorf("first execution did not observe %s; got %v", messageType, first) + } + } + + second := executeTestCode(t, ctx, shell, signer, sessionID, + "shared_value + 2", messages, errorsCh) + for _, messageType := range []string{"execute_result", "result:43", "execute_reply", "status:idle"} { + if !second[messageType] { + t.Errorf("second execution did not observe %s; got %v", messageType, second) + } + } + + shutdown := newTestRequest(t, "shutdown_request", sessionID, `{"restart":false}`) + sendTestMessage(t, control, signer, shutdown) + shutdownReply := receiveTestMessage(t, control, signer, ChannelControl) + if shutdownReply.header.MessageType != "shutdown_reply" { + t.Fatalf("control reply type = %q, want shutdown_reply", shutdownReply.header.MessageType) + } + + select { + case err := <-processDone: + if err != nil { + t.Fatalf("ipykernel exited after shutdown with error: %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("ipykernel did not exit after shutdown_request") + } +} + +func findIPyKernelPython(t *testing.T) string { + t.Helper() + if configured := strings.TrimSpace(os.Getenv("RUNME_TEST_PYTHON")); configured != "" { + if err := exec.Command(configured, "-c", "import ipykernel").Run(); err != nil { + t.Fatalf("RUNME_TEST_PYTHON does not provide ipykernel: %v", err) + } + return configured + } + python, err := exec.LookPath("python3") + if err != nil || exec.Command(python, "-c", "import ipykernel").Run() != nil { + t.Skip("ipykernel integration test requires RUNME_TEST_PYTHON or python3 with ipykernel") + } + return python +} + +func waitForConnectionFile(t *testing.T, ctx context.Context, path string, processDone <-chan error) ConnectionInfo { + t.Helper() + ticker := time.NewTicker(20 * time.Millisecond) + defer ticker.Stop() + var lastErr error + for { + connection, err := LoadConnectionFile(path) + if err == nil { + return connection + } + lastErr = err + select { + case <-ctx.Done(): + t.Fatalf("wait for ipykernel connection file: %v (last error: %v)", ctx.Err(), lastErr) + case err := <-processDone: + t.Fatalf("ipykernel exited before connection file was ready: %v (last error: %v)", err, lastErr) + case <-ticker.C: + } + } +} + +func openTestSocket( + t *testing.T, + ctx context.Context, + factory SocketFactory, + socketType SocketType, + identity []byte, + connection ConnectionInfo, + channel Channel, +) Socket { + t.Helper() + socket, err := factory.NewSocket(ctx, socketType, identity) + if err != nil { + t.Fatalf("create %s socket: %v", channel, err) + } + endpoint, err := connection.Endpoint(channel) + if err != nil { + _ = socket.Close() + t.Fatalf("resolve %s endpoint: %v", channel, err) + } + if err := socket.Dial(endpoint); err != nil { + _ = socket.Close() + t.Fatalf("dial %s endpoint: %v", channel, err) + } + return socket +} + +func newTestRequest(t *testing.T, messageType, sessionID, content string) Message { + t.Helper() + header, err := json.Marshal(map[string]string{ + "msg_id": randomHex(t, 16), + "username": "runme-test", + "session": sessionID, + "date": time.Now().UTC().Format(time.RFC3339Nano), + "msg_type": messageType, + "version": "5.3", + }) + if err != nil { + t.Fatal(err) + } + return Message{ + Header: header, + ParentHeader: []byte(`{}`), + Metadata: []byte(`{}`), + Content: []byte(content), + } +} + +func messageID(t *testing.T, message Message) string { + t.Helper() + var header testHeader + if err := json.Unmarshal(message.Header, &header); err != nil { + t.Fatal(err) + } + return header.MessageID +} + +func sendTestMessage(t *testing.T, socket Socket, signer *Signer, message Message) { + t.Helper() + frames, err := MarshalMultipart(message, signer, Limits{}) + if err != nil { + t.Fatalf("marshal Jupyter message: %v", err) + } + if err := socket.SendMultipart(frames); err != nil { + t.Fatalf("send Jupyter message: %v", err) + } +} + +func receiveTestMessage(t *testing.T, socket Socket, signer *Signer, channel Channel) receivedMessage { + t.Helper() + frames, err := socket.ReceiveMultipart() + if err != nil { + t.Fatalf("receive %s message: %v", channel, err) + } + message, err := decodeTestMessage(frames, signer, channel) + if err != nil { + t.Fatalf("decode %s message: %v", channel, err) + } + return message +} + +func decodeTestMessage(frames [][]byte, signer *Signer, channel Channel) (receivedMessage, error) { + wire, err := ParseMultipart(frames, signer, Limits{}) + if err != nil { + return receivedMessage{}, err + } + var header, parent testHeader + if err := json.Unmarshal(wire.Header, &header); err != nil { + return receivedMessage{}, fmt.Errorf("parse header: %w", err) + } + if err := json.Unmarshal(wire.ParentHeader, &parent); err != nil { + return receivedMessage{}, fmt.Errorf("parse parent header: %w", err) + } + return receivedMessage{channel: channel, wire: wire, header: header, parent: parent}, nil +} + +func readTestMessages( + ctx context.Context, + socket Socket, + signer *Signer, + channel Channel, + messages chan<- receivedMessage, + errorsCh chan<- error, +) { + for { + frames, err := socket.ReceiveMultipart() + if err != nil { + if ctx.Err() == nil { + select { + case errorsCh <- fmt.Errorf("receive %s: %w", channel, err): + case <-ctx.Done(): + } + } + return + } + message, err := decodeTestMessage(frames, signer, channel) + if err != nil { + select { + case errorsCh <- fmt.Errorf("decode %s: %w", channel, err): + case <-ctx.Done(): + } + return + } + select { + case messages <- message: + case <-ctx.Done(): + return + } + } +} + +func executeTestCode( + t *testing.T, + ctx context.Context, + shell Socket, + signer *Signer, + sessionID, code string, + messages <-chan receivedMessage, + errorsCh <-chan error, +) map[string]bool { + t.Helper() + content, err := json.Marshal(map[string]any{ + "code": code, + "silent": false, + "store_history": true, + "user_expressions": map[string]any{}, + "allow_stdin": false, + "stop_on_error": true, + }) + if err != nil { + t.Fatal(err) + } + request := newTestRequest(t, "execute_request", sessionID, string(content)) + requestID := messageID(t, request) + sendTestMessage(t, shell, signer, request) + + seen := make(map[string]bool) + for !(seen["execute_reply"] && seen["status:idle"]) { + select { + case <-ctx.Done(): + t.Fatalf("wait for execution %s: %v; seen %v", requestID, ctx.Err(), seen) + case err := <-errorsCh: + t.Fatalf("read kernel messages: %v; seen %v", err, seen) + case message := <-messages: + if message.parent.MessageID != requestID { + continue + } + switch message.header.MessageType { + case "status": + var content struct { + ExecutionState string `json:"execution_state"` + } + if err := json.Unmarshal(message.wire.Content, &content); err != nil { + t.Fatalf("parse status content: %v", err) + } + seen["status:"+content.ExecutionState] = true + case "stream": + var content struct { + Text string `json:"text"` + } + if err := json.Unmarshal(message.wire.Content, &content); err != nil { + t.Fatalf("parse stream content: %v", err) + } + if strings.Contains(content.Text, "runme-stream") { + seen["stream"] = true + } + case "execute_result": + var content struct { + Data map[string]any `json:"data"` + } + if err := json.Unmarshal(message.wire.Content, &content); err != nil { + t.Fatalf("parse execute_result content: %v", err) + } + if _, ok := content.Data["text/plain"]; ok { + seen["execute_result"] = true + seen["result:"+fmt.Sprint(content.Data["text/plain"])] = true + } + default: + seen[message.header.MessageType] = true + } + } + } + return seen +} + +func randomHex(t *testing.T, bytesCount int) string { + t.Helper() + raw := make([]byte, bytesCount) + if _, err := rand.Read(raw); err != nil { + t.Fatal(err) + } + return hex.EncodeToString(raw) +} + +func stopTestKernel(t *testing.T, command *exec.Cmd, processDone <-chan error) { + t.Helper() + if command.ProcessState != nil && command.ProcessState.Exited() { + return + } + select { + case <-processDone: + return + default: + } + if command.Process == nil { + return + } + if err := command.Process.Kill(); err != nil && !errors.Is(err, os.ErrProcessDone) { + t.Errorf("kill test ipykernel: %v", err) + } + select { + case <-processDone: + case <-time.After(5 * time.Second): + if runtime.GOOS != "windows" { + t.Error("test ipykernel was not reaped after kill") + } + } +} diff --git a/pkg/agent/jupyter/kernel_manager.go b/pkg/agent/jupyter/kernel_manager.go new file mode 100644 index 000000000..df8178540 --- /dev/null +++ b/pkg/agent/jupyter/kernel_manager.go @@ -0,0 +1,800 @@ +package jupyter + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "sort" + "strings" + "sync" + "time" + + "github.com/google/uuid" +) + +const connectionFilePlaceholder = "{connection_file}" + +// KernelState describes the lifecycle of a locally managed kernel. +type KernelState string + +const ( + KernelStateStarting KernelState = "starting" + KernelStateIdle KernelState = "idle" + KernelStateBusy KernelState = "busy" + KernelStateInterrupting KernelState = "interrupting" + KernelStateRestarting KernelState = "restarting" + KernelStateStopping KernelState = "stopping" + KernelStateDead KernelState = "dead" +) + +// KernelModel is compatible with the Jupyter Server kernel model fields used +// by Runme Web. +type KernelModel struct { + ID string `json:"id"` + Name string `json:"name"` + LastActivity string `json:"last_activity"` + ExecutionState string `json:"execution_state"` + Connections int `json:"connections"` +} + +// KernelLaunchSpec describes one client-requested kernel process. Argv follows +// the Jupyter kernelspec convention and must contain exactly one +// {connection_file} placeholder. The server replaces that placeholder with an +// owner-only connection file before launching the process directly, without a +// shell. +type KernelLaunchSpec struct { + Name string + Argv []string +} + +// KernelManagerConfig configures local kernel lifecycle boundaries. +type KernelManagerConfig struct { + RuntimeDir string + SocketFactory SocketFactory + StartupTimeout time.Duration + ReadinessTimeout time.Duration + GracefulTimeout time.Duration + TerminateTimeout time.Duration + KillTimeout time.Duration + DiagnosticsBytes int +} + +// KernelManager owns all local kernel processes and immutable connection +// generations independently of HTTP routing. +type KernelManager struct { + mu sync.RWMutex + config KernelManagerConfig + kernels map[string]*managedKernel + ctx context.Context + cancel context.CancelFunc + closed bool +} + +type managedKernel struct { + lifecycle sync.Mutex + id string + name string + state KernelState + lastActivity time.Time + connections int + generation uint64 + connection *kernelConnection + command *exec.Cmd + exited chan struct{} + exitErr error + runtimeDir string + diagnostics *boundedWriter + launchSpec KernelLaunchSpec +} + +type kernelConnection struct { + Generation uint64 + Info ConnectionInfo + Context context.Context + cancel context.CancelFunc +} + +// NewKernelManager creates a manager and prepares an owner-only runtime +// directory. +func NewKernelManager(config KernelManagerConfig) (*KernelManager, error) { + if strings.TrimSpace(config.RuntimeDir) == "" { + return nil, errors.New("jupyter kernel runtime directory is required") + } + if config.SocketFactory == nil { + config.SocketFactory = NewGoZeroMQFactory(Limits{}) + } + if config.StartupTimeout <= 0 { + config.StartupTimeout = 15 * time.Second + } + if config.ReadinessTimeout <= 0 { + config.ReadinessTimeout = 10 * time.Second + } + if config.GracefulTimeout <= 0 { + config.GracefulTimeout = 3 * time.Second + } + if config.TerminateTimeout <= 0 { + config.TerminateTimeout = 2 * time.Second + } + if config.KillTimeout <= 0 { + config.KillTimeout = 2 * time.Second + } + if config.DiagnosticsBytes <= 0 { + config.DiagnosticsBytes = 64 * 1024 + } + if err := os.MkdirAll(config.RuntimeDir, 0o700); err != nil { + return nil, fmt.Errorf("create Jupyter runtime directory: %w", err) + } + if err := os.Chmod(config.RuntimeDir, 0o700); err != nil { + return nil, fmt.Errorf("secure Jupyter runtime directory: %w", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + return &KernelManager{ + config: config, + kernels: make(map[string]*managedKernel), + ctx: ctx, + cancel: cancel, + }, nil +} + +func validateKernelLaunchSpec(spec KernelLaunchSpec) error { + if strings.TrimSpace(spec.Name) == "" { + return errors.New("jupyter kernel name is required") + } + if len(spec.Argv) == 0 { + return errors.New("jupyter kernel argv is required") + } + if len(spec.Argv) > 128 { + return errors.New("jupyter kernel argv exceeds 128 entries") + } + for i, arg := range spec.Argv { + if strings.IndexByte(arg, 0) >= 0 { + return fmt.Errorf("jupyter kernel argv[%d] contains a NUL byte", i) + } + if len(arg) > 8*1024 { + return fmt.Errorf("jupyter kernel argv[%d] exceeds 8192 bytes", i) + } + } + if strings.TrimSpace(spec.Argv[0]) == "" { + return errors.New("jupyter kernel executable is required") + } + if strings.Contains(spec.Argv[0], connectionFilePlaceholder) { + return errors.New("jupyter kernel executable cannot contain the connection-file placeholder") + } + placeholderCount := 0 + for _, arg := range spec.Argv[1:] { + placeholderCount += strings.Count(arg, connectionFilePlaceholder) + } + if placeholderCount != 1 { + return fmt.Errorf("jupyter kernel argv must contain exactly one %s placeholder", connectionFilePlaceholder) + } + return nil +} + +// Start launches a validated client-requested process and returns only after +// heartbeat and kernel_info readiness checks succeed. +func (m *KernelManager) Start(ctx context.Context, spec KernelLaunchSpec) (KernelModel, error) { + started := time.Now() + defer func() { observeLifecycle("start", started) }() + if err := validateKernelLaunchSpec(spec); err != nil { + return KernelModel{}, err + } + spec.Name = strings.TrimSpace(spec.Name) + spec.Argv = append([]string(nil), spec.Argv...) + kernel := &managedKernel{ + id: uuid.NewString(), + name: spec.Name, + state: KernelStateStarting, + lastActivity: time.Now().UTC(), + launchSpec: spec, + } + kernel.lifecycle.Lock() + defer kernel.lifecycle.Unlock() + + m.mu.Lock() + if m.closed { + m.mu.Unlock() + return KernelModel{}, errors.New("jupyter kernel manager is closed") + } + m.kernels[kernel.id] = kernel + addKernelState(kernel.state, 1) + m.mu.Unlock() + + if err := m.launch(ctx, kernel, spec); err != nil { + _ = m.stopProcess(context.Background(), kernel, false) + m.mu.Lock() + addKernelState(kernel.state, -1) + delete(m.kernels, kernel.id) + m.mu.Unlock() + return KernelModel{}, fmt.Errorf("start Jupyter kernel: %w", err) + } + return m.model(kernel.id) +} + +// List returns stable snapshots sorted by opaque kernel ID. +func (m *KernelManager) List() []KernelModel { + m.mu.RLock() + models := make([]KernelModel, 0, len(m.kernels)) + for _, kernel := range m.kernels { + models = append(models, kernelModel(kernel)) + } + m.mu.RUnlock() + sort.Slice(models, func(i, j int) bool { return models[i].ID < models[j].ID }) + return models +} + +// Get returns one kernel snapshot. +func (m *KernelManager) Get(id string) (KernelModel, error) { + return m.model(id) +} + +func (m *KernelManager) model(id string) (KernelModel, error) { + m.mu.RLock() + defer m.mu.RUnlock() + kernel, ok := m.kernels[id] + if !ok { + return KernelModel{}, fmt.Errorf("jupyter kernel %q not found", id) + } + return kernelModel(kernel), nil +} + +func kernelModel(kernel *managedKernel) KernelModel { + return KernelModel{ + ID: kernel.id, + Name: kernel.name, + LastActivity: kernel.lastActivity.Format(time.RFC3339Nano), + ExecutionState: string(kernel.state), + Connections: kernel.connections, + } +} + +func (m *KernelManager) launch(ctx context.Context, kernel *managedKernel, spec KernelLaunchSpec) error { + runtimeDir, err := os.MkdirTemp(m.config.RuntimeDir, "kernel-") + if err != nil { + return fmt.Errorf("create kernel runtime directory: %w", err) + } + if err := os.Chmod(runtimeDir, 0o700); err != nil { + _ = os.RemoveAll(runtimeDir) + return fmt.Errorf("secure kernel runtime directory: %w", err) + } + connectionPath := filepath.Join(runtimeDir, "connection.json") + argv := make([]string, len(spec.Argv)) + for i, arg := range spec.Argv { + argv[i] = strings.ReplaceAll(arg, connectionFilePlaceholder, connectionPath) + } + diagnostics := newBoundedWriter(m.config.DiagnosticsBytes) + command := exec.Command(argv[0], argv[1:]...) + command.Stdout = diagnostics + command.Stderr = diagnostics + configureManagedProcess(command) + if err := command.Start(); err != nil { + _ = os.RemoveAll(runtimeDir) + return fmt.Errorf("launch kernel %q: %w", spec.Name, err) + } + + exited := make(chan struct{}) + m.mu.Lock() + kernel.command = command + kernel.exited = exited + kernel.exitErr = nil + kernel.runtimeDir = runtimeDir + kernel.diagnostics = diagnostics + m.setKernelStateLocked(kernel, KernelStateStarting) + kernel.lastActivity = time.Now().UTC() + m.mu.Unlock() + go m.waitForProcess(kernel, command, exited) + + startupCtx, cancel := context.WithTimeout(ctx, m.config.StartupTimeout) + defer cancel() + connection, err := waitForManagedConnection(startupCtx, connectionPath, exited) + if err != nil { + return fmt.Errorf("wait for connection file: %w; diagnostics: %s", err, diagnostics.String()) + } + if err := os.Chmod(connectionPath, 0o600); err != nil { + return fmt.Errorf("secure connection file: %w", err) + } + + readinessCtx, readinessCancel := context.WithTimeout(startupCtx, m.config.ReadinessTimeout) + err = probeKernelReadiness(readinessCtx, m.config.SocketFactory, connection) + readinessCancel() + if err != nil { + return fmt.Errorf("kernel readiness failed: %w; diagnostics: %s", err, diagnostics.String()) + } + + generationCtx, generationCancel := context.WithCancel(m.ctx) + m.mu.Lock() + if kernel.command != command || kernel.state == KernelStateDead { + m.mu.Unlock() + generationCancel() + return errors.New("kernel exited during readiness") + } + kernel.generation++ + kernel.connection = &kernelConnection{ + Generation: kernel.generation, + Info: connection, + Context: generationCtx, + cancel: generationCancel, + } + m.setKernelStateLocked(kernel, KernelStateIdle) + kernel.lastActivity = time.Now().UTC() + m.mu.Unlock() + return nil +} + +func (m *KernelManager) waitForProcess(kernel *managedKernel, command *exec.Cmd, exited chan struct{}) { + err := command.Wait() + m.mu.Lock() + defer m.mu.Unlock() + if kernel.command != command { + close(exited) + return + } + kernel.exitErr = err + close(exited) + if kernel.state != KernelStateStopping && kernel.state != KernelStateRestarting { + m.setKernelStateLocked(kernel, KernelStateDead) + addCounter(jupyterUnexpectedExits) + kernel.lastActivity = time.Now().UTC() + if kernel.connection != nil { + kernel.connection.cancel() + kernel.connection = nil + } + } +} + +func waitForManagedConnection(ctx context.Context, path string, exited <-chan struct{}) (ConnectionInfo, error) { + ticker := time.NewTicker(20 * time.Millisecond) + defer ticker.Stop() + var lastErr error + for { + connection, err := LoadConnectionFile(path) + if err == nil { + return connection, nil + } + lastErr = err + select { + case <-ctx.Done(): + return ConnectionInfo{}, fmt.Errorf("%w (last connection error: %v)", ctx.Err(), lastErr) + case <-exited: + return ConnectionInfo{}, fmt.Errorf("kernel process exited (last connection error: %v)", lastErr) + case <-ticker.C: + } + } +} + +// Interrupt sends a process-group interrupt without destroying the kernel. +func (m *KernelManager) Interrupt(id string) error { + started := time.Now() + defer func() { observeLifecycle("interrupt", started) }() + m.mu.RLock() + kernel, ok := m.kernels[id] + m.mu.RUnlock() + if !ok { + return fmt.Errorf("jupyter kernel %q not found", id) + } + kernel.lifecycle.Lock() + defer kernel.lifecycle.Unlock() + + m.mu.Lock() + kernel, ok = m.kernels[id] + if !ok { + m.mu.Unlock() + return fmt.Errorf("jupyter kernel %q not found", id) + } + if kernel.command == nil || kernel.command.Process == nil || kernel.state == KernelStateDead { + m.mu.Unlock() + return fmt.Errorf("jupyter kernel %q is not running", id) + } + m.setKernelStateLocked(kernel, KernelStateInterrupting) + kernel.lastActivity = time.Now().UTC() + command := kernel.command + m.mu.Unlock() + + if err := interruptManagedProcess(command); err != nil { + return fmt.Errorf("interrupt Jupyter kernel: %w", err) + } + return nil +} + +// Restart retains the public kernel ID while replacing its process, ports, +// signing key, and immutable connection generation. +func (m *KernelManager) Restart(ctx context.Context, id string) (KernelModel, error) { + started := time.Now() + defer func() { observeLifecycle("restart", started) }() + m.mu.RLock() + kernel, ok := m.kernels[id] + m.mu.RUnlock() + if !ok { + return KernelModel{}, fmt.Errorf("jupyter kernel %q not found", id) + } + kernel.lifecycle.Lock() + defer kernel.lifecycle.Unlock() + + m.mu.Lock() + kernel, ok = m.kernels[id] + if !ok { + m.mu.Unlock() + return KernelModel{}, fmt.Errorf("jupyter kernel %q not found", id) + } + spec := kernel.launchSpec + m.setKernelStateLocked(kernel, KernelStateRestarting) + kernel.lastActivity = time.Now().UTC() + if kernel.connection != nil { + kernel.connection.cancel() + kernel.connection = nil + } + m.mu.Unlock() + + if err := m.stopProcess(ctx, kernel, true); err != nil { + return KernelModel{}, fmt.Errorf("stop old Jupyter kernel generation: %w", err) + } + if err := m.launch(ctx, kernel, spec); err != nil { + m.mu.Lock() + m.setKernelStateLocked(kernel, KernelStateDead) + kernel.lastActivity = time.Now().UTC() + m.mu.Unlock() + return KernelModel{}, fmt.Errorf("start new Jupyter kernel generation: %w", err) + } + return m.model(id) +} + +// Stop shuts down, reaps, and removes a managed kernel. +func (m *KernelManager) Stop(ctx context.Context, id string) error { + started := time.Now() + defer func() { observeLifecycle("stop", started) }() + m.mu.RLock() + kernel, ok := m.kernels[id] + m.mu.RUnlock() + if !ok { + return fmt.Errorf("jupyter kernel %q not found", id) + } + kernel.lifecycle.Lock() + defer kernel.lifecycle.Unlock() + + m.mu.Lock() + kernel, ok = m.kernels[id] + if !ok { + m.mu.Unlock() + return fmt.Errorf("jupyter kernel %q not found", id) + } + m.setKernelStateLocked(kernel, KernelStateStopping) + kernel.lastActivity = time.Now().UTC() + if kernel.connection != nil { + kernel.connection.cancel() + } + m.mu.Unlock() + + if err := m.stopProcess(ctx, kernel, true); err != nil { + return err + } + m.mu.Lock() + addKernelState(kernel.state, -1) + delete(m.kernels, id) + m.mu.Unlock() + return nil +} + +func (m *KernelManager) stopProcess(ctx context.Context, kernel *managedKernel, graceful bool) error { + m.mu.RLock() + command := kernel.command + exited := kernel.exited + connection := kernel.connection + runtimeDir := kernel.runtimeDir + m.mu.RUnlock() + if command == nil || command.Process == nil || exited == nil { + if runtimeDir != "" { + _ = os.RemoveAll(runtimeDir) + } + return nil + } + + if graceful && connection != nil { + shutdownCtx, cancel := context.WithTimeout(ctx, m.config.GracefulTimeout) + _ = requestKernelShutdown(shutdownCtx, m.config.SocketFactory, connection.Info) + cancel() + } + if graceful && waitForExit(exited, m.config.GracefulTimeout) { + _ = os.RemoveAll(runtimeDir) + return nil + } + select { + case <-exited: + _ = os.RemoveAll(runtimeDir) + return nil + default: + } + if err := terminateManagedProcess(command); err != nil && !isProcessDoneError(err) { + return fmt.Errorf("terminate Jupyter kernel: %w", err) + } + if waitForExit(exited, m.config.TerminateTimeout) { + _ = os.RemoveAll(runtimeDir) + return nil + } + if err := killManagedProcess(command); err != nil && !isProcessDoneError(err) { + return fmt.Errorf("kill Jupyter kernel: %w", err) + } + if !waitForExit(exited, m.config.KillTimeout) { + return errors.New("timed out reaping Jupyter kernel after kill") + } + _ = os.RemoveAll(runtimeDir) + return nil +} + +func waitForExit(exited <-chan struct{}, timeout time.Duration) bool { + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case <-exited: + return true + case <-timer.C: + return false + } +} + +// Connect obtains an immutable connection generation and increments the public +// connection count. The release function is idempotent. +func (m *KernelManager) Connect(id string) (kernelConnection, func(), error) { + m.mu.Lock() + kernel, ok := m.kernels[id] + if !ok || kernel.connection == nil || kernel.state == KernelStateDead || kernel.state == KernelStateStopping { + m.mu.Unlock() + return kernelConnection{}, nil, fmt.Errorf("jupyter kernel %q is not connectable", id) + } + connection := *kernel.connection + kernel.connections++ + kernel.lastActivity = time.Now().UTC() + m.mu.Unlock() + + var once sync.Once + release := func() { + once.Do(func() { + m.mu.Lock() + if current, exists := m.kernels[id]; exists && current.connections > 0 { + current.connections-- + current.lastActivity = time.Now().UTC() + } + m.mu.Unlock() + }) + } + return connection, release, nil +} + +// SetExecutionState records authenticated IOPub busy/idle transitions for the +// current immutable generation only. +func (m *KernelManager) SetExecutionState(id string, generation uint64, state KernelState) { + if state != KernelStateBusy && state != KernelStateIdle { + return + } + m.mu.Lock() + defer m.mu.Unlock() + kernel, ok := m.kernels[id] + if !ok || kernel.connection == nil || kernel.connection.Generation != generation { + return + } + m.setKernelStateLocked(kernel, state) + kernel.lastActivity = time.Now().UTC() +} + +func (m *KernelManager) setKernelStateLocked(kernel *managedKernel, state KernelState) { + if kernel.state == state { + return + } + addKernelState(kernel.state, -1) + kernel.state = state + addKernelState(kernel.state, 1) +} + +// Diagnostics returns bounded process output for local troubleshooting. +func (m *KernelManager) Diagnostics(id string) (string, error) { + m.mu.RLock() + defer m.mu.RUnlock() + kernel, ok := m.kernels[id] + if !ok { + return "", fmt.Errorf("jupyter kernel %q not found", id) + } + if kernel.diagnostics == nil { + return "", nil + } + return kernel.diagnostics.String(), nil +} + +// Close stops and reaps every managed kernel. It is safe to call repeatedly. +func (m *KernelManager) Close(ctx context.Context) error { + m.mu.Lock() + if m.closed { + m.mu.Unlock() + return nil + } + m.closed = true + ids := make([]string, 0, len(m.kernels)) + for id := range m.kernels { + ids = append(ids, id) + } + m.mu.Unlock() + + var errs []error + for _, id := range ids { + if err := m.Stop(ctx, id); err != nil { + errs = append(errs, err) + } + } + m.cancel() + return errors.Join(errs...) +} + +func probeKernelReadiness(ctx context.Context, factory SocketFactory, connection ConnectionInfo) error { + var lastErr error + for { + attemptCtx, cancel := context.WithTimeout(ctx, 750*time.Millisecond) + if err := heartbeatRoundTrip(attemptCtx, factory, connection); err != nil { + lastErr = err + } else if err := requestReply(attemptCtx, factory, connection, ChannelControl, "kernel_info_request", "kernel_info_reply", `{}`); err != nil { + lastErr = err + } else { + cancel() + return nil + } + cancel() + select { + case <-ctx.Done(): + return fmt.Errorf("jupyter readiness deadline exceeded: %w (last error: %v)", ctx.Err(), lastErr) + case <-time.After(100 * time.Millisecond): + } + } +} + +func requestKernelShutdown(ctx context.Context, factory SocketFactory, connection ConnectionInfo) error { + return requestReply(ctx, factory, connection, ChannelControl, "shutdown_request", "shutdown_reply", `{"restart":false}`) +} + +func heartbeatRoundTrip(ctx context.Context, factory SocketFactory, connection ConnectionInfo) error { + socket, err := factory.NewSocket(ctx, SocketTypeReq, nil) + if err != nil { + return err + } + defer socket.Close() + endpoint, err := connection.Endpoint(ChannelHeartbeat) + if err != nil { + return err + } + if err := socket.Dial(endpoint); err != nil { + return err + } + payload := []byte("runme-ready-" + randomIdentifier()) + if err := socket.SendMultipart([][]byte{payload}); err != nil { + return err + } + reply, err := socket.ReceiveMultipart() + if err != nil { + return err + } + if len(reply) != 1 || !bytes.Equal(reply[0], payload) { + return errors.New("unexpected Jupyter heartbeat reply") + } + return nil +} + +func requestReply( + ctx context.Context, + factory SocketFactory, + connection ConnectionInfo, + channel Channel, + requestType, replyType, content string, +) error { + identity := []byte("runme-" + randomIdentifier()) + socket, err := factory.NewSocket(ctx, SocketTypeDealer, identity) + if err != nil { + return err + } + defer socket.Close() + endpoint, err := connection.Endpoint(channel) + if err != nil { + return err + } + if err := socket.Dial(endpoint); err != nil { + return err + } + signer, err := NewSigner(connection.Key, connection.SignatureScheme) + if err != nil { + return err + } + request, requestID, err := newProtocolRequest(requestType, content) + if err != nil { + return err + } + frames, err := MarshalMultipart(request, signer, Limits{}) + if err != nil { + return err + } + if err := socket.SendMultipart(frames); err != nil { + return err + } + replyFrames, err := socket.ReceiveMultipart() + if err != nil { + return err + } + reply, err := ParseMultipart(replyFrames, signer, Limits{}) + if err != nil { + return err + } + var header, parent struct { + MessageID string `json:"msg_id"` + MessageType string `json:"msg_type"` + } + if err := json.Unmarshal(reply.Header, &header); err != nil { + return err + } + if err := json.Unmarshal(reply.ParentHeader, &parent); err != nil { + return err + } + if header.MessageType != replyType || parent.MessageID != requestID { + return fmt.Errorf("unexpected Jupyter reply type %q or parent %q", header.MessageType, parent.MessageID) + } + return nil +} + +func newProtocolRequest(messageType, content string) (Message, string, error) { + messageID := randomIdentifier() + header, err := json.Marshal(map[string]string{ + "msg_id": messageID, + "username": "runme", + "session": randomIdentifier(), + "date": time.Now().UTC().Format(time.RFC3339Nano), + "msg_type": messageType, + "version": "5.3", + }) + if err != nil { + return Message{}, "", err + } + message := Message{ + Header: header, + ParentHeader: []byte(`{}`), + Metadata: []byte(`{}`), + Content: []byte(content), + } + if err := validateJSONObject(message.Content); err != nil { + return Message{}, "", err + } + return message, messageID, nil +} + +func randomIdentifier() string { + raw := make([]byte, 16) + if _, err := rand.Read(raw); err != nil { + return strings.ReplaceAll(uuid.NewString(), "-", "") + } + return hex.EncodeToString(raw) +} + +type boundedWriter struct { + mu sync.Mutex + capacity int + data []byte +} + +func newBoundedWriter(capacity int) *boundedWriter { + return &boundedWriter{capacity: capacity} +} + +func (w *boundedWriter) Write(p []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + w.data = append(w.data, p...) + if len(w.data) > w.capacity { + w.data = append([]byte(nil), w.data[len(w.data)-w.capacity:]...) + } + return len(p), nil +} + +func (w *boundedWriter) String() string { + w.mu.Lock() + defer w.mu.Unlock() + return strings.TrimSpace(string(w.data)) +} diff --git a/pkg/agent/jupyter/kernel_manager_test.go b/pkg/agent/jupyter/kernel_manager_test.go new file mode 100644 index 000000000..4b99cecf5 --- /dev/null +++ b/pkg/agent/jupyter/kernel_manager_test.go @@ -0,0 +1,251 @@ +package jupyter + +import ( + "context" + "errors" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" +) + +func TestKernelManagerStartRestartStop(t *testing.T) { + manager := newTestKernelManager(t) + defer manager.Close(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + + started, err := manager.Start(ctx, testPythonLaunchSpec(t)) + if err != nil { + t.Fatalf("Start() error = %v", err) + } + if started.ExecutionState != string(KernelStateIdle) { + t.Fatalf("Start() state = %q, want idle", started.ExecutionState) + } + if started.ID == "" || started.Name != "python3" { + t.Fatalf("Start() model = %+v", started) + } + + firstConnection, release, err := manager.Connect(started.ID) + if err != nil { + t.Fatalf("Connect() error = %v", err) + } + if got, _ := manager.Get(started.ID); got.Connections != 1 { + t.Fatalf("connections = %d, want 1", got.Connections) + } + release() + release() + if got, _ := manager.Get(started.ID); got.Connections != 0 { + t.Fatalf("connections after idempotent release = %d, want 0", got.Connections) + } + + firstCommand, firstRuntimeDir := testKernelProcess(manager, started.ID) + assertOwnerOnlyPath(t, firstRuntimeDir, 0o700) + assertOwnerOnlyPath(t, filepath.Join(firstRuntimeDir, "connection.json"), 0o600) + + restarted, err := manager.Restart(ctx, started.ID) + if err != nil { + t.Fatalf("Restart() error = %v", err) + } + if restarted.ID != started.ID || restarted.ExecutionState != string(KernelStateIdle) { + t.Fatalf("Restart() model = %+v, want retained idle kernel", restarted) + } + if firstCommand.ProcessState == nil { + t.Fatal("old kernel process was not reaped during restart") + } + select { + case <-firstConnection.Context.Done(): + default: + t.Fatal("old connection generation was not canceled during restart") + } + if _, err := os.Stat(firstRuntimeDir); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("old runtime directory still exists or stat failed unexpectedly: %v", err) + } + + secondConnection, releaseSecond, err := manager.Connect(started.ID) + if err != nil { + t.Fatalf("Connect() after restart error = %v", err) + } + defer releaseSecond() + if secondConnection.Generation <= firstConnection.Generation { + t.Fatalf("generation = %d, want greater than %d", secondConnection.Generation, firstConnection.Generation) + } + if secondConnection.Info == firstConnection.Info { + t.Fatal("restart reused old ports and signing key") + } + secondCommand, secondRuntimeDir := testKernelProcess(manager, started.ID) + if secondRuntimeDir == firstRuntimeDir { + t.Fatal("restart reused old runtime directory") + } + + if err := manager.Stop(ctx, started.ID); err != nil { + t.Fatalf("Stop() error = %v", err) + } + if len(manager.List()) != 0 { + t.Fatalf("List() after stop = %+v, want empty", manager.List()) + } + if secondCommand.ProcessState == nil { + t.Fatal("restarted kernel process was not reaped during stop") + } + if _, err := os.Stat(secondRuntimeDir); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("runtime directory still exists after stop or stat failed unexpectedly: %v", err) + } +} + +func TestKernelManagerUnexpectedExitBecomesDead(t *testing.T) { + manager := newTestKernelManager(t) + defer manager.Close(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) + if err != nil { + t.Fatal(err) + } + connection, release, err := manager.Connect(model.ID) + if err != nil { + t.Fatal(err) + } + defer release() + command, _ := testKernelProcess(manager, model.ID) + if err := killManagedProcess(command); err != nil { + t.Fatalf("killManagedProcess() error = %v", err) + } + + waitForKernelState(t, manager, model.ID, KernelStateDead) + select { + case <-connection.Context.Done(): + case <-time.After(2 * time.Second): + t.Fatal("connection generation remained live after process crash") + } + if _, _, err := manager.Connect(model.ID); err == nil { + t.Fatal("Connect() accepted a dead kernel") + } + if err := manager.Stop(ctx, model.ID); err != nil { + t.Fatalf("Stop() dead kernel error = %v", err) + } +} + +func TestKernelManagerCloseReapsAllKernels(t *testing.T) { + manager := newTestKernelManager(t) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + + var commands []*exec.Cmd + var runtimeDirs []string + for range 2 { + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) + if err != nil { + t.Fatal(err) + } + command, runtimeDir := testKernelProcess(manager, model.ID) + commands = append(commands, command) + runtimeDirs = append(runtimeDirs, runtimeDir) + } + if err := manager.Close(ctx); err != nil { + t.Fatalf("Close() error = %v", err) + } + if len(manager.List()) != 0 { + t.Fatalf("List() after Close() = %+v, want empty", manager.List()) + } + for i, command := range commands { + if command.ProcessState == nil { + t.Errorf("kernel process %d was not reaped", i) + } + if _, err := os.Stat(runtimeDirs[i]); !errors.Is(err, os.ErrNotExist) { + t.Errorf("runtime directory %d still exists or stat failed unexpectedly: %v", i, err) + } + } + if _, err := manager.Start(ctx, testPythonLaunchSpec(t)); err == nil { + t.Fatal("Start() succeeded after Close()") + } +} + +func TestKernelManagerValidatesClientLaunchSpec(t *testing.T) { + manager := newTestKernelManager(t) + defer manager.Close(context.Background()) + _, err := manager.Start(context.Background(), KernelLaunchSpec{ + Name: "python3", + Argv: []string{"python3", "-c", "pass"}, + }) + if err == nil || !strings.Contains(err.Error(), connectionFilePlaceholder) { + t.Fatalf("Start() error = %v, want missing placeholder error", err) + } + if _, err := manager.Start(context.Background(), KernelLaunchSpec{ + Name: "python3", + Argv: []string{connectionFilePlaceholder}, + }); err == nil || !strings.Contains(err.Error(), "executable") { + t.Fatalf("Start() executable-placeholder error = %v", err) + } +} + +func TestBoundedWriterKeepsOnlyTail(t *testing.T) { + writer := newBoundedWriter(5) + if _, err := writer.Write([]byte("1234")); err != nil { + t.Fatal(err) + } + if _, err := writer.Write([]byte("5678")); err != nil { + t.Fatal(err) + } + if got, want := writer.String(), "45678"; got != want { + t.Fatalf("String() = %q, want %q", got, want) + } +} + +func newTestKernelManager(t *testing.T) *KernelManager { + t.Helper() + manager, err := NewKernelManager(KernelManagerConfig{ + RuntimeDir: filepath.Join(t.TempDir(), "runtime"), + StartupTimeout: 10 * time.Second, + ReadinessTimeout: 5 * time.Second, + GracefulTimeout: 2 * time.Second, + TerminateTimeout: 2 * time.Second, + KillTimeout: 2 * time.Second, + }) + if err != nil { + t.Fatalf("NewKernelManager() error = %v", err) + } + return manager +} + +func testPythonLaunchSpec(t *testing.T) KernelLaunchSpec { + t.Helper() + return KernelLaunchSpec{ + Name: "python3", + Argv: []string{findIPyKernelPython(t), "-m", "ipykernel_launcher", "-f", connectionFilePlaceholder}, + } +} + +func testKernelProcess(manager *KernelManager, id string) (*exec.Cmd, string) { + manager.mu.RLock() + defer manager.mu.RUnlock() + kernel := manager.kernels[id] + return kernel.command, kernel.runtimeDir +} + +func assertOwnerOnlyPath(t *testing.T, path string, want os.FileMode) { + t.Helper() + info, err := os.Stat(path) + if err != nil { + t.Fatalf("stat %s: %v", path, err) + } + if got := info.Mode().Perm(); got != want { + t.Fatalf("permissions for %s = %o, want %o", path, got, want) + } +} + +func waitForKernelState(t *testing.T, manager *KernelManager, id string, state KernelState) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + model, err := manager.Get(id) + if err == nil && model.ExecutionState == string(state) { + return + } + time.Sleep(10 * time.Millisecond) + } + model, err := manager.Get(id) + t.Fatalf("kernel state = %+v, error = %v; want %s", model, err, state) +} diff --git a/pkg/agent/jupyter/message.go b/pkg/agent/jupyter/message.go new file mode 100644 index 000000000..59d401028 --- /dev/null +++ b/pkg/agent/jupyter/message.go @@ -0,0 +1,246 @@ +package jupyter + +import ( + "bytes" + "crypto/hmac" + "crypto/sha256" + "crypto/subtle" + "encoding/hex" + "encoding/json" + "errors" + "fmt" +) + +const ( + // SignatureSchemeHMACSHA256 is the only signing scheme supported by the + // Jupyter messaging specification and the direct bridge. + SignatureSchemeHMACSHA256 = "hmac-sha256" + + defaultMaxFrameCount = 64 + defaultMaxFrameBytes = 16 * 1024 * 1024 + defaultMaxMessageBytes = 32 * 1024 * 1024 +) + +var delimiter = []byte("") + +// Limits bounds untrusted multipart traffic. Zero values select conservative +// defaults. +type Limits struct { + MaxFrameCount int + MaxFrameBytes int + MaxMessageBytes int +} + +func (l Limits) normalized() Limits { + if l.MaxFrameCount <= 0 { + l.MaxFrameCount = defaultMaxFrameCount + } + if l.MaxFrameBytes <= 0 { + l.MaxFrameBytes = defaultMaxFrameBytes + } + if l.MaxMessageBytes <= 0 { + l.MaxMessageBytes = defaultMaxMessageBytes + } + return l +} + +// Message is a Jupyter wire message. JSON dictionaries remain raw so unknown +// fields and message types survive a bridge round trip unchanged. +type Message struct { + RoutingPrefix [][]byte + Signature []byte + Header json.RawMessage + ParentHeader json.RawMessage + Metadata json.RawMessage + Content json.RawMessage + Buffers [][]byte +} + +// Signer signs and verifies the four JSON frames in a Jupyter wire message. +type Signer struct { + key []byte +} + +// NewSigner constructs a signer for a validated Jupyter signature scheme. An +// empty key is valid and represents authentication-disabled Jupyter messages. +func NewSigner(key, signatureScheme string) (*Signer, error) { + if signatureScheme != SignatureSchemeHMACSHA256 { + return nil, fmt.Errorf("unsupported Jupyter signature scheme %q", signatureScheme) + } + return &Signer{key: []byte(key)}, nil +} + +// Sign returns the lowercase hexadecimal HMAC for the four serialized JSON +// frames. It returns an empty signature when the connection key is empty. +func (s *Signer) Sign(parts ...[]byte) []byte { + if s == nil || len(s.key) == 0 { + return []byte{} + } + mac := hmac.New(sha256.New, s.key) + for _, part := range parts { + _, _ = mac.Write(part) + } + encoded := make([]byte, hex.EncodedLen(sha256.Size)) + hex.Encode(encoded, mac.Sum(nil)) + return encoded +} + +// Verify checks a wire signature without timing-sensitive byte comparison. +func (s *Signer) Verify(signature []byte, parts ...[]byte) error { + if s == nil { + return errors.New("jupyter signer is required") + } + if len(s.key) == 0 { + if len(signature) != 0 { + return errors.New("unexpected signature for empty Jupyter key") + } + return nil + } + if len(signature) != hex.EncodedLen(sha256.Size) { + return errors.New("invalid Jupyter signature length") + } + provided := make([]byte, sha256.Size) + if _, err := hex.Decode(provided, signature); err != nil { + return errors.New("invalid Jupyter signature encoding") + } + + mac := hmac.New(sha256.New, s.key) + for _, part := range parts { + _, _ = mac.Write(part) + } + if subtle.ConstantTimeCompare(provided, mac.Sum(nil)) != 1 { + return errors.New("invalid Jupyter message signature") + } + return nil +} + +// ParseMultipart validates and decodes a Jupyter ZeroMQ multipart message. It +// searches for the delimiter because ROUTER and IOPub prefixes are variable. +func ParseMultipart(frames [][]byte, signer *Signer, limits Limits) (Message, error) { + if err := validateMultipartLimits(frames, limits); err != nil { + return Message{}, err + } + + delimiterIndex := -1 + for i, frame := range frames { + if bytes.Equal(frame, delimiter) { + delimiterIndex = i + break + } + } + if delimiterIndex < 0 { + return Message{}, errors.New("jupyter message delimiter not found") + } + if len(frames)-delimiterIndex-1 < 5 { + return Message{}, errors.New("jupyter message is missing signed JSON frames") + } + + signature := frames[delimiterIndex+1] + header := frames[delimiterIndex+2] + parentHeader := frames[delimiterIndex+3] + metadata := frames[delimiterIndex+4] + content := frames[delimiterIndex+5] + // Authenticate exact bytes before attempting to parse attacker-controlled + // JSON. Callers can therefore rely on ParseMultipart never exposing an + // unauthenticated logical message. + if err := signer.Verify(signature, header, parentHeader, metadata, content); err != nil { + return Message{}, err + } + for name, raw := range map[string][]byte{ + "header": header, + "parent_header": parentHeader, + "metadata": metadata, + "content": content, + } { + if err := validateJSONObject(raw); err != nil { + return Message{}, fmt.Errorf("invalid Jupyter %s: %w", name, err) + } + } + + return Message{ + RoutingPrefix: cloneFrames(frames[:delimiterIndex]), + Signature: bytes.Clone(signature), + Header: bytes.Clone(header), + ParentHeader: bytes.Clone(parentHeader), + Metadata: bytes.Clone(metadata), + Content: bytes.Clone(content), + Buffers: cloneFrames(frames[delimiterIndex+6:]), + }, nil +} + +// MarshalMultipart validates and serializes a Jupyter message into ZeroMQ +// multipart frames, recomputing the signature over the exact JSON bytes. +func MarshalMultipart(message Message, signer *Signer, limits Limits) ([][]byte, error) { + for name, raw := range map[string][]byte{ + "header": message.Header, + "parent_header": message.ParentHeader, + "metadata": message.Metadata, + "content": message.Content, + } { + if err := validateJSONObject(raw); err != nil { + return nil, fmt.Errorf("invalid Jupyter %s: %w", name, err) + } + } + if signer == nil { + return nil, errors.New("jupyter signer is required") + } + + frames := make([][]byte, 0, len(message.RoutingPrefix)+6+len(message.Buffers)) + frames = append(frames, cloneFrames(message.RoutingPrefix)...) + frames = append(frames, + bytes.Clone(delimiter), + signer.Sign(message.Header, message.ParentHeader, message.Metadata, message.Content), + bytes.Clone(message.Header), + bytes.Clone(message.ParentHeader), + bytes.Clone(message.Metadata), + bytes.Clone(message.Content), + ) + frames = append(frames, cloneFrames(message.Buffers)...) + if err := validateMultipartLimits(frames, limits); err != nil { + return nil, err + } + return frames, nil +} + +func validateJSONObject(raw []byte) error { + if len(raw) == 0 || !json.Valid(raw) { + return errors.New("must be valid JSON") + } + var value map[string]json.RawMessage + if err := json.Unmarshal(raw, &value); err != nil || value == nil { + return errors.New("must be a JSON object") + } + return nil +} + +func validateMultipartLimits(frames [][]byte, limits Limits) error { + limits = limits.normalized() + if len(frames) == 0 { + return errors.New("jupyter multipart message is empty") + } + if len(frames) > limits.MaxFrameCount { + return fmt.Errorf("jupyter multipart frame count %d exceeds limit %d", len(frames), limits.MaxFrameCount) + } + total := 0 + for _, frame := range frames { + if len(frame) > limits.MaxFrameBytes { + return fmt.Errorf("jupyter frame size %d exceeds limit %d", len(frame), limits.MaxFrameBytes) + } + if len(frame) > limits.MaxMessageBytes-total { + return fmt.Errorf("jupyter multipart payload exceeds limit %d", limits.MaxMessageBytes) + } + total += len(frame) + } + return nil +} + +func cloneFrames(frames [][]byte) [][]byte { + if len(frames) == 0 { + return nil + } + cloned := make([][]byte, len(frames)) + for i, frame := range frames { + cloned[i] = bytes.Clone(frame) + } + return cloned +} diff --git a/pkg/agent/jupyter/message_test.go b/pkg/agent/jupyter/message_test.go new file mode 100644 index 000000000..a828c5e38 --- /dev/null +++ b/pkg/agent/jupyter/message_test.go @@ -0,0 +1,156 @@ +package jupyter + +import ( + "bytes" + "reflect" + "testing" +) + +func TestMarshalMultipartGoldenAndRoundTrip(t *testing.T) { + signer, err := NewSigner("secret", SignatureSchemeHMACSHA256) + if err != nil { + t.Fatal(err) + } + message := Message{ + RoutingPrefix: [][]byte{[]byte("client-a"), []byte("route-b")}, + Header: []byte(`{"msg_id":"abc","msg_type":"execute_request","session":"s","username":"runme","version":"5.3"}`), + ParentHeader: []byte(`{}`), + Metadata: []byte(`{"x":1}`), + Content: []byte(`{"code":"print(1)","silent":false}`), + Buffers: [][]byte{{0x00, 0x01, 0xfe, 0xff}}, + } + + frames, err := MarshalMultipart(message, signer, Limits{}) + if err != nil { + t.Fatalf("MarshalMultipart() error = %v", err) + } + want := [][]byte{ + []byte("client-a"), + []byte("route-b"), + []byte(""), + []byte("032afa1cd3332a9794de5989bf088fbe4ca39b6ca5d28c4d9e61c5271e6fda3c"), + message.Header, + message.ParentHeader, + message.Metadata, + message.Content, + {0x00, 0x01, 0xfe, 0xff}, + } + if !reflect.DeepEqual(frames, want) { + t.Fatalf("MarshalMultipart() frames differ\n got: %q\nwant: %q", frames, want) + } + + parsed, err := ParseMultipart(frames, signer, Limits{}) + if err != nil { + t.Fatalf("ParseMultipart() error = %v", err) + } + if !reflect.DeepEqual(parsed.RoutingPrefix, message.RoutingPrefix) { + t.Fatalf("routing prefix = %q, want %q", parsed.RoutingPrefix, message.RoutingPrefix) + } + if !reflect.DeepEqual(parsed.Buffers, message.Buffers) { + t.Fatalf("buffers = %v, want %v", parsed.Buffers, message.Buffers) + } +} + +func TestParseMultipartRejectsInvalidTraffic(t *testing.T) { + signer, err := NewSigner("secret", SignatureSchemeHMACSHA256) + if err != nil { + t.Fatal(err) + } + valid, err := MarshalMultipart(Message{ + Header: []byte(`{"msg_id":"abc"}`), + ParentHeader: []byte(`{}`), + Metadata: []byte(`{}`), + Content: []byte(`{"code":"1"}`), + }, signer, Limits{}) + if err != nil { + t.Fatal(err) + } + + tests := map[string]struct { + frames [][]byte + limits Limits + }{ + "missing delimiter": {frames: valid[1:]}, + "missing content": {frames: valid[:len(valid)-1]}, + "bad signature": { + frames: func() [][]byte { + frames := cloneFrames(valid) + frames[1] = bytes.Repeat([]byte("0"), 64) + return frames + }(), + }, + "malformed header": { + frames: func() [][]byte { + frames := cloneFrames(valid) + frames[2] = []byte(`[]`) + frames[1] = signer.Sign(frames[2], frames[3], frames[4], frames[5]) + return frames + }(), + }, + "too many frames": { + frames: valid, + limits: Limits{MaxFrameCount: len(valid) - 1}, + }, + "oversized frame": { + frames: valid, + limits: Limits{MaxFrameBytes: 4}, + }, + "oversized message": { + frames: valid, + limits: Limits{MaxFrameBytes: 1024, MaxMessageBytes: 8}, + }, + } + for name, test := range tests { + t.Run(name, func(t *testing.T) { + if _, err := ParseMultipart(test.frames, signer, test.limits); err == nil { + t.Fatal("ParseMultipart() unexpectedly succeeded") + } + }) + } +} + +func TestEmptyKeyRequiresEmptySignature(t *testing.T) { + signer, err := NewSigner("", SignatureSchemeHMACSHA256) + if err != nil { + t.Fatal(err) + } + frames, err := MarshalMultipart(Message{ + Header: []byte(`{}`), + ParentHeader: []byte(`{}`), + Metadata: []byte(`{}`), + Content: []byte(`{}`), + }, signer, Limits{}) + if err != nil { + t.Fatal(err) + } + if len(frames[1]) != 0 { + t.Fatalf("signature = %q, want empty", frames[1]) + } + frames[1] = []byte("unexpected") + if _, err := ParseMultipart(frames, signer, Limits{}); err == nil { + t.Fatal("ParseMultipart() accepted a signature with an empty key") + } +} + +func FuzzParseMultipart(f *testing.F) { + f.Add([]byte("\x00{}{}{}{}"), uint8(0)) + f.Add([]byte{0xff, 0x00, 0x01}, uint8(3)) + f.Fuzz(func(t *testing.T, raw []byte, splitSeed uint8) { + signer, err := NewSigner("fuzz-key", SignatureSchemeHMACSHA256) + if err != nil { + t.Fatal(err) + } + frameCount := int(splitSeed%16) + 1 + frames := make([][]byte, 0, frameCount) + for i := 0; i < frameCount; i++ { + start := len(raw) * i / frameCount + end := len(raw) * (i + 1) / frameCount + frames = append(frames, raw[start:end]) + } + _, _ = ParseMultipart(frames, signer, Limits{ + MaxFrameCount: 32, + MaxFrameBytes: 64 * 1024, + MaxMessageBytes: 128 * 1024, + }) + }) +} diff --git a/pkg/agent/jupyter/metrics.go b/pkg/agent/jupyter/metrics.go new file mode 100644 index 000000000..d0bb5c770 --- /dev/null +++ b/pkg/agent/jupyter/metrics.go @@ -0,0 +1,72 @@ +package jupyter + +import ( + "context" + "errors" + "strings" + "time" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + otelmetric "go.opentelemetry.io/otel/metric" +) + +var ( + jupyterMeter = otel.Meter("github.com/runmedev/runme/v3/pkg/agent/jupyter") + jupyterKernelStates = mustMetric(jupyterMeter.Int64UpDownCounter("runme.jupyter.kernels", otelmetric.WithDescription("Current directly managed Jupyter kernels by lifecycle state."), otelmetric.WithUnit("{kernel}"))) + jupyterLifecycleLatency = mustMetric(jupyterMeter.Float64Histogram("runme.jupyter.lifecycle.duration", otelmetric.WithDescription("Duration of direct Jupyter kernel lifecycle operations."), otelmetric.WithUnit("s"))) + jupyterActiveBridges = mustMetric(jupyterMeter.Int64UpDownCounter("runme.jupyter.bridges.active", otelmetric.WithDescription("Current direct browser-to-kernel Jupyter bridges."), otelmetric.WithUnit("{bridge}"))) + jupyterChannelMessages = mustMetric(jupyterMeter.Int64Counter("runme.jupyter.channel.messages", otelmetric.WithDescription("Direct Jupyter channel messages by direction and channel."), otelmetric.WithUnit("{message}"))) + jupyterChannelBytes = mustMetric(jupyterMeter.Int64Counter("runme.jupyter.channel.bytes", otelmetric.WithDescription("Direct Jupyter channel payload bytes by direction and channel."), otelmetric.WithUnit("By"))) + jupyterProtocolErrors = mustMetric(jupyterMeter.Int64Counter("runme.jupyter.protocol.errors", otelmetric.WithDescription("Rejected direct Jupyter protocol messages by bounded reason."), otelmetric.WithUnit("{error}"))) + jupyterRateLimits = mustMetric(jupyterMeter.Int64Counter("runme.jupyter.rate_limits", otelmetric.WithDescription("Direct Jupyter bridges closed by IOPub rate limits."), otelmetric.WithUnit("{event}"))) + jupyterHeartbeatMisses = mustMetric(jupyterMeter.Int64Counter("runme.jupyter.heartbeat.misses", otelmetric.WithDescription("Missed direct Jupyter kernel heartbeat probes."), otelmetric.WithUnit("{event}"))) + jupyterUnexpectedExits = mustMetric(jupyterMeter.Int64Counter("runme.jupyter.kernel.unexpected_exits", otelmetric.WithDescription("Unexpected exits of directly managed Jupyter kernels."), otelmetric.WithUnit("{event}"))) + jupyterBridgeCloses = mustMetric(jupyterMeter.Int64Counter("runme.jupyter.bridge.closes", otelmetric.WithDescription("Direct Jupyter bridge closes by bounded reason."), otelmetric.WithUnit("{event}"))) +) + +func mustMetric[T any](instrument T, err error) T { + if err != nil { + panic(err) + } + return instrument +} + +func observeLifecycle(operation string, started time.Time) { + jupyterLifecycleLatency.Record(context.Background(), time.Since(started).Seconds(), otelmetric.WithAttributes(attribute.String("operation", operation))) +} + +func addKernelState(state KernelState, delta int64) { + jupyterKernelStates.Add(context.Background(), delta, otelmetric.WithAttributes(attribute.String("state", string(state)))) +} + +func addActiveBridge(delta int64) { + jupyterActiveBridges.Add(context.Background(), delta) +} + +func observeChannelMessage(direction string, channel Channel, payloadBytes int) { + attributes := otelmetric.WithAttributes( + attribute.String("direction", direction), + attribute.String("channel", string(channel)), + ) + jupyterChannelMessages.Add(context.Background(), 1, attributes) + jupyterChannelBytes.Add(context.Background(), int64(payloadBytes), attributes) +} + +func observeBridgeClose(reason string) { + addCounter(jupyterBridgeCloses, attribute.String("reason", reason)) +} + +func addCounter(counter otelmetric.Int64Counter, attributes ...attribute.KeyValue) { + counter.Add(context.Background(), 1, otelmetric.WithAttributes(attributes...)) +} + +func observeProtocolError(err error) { + reason := "malformed_message" + if errors.Is(err, ErrBinaryBuffersUnsupported) { + reason = "unsupported_binary_buffers" + } else if err != nil && strings.Contains(strings.ToLower(err.Error()), "signature") { + reason = "bad_signature" + } + addCounter(jupyterProtocolErrors, attribute.String("reason", reason)) +} diff --git a/pkg/agent/jupyter/process_unix.go b/pkg/agent/jupyter/process_unix.go new file mode 100644 index 000000000..55aa41a72 --- /dev/null +++ b/pkg/agent/jupyter/process_unix.go @@ -0,0 +1,36 @@ +//go:build !windows + +package jupyter + +import ( + "errors" + "os/exec" + "syscall" +) + +func configureManagedProcess(command *exec.Cmd) { + command.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} +} + +func interruptManagedProcess(command *exec.Cmd) error { + return signalManagedProcessGroup(command, syscall.SIGINT) +} + +func terminateManagedProcess(command *exec.Cmd) error { + return signalManagedProcessGroup(command, syscall.SIGTERM) +} + +func killManagedProcess(command *exec.Cmd) error { + return signalManagedProcessGroup(command, syscall.SIGKILL) +} + +func signalManagedProcessGroup(command *exec.Cmd, signal syscall.Signal) error { + if command == nil || command.Process == nil { + return syscall.ESRCH + } + return syscall.Kill(-command.Process.Pid, signal) +} + +func isProcessDoneError(err error) bool { + return errors.Is(err, syscall.ESRCH) +} diff --git a/pkg/agent/jupyter/process_windows.go b/pkg/agent/jupyter/process_windows.go new file mode 100644 index 000000000..9b3a17235 --- /dev/null +++ b/pkg/agent/jupyter/process_windows.go @@ -0,0 +1,38 @@ +//go:build windows + +package jupyter + +import ( + "errors" + "os" + "os/exec" +) + +// Windows process-group control remains isolated here so it can be upgraded to +// console control events without changing KernelManager. +func configureManagedProcess(_ *exec.Cmd) {} + +func interruptManagedProcess(command *exec.Cmd) error { + if command == nil || command.Process == nil { + return os.ErrProcessDone + } + return command.Process.Signal(os.Interrupt) +} + +func terminateManagedProcess(command *exec.Cmd) error { + if command == nil || command.Process == nil { + return os.ErrProcessDone + } + return command.Process.Kill() +} + +func killManagedProcess(command *exec.Cmd) error { + if command == nil || command.Process == nil { + return os.ErrProcessDone + } + return command.Process.Kill() +} + +func isProcessDoneError(err error) bool { + return errors.Is(err, os.ErrProcessDone) +} diff --git a/pkg/agent/jupyter/socket.go b/pkg/agent/jupyter/socket.go new file mode 100644 index 000000000..12cc9718d --- /dev/null +++ b/pkg/agent/jupyter/socket.go @@ -0,0 +1,204 @@ +package jupyter + +import ( + "context" + "errors" + "fmt" + "io" + "log" + "net" + "net/url" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/go-zeromq/zmq4" +) + +// SocketType is one of the three ZeroMQ patterns used by Jupyter clients. +type SocketType string + +const ( + SocketTypeDealer SocketType = "dealer" + SocketTypeSub SocketType = "sub" + SocketTypeReq SocketType = "req" +) + +// Socket is the narrow multipart interface used by the kernel manager and +// channels bridge. It intentionally hides the selected ZeroMQ implementation. +type Socket interface { + Dial(endpoint string) error + SendMultipart(frames [][]byte) error + ReceiveMultipart() ([][]byte, error) + Close() error +} + +// SocketFactory creates context-bound Jupyter sockets. DEALER identities are +// supplied by the bridge so shell, control, and stdin can share one identity. +type SocketFactory interface { + NewSocket(ctx context.Context, socketType SocketType, identity []byte) (Socket, error) +} + +// GoZeroMQFactory adapts github.com/go-zeromq/zmq4 to SocketFactory. +type GoZeroMQFactory struct { + Limits Limits + DialTimeout time.Duration + DialRetry time.Duration + MaxDialRetries int +} + +// NewGoZeroMQFactory returns a pure-Go ZeroMQ socket factory with bounded dial +// behavior and multipart limits. +func NewGoZeroMQFactory(limits Limits) *GoZeroMQFactory { + return &GoZeroMQFactory{ + Limits: limits.normalized(), + DialTimeout: 5 * time.Second, + DialRetry: 100 * time.Millisecond, + MaxDialRetries: 2, + } +} + +// NewSocket creates a DEALER, subscribe-all SUB, or REQ socket. +func (f *GoZeroMQFactory) NewSocket(ctx context.Context, socketType SocketType, identity []byte) (Socket, error) { + if ctx == nil { + return nil, errors.New("jupyter socket context is required") + } + if f == nil { + return nil, errors.New("jupyter socket factory is required") + } + if socketType == SocketTypeDealer && len(identity) == 0 { + return nil, errors.New("jupyter DEALER identity is required") + } + + dialTimeout := f.DialTimeout + if dialTimeout <= 0 { + dialTimeout = 5 * time.Second + } + dialRetry := f.DialRetry + if dialRetry <= 0 { + dialRetry = 100 * time.Millisecond + } + maxDialRetries := f.MaxDialRetries + if maxDialRetries < 0 { + return nil, errors.New("unbounded Jupyter dial retries are not allowed") + } + + opts := []zmq4.Option{ + zmq4.WithDialerTimeout(dialTimeout), + zmq4.WithDialerRetry(dialRetry), + zmq4.WithDialerMaxRetries(maxDialRetries), + zmq4.WithAutomaticReconnect(false), + zmq4.WithLogger(log.New(io.Discard, "", 0)), + } + if len(identity) > 0 { + opts = append(opts, zmq4.WithID(zmq4.SocketIdentity(identity))) + } + + var raw zmq4.Socket + switch socketType { + case SocketTypeDealer: + raw = zmq4.NewDealer(ctx, opts...) + case SocketTypeSub: + raw = zmq4.NewSub(ctx, opts...) + if err := raw.SetOption(zmq4.OptionSubscribe, ""); err != nil { + _ = raw.Close() + return nil, fmt.Errorf("subscribe to all Jupyter IOPub topics: %w", err) + } + case SocketTypeReq: + raw = zmq4.NewReq(ctx, opts...) + default: + return nil, fmt.Errorf("unsupported Jupyter socket type %q", socketType) + } + + return &goZeroMQSocket{ + raw: raw, + limits: f.Limits.normalized(), + }, nil +} + +type goZeroMQSocket struct { + raw zmq4.Socket + limits Limits + closed atomic.Bool + closeOnce sync.Once + closeErr error +} + +func (s *goZeroMQSocket) Dial(endpoint string) error { + if s == nil || s.raw == nil || s.closed.Load() { + return errors.New("jupyter socket is closed") + } + if err := validateLoopbackTCPEndpoint(endpoint); err != nil { + return err + } + if err := s.raw.Dial(endpoint); err != nil { + return fmt.Errorf("dial Jupyter socket: %w", err) + } + return nil +} + +func (s *goZeroMQSocket) SendMultipart(frames [][]byte) error { + if s == nil || s.raw == nil || s.closed.Load() { + return errors.New("jupyter socket is closed") + } + if err := validateMultipartLimits(frames, s.limits); err != nil { + return err + } + if err := s.raw.SendMulti(zmq4.NewMsgFrom(cloneFrames(frames)...)); err != nil { + return fmt.Errorf("send Jupyter multipart message: %w", err) + } + return nil +} + +func (s *goZeroMQSocket) ReceiveMultipart() ([][]byte, error) { + if s == nil || s.raw == nil || s.closed.Load() { + return nil, errors.New("jupyter socket is closed") + } + message, err := s.raw.Recv() + if err != nil { + return nil, fmt.Errorf("receive Jupyter multipart message: %w", err) + } + if err := message.Err(); err != nil { + return nil, fmt.Errorf("receive Jupyter multipart message: %w", err) + } + if err := validateMultipartLimits(message.Frames, s.limits); err != nil { + return nil, err + } + return cloneFrames(message.Frames), nil +} + +func (s *goZeroMQSocket) Close() error { + if s == nil { + return nil + } + s.closeOnce.Do(func() { + if s.raw != nil { + s.closed.Store(true) + s.closeErr = s.raw.Close() + } + }) + return s.closeErr +} + +func validateLoopbackTCPEndpoint(endpoint string) error { + parsed, err := url.Parse(endpoint) + if err != nil { + return fmt.Errorf("invalid Jupyter endpoint: %w", err) + } + if parsed.Scheme != "tcp" { + return fmt.Errorf("unsupported Jupyter endpoint scheme %q", parsed.Scheme) + } + if parsed.User != nil || parsed.Path != "" || parsed.RawQuery != "" || parsed.Fragment != "" { + return errors.New("invalid Jupyter TCP endpoint") + } + host, port, err := net.SplitHostPort(parsed.Host) + if err != nil || strings.TrimSpace(port) == "" { + return errors.New("invalid Jupyter TCP endpoint address") + } + ip := net.ParseIP(host) + if ip == nil || !ip.IsLoopback() { + return errors.New("jupyter TCP endpoint must use a loopback IP") + } + return nil +} diff --git a/pkg/agent/jupyter/socket_test.go b/pkg/agent/jupyter/socket_test.go new file mode 100644 index 000000000..358f5bd93 --- /dev/null +++ b/pkg/agent/jupyter/socket_test.go @@ -0,0 +1,77 @@ +package jupyter + +import ( + "context" + "errors" + "testing" + "time" +) + +func TestGoZeroMQFactoryValidatesSocketConfiguration(t *testing.T) { + factory := NewGoZeroMQFactory(Limits{}) + ctx := context.Background() + + if _, err := factory.NewSocket(ctx, SocketTypeDealer, nil); err == nil { + t.Fatal("NewSocket() accepted an empty DEALER identity") + } + if _, err := factory.NewSocket(ctx, SocketType("push"), []byte("client")); err == nil { + t.Fatal("NewSocket() accepted an unsupported socket type") + } + + dealer, err := factory.NewSocket(ctx, SocketTypeDealer, []byte("client")) + if err != nil { + t.Fatalf("NewSocket() error = %v", err) + } + defer dealer.Close() + if err := dealer.Dial("tcp://192.0.2.10:1234"); err == nil { + t.Fatal("Dial() accepted a non-loopback endpoint") + } + if err := dealer.Dial("ipc:///tmp/kernel.sock"); err == nil { + t.Fatal("Dial() accepted a non-TCP endpoint") + } +} + +func TestGoZeroMQSocketEnforcesMultipartLimits(t *testing.T) { + factory := NewGoZeroMQFactory(Limits{ + MaxFrameCount: 2, + MaxFrameBytes: 4, + MaxMessageBytes: 8, + }) + socket, err := factory.NewSocket(context.Background(), SocketTypeReq, nil) + if err != nil { + t.Fatal(err) + } + defer socket.Close() + + if err := socket.SendMultipart([][]byte{[]byte("12345")}); err == nil { + t.Fatal("SendMultipart() accepted an oversized frame") + } + if err := socket.SendMultipart([][]byte{[]byte("1"), []byte("2"), []byte("3")}); err == nil { + t.Fatal("SendMultipart() accepted too many frames") + } +} + +func TestGoZeroMQSocketReceiveHonorsContextCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + factory := NewGoZeroMQFactory(Limits{}) + socket, err := factory.NewSocket(ctx, SocketTypeSub, nil) + if err != nil { + t.Fatal(err) + } + defer socket.Close() + cancel() + + done := make(chan error, 1) + go func() { + _, err := socket.ReceiveMultipart() + done <- err + }() + select { + case err := <-done: + if err == nil || !errors.Is(err, context.Canceled) { + t.Fatalf("ReceiveMultipart() error = %v, want context cancellation", err) + } + case <-time.After(2 * time.Second): + t.Fatal("ReceiveMultipart() did not unblock after context cancellation") + } +} diff --git a/pkg/agent/server/jupyter_kernels.go b/pkg/agent/server/jupyter_kernels.go new file mode 100644 index 000000000..f0f9645c6 --- /dev/null +++ b/pkg/agent/server/jupyter_kernels.go @@ -0,0 +1,341 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "sync" + "time" + + "github.com/gorilla/websocket" + + jupyterbridge "github.com/runmedev/runme/v3/pkg/agent/jupyter" + "github.com/runmedev/runme/v3/pkg/agent/logs" +) + +const ( + jupyterKernelsRoute = "/v1/jupyter/kernels" + maxKernelRequestBody = 64 * 1024 + + // Keep idle Jupyter channels sockets alive across intermediaries. Long-running + // cells can be quiet, so transport liveness cannot depend on kernel traffic. + jupyterChannelsPingInterval = 20 * time.Second + jupyterChannelsPongWait = 60 * time.Second + jupyterChannelsPingWriteTimeout = 10 * time.Second +) + +type jupyterKernelsHandler struct { + manager *jupyterbridge.KernelManager + bridge *jupyterbridge.KernelChannelsBridge + upgrader websocket.Upgrader +} + +type kernelCreateRequest struct { + Name string `json:"name"` + Argv []string `json:"argv"` +} + +func newJupyterKernelsHandler(manager *jupyterbridge.KernelManager) (*jupyterKernelsHandler, error) { + if manager == nil { + return nil, errors.New("jupyter kernel manager is required") + } + bridge, err := jupyterbridge.NewKernelChannelsBridge(manager, nil, jupyterbridge.BridgeConfig{}) + if err != nil { + return nil, err + } + return &jupyterKernelsHandler{ + manager: manager, + bridge: bridge, + upgrader: websocket.Upgrader{ + ReadBufferSize: 1024, + WriteBufferSize: 1024, + CheckOrigin: func(_ *http.Request) bool { + return true + }, + }, + }, nil +} + +func (h *jupyterKernelsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == jupyterKernelsRoute { + h.handleCollection(w, r) + return + } + if !strings.HasPrefix(r.URL.Path, jupyterKernelsRoute+"/") { + writeHTTPError(w, http.StatusNotFound, "not found") + return + } + + trimmed := strings.TrimPrefix(r.URL.Path, jupyterKernelsRoute+"/") + parts := strings.Split(trimmed, "/") + if len(parts) < 1 || len(parts) > 2 || strings.TrimSpace(parts[0]) == "" { + writeHTTPError(w, http.StatusNotFound, "not found") + return + } + kernelID, err := url.PathUnescape(parts[0]) + if err != nil || strings.TrimSpace(kernelID) == "" || strings.Contains(kernelID, "/") { + writeHTTPError(w, http.StatusBadRequest, "invalid kernel id") + return + } + if _, err := h.manager.Get(kernelID); err != nil { + writeHTTPError(w, http.StatusNotFound, "kernel not found") + return + } + + if len(parts) == 1 { + h.handleKernel(w, r, kernelID) + return + } + switch parts[1] { + case "interrupt": + h.handleInterrupt(w, r, kernelID) + case "restart": + h.handleRestart(w, r, kernelID) + case "channels": + h.handleChannels(w, r, kernelID) + default: + writeHTTPError(w, http.StatusNotFound, "not found") + } +} + +func (h *jupyterKernelsHandler) handleCollection(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodGet: + writeJSON(w, http.StatusOK, h.manager.List()) + case http.MethodPost: + var request kernelCreateRequest + decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, maxKernelRequestBody)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&request); err != nil && !errors.Is(err, io.EOF) { + writeHTTPError(w, http.StatusBadRequest, "invalid kernel request") + return + } + if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) { + writeHTTPError(w, http.StatusBadRequest, "invalid kernel request") + return + } + name := strings.TrimSpace(request.Name) + kernel, err := h.manager.Start(r.Context(), jupyterbridge.KernelLaunchSpec{ + Name: name, + Argv: request.Argv, + }) + if err != nil { + logs.FromContext(r.Context()).Error(err, "failed to start Jupyter kernel", "kernel_name", name) + writeHTTPError(w, http.StatusBadRequest, "failed to start kernel") + return + } + logs.FromContext(r.Context()).Info("started Jupyter kernel", "kernel_id", kernel.ID, "kernel_name", name) + writeJSON(w, http.StatusCreated, kernel) + default: + writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") + } +} + +func (h *jupyterKernelsHandler) handleKernel(w http.ResponseWriter, r *http.Request, kernelID string) { + switch r.Method { + case http.MethodGet: + kernel, err := h.manager.Get(kernelID) + if err != nil { + writeHTTPError(w, http.StatusNotFound, "kernel not found") + return + } + writeJSON(w, http.StatusOK, kernel) + case http.MethodDelete: + logs.FromContext(r.Context()).Info("stopping Jupyter kernel", "kernel_id", kernelID) + if err := h.manager.Stop(r.Context(), kernelID); err != nil { + writeKernelOperationError(w, r, "stop", err) + return + } + w.WriteHeader(http.StatusNoContent) + default: + writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") + } +} + +func (h *jupyterKernelsHandler) handleInterrupt(w http.ResponseWriter, r *http.Request, kernelID string) { + if r.Method != http.MethodPost { + writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + if err := h.manager.Interrupt(kernelID); err != nil { + writeKernelOperationError(w, r, "interrupt", err) + return + } + logs.FromContext(r.Context()).Info("interrupted Jupyter kernel", "kernel_id", kernelID) + w.WriteHeader(http.StatusNoContent) +} + +func (h *jupyterKernelsHandler) handleRestart(w http.ResponseWriter, r *http.Request, kernelID string) { + if r.Method != http.MethodPost { + writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + kernel, err := h.manager.Restart(r.Context(), kernelID) + if err != nil { + writeKernelOperationError(w, r, "restart", err) + return + } + logs.FromContext(r.Context()).Info("restarted Jupyter kernel", "kernel_id", kernelID) + writeJSON(w, http.StatusOK, kernel) +} + +func writeKernelOperationError(w http.ResponseWriter, r *http.Request, operation string, err error) { + logs.FromContext(r.Context()).Error(err, "Jupyter kernel operation failed", "operation", operation) + writeHTTPError(w, http.StatusInternalServerError, "kernel operation failed") +} + +func (h *jupyterKernelsHandler) handleChannels(w http.ResponseWriter, r *http.Request, kernelID string) { + if r.Method != http.MethodGet { + writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + for _, protocol := range websocket.Subprotocols(r) { + if protocol == "v1.kernel.websocket.jupyter.org" { + writeHTTPError(w, http.StatusBadRequest, "binary Jupyter WebSocket protocol is not supported") + return + } + } + conn, err := h.upgrader.Upgrade(w, r, nil) + if err != nil { + logs.FromContext(r.Context()).Error(err, "failed to upgrade Jupyter channels WebSocket") + return + } + client := newJupyterChannelsClient(conn) + if err := configureWebSocketPongKeepalive(conn, jupyterChannelsPongWait); err != nil { + _ = client.Close(websocket.CloseInternalServerErr, "failed to configure WebSocket keepalive") + return + } + + keepaliveErr := make(chan error, 1) + stopKeepalive := make(chan struct{}) + go sendWebSocketKeepalivePings(conn, keepaliveErr, stopKeepalive) + bridgeErr := h.bridge.Bridge(r.Context(), kernelID, client) + close(stopKeepalive) + select { + case err := <-keepaliveErr: + if bridgeErr == nil { + bridgeErr = err + } + default: + } + if bridgeErr != nil && !errors.Is(bridgeErr, context.Canceled) { + logs.FromContext(r.Context()).Error(bridgeErr, "Jupyter channels bridge closed", "kernel_id", kernelID) + } + _ = client.Close(websocket.CloseNormalClosure, "Jupyter channels closed") +} + +type jupyterChannelsClient struct { + conn *websocket.Conn + closeOnce sync.Once +} + +func newJupyterChannelsClient(conn *websocket.Conn) *jupyterChannelsClient { + return &jupyterChannelsClient{conn: conn} +} + +func (c *jupyterChannelsClient) Read(ctx context.Context) ([]byte, error) { + done := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + _ = c.conn.SetReadDeadline(time.Now()) + case <-done: + } + }() + messageType, payload, err := c.conn.ReadMessage() + close(done) + if err != nil { + return nil, err + } + if messageType != websocket.TextMessage { + return nil, jupyterbridge.ErrBinaryBuffersUnsupported + } + return payload, nil +} + +func (c *jupyterChannelsClient) Write(ctx context.Context, payload []byte) error { + deadline := time.Now().Add(jupyterChannelsPingWriteTimeout) + if ctxDeadline, ok := ctx.Deadline(); ok && ctxDeadline.Before(deadline) { + deadline = ctxDeadline + } + if err := c.conn.SetWriteDeadline(deadline); err != nil { + return err + } + return c.conn.WriteMessage(websocket.TextMessage, payload) +} + +func (c *jupyterChannelsClient) Close(code int, reason string) error { + var closeErr error + c.closeOnce.Do(func() { + closeErr = c.conn.WriteControl( + websocket.CloseMessage, + websocket.FormatCloseMessage(code, reason), + time.Now().Add(jupyterChannelsPingWriteTimeout), + ) + if err := c.conn.Close(); closeErr == nil { + closeErr = err + } + }) + return closeErr +} + +func configureWebSocketPongKeepalive(conn *websocket.Conn, pongWait time.Duration) error { + if conn == nil { + return errors.New("websocket connection is required") + } + if pongWait <= 0 { + return errors.New("pong wait must be positive") + } + if err := conn.SetReadDeadline(time.Now().Add(pongWait)); err != nil { + return err + } + conn.SetPongHandler(func(_ string) error { + return conn.SetReadDeadline(time.Now().Add(pongWait)) + }) + return nil +} + +func sendWebSocketKeepalivePings(conn *websocket.Conn, errCh chan<- error, stop <-chan struct{}) { + ticker := time.NewTicker(jupyterChannelsPingInterval) + defer ticker.Stop() + for { + select { + case <-stop: + return + case <-ticker.C: + if err := conn.WriteControl( + websocket.PingMessage, + []byte("runme-keepalive"), + time.Now().Add(jupyterChannelsPingWriteTimeout), + ); err != nil { + select { + case errCh <- fmt.Errorf("write Jupyter channels keepalive: %w", err): + default: + } + return + } + } + } +} + +func writeJSON(w http.ResponseWriter, statusCode int, value any) { + data, err := json.Marshal(value) + if err != nil { + writeHTTPError(w, http.StatusInternalServerError, "failed to serialize response") + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(statusCode) + _, _ = w.Write(data) +} + +func writeHTTPError(w http.ResponseWriter, statusCode int, message string) { + writeJSON(w, statusCode, map[string]any{"error": message}) +} + +var _ http.Handler = (*jupyterKernelsHandler)(nil) diff --git a/pkg/agent/server/jupyter_kernels_test.go b/pkg/agent/server/jupyter_kernels_test.go new file mode 100644 index 000000000..80d7bcf15 --- /dev/null +++ b/pkg/agent/server/jupyter_kernels_test.go @@ -0,0 +1,367 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "net/url" + "os" + "os/exec" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" + + "github.com/runmedev/runme/v3/pkg/agent/config" + jupyterbridge "github.com/runmedev/runme/v3/pkg/agent/jupyter" +) + +func newTestJupyterKernelsHandler(t *testing.T) (*jupyterKernelsHandler, *jupyterbridge.KernelManager) { + t.Helper() + manager, err := jupyterbridge.NewKernelManager(jupyterbridge.KernelManagerConfig{ + RuntimeDir: t.TempDir(), + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = manager.Close(context.Background()) }) + handler, err := newJupyterKernelsHandler(manager) + if err != nil { + t.Fatal(err) + } + return handler, manager +} + +func TestJupyterKernelsRouteContract(t *testing.T) { + handler, _ := newTestJupyterKernelsHandler(t) + + tests := []struct { + name string + method string + path string + body string + status int + }{ + {name: "list", method: http.MethodGet, path: jupyterKernelsRoute, status: http.StatusOK}, + {name: "collection method", method: http.MethodDelete, path: jupyterKernelsRoute, status: http.StatusMethodNotAllowed}, + {name: "require argv", method: http.MethodPost, path: jupyterKernelsRoute, body: `{"name":"python3"}`, status: http.StatusBadRequest}, + {name: "reject custom path", method: http.MethodPost, path: jupyterKernelsRoute, body: `{"name":"python3","argv":["python3","{connection_file}"],"path":"/tmp/connection.json"}`, status: http.StatusBadRequest}, + {name: "reject multiple objects", method: http.MethodPost, path: jupyterKernelsRoute, body: `{ } { }`, status: http.StatusBadRequest}, + {name: "unknown kernel", method: http.MethodGet, path: jupyterKernelsRoute + "/missing", status: http.StatusNotFound}, + {name: "unknown action", method: http.MethodPost, path: jupyterKernelsRoute + "/missing/unknown", status: http.StatusNotFound}, + {name: "old server registry removed", method: http.MethodGet, path: "/v1/jupyter/servers", status: http.StatusNotFound}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + response := httptest.NewRecorder() + request := httptest.NewRequest(tt.method, tt.path, strings.NewReader(tt.body)) + handler.ServeHTTP(response, request) + if response.Code != tt.status { + t.Fatalf("status = %d, want %d; body: %s", response.Code, tt.status, response.Body.String()) + } + }) + } + + response := httptest.NewRecorder() + handler.ServeHTTP(response, httptest.NewRequest(http.MethodGet, jupyterKernelsRoute, nil)) + var kernels []jupyterbridge.KernelModel + if err := json.Unmarshal(response.Body.Bytes(), &kernels); err != nil { + t.Fatal(err) + } + if len(kernels) != 0 { + t.Fatalf("kernels = %#v, want empty list", kernels) + } +} + +func TestNewJupyterKernelsHandlerRequiresManager(t *testing.T) { + if _, err := newJupyterKernelsHandler(nil); err == nil { + t.Fatal("newJupyterKernelsHandler(nil) succeeded") + } +} + +func TestServerDoesNotRegisterJupyterRoutesByDefault(t *testing.T) { + runmeServer, err := NewServer(Options{ + Server: &config.AssistantServerConfig{ + ParserService: true, + }, + ConfigDir: t.TempDir(), + }) + if err != nil { + t.Fatal(err) + } + handler, err := runmeServer.BuildHandler() + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if runmeServer.jupyterManager != nil { + _ = runmeServer.jupyterManager.Close(context.Background()) + } + }) + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + + response, err := http.Get(server.URL + jupyterKernelsRoute) + if err != nil { + t.Fatal(err) + } + response.Body.Close() + if response.StatusCode != http.StatusNotFound { + t.Fatalf("kernel route status = %d, want 404", response.StatusCode) + } + if runmeServer.jupyterManager != nil { + t.Fatal("Jupyter kernel manager initialized while service is disabled") + } + if runmeServer.runner != nil { + t.Fatal("runner initialized while runner service is disabled") + } + if runmeServer.serverConfig.JupyterService { + t.Fatal("Jupyter service unexpectedly enabled by default") + } + if runmeServer.serverConfig.RunnerService { + t.Fatal("runner service unexpectedly enabled by default") + } + + legacyResponse, err := http.Get(server.URL + "/v1/jupyter/servers") + if err != nil { + t.Fatal(err) + } + legacyResponse.Body.Close() + if legacyResponse.StatusCode != http.StatusNotFound { + t.Fatalf("legacy server route status = %d, want 404", legacyResponse.StatusCode) + } +} + +func TestServerRegistersOnlyKernelScopedJupyterRoutesWhenEnabled(t *testing.T) { + runmeServer, err := NewServer(Options{ + Server: &config.AssistantServerConfig{ + RunnerService: true, + JupyterService: true, + }, + ConfigDir: t.TempDir(), + }) + if err != nil { + t.Fatal(err) + } + handler, err := runmeServer.BuildHandler() + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if runmeServer.jupyterManager != nil { + _ = runmeServer.jupyterManager.Close(context.Background()) + } + }) + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + + response, err := http.Get(server.URL + jupyterKernelsRoute) + if err != nil { + t.Fatal(err) + } + response.Body.Close() + if response.StatusCode != http.StatusOK { + t.Fatalf("kernel route status = %d, want 200", response.StatusCode) + } + + legacyResponse, err := http.Get(server.URL + "/v1/jupyter/servers") + if err != nil { + t.Fatal(err) + } + legacyResponse.Body.Close() + if legacyResponse.StatusCode != http.StatusNotFound { + t.Fatalf("legacy server route status = %d, want 404", legacyResponse.StatusCode) + } +} + +func TestNewServerRejectsJupyterServiceWithoutRunnerService(t *testing.T) { + _, err := NewServer(Options{ + Server: &config.AssistantServerConfig{ + ParserService: true, + JupyterService: true, + }, + }) + if err == nil { + t.Fatal("NewServer succeeded with Jupyter service enabled and runner service disabled") + } + if err.Error() != "jupyter service requires runner service" { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestDirectJupyterKernelAPIAndWebSocket(t *testing.T) { + python := os.Getenv("RUNME_TEST_PYTHON") + if python == "" { + var err error + python, err = exec.LookPath("python3") + if err != nil || exec.Command(python, "-c", "import ipykernel").Run() != nil { + t.Skip("ipykernel integration test requires RUNME_TEST_PYTHON or python3 with ipykernel") + } + } else if err := exec.Command(python, "-c", "import ipykernel").Run(); err != nil { + t.Fatalf("RUNME_TEST_PYTHON does not provide ipykernel: %v", err) + } + manager, err := jupyterbridge.NewKernelManager(jupyterbridge.KernelManagerConfig{ + RuntimeDir: t.TempDir(), + StartupTimeout: 15 * time.Second, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = manager.Close(context.Background()) }) + handler, err := newJupyterKernelsHandler(manager) + if err != nil { + t.Fatal(err) + } + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + + createBody, err := json.Marshal(kernelCreateRequest{ + Name: "python3", + Argv: []string{python, "-m", "ipykernel_launcher", "-f", "{connection_file}"}, + }) + if err != nil { + t.Fatal(err) + } + createResponse, err := http.Post( + server.URL+jupyterKernelsRoute, + "application/json", + strings.NewReader(string(createBody)), + ) + if err != nil { + t.Fatal(err) + } + defer createResponse.Body.Close() + if createResponse.StatusCode != http.StatusCreated { + body, _ := io.ReadAll(createResponse.Body) + t.Fatalf("create status = %d; body: %s", createResponse.StatusCode, body) + } + var kernel jupyterbridge.KernelModel + if err := json.NewDecoder(createResponse.Body).Decode(&kernel); err != nil { + t.Fatal(err) + } + + getResponse, err := http.Get(server.URL + jupyterKernelsRoute + "/" + url.PathEscape(kernel.ID)) + if err != nil { + t.Fatal(err) + } + getResponse.Body.Close() + if getResponse.StatusCode != http.StatusOK { + t.Fatalf("get status = %d, want 200", getResponse.StatusCode) + } + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + jupyterKernelsRoute + "/" + url.PathEscape(kernel.ID) + "/channels" + binaryDialer := websocket.Dialer{Subprotocols: []string{"v1.kernel.websocket.jupyter.org"}} + unsupported, response, err := binaryDialer.Dial(wsURL, nil) + if unsupported != nil { + unsupported.Close() + } + if err == nil || response == nil || response.StatusCode != http.StatusBadRequest { + t.Fatalf("binary subprotocol dial error = %v, response = %#v; want HTTP 400", err, response) + } + response.Body.Close() + + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + if err := conn.SetReadDeadline(time.Now().Add(20 * time.Second)); err != nil { + t.Fatal(err) + } + const requestID = "runme-direct-http-test" + request := map[string]any{ + "channel": "shell", + "header": map[string]any{ + "msg_id": requestID, "msg_type": "execute_request", "username": "runme", "session": "test-session", "version": "5.3", + }, + "parent_header": map[string]any{}, + "metadata": map[string]any{}, + "content": map[string]any{ + "code": "shared_http_value = 21\nshared_http_value * 2", "silent": false, "store_history": true, "user_expressions": map[string]any{}, "allow_stdin": false, "stop_on_error": true, + }, + "buffers": []any{}, + } + if err := conn.WriteJSON(request); err != nil { + t.Fatal(err) + } + seenReply, seenResult, seenIdle := false, false, false + for !(seenReply && seenResult && seenIdle) { + var message struct { + Channel string `json:"channel"` + Header map[string]any `json:"header"` + ParentHeader map[string]any `json:"parent_header"` + Content map[string]any `json:"content"` + } + if err := conn.ReadJSON(&message); err != nil { + t.Fatal(err) + } + if message.ParentHeader["msg_id"] != requestID { + continue + } + switch message.Header["msg_type"] { + case "execute_reply": + seenReply = message.Channel == "shell" && message.Content["status"] == "ok" + case "execute_result": + if data, ok := message.Content["data"].(map[string]any); ok { + seenResult = data["text/plain"] == "42" + } + case "status": + seenIdle = message.Content["execution_state"] == "idle" + } + } + request["buffers"] = []string{"AA=="} + request["header"].(map[string]any)["msg_id"] = "runme-binary-buffer-test" + if err := conn.WriteJSON(request); err != nil { + t.Fatal(err) + } + _, _, err = conn.ReadMessage() + var closeErr *websocket.CloseError + if !errors.As(err, &closeErr) || closeErr.Code != websocket.CloseUnsupportedData || closeErr.Text != "unsupported_binary_buffers" { + t.Fatalf("buffered JSON close error = %v, want code 1003 and unsupported_binary_buffers", err) + } + _ = conn.Close() + + restartResponse, err := http.Post( + server.URL+jupyterKernelsRoute+"/"+url.PathEscape(kernel.ID)+"/restart", + "application/json", + nil, + ) + if err != nil { + t.Fatal(err) + } + restartResponse.Body.Close() + if restartResponse.StatusCode != http.StatusOK { + t.Fatalf("restart status = %d, want 200", restartResponse.StatusCode) + } + + interruptResponse, err := http.Post( + server.URL+jupyterKernelsRoute+"/"+url.PathEscape(kernel.ID)+"/interrupt", + "application/json", + nil, + ) + if err != nil { + t.Fatal(err) + } + interruptResponse.Body.Close() + if interruptResponse.StatusCode != http.StatusNoContent { + t.Fatalf("interrupt status = %d, want 204", interruptResponse.StatusCode) + } + + deleteRequest, err := http.NewRequest(http.MethodDelete, server.URL+jupyterKernelsRoute+"/"+url.PathEscape(kernel.ID), nil) + if err != nil { + t.Fatal(err) + } + deleteResponse, err := http.DefaultClient.Do(deleteRequest) + if err != nil { + t.Fatal(err) + } + deleteResponse.Body.Close() + if deleteResponse.StatusCode != http.StatusNoContent { + t.Fatalf("delete status = %d, want 204", deleteResponse.StatusCode) + } +} diff --git a/pkg/agent/server/jupyter_proxy.go b/pkg/agent/server/jupyter_proxy.go deleted file mode 100644 index 55bc47695..000000000 --- a/pkg/agent/server/jupyter_proxy.go +++ /dev/null @@ -1,827 +0,0 @@ -package server - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net" - "net/http" - "net/url" - "os" - "path/filepath" - "sort" - "strings" - "sync" - "time" - - "github.com/go-logr/logr" - "github.com/gorilla/websocket" - "go.uber.org/zap" - - "github.com/runmedev/runme/v3/pkg/agent/logs" -) - -const ( - jupyterServersRoute = "/v1/jupyter/servers" - jupyterConfigDir = "jupyter" - - // Keep idle Jupyter channels sockets alive across intermediaries. - // Long-running cells can be quiet for extended periods, so liveness cannot - // depend on a steady stream of Jupyter messages. - jupyterChannelsPingInterval = 20 * time.Second - jupyterChannelsPongWait = 60 * time.Second - jupyterChannelsPingWriteTimeout = 10 * time.Second -) - -type jupyterProxyHandler struct { - registry *jupyterServerRegistry - httpClient *http.Client - upgrader websocket.Upgrader -} - -func newJupyterProxyHandler(configDir string) (*jupyterProxyHandler, error) { - registry, err := newJupyterServerRegistry(configDir) - if err != nil { - return nil, err - } - - return &jupyterProxyHandler{ - registry: registry, - httpClient: &http.Client{ - Timeout: 120 * time.Second, - }, - upgrader: websocket.Upgrader{ - ReadBufferSize: 1024, - WriteBufferSize: 1024, - CheckOrigin: func(_ *http.Request) bool { - return true - }, - }, - }, nil -} - -func (h *jupyterProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { - path := r.URL.Path - switch { - case path == jupyterServersRoute: - if r.Method != http.MethodGet { - writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") - return - } - h.handleListServers(w, r) - return - case strings.HasPrefix(path, jupyterServersRoute+"/"): - h.handleServerSubroute(w, r) - return - default: - writeHTTPError(w, http.StatusNotFound, "not found") - return - } -} - -func (h *jupyterProxyHandler) handleListServers(w http.ResponseWriter, r *http.Request) { - records := h.registry.List() - writeJSON(w, http.StatusOK, records) -} - -func (h *jupyterProxyHandler) handleServerSubroute(w http.ResponseWriter, r *http.Request) { - trimmed := strings.TrimPrefix(r.URL.Path, jupyterServersRoute+"/") - parts := strings.Split(trimmed, "/") - if len(parts) < 2 { - writeHTTPError(w, http.StatusNotFound, "not found") - return - } - - serverName, err := url.PathUnescape(parts[0]) - if err != nil || strings.TrimSpace(serverName) == "" { - writeHTTPError(w, http.StatusBadRequest, "invalid server name") - return - } - - if parts[1] != "kernels" { - writeHTTPError(w, http.StatusNotFound, "not found") - return - } - - switch { - case len(parts) == 2: - switch r.Method { - case http.MethodGet, http.MethodPost: - h.forwardHTTP(w, r, serverName, "/api/kernels") - default: - writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") - } - return - case len(parts) == 3: - kernelSegment, err := url.PathUnescape(parts[2]) - if err != nil || strings.TrimSpace(kernelSegment) == "" { - writeHTTPError(w, http.StatusBadRequest, "invalid kernel id") - return - } - switch { - case strings.HasSuffix(kernelSegment, ":interrupt"): - if r.Method != http.MethodPost { - writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") - return - } - kernelID := strings.TrimSuffix(kernelSegment, ":interrupt") - if strings.TrimSpace(kernelID) == "" { - writeHTTPError(w, http.StatusBadRequest, "invalid kernel id") - return - } - h.forwardHTTP(w, r, serverName, "/api/kernels/"+url.PathEscape(kernelID)+"/interrupt") - return - case strings.HasSuffix(kernelSegment, ":restart"): - if r.Method != http.MethodPost { - writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") - return - } - kernelID := strings.TrimSuffix(kernelSegment, ":restart") - if strings.TrimSpace(kernelID) == "" { - writeHTTPError(w, http.StatusBadRequest, "invalid kernel id") - return - } - h.forwardHTTP(w, r, serverName, "/api/kernels/"+url.PathEscape(kernelID)+"/restart") - return - default: - if r.Method != http.MethodDelete { - writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") - return - } - h.forwardHTTP(w, r, serverName, "/api/kernels/"+url.PathEscape(kernelSegment)) - return - } - case len(parts) == 4 && parts[3] == "channels": - kernelID, err := url.PathUnescape(parts[2]) - if err != nil || strings.TrimSpace(kernelID) == "" { - writeHTTPError(w, http.StatusBadRequest, "invalid kernel id") - return - } - if r.Method != http.MethodGet { - writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") - return - } - h.forwardChannelsWebSocket(w, r, serverName, kernelID) - return - default: - writeHTTPError(w, http.StatusNotFound, "not found") - return - } -} - -func (h *jupyterProxyHandler) forwardHTTP(w http.ResponseWriter, r *http.Request, serverName, upstreamPath string) { - ctx := r.Context() - log := logs.FromContext(ctx) - - server, err := h.registry.Resolve(serverName) - if err != nil { - writeHTTPError(w, http.StatusNotFound, err.Error()) - return - } - if strings.TrimSpace(server.Token) == "" { - writeHTTPError(w, http.StatusBadRequest, "server is missing jupyter token") - return - } - - requestBody, err := readProxyRequestBody(r) - if err != nil { - log.Error(err, "failed to read jupyter proxy request body", "server", serverName) - writeHTTPError(w, http.StatusBadRequest, "failed to read request body") - return - } - - resp, err := h.doJupyterHTTPRequest(ctx, r, server, upstreamPath, requestBody) - if err != nil { - log.Error(err, "Jupyter upstream request failed", "server", serverName) - writeHTTPError(w, http.StatusBadGateway, "failed to contact jupyter server") - return - } - - if resp.StatusCode == http.StatusForbidden { - refreshedServer, refreshErr := h.registry.ResolveRefreshed(serverName) - if refreshErr != nil { - log.Error(refreshErr, "failed to refresh jupyter server config after forbidden response", "server", serverName) - } else if shouldRetryJupyterRequestWithRefreshedToken(server, refreshedServer) { - log.Info("retrying jupyter request with refreshed token", "server", serverName, "method", r.Method) - resp.Body.Close() - resp, err = h.doJupyterHTTPRequest(ctx, r, refreshedServer, upstreamPath, requestBody) - if err != nil { - log.Error(err, "Jupyter upstream retry failed", "server", serverName) - writeHTTPError(w, http.StatusBadGateway, "failed to contact jupyter server") - return - } - } - } - defer resp.Body.Close() - - copyProxyResponseHeaders(w.Header(), resp.Header) - w.WriteHeader(resp.StatusCode) - if _, err := io.Copy(w, resp.Body); err != nil { - log.Error(err, "failed to copy jupyter response body") - } -} - -func (h *jupyterProxyHandler) doJupyterHTTPRequest(ctx context.Context, r *http.Request, server jupyterServerRecord, upstreamPath string, requestBody []byte) (*http.Response, error) { - upstreamURL, err := buildUpstreamURL(server.BaseURL, upstreamPath, r.URL.RawQuery) - if err != nil { - return nil, errors.New("failed to construct upstream url") - } - setUpstreamAuthToken(upstreamURL, server.Token) - - var body io.Reader - if len(requestBody) > 0 { - body = bytes.NewReader(requestBody) - } - upstreamReq, err := http.NewRequestWithContext(ctx, r.Method, upstreamURL.String(), body) - if err != nil { - return nil, errors.New("failed to build upstream request") - } - copyProxyRequestHeaders(upstreamReq.Header, r.Header) - upstreamReq.Header.Set("Authorization", "token "+server.Token) - - resp, err := h.httpClient.Do(upstreamReq) - if err != nil { - return nil, err - } - return resp, nil -} - -func readProxyRequestBody(r *http.Request) ([]byte, error) { - if r.Body == nil { - return nil, nil - } - return io.ReadAll(r.Body) -} - -func shouldRetryJupyterRequestWithRefreshedToken(oldServer, newServer jupyterServerRecord) bool { - oldToken := strings.TrimSpace(oldServer.Token) - newToken := strings.TrimSpace(newServer.Token) - return newToken != "" && newToken != oldToken -} - -func (h *jupyterProxyHandler) forwardChannelsWebSocket(w http.ResponseWriter, r *http.Request, serverName, kernelID string) { - ctx := r.Context() - log := logs.FromContext(ctx) - - server, err := h.registry.Resolve(serverName) - if err != nil { - writeHTTPError(w, http.StatusNotFound, err.Error()) - return - } - if strings.TrimSpace(server.Token) == "" { - writeHTTPError(w, http.StatusBadRequest, "server is missing jupyter token") - return - } - - upstreamURL, err := buildUpstreamURL( - server.BaseURL, - "/api/kernels/"+url.PathEscape(kernelID)+"/channels", - sanitizeClientQueryForUpstream(r.URL.RawQuery), - ) - if err != nil { - writeHTTPError(w, http.StatusInternalServerError, "failed to construct upstream url") - return - } - setUpstreamAuthToken(upstreamURL, server.Token) - switch upstreamURL.Scheme { - case "http": - upstreamURL.Scheme = "ws" - case "https": - upstreamURL.Scheme = "wss" - default: - writeHTTPError(w, http.StatusInternalServerError, "invalid upstream url scheme") - return - } - - upstreamHeader := http.Header{} - copyWebSocketDialHeaders(upstreamHeader, r.Header) - upstreamHeader.Set("Authorization", "token "+server.Token) - - upstreamConn, upstreamResp, err := websocket.DefaultDialer.DialContext(ctx, upstreamURL.String(), upstreamHeader) - if err != nil { - statusCode := http.StatusBadGateway - if upstreamResp != nil && upstreamResp.StatusCode > 0 { - statusCode = upstreamResp.StatusCode - } - log.Error(err, "failed to connect to jupyter channels websocket", "server", serverName, "kernel_id", kernelID) - writeHTTPError(w, statusCode, "failed to connect to jupyter kernel channels") - return - } - defer upstreamConn.Close() - - clientConn, err := h.upgrader.Upgrade(w, r, nil) - if err != nil { - log.Error(err, "failed to upgrade client websocket for jupyter channels") - return - } - defer clientConn.Close() - - if err := configureWebSocketPongKeepalive(clientConn, jupyterChannelsPongWait); err != nil { - log.Error(err, "failed to configure jupyter client websocket keepalive") - } - - errCh := make(chan error, 3) - stopKeepalive := make(chan struct{}) - defer close(stopKeepalive) - - go bridgeWebSocketMessagesWithContext(ctx, clientConn, upstreamConn, "upstream_to_client", errCh) - go bridgeWebSocketMessagesWithContext(ctx, upstreamConn, clientConn, "client_to_upstream", errCh) - go sendWebSocketKeepalivePings(clientConn, "runme_to_client_keepalive", errCh, stopKeepalive) - - if err := <-errCh; err != nil { - logJupyterWebSocketBridgeClose(log, serverName, kernelID, err) - } -} - -type jupyterServerRegistry struct { - configDir string - mu sync.RWMutex - servers map[string]jupyterServerRecord -} - -type jupyterServerRecord struct { - Name string - Runner string - BaseURL *url.URL - Token string -} - -type jupyterServerFile struct { - BaseURLSnake string `json:"base_url"` - BaseURLCamel string `json:"baseUrl"` - Token string `json:"token"` - Runner string `json:"runner"` -} - -type jupyterServerPublicRecord struct { - Name string `json:"name"` - Runner string `json:"runner"` - BaseURL string `json:"base_url"` - HasToken bool `json:"has_token"` -} - -func newJupyterServerRegistry(configDir string) (*jupyterServerRegistry, error) { - if strings.TrimSpace(configDir) == "" { - return nil, errors.New("config directory is required") - } - registry := &jupyterServerRegistry{ - configDir: filepath.Join(configDir, jupyterConfigDir), - servers: make(map[string]jupyterServerRecord), - } - if err := registry.reload(); err != nil { - return nil, err - } - return registry, nil -} - -func (r *jupyterServerRegistry) List() []jupyterServerPublicRecord { - r.mu.RLock() - defer r.mu.RUnlock() - - records := make([]jupyterServerPublicRecord, 0, len(r.servers)) - for _, server := range r.servers { - records = append(records, jupyterServerPublicRecord{ - Name: server.Name, - Runner: server.Runner, - BaseURL: server.BaseURL.String(), - HasToken: strings.TrimSpace(server.Token) != "", - }) - } - sort.Slice(records, func(i, j int) bool { - return records[i].Name < records[j].Name - }) - return records -} - -func (r *jupyterServerRegistry) Resolve(name string) (jupyterServerRecord, error) { - if strings.TrimSpace(name) == "" { - return jupyterServerRecord{}, errors.New("server name is required") - } - - r.mu.RLock() - server, ok := r.servers[name] - r.mu.RUnlock() - if ok { - return server, nil - } - - if err := r.reload(); err != nil { - return jupyterServerRecord{}, err - } - - r.mu.RLock() - server, ok = r.servers[name] - r.mu.RUnlock() - if !ok { - return jupyterServerRecord{}, fmt.Errorf("jupyter server %q not found", name) - } - return server, nil -} - -func (r *jupyterServerRegistry) ResolveRefreshed(name string) (jupyterServerRecord, error) { - if strings.TrimSpace(name) == "" { - return jupyterServerRecord{}, errors.New("server name is required") - } - if err := r.reload(); err != nil { - return jupyterServerRecord{}, err - } - - r.mu.RLock() - server, ok := r.servers[name] - r.mu.RUnlock() - if !ok { - return jupyterServerRecord{}, fmt.Errorf("jupyter server %q not found", name) - } - return server, nil -} - -func (r *jupyterServerRegistry) reload() error { - entries, err := os.ReadDir(r.configDir) - if err != nil { - if errors.Is(err, os.ErrNotExist) { - r.mu.Lock() - r.servers = make(map[string]jupyterServerRecord) - r.mu.Unlock() - return nil - } - return err - } - - next := make(map[string]jupyterServerRecord) - for _, entry := range entries { - if entry.IsDir() { - continue - } - if !strings.HasSuffix(entry.Name(), ".json") { - continue - } - serverName := strings.TrimSuffix(entry.Name(), ".json") - fullPath := filepath.Join(r.configDir, entry.Name()) - server, err := loadJupyterServerFile(serverName, fullPath) - if err != nil { - zap.L().Warn( - "Skipping invalid jupyter server config", - zap.String("path", fullPath), - zap.Error(err), - ) - continue - } - next[serverName] = server - } - - r.mu.Lock() - r.servers = next - r.mu.Unlock() - return nil -} - -func loadJupyterServerFile(name, path string) (jupyterServerRecord, error) { - raw, err := os.ReadFile(path) - if err != nil { - return jupyterServerRecord{}, err - } - - var payload jupyterServerFile - if err := json.Unmarshal(raw, &payload); err != nil { - return jupyterServerRecord{}, err - } - - baseURLRaw := strings.TrimSpace(payload.BaseURLSnake) - if baseURLRaw == "" { - baseURLRaw = strings.TrimSpace(payload.BaseURLCamel) - } - if baseURLRaw == "" { - return jupyterServerRecord{}, errors.New("base_url is required") - } - - parsed, err := url.Parse(baseURLRaw) - if err != nil { - return jupyterServerRecord{}, errors.New("base_url is invalid") - } - if parsed.Scheme == "" || parsed.Host == "" { - return jupyterServerRecord{}, errors.New("base_url must include scheme and host") - } - if parsed.Path == "" { - parsed.Path = "/" - } - if !strings.HasSuffix(parsed.Path, "/") { - parsed.Path += "/" - } - - return jupyterServerRecord{ - Name: name, - Runner: strings.TrimSpace(payload.Runner), - BaseURL: parsed, - Token: strings.TrimSpace(payload.Token), - }, nil -} - -func buildUpstreamURL(base *url.URL, upstreamPath, rawQuery string) (*url.URL, error) { - if base == nil { - return nil, errors.New("base url is required") - } - if strings.TrimSpace(upstreamPath) == "" { - return nil, errors.New("upstream path is required") - } - ref := &url.URL{ - Path: upstreamPath, - RawQuery: rawQuery, - } - return base.ResolveReference(ref), nil -} - -func sanitizeClientQueryForUpstream(rawQuery string) string { - if strings.TrimSpace(rawQuery) == "" { - return "" - } - - queryValues, err := url.ParseQuery(rawQuery) - if err != nil { - return rawQuery - } - - for key := range queryValues { - if strings.EqualFold(strings.TrimSpace(key), "authorization") { - delete(queryValues, key) - } - } - - return queryValues.Encode() -} - -func setUpstreamAuthToken(upstreamURL *url.URL, token string) { - if upstreamURL == nil { - return - } - trimmedToken := strings.TrimSpace(token) - if trimmedToken == "" { - return - } - query := upstreamURL.Query() - query.Set("token", trimmedToken) - upstreamURL.RawQuery = query.Encode() -} - -func copyProxyRequestHeaders(dst, src http.Header) { - for k, values := range src { - if isHopByHopHeader(k) || isRestrictedProxyRequestHeader(k) { - continue - } - for _, value := range values { - dst.Add(k, value) - } - } -} - -func copyProxyResponseHeaders(dst, src http.Header) { - for k, values := range src { - if isHopByHopHeader(k) || isCORSResponseHeader(k) { - continue - } - for _, value := range values { - dst.Add(k, value) - } - } -} - -func copyWebSocketDialHeaders(dst, src http.Header) { - const secWebSocketProtocol = "Sec-WebSocket-Protocol" - for _, value := range src.Values(secWebSocketProtocol) { - dst.Add(secWebSocketProtocol, value) - } -} - -func isHopByHopHeader(key string) bool { - switch strings.ToLower(strings.TrimSpace(key)) { - case "connection", - "keep-alive", - "proxy-authenticate", - "proxy-authorization", - "te", - "trailers", - "transfer-encoding", - "upgrade": - return true - default: - return false - } -} - -func isCORSResponseHeader(key string) bool { - switch strings.ToLower(strings.TrimSpace(key)) { - case "access-control-allow-origin", - "access-control-allow-credentials", - "access-control-allow-methods", - "access-control-allow-headers", - "access-control-expose-headers", - "access-control-max-age": - return true - default: - return false - } -} - -func isRestrictedProxyRequestHeader(key string) bool { - switch strings.ToLower(strings.TrimSpace(key)) { - case "authorization", - "origin", - "referer", - "cookie", - "x-xsrftoken": - return true - default: - return false - } -} - -type webSocketBridgeError struct { - err error - direction string - operation string - srcRemote string - dstRemote string - messageType int -} - -func (e *webSocketBridgeError) Error() string { - if e == nil || e.err == nil { - return "" - } - return e.err.Error() -} - -func (e *webSocketBridgeError) Unwrap() error { - if e == nil { - return nil - } - return e.err -} - -func remoteAddr(conn *websocket.Conn) string { - if conn == nil || conn.UnderlyingConn() == nil { - return "" - } - return conn.UnderlyingConn().RemoteAddr().String() -} - -func logJupyterWebSocketBridgeClose(log logr.Logger, serverName, kernelID string, err error) { - if err == nil { - return - } - - fields := []any{ - "server", serverName, - "kernel_id", kernelID, - "error", err.Error(), - } - - var bridgeErr *webSocketBridgeError - if errors.As(err, &bridgeErr) { - fields = append( - fields, - "direction", bridgeErr.direction, - "operation", bridgeErr.operation, - "src_remote_addr", bridgeErr.srcRemote, - "dst_remote_addr", bridgeErr.dstRemote, - ) - if bridgeErr.messageType > 0 { - fields = append(fields, "message_type", bridgeErr.messageType) - } - } - - var closeErr *websocket.CloseError - if errors.As(err, &closeErr) { - fields = append( - fields, - "is_websocket_close", true, - "close_code", closeErr.Code, - "close_text", closeErr.Text, - ) - } else { - fields = append(fields, "is_websocket_close", false) - } - - var netErr net.Error - if errors.As(err, &netErr) { - fields = append( - fields, - "is_timeout", netErr.Timeout(), - "is_temporary", netErr.Temporary(), - ) - } - - log.Info("jupyter channels websocket bridge closed", fields...) -} - -func bridgeWebSocketMessagesWithContext( - ctx context.Context, - dst, src *websocket.Conn, - direction string, - errCh chan<- error, -) { - done := make(chan struct{}) - defer close(done) - - go func() { - select { - case <-ctx.Done(): - // Force src.ReadMessage to return so the bridge can exit promptly. - _ = src.SetReadDeadline(time.Now()) - case <-done: - } - }() - - for { - msgType, payload, err := src.ReadMessage() - if err != nil { - if ctx.Err() != nil { - err = ctx.Err() - } - errCh <- &webSocketBridgeError{ - err: err, - direction: direction, - operation: "read", - srcRemote: remoteAddr(src), - dstRemote: remoteAddr(dst), - } - return - } - if err := dst.WriteMessage(msgType, payload); err != nil { - errCh <- &webSocketBridgeError{ - err: err, - direction: direction, - operation: "write", - srcRemote: remoteAddr(src), - dstRemote: remoteAddr(dst), - messageType: msgType, - } - return - } - } -} - -func configureWebSocketPongKeepalive(conn *websocket.Conn, pongWait time.Duration) error { - if conn == nil { - return errors.New("websocket connection is required") - } - if pongWait <= 0 { - return errors.New("pong wait must be positive") - } - if err := conn.SetReadDeadline(time.Now().Add(pongWait)); err != nil { - return err - } - conn.SetPongHandler(func(_ string) error { - return conn.SetReadDeadline(time.Now().Add(pongWait)) - }) - return nil -} - -func sendWebSocketKeepalivePings(conn *websocket.Conn, direction string, errCh chan<- error, stop <-chan struct{}) { - ticker := time.NewTicker(jupyterChannelsPingInterval) - defer ticker.Stop() - - for { - select { - case <-stop: - return - case <-ticker.C: - if err := conn.WriteControl( - websocket.PingMessage, - []byte("runme-keepalive"), - time.Now().Add(jupyterChannelsPingWriteTimeout), - ); err != nil { - errCh <- &webSocketBridgeError{ - err: err, - direction: direction, - operation: "write_ping", - srcRemote: "runme", - dstRemote: remoteAddr(conn), - messageType: websocket.PingMessage, - } - return - } - } - } -} - -func writeJSON(w http.ResponseWriter, statusCode int, value any) { - data, err := json.Marshal(value) - if err != nil { - writeHTTPError(w, http.StatusInternalServerError, "failed to serialize response") - return - } - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(statusCode) - _, _ = w.Write(data) -} - -func writeHTTPError(w http.ResponseWriter, statusCode int, message string) { - writeJSON(w, statusCode, map[string]any{ - "error": message, - }) -} - -var _ http.Handler = (*jupyterProxyHandler)(nil) diff --git a/pkg/agent/server/jupyter_proxy_test.go b/pkg/agent/server/jupyter_proxy_test.go deleted file mode 100644 index 98a543e28..000000000 --- a/pkg/agent/server/jupyter_proxy_test.go +++ /dev/null @@ -1,209 +0,0 @@ -package server - -import ( - "encoding/json" - "io" - "net/http" - "net/http/httptest" - "net/url" - "os" - "path/filepath" - "reflect" - "strings" - "testing" -) - -func TestCopyProxyRequestHeaders_StripsSensitiveHeaders(t *testing.T) { - src := http.Header{} - src.Set("Authorization", "Bearer test") - src.Set("Origin", "http://localhost:5173") - src.Set("Referer", "http://localhost:5173/") - src.Set("Cookie", "session=abc123") - src.Set("X-XSRFToken", "token") - src.Set("Content-Type", "application/json") - src.Set("X-Test", "ok") - - dst := http.Header{} - copyProxyRequestHeaders(dst, src) - - for _, key := range []string{"Authorization", "Origin", "Referer", "Cookie", "X-XSRFToken"} { - if got := dst.Get(key); got != "" { - t.Fatalf("expected %s to be filtered, got %q", key, got) - } - } - - if got := dst.Get("Content-Type"); got != "application/json" { - t.Fatalf("expected Content-Type to be copied, got %q", got) - } - if got := dst.Get("X-Test"); got != "ok" { - t.Fatalf("expected X-Test to be copied, got %q", got) - } -} - -func TestCopyProxyResponseHeaders_StripsCORSHeaders(t *testing.T) { - src := http.Header{} - src.Set("Access-Control-Allow-Origin", "") - src.Set("Access-Control-Allow-Credentials", "true") - src.Set("Access-Control-Allow-Headers", "authorization") - src.Set("Access-Control-Allow-Methods", "GET") - src.Set("Access-Control-Expose-Headers", "x-custom") - src.Set("Access-Control-Max-Age", "7200") - src.Set("Content-Type", "application/json") - src.Set("X-Test", "ok") - - dst := http.Header{} - copyProxyResponseHeaders(dst, src) - - if got := dst.Get("Access-Control-Allow-Origin"); got != "" { - t.Fatalf("expected Access-Control-Allow-Origin to be filtered, got %q", got) - } - if got := dst.Get("Access-Control-Allow-Credentials"); got != "" { - t.Fatalf("expected Access-Control-Allow-Credentials to be filtered, got %q", got) - } - if got := dst.Get("Access-Control-Allow-Headers"); got != "" { - t.Fatalf("expected Access-Control-Allow-Headers to be filtered, got %q", got) - } - if got := dst.Get("Access-Control-Allow-Methods"); got != "" { - t.Fatalf("expected Access-Control-Allow-Methods to be filtered, got %q", got) - } - if got := dst.Get("Access-Control-Expose-Headers"); got != "" { - t.Fatalf("expected Access-Control-Expose-Headers to be filtered, got %q", got) - } - if got := dst.Get("Access-Control-Max-Age"); got != "" { - t.Fatalf("expected Access-Control-Max-Age to be filtered, got %q", got) - } - - if got := dst.Get("Content-Type"); got != "application/json" { - t.Fatalf("expected Content-Type to be copied, got %q", got) - } - if got := dst.Get("X-Test"); got != "ok" { - t.Fatalf("expected X-Test to be copied, got %q", got) - } -} - -func TestSetUpstreamAuthToken(t *testing.T) { - upstreamURL, err := url.Parse("http://localhost:8888/api/kernels?foo=bar") - if err != nil { - t.Fatalf("failed to parse url: %v", err) - } - - setUpstreamAuthToken(upstreamURL, "abc123") - got := upstreamURL.Query().Get("token") - if got != "abc123" { - t.Fatalf("expected token query param to be set, got %q", got) - } - if got := upstreamURL.Query().Get("foo"); got != "bar" { - t.Fatalf("expected existing query params to be preserved, got %q", got) - } -} - -func TestSanitizeClientQueryForUpstream_RemovesAuthorization(t *testing.T) { - query := "authorization=Bearer+abc123&session_id=s1&Authorization=Bearer+ignored" - got := sanitizeClientQueryForUpstream(query) - parsed, err := url.ParseQuery(got) - if err != nil { - t.Fatalf("failed to parse sanitized query: %v", err) - } - - if auth := parsed.Get("authorization"); auth != "" { - t.Fatalf("expected authorization query to be removed, got %q", auth) - } - if auth := parsed.Get("Authorization"); auth != "" { - t.Fatalf("expected Authorization query to be removed, got %q", auth) - } - if sessionID := parsed.Get("session_id"); sessionID != "s1" { - t.Fatalf("expected session_id to be preserved, got %q", sessionID) - } -} - -func TestSanitizeClientQueryForUpstream_InvalidQueryPassthrough(t *testing.T) { - const raw = "%zz=1" - if got := sanitizeClientQueryForUpstream(raw); got != raw { - t.Fatalf("expected invalid query to pass through, got %q", got) - } -} - -func TestForwardHTTP_ReloadsTokenAfterForbidden(t *testing.T) { - var tokens []string - upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - token := r.URL.Query().Get("token") - tokens = append(tokens, token) - - if got := r.Header.Get("Authorization"); got != "token "+token { - t.Errorf("Authorization header = %q, want %q", got, "token "+token) - } - if token == "old-token" { - http.Error(w, "forbidden", http.StatusForbidden) - return - } - - if token != "new-token" { - t.Errorf("unexpected token %q", token) - http.Error(w, "unexpected token", http.StatusBadRequest) - return - } - if got := r.URL.Path; got != "/api/kernels" { - t.Errorf("upstream path = %q, want /api/kernels", got) - } - body, err := io.ReadAll(r.Body) - if err != nil { - t.Errorf("failed to read upstream request body: %v", err) - } - if got := strings.TrimSpace(string(body)); got != `{"name":"python3"}` { - t.Errorf("upstream request body = %q, want %q", got, `{"name":"python3"}`) - } - - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusCreated) - _, _ = w.Write([]byte(`{"id":"kernel-1","name":"python3"}`)) - })) - defer upstream.Close() - - configDir := t.TempDir() - writeJupyterServerConfig(t, configDir, "port-8890", upstream.URL, "old-token") - handler, err := newJupyterProxyHandler(configDir) - if err != nil { - t.Fatalf("newJupyterProxyHandler failed: %v", err) - } - writeJupyterServerConfig(t, configDir, "port-8890", upstream.URL, "new-token") - - req := httptest.NewRequest( - http.MethodPost, - "/v1/jupyter/servers/port-8890/kernels", - strings.NewReader(`{"name":"python3"}`), - ) - req.Header.Set("Content-Type", "application/json") - rec := httptest.NewRecorder() - - handler.ServeHTTP(rec, req) - - if rec.Code != http.StatusCreated { - t.Fatalf("status = %d, want %d; body: %s", rec.Code, http.StatusCreated, rec.Body.String()) - } - if got := strings.TrimSpace(rec.Body.String()); got != `{"id":"kernel-1","name":"python3"}` { - t.Fatalf("body = %q, want created kernel JSON", got) - } - wantTokens := []string{"old-token", "new-token"} - if !reflect.DeepEqual(tokens, wantTokens) { - t.Fatalf("upstream tokens = %v, want %v", tokens, wantTokens) - } -} - -func writeJupyterServerConfig(t *testing.T, configDir, name, baseURL, token string) { - t.Helper() - jupyterDir := filepath.Join(configDir, jupyterConfigDir) - if err := os.MkdirAll(jupyterDir, 0o755); err != nil { - t.Fatalf("failed to create jupyter config dir: %v", err) - } - payload := map[string]string{ - "base_url": baseURL, - "token": token, - } - raw, err := json.Marshal(payload) - if err != nil { - t.Fatalf("failed to marshal jupyter config: %v", err) - } - if err := os.WriteFile(filepath.Join(jupyterDir, name+".json"), raw, 0o600); err != nil { - t.Fatalf("failed to write jupyter config: %v", err) - } -} diff --git a/pkg/agent/server/server.go b/pkg/agent/server/server.go index 0f93c0e43..f3ac49eed 100644 --- a/pkg/agent/server/server.go +++ b/pkg/agent/server/server.go @@ -8,6 +8,7 @@ import ( "net/http" "os" "os/signal" + "path/filepath" "syscall" "time" @@ -20,6 +21,7 @@ import ( "github.com/runmedev/runme/v3/pkg/agent/api" "github.com/runmedev/runme/v3/pkg/agent/iam" + "github.com/runmedev/runme/v3/pkg/agent/jupyter" "github.com/runmedev/runme/v3/pkg/agent/runme" "github.com/runmedev/runme/v3/pkg/agent/runme/stream" @@ -55,6 +57,7 @@ type Server struct { registerHandlers RegisterHandlers assetsFS fs.FS wsHandler *stream.WebSocketHandler + jupyterManager *jupyter.KernelManager tapFactory stream.TapFactory } @@ -80,6 +83,9 @@ func NewServer(opts Options) (*Server, error) { if opts.Server == nil { opts.Server = &config.AssistantServerConfig{} } + if opts.Server.JupyterService && !opts.Server.RunnerService { + return nil, errors.New("jupyter service requires runner service") + } assetsFS, err := loadStaticAssets(opts.Server.StaticAssets) if err != nil { @@ -327,22 +333,35 @@ func (s *Server) registerServices() error { mux.HandleProtected(runnerSvcPath, runnerSvcHandler, s.checker, api.RunnerUserRole) } - jupyterProxy, err := newJupyterProxyHandler(s.configDir) - if err != nil { - return errors.Wrap(err, "failed to initialize jupyter proxy handler") + if s.serverConfig.JupyterService { + if s.jupyterManager == nil { + s.jupyterManager, err = jupyter.NewKernelManager(jupyter.KernelManagerConfig{ + RuntimeDir: filepath.Join(s.configDir, "jupyter-kernels"), + }) + if err != nil { + return errors.Wrap(err, "failed to initialize direct Jupyter kernel manager") + } + } + jupyterKernels, err := newJupyterKernelsHandler(s.jupyterManager) + if err != nil { + return errors.Wrap(err, "failed to initialize Jupyter kernels handler") + } + mux.HandleProtected( + jupyterKernelsRoute, + otelhttp.NewHandler(jupyterKernels, jupyterKernelsRoute), + s.checker, + api.RunnerUserRole, + ) + mux.HandleProtected( + jupyterKernelsRoute+"/", + otelhttp.NewHandler(jupyterKernels, jupyterKernelsRoute+"/"), + s.checker, + api.RunnerUserRole, + ) + log.Info("Setting up Jupyter kernel service", "path", jupyterKernelsRoute) + } else { + log.Info("Jupyter kernel service is disabled") } - mux.HandleProtected( - "/v1/jupyter/servers", - otelhttp.NewHandler(jupyterProxy, "/v1/jupyter/servers"), - s.checker, - api.RunnerUserRole, - ) - mux.HandleProtected( - "/v1/jupyter/servers/", - otelhttp.NewHandler(jupyterProxy, "/v1/jupyter/servers/"), - s.checker, - api.RunnerUserRole, - ) // Health check should be public checker := grpchealth.NewStaticChecker() @@ -406,6 +425,14 @@ func (s *Server) shutdown() { } log.Info("WebSocket handler shutdown complete") } + if s.jupyterManager != nil { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + if err := s.jupyterManager.Close(ctx); err != nil { + log.Error(err, "Error shutting down Jupyter kernel manager") + } + log.Info("Jupyter kernel manager shutdown complete") + } if s.hServer != nil { ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) defer cancel()