From 399c098e69bbbb83bbc5bd5d99e955f9f901fe3b Mon Sep 17 00:00:00 2001 From: Jeremy Lewi Date: Tue, 4 Aug 2026 11:50:43 +0000 Subject: [PATCH 1/8] Add direct Jupyter kernel bridge Signed-off-by: Jeremy Lewi --- go.mod | 2 + go.sum | 4 + pkg/agent/config/config.go | 11 + pkg/agent/config/config_test.go | 24 + pkg/agent/jupyter/bridge.go | 677 +++++++++++++++ pkg/agent/jupyter/bridge_integration_test.go | 557 ++++++++++++ pkg/agent/jupyter/channel_json.go | 119 +++ pkg/agent/jupyter/channel_json_test.go | 87 ++ pkg/agent/jupyter/connection.go | 144 ++++ pkg/agent/jupyter/connection_test.go | 60 ++ .../jupyter/ipykernel_integration_test.go | 405 +++++++++ pkg/agent/jupyter/kernel_manager.go | 804 ++++++++++++++++++ pkg/agent/jupyter/kernel_manager_test.go | 247 ++++++ pkg/agent/jupyter/message.go | 246 ++++++ pkg/agent/jupyter/message_test.go | 156 ++++ pkg/agent/jupyter/metrics.go | 37 + pkg/agent/jupyter/process_unix.go | 36 + pkg/agent/jupyter/process_windows.go | 38 + pkg/agent/jupyter/socket.go | 204 +++++ pkg/agent/jupyter/socket_test.go | 77 ++ pkg/agent/server/jupyter_proxy.go | 282 +++++- pkg/agent/server/jupyter_proxy_test.go | 161 ++++ pkg/agent/server/server.go | 23 +- 23 files changed, 4385 insertions(+), 16 deletions(-) create mode 100644 pkg/agent/jupyter/bridge.go create mode 100644 pkg/agent/jupyter/bridge_integration_test.go create mode 100644 pkg/agent/jupyter/channel_json.go create mode 100644 pkg/agent/jupyter/channel_json_test.go create mode 100644 pkg/agent/jupyter/connection.go create mode 100644 pkg/agent/jupyter/connection_test.go create mode 100644 pkg/agent/jupyter/ipykernel_integration_test.go create mode 100644 pkg/agent/jupyter/kernel_manager.go create mode 100644 pkg/agent/jupyter/kernel_manager_test.go create mode 100644 pkg/agent/jupyter/message.go create mode 100644 pkg/agent/jupyter/message_test.go create mode 100644 pkg/agent/jupyter/metrics.go create mode 100644 pkg/agent/jupyter/process_unix.go create mode 100644 pkg/agent/jupyter/process_windows.go create mode 100644 pkg/agent/jupyter/socket.go create mode 100644 pkg/agent/jupyter/socket_test.go diff --git a/go.mod b/go.mod index b5e29061a..f201ae2bb 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 @@ -158,6 +159,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 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/config/config.go b/pkg/agent/config/config.go index 380d0d7f4..ebe14caf9 100644 --- a/pkg/agent/config/config.go +++ b/pkg/agent/config/config.go @@ -478,6 +478,11 @@ type AssistantServerConfig struct { // ParserService starts the Runme parser service if true otherwise it doesn't start the parser service. ParserService bool `json:"parserService" yaml:"parserService"` + // Jupyter configures the optional direct ipykernel provider. It is disabled + // unless explicitly enabled so the existing Jupyter Server proxy remains the + // default execution path during rollout. + Jupyter *JupyterConfig `json:"jupyter,omitempty" yaml:"jupyter,omitempty"` + // OIDC configuration OIDC *OIDCConfig `json:"oidc,omitempty" yaml:"oidc,omitempty"` @@ -489,6 +494,12 @@ type AssistantServerConfig struct { WebAppURL string `json:"webAppURL" yaml:"webAppURL"` } +// JupyterConfig controls the feature-flagged direct ZeroMQ kernel provider. +type JupyterConfig struct { + DirectZMQ bool `json:"directZMQ" yaml:"directZMQ"` + PythonCommand string `json:"pythonCommand,omitempty" yaml:"pythonCommand,omitempty"` +} + // OIDCConfig contains configuration for OIDC authentication type OIDCConfig struct { // Google contains Google-specific OIDC configuration diff --git a/pkg/agent/config/config_test.go b/pkg/agent/config/config_test.go index 6c8b0867f..d34c64359 100644 --- a/pkg/agent/config/config_test.go +++ b/pkg/agent/config/config_test.go @@ -142,6 +142,30 @@ func Test_AssistantServerConfig_GetAgentService(t *testing.T) { } } +func Test_DirectJupyterConfiguration(t *testing.T) { + v := viper.New() + v.SetConfigType("yaml") + if err := v.ReadConfig(strings.NewReader(` +assistantServer: + runnerService: true + jupyter: + directZMQ: true + pythonCommand: /opt/runme/python +`)); err != nil { + t.Fatal(err) + } + cfg, err := getConfigFromViper(v) + if err != nil { + t.Fatal(err) + } + if cfg.AssistantServer == nil || cfg.AssistantServer.Jupyter == nil { + t.Fatal("direct Jupyter configuration was not loaded") + } + if !cfg.AssistantServer.Jupyter.DirectZMQ || cfg.AssistantServer.Jupyter.PythonCommand != "/opt/runme/python" { + t.Fatalf("Jupyter config = %#v", cfg.AssistantServer.Jupyter) + } +} + // setEnvVars sets the provided environment variables and returns a function to restore the previous state. func setEnvVars(env map[string]string) func() { if env == nil { diff --git a/pkg/agent/jupyter/bridge.go b/pkg/agent/jupyter/bridge.go new file mode 100644 index 000000000..6b72b2fc4 --- /dev/null +++ b/pkg/agent/jupyter/bridge.go @@ -0,0 +1,677 @@ +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") + } + jupyterActiveBridges.Inc() + defer jupyterActiveBridges.Dec() + 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, 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 isClosed(generationChanged): + closeMetricReason = "kernel_restart" + } + jupyterBridgeCloses.WithLabelValues(closeMetricReason).Inc() + _ = 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() + + 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 + } + infoReply := make(chan struct{}, 1) + 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" && parentID == requestID { + select { + case infoReply <- struct{}{}: + 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 + send := func() error { + 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 ctx.Err() + case err := <-errorsCh: + return err + case <-infoReply: + 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 + } + jupyterChannelMessages.WithLabelValues("browser_to_kernel", string(channel)).Inc() + jupyterChannelBytes.WithLabelValues("browser_to_kernel", string(channel)).Add(float64(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 { + jupyterRateLimits.Inc() + return err + } + b.observeKernelState(kernelID, generation, message) + } + if err := queue.Enqueue(ctx, payload); err != nil { + return err + } + jupyterChannelMessages.WithLabelValues("kernel_to_browser", string(channel)).Inc() + jupyterChannelBytes.WithLabelValues("kernel_to_browser", string(channel)).Add(float64(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 { + jupyterHeartbeatMisses.Inc() + _ = current.Close() + replacement, err := b.openSocket(ctx, SocketTypeReq, nil, connection, ChannelHeartbeat) + if err != nil { + continue + } + current = replacement + defer current.Close() + } + continue + } + jupyterHeartbeatMisses.Inc() + 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..ccb350019 --- /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, "python3") + 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, "python3") + 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, "python3") + 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, "python3") + 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, "python3") + 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, "python3") + 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, "python3") + 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..a0708e06f --- /dev/null +++ b/pkg/agent/jupyter/kernel_manager.go @@ -0,0 +1,804 @@ +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"` +} + +// LaunchProfile is a server-side allowlisted kernel command. Browser requests +// select only Name; they cannot supply Command, Args, or a connection path. +type LaunchProfile struct { + Name string + Command string + Args []string +} + +// PythonLaunchProfile returns the standard ipykernel launch profile. +func PythonLaunchProfile(command string) LaunchProfile { + if strings.TrimSpace(command) == "" { + command = "python3" + } + return LaunchProfile{ + Name: "python3", + Command: command, + Args: []string{"-m", "ipykernel_launcher", "-f", connectionFilePlaceholder}, + } +} + +// KernelManagerConfig configures local kernel lifecycle boundaries. +type KernelManagerConfig struct { + RuntimeDir string + Profiles map[string]LaunchProfile + 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 +} + +type kernelConnection struct { + Generation uint64 + Info ConnectionInfo + Context context.Context + cancel context.CancelFunc +} + +// NewKernelManager creates a manager after validating all launch profiles and +// preparing 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 len(config.Profiles) == 0 { + profile := PythonLaunchProfile("python3") + config.Profiles = map[string]LaunchProfile{profile.Name: profile} + } + profiles := make(map[string]LaunchProfile, len(config.Profiles)) + for key, profile := range config.Profiles { + if err := validateLaunchProfile(key, profile); err != nil { + return nil, err + } + profile.Args = append([]string(nil), profile.Args...) + profiles[key] = profile + } + config.Profiles = profiles + 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 validateLaunchProfile(key string, profile LaunchProfile) error { + if strings.TrimSpace(key) == "" || key != profile.Name { + return fmt.Errorf("jupyter launch profile key %q must match its non-empty name", key) + } + if strings.TrimSpace(profile.Command) == "" { + return fmt.Errorf("jupyter launch profile %q has no command", key) + } + placeholderCount := 0 + for _, arg := range profile.Args { + placeholderCount += strings.Count(arg, connectionFilePlaceholder) + } + if placeholderCount != 1 { + return fmt.Errorf("jupyter launch profile %q must contain exactly one %s placeholder", key, connectionFilePlaceholder) + } + return nil +} + +// Start launches an allowlisted profile and returns only after heartbeat and +// kernel_info readiness checks succeed. +func (m *KernelManager) Start(ctx context.Context, profileName string) (KernelModel, error) { + started := time.Now() + defer func() { observeLifecycle("start", started) }() + profile, ok := m.config.Profiles[profileName] + if !ok { + return KernelModel{}, fmt.Errorf("jupyter launch profile %q is not allowed", profileName) + } + kernel := &managedKernel{ + id: uuid.NewString(), + name: profile.Name, + state: KernelStateStarting, + lastActivity: time.Now().UTC(), + } + 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 + jupyterKernelStates.WithLabelValues(string(kernel.state)).Inc() + m.mu.Unlock() + + if err := m.launch(ctx, kernel, profile); err != nil { + _ = m.stopProcess(context.Background(), kernel, false) + m.mu.Lock() + jupyterKernelStates.WithLabelValues(string(kernel.state)).Dec() + 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, profile LaunchProfile) 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") + args := make([]string, len(profile.Args)) + for i, arg := range profile.Args { + args[i] = strings.ReplaceAll(arg, connectionFilePlaceholder, connectionPath) + } + diagnostics := newBoundedWriter(m.config.DiagnosticsBytes) + command := exec.Command(profile.Command, args...) + command.Stdout = diagnostics + command.Stderr = diagnostics + configureManagedProcess(command) + if err := command.Start(); err != nil { + _ = os.RemoveAll(runtimeDir) + return fmt.Errorf("launch profile %q: %w", profile.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) + jupyterUnexpectedExits.Inc() + 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) + } + profile := m.config.Profiles[kernel.name] + 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, profile); 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() + jupyterKernelStates.WithLabelValues(string(kernel.state)).Dec() + 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 + } + jupyterKernelStates.WithLabelValues(string(kernel.state)).Dec() + kernel.state = state + jupyterKernelStates.WithLabelValues(string(kernel.state)).Inc() +} + +// 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..c13ddb8b5 --- /dev/null +++ b/pkg/agent/jupyter/kernel_manager_test.go @@ -0,0 +1,247 @@ +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, "python3") + 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, "python3") + 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, "python3") + 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, "python3"); err == nil { + t.Fatal("Start() succeeded after Close()") + } +} + +func TestKernelManagerRejectsArbitraryLaunchCommands(t *testing.T) { + python := findIPyKernelPython(t) + _, err := NewKernelManager(KernelManagerConfig{ + RuntimeDir: t.TempDir(), + Profiles: map[string]LaunchProfile{ + "python3": {Name: "python3", Command: python, Args: []string{"-c", "pass"}}, + }, + }) + if err == nil || !strings.Contains(err.Error(), connectionFilePlaceholder) { + t.Fatalf("NewKernelManager() error = %v, want missing placeholder error", err) + } + + manager := newTestKernelManager(t) + defer manager.Close(context.Background()) + if _, err := manager.Start(context.Background(), "browser-supplied-command"); err == nil { + t.Fatal("Start() accepted a non-allowlisted profile") + } +} + +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() + python := findIPyKernelPython(t) + profile := PythonLaunchProfile(python) + manager, err := NewKernelManager(KernelManagerConfig{ + RuntimeDir: filepath.Join(t.TempDir(), "runtime"), + Profiles: map[string]LaunchProfile{profile.Name: profile}, + 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 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..abed5afd4 --- /dev/null +++ b/pkg/agent/jupyter/metrics.go @@ -0,0 +1,37 @@ +package jupyter + +import ( + "strings" + "time" + + "github.com/prometheus/client_golang/prometheus" +) + +var ( + jupyterKernelStates = prometheus.NewGaugeVec(prometheus.GaugeOpts{Name: "runme_jupyter_kernels", Help: "Current directly managed Jupyter kernels by lifecycle state."}, []string{"state"}) + jupyterLifecycleLatency = prometheus.NewHistogramVec(prometheus.HistogramOpts{Name: "runme_jupyter_lifecycle_duration_seconds", Help: "Duration of direct Jupyter kernel lifecycle operations.", Buckets: prometheus.DefBuckets}, []string{"operation"}) + jupyterActiveBridges = prometheus.NewGauge(prometheus.GaugeOpts{Name: "runme_jupyter_active_bridges", Help: "Current direct browser-to-kernel Jupyter bridges."}) + jupyterChannelMessages = prometheus.NewCounterVec(prometheus.CounterOpts{Name: "runme_jupyter_channel_messages_total", Help: "Direct Jupyter channel messages by direction and channel."}, []string{"direction", "channel"}) + jupyterChannelBytes = prometheus.NewCounterVec(prometheus.CounterOpts{Name: "runme_jupyter_channel_bytes_total", Help: "Direct Jupyter channel payload bytes by direction and channel."}, []string{"direction", "channel"}) + jupyterProtocolErrors = prometheus.NewCounterVec(prometheus.CounterOpts{Name: "runme_jupyter_protocol_errors_total", Help: "Rejected direct Jupyter protocol messages by bounded reason."}, []string{"reason"}) + jupyterRateLimits = prometheus.NewCounter(prometheus.CounterOpts{Name: "runme_jupyter_rate_limits_total", Help: "Direct Jupyter bridges closed by IOPub rate limits."}) + jupyterHeartbeatMisses = prometheus.NewCounter(prometheus.CounterOpts{Name: "runme_jupyter_heartbeat_misses_total", Help: "Missed direct Jupyter kernel heartbeat probes."}) + jupyterUnexpectedExits = prometheus.NewCounter(prometheus.CounterOpts{Name: "runme_jupyter_unexpected_exits_total", Help: "Unexpected exits of directly managed Jupyter kernels."}) + jupyterBridgeCloses = prometheus.NewCounterVec(prometheus.CounterOpts{Name: "runme_jupyter_bridge_closes_total", Help: "Direct Jupyter bridge closes by bounded reason."}, []string{"reason"}) +) + +func init() { + prometheus.MustRegister(jupyterKernelStates, jupyterLifecycleLatency, jupyterActiveBridges, jupyterChannelMessages, jupyterChannelBytes, jupyterProtocolErrors, jupyterRateLimits, jupyterHeartbeatMisses, jupyterUnexpectedExits, jupyterBridgeCloses) +} + +func observeLifecycle(operation string, started time.Time) { + jupyterLifecycleLatency.WithLabelValues(operation).Observe(time.Since(started).Seconds()) +} + +func observeProtocolError(err error) { + reason := "malformed_message" + if err != nil && strings.Contains(strings.ToLower(err.Error()), "signature") { + reason = "bad_signature" + } + jupyterProtocolErrors.WithLabelValues(reason).Inc() +} 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_proxy.go b/pkg/agent/server/jupyter_proxy.go index 55bc47695..4361ade52 100644 --- a/pkg/agent/server/jupyter_proxy.go +++ b/pkg/agent/server/jupyter_proxy.go @@ -21,12 +21,15 @@ import ( "github.com/gorilla/websocket" "go.uber.org/zap" + jupyterbridge "github.com/runmedev/runme/v3/pkg/agent/jupyter" "github.com/runmedev/runme/v3/pkg/agent/logs" ) const ( - jupyterServersRoute = "/v1/jupyter/servers" - jupyterConfigDir = "jupyter" + jupyterServersRoute = "/v1/jupyter/servers" + jupyterConfigDir = "jupyter" + directZMQProvider = "direct-zmq" + maxDirectRequestBody = 64 * 1024 // Keep idle Jupyter channels sockets alive across intermediaries. // Long-running cells can be quiet for extended periods, so liveness cannot @@ -37,19 +40,250 @@ const ( ) type jupyterProxyHandler struct { - registry *jupyterServerRegistry - httpClient *http.Client - upgrader websocket.Upgrader + registry *jupyterServerRegistry + directManager *jupyterbridge.KernelManager + directBridge *jupyterbridge.KernelChannelsBridge + httpClient *http.Client + upgrader websocket.Upgrader } -func newJupyterProxyHandler(configDir string) (*jupyterProxyHandler, error) { +type directKernelCreateRequest struct { + Name string `json:"name"` + Path string `json:"path"` +} + +func (h *jupyterProxyHandler) handleDirectSubroute(w http.ResponseWriter, r *http.Request, parts []string) { + switch { + case len(parts) == 2: + switch r.Method { + case http.MethodGet: + writeJSON(w, http.StatusOK, h.directManager.List()) + case http.MethodPost: + var request directKernelCreateRequest + decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, maxDirectRequestBody)) + if err := decoder.Decode(&request); err != nil { + writeHTTPError(w, http.StatusBadRequest, "invalid kernel request") + return + } + if strings.TrimSpace(request.Path) != "" { + writeHTTPError(w, http.StatusBadRequest, "custom kernel paths are not supported") + return + } + profile := strings.TrimSpace(request.Name) + if profile == "" { + profile = "python3" + } + kernel, err := h.directManager.Start(r.Context(), profile) + if err != nil { + logs.FromContext(r.Context()).Error(err, "failed to start direct Jupyter kernel", "profile", profile) + writeHTTPError(w, http.StatusBadRequest, "failed to start kernel") + return + } + logs.FromContext(r.Context()).Info("started direct Jupyter kernel", "kernel_id", kernel.ID, "profile", profile) + writeJSON(w, http.StatusCreated, kernel) + 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 + } + action := "delete" + kernelID := kernelSegment + if strings.HasSuffix(kernelSegment, ":interrupt") { + action = "interrupt" + kernelID = strings.TrimSuffix(kernelSegment, ":interrupt") + } else if strings.HasSuffix(kernelSegment, ":restart") { + action = "restart" + kernelID = strings.TrimSuffix(kernelSegment, ":restart") + } + if strings.TrimSpace(kernelID) == "" { + writeHTTPError(w, http.StatusBadRequest, "invalid kernel id") + return + } + if _, err := h.directManager.Get(kernelID); err != nil { + writeHTTPError(w, http.StatusNotFound, "kernel not found") + return + } + switch action { + case "interrupt": + if r.Method != http.MethodPost { + writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + if err := h.directManager.Interrupt(kernelID); err != nil { + writeDirectKernelOperationError(w, r, "interrupt", err) + return + } + logs.FromContext(r.Context()).Info("interrupted direct Jupyter kernel", "kernel_id", kernelID) + w.WriteHeader(http.StatusNoContent) + case "restart": + if r.Method != http.MethodPost { + writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + kernel, err := h.directManager.Restart(r.Context(), kernelID) + if err != nil { + writeDirectKernelOperationError(w, r, "restart", err) + return + } + logs.FromContext(r.Context()).Info("restarted direct Jupyter kernel", "kernel_id", kernelID) + writeJSON(w, http.StatusOK, kernel) + default: + if r.Method != http.MethodDelete { + writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + logs.FromContext(r.Context()).Info("stopping direct Jupyter kernel", "kernel_id", kernelID) + if err := h.directManager.Stop(r.Context(), kernelID); err != nil { + writeDirectKernelOperationError(w, r, "stop", err) + return + } + logs.FromContext(r.Context()).Info("stopped direct Jupyter kernel", "kernel_id", kernelID) + w.WriteHeader(http.StatusNoContent) + } + return + case len(parts) == 4 && parts[3] == "channels": + if r.Method != http.MethodGet { + writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + kernelID, err := url.PathUnescape(parts[2]) + if err != nil || strings.TrimSpace(kernelID) == "" { + writeHTTPError(w, http.StatusBadRequest, "invalid kernel id") + return + } + if _, err := h.directManager.Get(kernelID); err != nil { + writeHTTPError(w, http.StatusNotFound, "kernel not found") + return + } + h.handleDirectChannels(w, r, kernelID) + return + default: + writeHTTPError(w, http.StatusNotFound, "not found") + } +} + +func writeDirectKernelOperationError(w http.ResponseWriter, r *http.Request, operation string, err error) { + logs.FromContext(r.Context()).Error(err, "direct Jupyter kernel operation failed", "operation", operation) + writeHTTPError(w, http.StatusInternalServerError, "kernel operation failed") +} + +func (h *jupyterProxyHandler) handleDirectChannels(w http.ResponseWriter, r *http.Request, kernelID string) { + 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 direct Jupyter channels WebSocket") + return + } + client := newDirectChannelsClient(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, "runme_to_direct_client_keepalive", keepaliveErr, stopKeepalive) + bridgeErr := h.directBridge.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, "direct Jupyter channels bridge closed", "kernel_id", kernelID) + } + _ = client.Close(websocket.CloseNormalClosure, "Jupyter channels closed") +} + +type directChannelsClient struct { + conn *websocket.Conn + closeOnce sync.Once +} + +func newDirectChannelsClient(conn *websocket.Conn) *directChannelsClient { + return &directChannelsClient{conn: conn} +} + +func (c *directChannelsClient) 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 *directChannelsClient) 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 *directChannelsClient) 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 newJupyterProxyHandler(configDir string, managers ...*jupyterbridge.KernelManager) (*jupyterProxyHandler, error) { registry, err := newJupyterServerRegistry(configDir) if err != nil { return nil, err } + var directManager *jupyterbridge.KernelManager + if len(managers) > 0 { + directManager = managers[0] + } + var directBridge *jupyterbridge.KernelChannelsBridge + if directManager != nil { + directBridge, err = jupyterbridge.NewKernelChannelsBridge(directManager, nil, jupyterbridge.BridgeConfig{}) + if err != nil { + return nil, err + } + } return &jupyterProxyHandler{ - registry: registry, + registry: registry, + directManager: directManager, + directBridge: directBridge, httpClient: &http.Client{ Timeout: 120 * time.Second, }, @@ -84,6 +318,16 @@ func (h *jupyterProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) func (h *jupyterProxyHandler) handleListServers(w http.ResponseWriter, r *http.Request) { records := h.registry.List() + if h.directManager != nil { + records = append(records, jupyterServerPublicRecord{ + Name: directZMQProvider, + BaseURL: "", + HasToken: false, + Backend: directZMQProvider, + Capabilities: []string{"kernels", "channels", "interrupt", "restart"}, + }) + sort.Slice(records, func(i, j int) bool { return records[i].Name < records[j].Name }) + } writeJSON(w, http.StatusOK, records) } @@ -105,6 +349,10 @@ func (h *jupyterProxyHandler) handleServerSubroute(w http.ResponseWriter, r *htt writeHTTPError(w, http.StatusNotFound, "not found") return } + if serverName == directZMQProvider && h.directManager != nil { + h.handleDirectSubroute(w, r, parts) + return + } switch { case len(parts) == 2: @@ -357,10 +605,12 @@ type jupyterServerFile struct { } type jupyterServerPublicRecord struct { - Name string `json:"name"` - Runner string `json:"runner"` - BaseURL string `json:"base_url"` - HasToken bool `json:"has_token"` + Name string `json:"name"` + Runner string `json:"runner"` + BaseURL string `json:"base_url"` + HasToken bool `json:"has_token"` + Backend string `json:"backend"` + Capabilities []string `json:"capabilities"` } func newJupyterServerRegistry(configDir string) (*jupyterServerRegistry, error) { @@ -384,10 +634,12 @@ func (r *jupyterServerRegistry) List() []jupyterServerPublicRecord { 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) != "", + Name: server.Name, + Runner: server.Runner, + BaseURL: server.BaseURL.String(), + HasToken: strings.TrimSpace(server.Token) != "", + Backend: "jupyter-server", + Capabilities: []string{"kernels", "channels", "interrupt", "restart"}, }) } sort.Slice(records, func(i, j int) bool { diff --git a/pkg/agent/server/jupyter_proxy_test.go b/pkg/agent/server/jupyter_proxy_test.go index 98a543e28..ab5a869df 100644 --- a/pkg/agent/server/jupyter_proxy_test.go +++ b/pkg/agent/server/jupyter_proxy_test.go @@ -1,6 +1,7 @@ package server import ( + "context" "encoding/json" "io" "net/http" @@ -11,8 +12,168 @@ import ( "reflect" "strings" "testing" + "time" + + "github.com/gorilla/websocket" + + jupyterbridge "github.com/runmedev/runme/v3/pkg/agent/jupyter" ) +func TestJupyterProviderDiscoveryFeatureFlag(t *testing.T) { + configDir := t.TempDir() + disabled, err := newJupyterProxyHandler(configDir) + if err != nil { + t.Fatal(err) + } + disabledResponse := httptest.NewRecorder() + disabled.ServeHTTP(disabledResponse, httptest.NewRequest(http.MethodGet, jupyterServersRoute, nil)) + if strings.Contains(disabledResponse.Body.String(), directZMQProvider) { + t.Fatalf("disabled provider was advertised: %s", disabledResponse.Body.String()) + } + + profile := jupyterbridge.PythonLaunchProfile("python3") + manager, err := jupyterbridge.NewKernelManager(jupyterbridge.KernelManagerConfig{ + RuntimeDir: t.TempDir(), + Profiles: map[string]jupyterbridge.LaunchProfile{profile.Name: profile}, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = manager.Close(context.Background()) }) + enabled, err := newJupyterProxyHandler(configDir, manager) + if err != nil { + t.Fatal(err) + } + enabledResponse := httptest.NewRecorder() + enabled.ServeHTTP(enabledResponse, httptest.NewRequest(http.MethodGet, jupyterServersRoute, nil)) + if enabledResponse.Code != http.StatusOK { + t.Fatalf("status = %d; body: %s", enabledResponse.Code, enabledResponse.Body.String()) + } + var records []jupyterServerPublicRecord + if err := json.Unmarshal(enabledResponse.Body.Bytes(), &records); err != nil { + t.Fatal(err) + } + if len(records) != 1 || records[0].Name != directZMQProvider || records[0].Backend != directZMQProvider { + t.Fatalf("records = %#v, want direct provider", records) + } + if records[0].HasToken || len(records[0].Capabilities) != 4 { + t.Fatalf("direct provider metadata = %#v", records[0]) + } +} + +func TestDirectJupyterHTTPAndWebSocketCompatibility(t *testing.T) { + python := os.Getenv("RUNME_TEST_PYTHON") + if python == "" { + t.Skip("set RUNME_TEST_PYTHON to an ipykernel-capable Python interpreter") + } + profile := jupyterbridge.PythonLaunchProfile(python) + manager, err := jupyterbridge.NewKernelManager(jupyterbridge.KernelManagerConfig{ + RuntimeDir: t.TempDir(), + Profiles: map[string]jupyterbridge.LaunchProfile{profile.Name: profile}, + StartupTimeout: 15 * time.Second, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = manager.Close(context.Background()) }) + handler, err := newJupyterProxyHandler(t.TempDir(), manager) + if err != nil { + t.Fatal(err) + } + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + + createResponse, err := http.Post( + server.URL+jupyterServersRoute+"/"+directZMQProvider+"/kernels", + "application/json", + strings.NewReader(`{"name":"python3"}`), + ) + 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) + } + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + jupyterServersRoute + "/" + directZMQProvider + "/kernels/" + url.PathEscape(kernel.ID) + "/channels" + 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": "21 * 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" + } + } + + restartResponse, err := http.Post(server.URL+jupyterServersRoute+"/"+directZMQProvider+"/kernels/"+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) + } + deleteRequest, err := http.NewRequest(http.MethodDelete, server.URL+jupyterServersRoute+"/"+directZMQProvider+"/kernels/"+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) + } +} + func TestCopyProxyRequestHeaders_StripsSensitiveHeaders(t *testing.T) { src := http.Header{} src.Set("Authorization", "Bearer test") diff --git a/pkg/agent/server/server.go b/pkg/agent/server/server.go index 0f93c0e43..ada0fcd99 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 } @@ -327,7 +330,17 @@ func (s *Server) registerServices() error { mux.HandleProtected(runnerSvcPath, runnerSvcHandler, s.checker, api.RunnerUserRole) } - jupyterProxy, err := newJupyterProxyHandler(s.configDir) + if s.serverConfig.Jupyter != nil && s.serverConfig.Jupyter.DirectZMQ && s.jupyterManager == nil { + pythonProfile := jupyter.PythonLaunchProfile(s.serverConfig.Jupyter.PythonCommand) + s.jupyterManager, err = jupyter.NewKernelManager(jupyter.KernelManagerConfig{ + RuntimeDir: filepath.Join(s.configDir, "jupyter-kernels"), + Profiles: map[string]jupyter.LaunchProfile{pythonProfile.Name: pythonProfile}, + }) + if err != nil { + return errors.Wrap(err, "failed to initialize direct Jupyter kernel manager") + } + } + jupyterProxy, err := newJupyterProxyHandler(s.configDir, s.jupyterManager) if err != nil { return errors.Wrap(err, "failed to initialize jupyter proxy handler") } @@ -406,6 +419,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() From 1cbea03ed22a7c43374d5deba2f0d85cfd8aa38d Mon Sep 17 00:00:00 2001 From: Jeremy lewi Date: Wed, 5 Aug 2026 13:39:38 +0200 Subject: [PATCH 2/8] Replace Jupyter proxy with direct kernel API Signed-off-by: Jeremy lewi --- pkg/agent/config/config.go | 8 +- pkg/agent/config/config_test.go | 7 +- pkg/agent/jupyter/bridge.go | 4 + pkg/agent/jupyter/metrics.go | 5 +- pkg/agent/server/jupyter_kernels.go | 344 +++++++ pkg/agent/server/jupyter_kernels_test.go | 294 ++++++ pkg/agent/server/jupyter_proxy.go | 1079 ---------------------- pkg/agent/server/jupyter_proxy_test.go | 370 -------- pkg/agent/server/server.go | 20 +- 9 files changed, 664 insertions(+), 1467 deletions(-) create mode 100644 pkg/agent/server/jupyter_kernels.go create mode 100644 pkg/agent/server/jupyter_kernels_test.go delete mode 100644 pkg/agent/server/jupyter_proxy.go delete mode 100644 pkg/agent/server/jupyter_proxy_test.go diff --git a/pkg/agent/config/config.go b/pkg/agent/config/config.go index ebe14caf9..a6e880840 100644 --- a/pkg/agent/config/config.go +++ b/pkg/agent/config/config.go @@ -478,9 +478,8 @@ type AssistantServerConfig struct { // ParserService starts the Runme parser service if true otherwise it doesn't start the parser service. ParserService bool `json:"parserService" yaml:"parserService"` - // Jupyter configures the optional direct ipykernel provider. It is disabled - // unless explicitly enabled so the existing Jupyter Server proxy remains the - // default execution path during rollout. + // Jupyter configures Runme-managed kernels. Runme always exposes the direct + // kernel API; this block only overrides launch-profile settings. Jupyter *JupyterConfig `json:"jupyter,omitempty" yaml:"jupyter,omitempty"` // OIDC configuration @@ -494,9 +493,8 @@ type AssistantServerConfig struct { WebAppURL string `json:"webAppURL" yaml:"webAppURL"` } -// JupyterConfig controls the feature-flagged direct ZeroMQ kernel provider. +// JupyterConfig controls the built-in Python kernel launch profile. type JupyterConfig struct { - DirectZMQ bool `json:"directZMQ" yaml:"directZMQ"` PythonCommand string `json:"pythonCommand,omitempty" yaml:"pythonCommand,omitempty"` } diff --git a/pkg/agent/config/config_test.go b/pkg/agent/config/config_test.go index d34c64359..f12ff2c63 100644 --- a/pkg/agent/config/config_test.go +++ b/pkg/agent/config/config_test.go @@ -142,14 +142,13 @@ func Test_AssistantServerConfig_GetAgentService(t *testing.T) { } } -func Test_DirectJupyterConfiguration(t *testing.T) { +func Test_JupyterConfiguration(t *testing.T) { v := viper.New() v.SetConfigType("yaml") if err := v.ReadConfig(strings.NewReader(` assistantServer: runnerService: true jupyter: - directZMQ: true pythonCommand: /opt/runme/python `)); err != nil { t.Fatal(err) @@ -159,9 +158,9 @@ assistantServer: t.Fatal(err) } if cfg.AssistantServer == nil || cfg.AssistantServer.Jupyter == nil { - t.Fatal("direct Jupyter configuration was not loaded") + t.Fatal("Jupyter configuration was not loaded") } - if !cfg.AssistantServer.Jupyter.DirectZMQ || cfg.AssistantServer.Jupyter.PythonCommand != "/opt/runme/python" { + if cfg.AssistantServer.Jupyter.PythonCommand != "/opt/runme/python" { t.Fatalf("Jupyter config = %#v", cfg.AssistantServer.Jupyter) } } diff --git a/pkg/agent/jupyter/bridge.go b/pkg/agent/jupyter/bridge.go index 6b72b2fc4..b964d6ce2 100644 --- a/pkg/agent/jupyter/bridge.go +++ b/pkg/agent/jupyter/bridge.go @@ -191,6 +191,8 @@ func (b *KernelChannelsBridge) Bridge(ctx context.Context, kernelID string, clie 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) { @@ -204,6 +206,8 @@ func (b *KernelChannelsBridge) Bridge(ctx context.Context, kernelID string, clie 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" } diff --git a/pkg/agent/jupyter/metrics.go b/pkg/agent/jupyter/metrics.go index abed5afd4..f00a8beb6 100644 --- a/pkg/agent/jupyter/metrics.go +++ b/pkg/agent/jupyter/metrics.go @@ -1,6 +1,7 @@ package jupyter import ( + "errors" "strings" "time" @@ -30,7 +31,9 @@ func observeLifecycle(operation string, started time.Time) { func observeProtocolError(err error) { reason := "malformed_message" - if err != nil && strings.Contains(strings.ToLower(err.Error()), "signature") { + if errors.Is(err, ErrBinaryBuffersUnsupported) { + reason = "unsupported_binary_buffers" + } else if err != nil && strings.Contains(strings.ToLower(err.Error()), "signature") { reason = "bad_signature" } jupyterProtocolErrors.WithLabelValues(reason).Inc() diff --git a/pkg/agent/server/jupyter_kernels.go b/pkg/agent/server/jupyter_kernels.go new file mode 100644 index 000000000..d74b10615 --- /dev/null +++ b/pkg/agent/server/jupyter_kernels.go @@ -0,0 +1,344 @@ +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"` + Path string `json:"path"` +} + +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)) + 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 + } + if strings.TrimSpace(request.Path) != "" { + writeHTTPError(w, http.StatusBadRequest, "custom kernel paths are not supported") + return + } + profile := strings.TrimSpace(request.Name) + if profile == "" { + profile = "python3" + } + kernel, err := h.manager.Start(r.Context(), profile) + if err != nil { + logs.FromContext(r.Context()).Error(err, "failed to start Jupyter kernel", "profile", profile) + writeHTTPError(w, http.StatusBadRequest, "failed to start kernel") + return + } + logs.FromContext(r.Context()).Info("started Jupyter kernel", "kernel_id", kernel.ID, "profile", profile) + 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..34ab24dc2 --- /dev/null +++ b/pkg/agent/server/jupyter_kernels_test.go @@ -0,0 +1,294 @@ +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() + profile := jupyterbridge.PythonLaunchProfile("python3") + manager, err := jupyterbridge.NewKernelManager(jupyterbridge.KernelManagerConfig{ + RuntimeDir: t.TempDir(), + Profiles: map[string]jupyterbridge.LaunchProfile{profile.Name: profile}, + }) + 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: "reject custom path", method: http.MethodPost, path: jupyterKernelsRoute, body: `{"name":"python3","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 TestServerRegistersOnlyKernelScopedJupyterRoutes(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.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 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) + } + profile := jupyterbridge.PythonLaunchProfile(python) + manager, err := jupyterbridge.NewKernelManager(jupyterbridge.KernelManagerConfig{ + RuntimeDir: t.TempDir(), + Profiles: map[string]jupyterbridge.LaunchProfile{profile.Name: profile}, + 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) + + createResponse, err := http.Post( + server.URL+jupyterKernelsRoute, + "application/json", + strings.NewReader(`{"name":"python3"}`), + ) + 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 4361ade52..000000000 --- a/pkg/agent/server/jupyter_proxy.go +++ /dev/null @@ -1,1079 +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" - - jupyterbridge "github.com/runmedev/runme/v3/pkg/agent/jupyter" - "github.com/runmedev/runme/v3/pkg/agent/logs" -) - -const ( - jupyterServersRoute = "/v1/jupyter/servers" - jupyterConfigDir = "jupyter" - directZMQProvider = "direct-zmq" - maxDirectRequestBody = 64 * 1024 - - // 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 - directManager *jupyterbridge.KernelManager - directBridge *jupyterbridge.KernelChannelsBridge - httpClient *http.Client - upgrader websocket.Upgrader -} - -type directKernelCreateRequest struct { - Name string `json:"name"` - Path string `json:"path"` -} - -func (h *jupyterProxyHandler) handleDirectSubroute(w http.ResponseWriter, r *http.Request, parts []string) { - switch { - case len(parts) == 2: - switch r.Method { - case http.MethodGet: - writeJSON(w, http.StatusOK, h.directManager.List()) - case http.MethodPost: - var request directKernelCreateRequest - decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, maxDirectRequestBody)) - if err := decoder.Decode(&request); err != nil { - writeHTTPError(w, http.StatusBadRequest, "invalid kernel request") - return - } - if strings.TrimSpace(request.Path) != "" { - writeHTTPError(w, http.StatusBadRequest, "custom kernel paths are not supported") - return - } - profile := strings.TrimSpace(request.Name) - if profile == "" { - profile = "python3" - } - kernel, err := h.directManager.Start(r.Context(), profile) - if err != nil { - logs.FromContext(r.Context()).Error(err, "failed to start direct Jupyter kernel", "profile", profile) - writeHTTPError(w, http.StatusBadRequest, "failed to start kernel") - return - } - logs.FromContext(r.Context()).Info("started direct Jupyter kernel", "kernel_id", kernel.ID, "profile", profile) - writeJSON(w, http.StatusCreated, kernel) - 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 - } - action := "delete" - kernelID := kernelSegment - if strings.HasSuffix(kernelSegment, ":interrupt") { - action = "interrupt" - kernelID = strings.TrimSuffix(kernelSegment, ":interrupt") - } else if strings.HasSuffix(kernelSegment, ":restart") { - action = "restart" - kernelID = strings.TrimSuffix(kernelSegment, ":restart") - } - if strings.TrimSpace(kernelID) == "" { - writeHTTPError(w, http.StatusBadRequest, "invalid kernel id") - return - } - if _, err := h.directManager.Get(kernelID); err != nil { - writeHTTPError(w, http.StatusNotFound, "kernel not found") - return - } - switch action { - case "interrupt": - if r.Method != http.MethodPost { - writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") - return - } - if err := h.directManager.Interrupt(kernelID); err != nil { - writeDirectKernelOperationError(w, r, "interrupt", err) - return - } - logs.FromContext(r.Context()).Info("interrupted direct Jupyter kernel", "kernel_id", kernelID) - w.WriteHeader(http.StatusNoContent) - case "restart": - if r.Method != http.MethodPost { - writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") - return - } - kernel, err := h.directManager.Restart(r.Context(), kernelID) - if err != nil { - writeDirectKernelOperationError(w, r, "restart", err) - return - } - logs.FromContext(r.Context()).Info("restarted direct Jupyter kernel", "kernel_id", kernelID) - writeJSON(w, http.StatusOK, kernel) - default: - if r.Method != http.MethodDelete { - writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") - return - } - logs.FromContext(r.Context()).Info("stopping direct Jupyter kernel", "kernel_id", kernelID) - if err := h.directManager.Stop(r.Context(), kernelID); err != nil { - writeDirectKernelOperationError(w, r, "stop", err) - return - } - logs.FromContext(r.Context()).Info("stopped direct Jupyter kernel", "kernel_id", kernelID) - w.WriteHeader(http.StatusNoContent) - } - return - case len(parts) == 4 && parts[3] == "channels": - if r.Method != http.MethodGet { - writeHTTPError(w, http.StatusMethodNotAllowed, "method not allowed") - return - } - kernelID, err := url.PathUnescape(parts[2]) - if err != nil || strings.TrimSpace(kernelID) == "" { - writeHTTPError(w, http.StatusBadRequest, "invalid kernel id") - return - } - if _, err := h.directManager.Get(kernelID); err != nil { - writeHTTPError(w, http.StatusNotFound, "kernel not found") - return - } - h.handleDirectChannels(w, r, kernelID) - return - default: - writeHTTPError(w, http.StatusNotFound, "not found") - } -} - -func writeDirectKernelOperationError(w http.ResponseWriter, r *http.Request, operation string, err error) { - logs.FromContext(r.Context()).Error(err, "direct Jupyter kernel operation failed", "operation", operation) - writeHTTPError(w, http.StatusInternalServerError, "kernel operation failed") -} - -func (h *jupyterProxyHandler) handleDirectChannels(w http.ResponseWriter, r *http.Request, kernelID string) { - 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 direct Jupyter channels WebSocket") - return - } - client := newDirectChannelsClient(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, "runme_to_direct_client_keepalive", keepaliveErr, stopKeepalive) - bridgeErr := h.directBridge.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, "direct Jupyter channels bridge closed", "kernel_id", kernelID) - } - _ = client.Close(websocket.CloseNormalClosure, "Jupyter channels closed") -} - -type directChannelsClient struct { - conn *websocket.Conn - closeOnce sync.Once -} - -func newDirectChannelsClient(conn *websocket.Conn) *directChannelsClient { - return &directChannelsClient{conn: conn} -} - -func (c *directChannelsClient) 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 *directChannelsClient) 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 *directChannelsClient) 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 newJupyterProxyHandler(configDir string, managers ...*jupyterbridge.KernelManager) (*jupyterProxyHandler, error) { - registry, err := newJupyterServerRegistry(configDir) - if err != nil { - return nil, err - } - - var directManager *jupyterbridge.KernelManager - if len(managers) > 0 { - directManager = managers[0] - } - var directBridge *jupyterbridge.KernelChannelsBridge - if directManager != nil { - directBridge, err = jupyterbridge.NewKernelChannelsBridge(directManager, nil, jupyterbridge.BridgeConfig{}) - if err != nil { - return nil, err - } - } - return &jupyterProxyHandler{ - registry: registry, - directManager: directManager, - directBridge: directBridge, - 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() - if h.directManager != nil { - records = append(records, jupyterServerPublicRecord{ - Name: directZMQProvider, - BaseURL: "", - HasToken: false, - Backend: directZMQProvider, - Capabilities: []string{"kernels", "channels", "interrupt", "restart"}, - }) - sort.Slice(records, func(i, j int) bool { return records[i].Name < records[j].Name }) - } - 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 - } - if serverName == directZMQProvider && h.directManager != nil { - h.handleDirectSubroute(w, r, parts) - 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"` - Backend string `json:"backend"` - Capabilities []string `json:"capabilities"` -} - -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) != "", - Backend: "jupyter-server", - Capabilities: []string{"kernels", "channels", "interrupt", "restart"}, - }) - } - 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 ab5a869df..000000000 --- a/pkg/agent/server/jupyter_proxy_test.go +++ /dev/null @@ -1,370 +0,0 @@ -package server - -import ( - "context" - "encoding/json" - "io" - "net/http" - "net/http/httptest" - "net/url" - "os" - "path/filepath" - "reflect" - "strings" - "testing" - "time" - - "github.com/gorilla/websocket" - - jupyterbridge "github.com/runmedev/runme/v3/pkg/agent/jupyter" -) - -func TestJupyterProviderDiscoveryFeatureFlag(t *testing.T) { - configDir := t.TempDir() - disabled, err := newJupyterProxyHandler(configDir) - if err != nil { - t.Fatal(err) - } - disabledResponse := httptest.NewRecorder() - disabled.ServeHTTP(disabledResponse, httptest.NewRequest(http.MethodGet, jupyterServersRoute, nil)) - if strings.Contains(disabledResponse.Body.String(), directZMQProvider) { - t.Fatalf("disabled provider was advertised: %s", disabledResponse.Body.String()) - } - - profile := jupyterbridge.PythonLaunchProfile("python3") - manager, err := jupyterbridge.NewKernelManager(jupyterbridge.KernelManagerConfig{ - RuntimeDir: t.TempDir(), - Profiles: map[string]jupyterbridge.LaunchProfile{profile.Name: profile}, - }) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _ = manager.Close(context.Background()) }) - enabled, err := newJupyterProxyHandler(configDir, manager) - if err != nil { - t.Fatal(err) - } - enabledResponse := httptest.NewRecorder() - enabled.ServeHTTP(enabledResponse, httptest.NewRequest(http.MethodGet, jupyterServersRoute, nil)) - if enabledResponse.Code != http.StatusOK { - t.Fatalf("status = %d; body: %s", enabledResponse.Code, enabledResponse.Body.String()) - } - var records []jupyterServerPublicRecord - if err := json.Unmarshal(enabledResponse.Body.Bytes(), &records); err != nil { - t.Fatal(err) - } - if len(records) != 1 || records[0].Name != directZMQProvider || records[0].Backend != directZMQProvider { - t.Fatalf("records = %#v, want direct provider", records) - } - if records[0].HasToken || len(records[0].Capabilities) != 4 { - t.Fatalf("direct provider metadata = %#v", records[0]) - } -} - -func TestDirectJupyterHTTPAndWebSocketCompatibility(t *testing.T) { - python := os.Getenv("RUNME_TEST_PYTHON") - if python == "" { - t.Skip("set RUNME_TEST_PYTHON to an ipykernel-capable Python interpreter") - } - profile := jupyterbridge.PythonLaunchProfile(python) - manager, err := jupyterbridge.NewKernelManager(jupyterbridge.KernelManagerConfig{ - RuntimeDir: t.TempDir(), - Profiles: map[string]jupyterbridge.LaunchProfile{profile.Name: profile}, - StartupTimeout: 15 * time.Second, - }) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { _ = manager.Close(context.Background()) }) - handler, err := newJupyterProxyHandler(t.TempDir(), manager) - if err != nil { - t.Fatal(err) - } - server := httptest.NewServer(handler) - t.Cleanup(server.Close) - - createResponse, err := http.Post( - server.URL+jupyterServersRoute+"/"+directZMQProvider+"/kernels", - "application/json", - strings.NewReader(`{"name":"python3"}`), - ) - 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) - } - - wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + jupyterServersRoute + "/" + directZMQProvider + "/kernels/" + url.PathEscape(kernel.ID) + "/channels" - 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": "21 * 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" - } - } - - restartResponse, err := http.Post(server.URL+jupyterServersRoute+"/"+directZMQProvider+"/kernels/"+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) - } - deleteRequest, err := http.NewRequest(http.MethodDelete, server.URL+jupyterServersRoute+"/"+directZMQProvider+"/kernels/"+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) - } -} - -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 ada0fcd99..03ac8f4de 100644 --- a/pkg/agent/server/server.go +++ b/pkg/agent/server/server.go @@ -330,8 +330,12 @@ func (s *Server) registerServices() error { mux.HandleProtected(runnerSvcPath, runnerSvcHandler, s.checker, api.RunnerUserRole) } - if s.serverConfig.Jupyter != nil && s.serverConfig.Jupyter.DirectZMQ && s.jupyterManager == nil { - pythonProfile := jupyter.PythonLaunchProfile(s.serverConfig.Jupyter.PythonCommand) + if s.jupyterManager == nil { + pythonCommand := "" + if s.serverConfig.Jupyter != nil { + pythonCommand = s.serverConfig.Jupyter.PythonCommand + } + pythonProfile := jupyter.PythonLaunchProfile(pythonCommand) s.jupyterManager, err = jupyter.NewKernelManager(jupyter.KernelManagerConfig{ RuntimeDir: filepath.Join(s.configDir, "jupyter-kernels"), Profiles: map[string]jupyter.LaunchProfile{pythonProfile.Name: pythonProfile}, @@ -340,19 +344,19 @@ func (s *Server) registerServices() error { return errors.Wrap(err, "failed to initialize direct Jupyter kernel manager") } } - jupyterProxy, err := newJupyterProxyHandler(s.configDir, s.jupyterManager) + jupyterKernels, err := newJupyterKernelsHandler(s.jupyterManager) if err != nil { - return errors.Wrap(err, "failed to initialize jupyter proxy handler") + return errors.Wrap(err, "failed to initialize Jupyter kernels handler") } mux.HandleProtected( - "/v1/jupyter/servers", - otelhttp.NewHandler(jupyterProxy, "/v1/jupyter/servers"), + jupyterKernelsRoute, + otelhttp.NewHandler(jupyterKernels, jupyterKernelsRoute), s.checker, api.RunnerUserRole, ) mux.HandleProtected( - "/v1/jupyter/servers/", - otelhttp.NewHandler(jupyterProxy, "/v1/jupyter/servers/"), + jupyterKernelsRoute+"/", + otelhttp.NewHandler(jupyterKernels, jupyterKernelsRoute+"/"), s.checker, api.RunnerUserRole, ) From d933a23f34bfcffdf4760a89993da85020822ae6 Mon Sep 17 00:00:00 2001 From: Jeremy lewi Date: Fri, 14 Aug 2026 07:08:26 +0200 Subject: [PATCH 3/8] Allow clients to select Jupyter kernel commands Signed-off-by: Jeremy lewi --- pkg/agent/README.md | 9 +- pkg/agent/config/config.go | 9 -- pkg/agent/config/config_test.go | 23 ---- pkg/agent/iam/oidc.go | 7 +- pkg/agent/jupyter/bridge_integration_test.go | 14 +-- pkg/agent/jupyter/kernel_manager.go | 113 ++++++++++--------- pkg/agent/jupyter/kernel_manager_test.go | 41 +++---- pkg/agent/server/jupyter_kernels.go | 25 ++-- pkg/agent/server/jupyter_kernels_test.go | 16 ++- pkg/agent/server/server.go | 6 - 10 files changed, 117 insertions(+), 146 deletions(-) diff --git a/pkg/agent/README.md b/pkg/agent/README.md index 767ed0c15..0a5bdb42e 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,9 @@ 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. 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 diff --git a/pkg/agent/config/config.go b/pkg/agent/config/config.go index a6e880840..380d0d7f4 100644 --- a/pkg/agent/config/config.go +++ b/pkg/agent/config/config.go @@ -478,10 +478,6 @@ type AssistantServerConfig struct { // ParserService starts the Runme parser service if true otherwise it doesn't start the parser service. ParserService bool `json:"parserService" yaml:"parserService"` - // Jupyter configures Runme-managed kernels. Runme always exposes the direct - // kernel API; this block only overrides launch-profile settings. - Jupyter *JupyterConfig `json:"jupyter,omitempty" yaml:"jupyter,omitempty"` - // OIDC configuration OIDC *OIDCConfig `json:"oidc,omitempty" yaml:"oidc,omitempty"` @@ -493,11 +489,6 @@ type AssistantServerConfig struct { WebAppURL string `json:"webAppURL" yaml:"webAppURL"` } -// JupyterConfig controls the built-in Python kernel launch profile. -type JupyterConfig struct { - PythonCommand string `json:"pythonCommand,omitempty" yaml:"pythonCommand,omitempty"` -} - // OIDCConfig contains configuration for OIDC authentication type OIDCConfig struct { // Google contains Google-specific OIDC configuration diff --git a/pkg/agent/config/config_test.go b/pkg/agent/config/config_test.go index f12ff2c63..6c8b0867f 100644 --- a/pkg/agent/config/config_test.go +++ b/pkg/agent/config/config_test.go @@ -142,29 +142,6 @@ func Test_AssistantServerConfig_GetAgentService(t *testing.T) { } } -func Test_JupyterConfiguration(t *testing.T) { - v := viper.New() - v.SetConfigType("yaml") - if err := v.ReadConfig(strings.NewReader(` -assistantServer: - runnerService: true - jupyter: - pythonCommand: /opt/runme/python -`)); err != nil { - t.Fatal(err) - } - cfg, err := getConfigFromViper(v) - if err != nil { - t.Fatal(err) - } - if cfg.AssistantServer == nil || cfg.AssistantServer.Jupyter == nil { - t.Fatal("Jupyter configuration was not loaded") - } - if cfg.AssistantServer.Jupyter.PythonCommand != "/opt/runme/python" { - t.Fatalf("Jupyter config = %#v", cfg.AssistantServer.Jupyter) - } -} - // setEnvVars sets the provided environment variables and returns a function to restore the previous state. func setEnvVars(env map[string]string) func() { if env == nil { 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_integration_test.go b/pkg/agent/jupyter/bridge_integration_test.go index ccb350019..2b47607ce 100644 --- a/pkg/agent/jupyter/bridge_integration_test.go +++ b/pkg/agent/jupyter/bridge_integration_test.go @@ -57,7 +57,7 @@ func TestKernelChannelsBridgeSharedStateAndTwoClients(t *testing.T) { defer manager.Close(context.Background()) ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) defer cancel() - model, err := manager.Start(ctx, "python3") + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) if err != nil { t.Fatal(err) } @@ -92,7 +92,7 @@ func TestKernelChannelsBridgeConcurrentExecutionsRouteByParent(t *testing.T) { defer manager.Close(context.Background()) ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) defer cancel() - model, err := manager.Start(ctx, "python3") + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) if err != nil { t.Fatal(err) } @@ -124,7 +124,7 @@ func TestKernelChannelsBridgeInterruptAndReuse(t *testing.T) { defer manager.Close(context.Background()) ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) defer cancel() - model, err := manager.Start(ctx, "python3") + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) if err != nil { t.Fatal(err) } @@ -157,7 +157,7 @@ func TestKernelChannelsBridgeRestartClosesOldGeneration(t *testing.T) { defer manager.Close(context.Background()) ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) defer cancel() - model, err := manager.Start(ctx, "python3") + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) if err != nil { t.Fatal(err) } @@ -202,7 +202,7 @@ func TestKernelChannelsBridgeRichOutputsAndQuietExecution(t *testing.T) { defer manager.Close(context.Background()) ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() - model, err := manager.Start(ctx, "python3") + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) if err != nil { t.Fatal(err) } @@ -244,7 +244,7 @@ func TestKernelChannelsBridgeStdinAndControlRouting(t *testing.T) { defer manager.Close(context.Background()) ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) defer cancel() - model, err := manager.Start(ctx, "python3") + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) if err != nil { t.Fatal(err) } @@ -307,7 +307,7 @@ func TestKernelChannelsBridgeReconnectSoak(t *testing.T) { defer manager.Close(context.Background()) ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second) defer cancel() - model, err := manager.Start(ctx, "python3") + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) if err != nil { t.Fatal(err) } diff --git a/pkg/agent/jupyter/kernel_manager.go b/pkg/agent/jupyter/kernel_manager.go index a0708e06f..047cd9625 100644 --- a/pkg/agent/jupyter/kernel_manager.go +++ b/pkg/agent/jupyter/kernel_manager.go @@ -44,30 +44,30 @@ type KernelModel struct { Connections int `json:"connections"` } -// LaunchProfile is a server-side allowlisted kernel command. Browser requests -// select only Name; they cannot supply Command, Args, or a connection path. -type LaunchProfile struct { - Name string - Command string - Args []string -} - -// PythonLaunchProfile returns the standard ipykernel launch profile. -func PythonLaunchProfile(command string) LaunchProfile { +// 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 +} + +// PythonKernelLaunchSpec returns the standard ipykernel launch specification. +func PythonKernelLaunchSpec(command string) KernelLaunchSpec { if strings.TrimSpace(command) == "" { command = "python3" } - return LaunchProfile{ - Name: "python3", - Command: command, - Args: []string{"-m", "ipykernel_launcher", "-f", connectionFilePlaceholder}, + return KernelLaunchSpec{ + Name: "python3", + Argv: []string{command, "-m", "ipykernel_launcher", "-f", connectionFilePlaceholder}, } } // KernelManagerConfig configures local kernel lifecycle boundaries. type KernelManagerConfig struct { RuntimeDir string - Profiles map[string]LaunchProfile SocketFactory SocketFactory StartupTimeout time.Duration ReadinessTimeout time.Duration @@ -102,6 +102,7 @@ type managedKernel struct { exitErr error runtimeDir string diagnostics *boundedWriter + launchSpec KernelLaunchSpec } type kernelConnection struct { @@ -111,25 +112,12 @@ type kernelConnection struct { cancel context.CancelFunc } -// NewKernelManager creates a manager after validating all launch profiles and -// preparing an owner-only runtime directory. +// 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 len(config.Profiles) == 0 { - profile := PythonLaunchProfile("python3") - config.Profiles = map[string]LaunchProfile{profile.Name: profile} - } - profiles := make(map[string]LaunchProfile, len(config.Profiles)) - for key, profile := range config.Profiles { - if err := validateLaunchProfile(key, profile); err != nil { - return nil, err - } - profile.Args = append([]string(nil), profile.Args...) - profiles[key] = profile - } - config.Profiles = profiles if config.SocketFactory == nil { config.SocketFactory = NewGoZeroMQFactory(Limits{}) } @@ -167,37 +155,56 @@ func NewKernelManager(config KernelManagerConfig) (*KernelManager, error) { }, nil } -func validateLaunchProfile(key string, profile LaunchProfile) error { - if strings.TrimSpace(key) == "" || key != profile.Name { - return fmt.Errorf("jupyter launch profile key %q must match its non-empty name", key) +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 strings.TrimSpace(profile.Command) == "" { - return fmt.Errorf("jupyter launch profile %q has no command", key) + 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 profile.Args { + for _, arg := range spec.Argv[1:] { placeholderCount += strings.Count(arg, connectionFilePlaceholder) } if placeholderCount != 1 { - return fmt.Errorf("jupyter launch profile %q must contain exactly one %s placeholder", key, connectionFilePlaceholder) + return fmt.Errorf("jupyter kernel argv must contain exactly one %s placeholder", connectionFilePlaceholder) } return nil } -// Start launches an allowlisted profile and returns only after heartbeat and -// kernel_info readiness checks succeed. -func (m *KernelManager) Start(ctx context.Context, profileName string) (KernelModel, error) { +// 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) }() - profile, ok := m.config.Profiles[profileName] - if !ok { - return KernelModel{}, fmt.Errorf("jupyter launch profile %q is not allowed", profileName) + 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: profile.Name, + name: spec.Name, state: KernelStateStarting, lastActivity: time.Now().UTC(), + launchSpec: spec, } kernel.lifecycle.Lock() defer kernel.lifecycle.Unlock() @@ -211,7 +218,7 @@ func (m *KernelManager) Start(ctx context.Context, profileName string) (KernelMo jupyterKernelStates.WithLabelValues(string(kernel.state)).Inc() m.mu.Unlock() - if err := m.launch(ctx, kernel, profile); err != nil { + if err := m.launch(ctx, kernel, spec); err != nil { _ = m.stopProcess(context.Background(), kernel, false) m.mu.Lock() jupyterKernelStates.WithLabelValues(string(kernel.state)).Dec() @@ -259,7 +266,7 @@ func kernelModel(kernel *managedKernel) KernelModel { } } -func (m *KernelManager) launch(ctx context.Context, kernel *managedKernel, profile LaunchProfile) error { +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) @@ -269,18 +276,18 @@ func (m *KernelManager) launch(ctx context.Context, kernel *managedKernel, profi return fmt.Errorf("secure kernel runtime directory: %w", err) } connectionPath := filepath.Join(runtimeDir, "connection.json") - args := make([]string, len(profile.Args)) - for i, arg := range profile.Args { - args[i] = strings.ReplaceAll(arg, connectionFilePlaceholder, connectionPath) + 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(profile.Command, args...) + 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 profile %q: %w", profile.Name, err) + return fmt.Errorf("launch kernel %q: %w", spec.Name, err) } exited := make(chan struct{}) @@ -427,7 +434,7 @@ func (m *KernelManager) Restart(ctx context.Context, id string) (KernelModel, er m.mu.Unlock() return KernelModel{}, fmt.Errorf("jupyter kernel %q not found", id) } - profile := m.config.Profiles[kernel.name] + spec := kernel.launchSpec m.setKernelStateLocked(kernel, KernelStateRestarting) kernel.lastActivity = time.Now().UTC() if kernel.connection != nil { @@ -439,7 +446,7 @@ func (m *KernelManager) Restart(ctx context.Context, id string) (KernelModel, er 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, profile); err != nil { + if err := m.launch(ctx, kernel, spec); err != nil { m.mu.Lock() m.setKernelStateLocked(kernel, KernelStateDead) kernel.lastActivity = time.Now().UTC() diff --git a/pkg/agent/jupyter/kernel_manager_test.go b/pkg/agent/jupyter/kernel_manager_test.go index c13ddb8b5..27f10ea9d 100644 --- a/pkg/agent/jupyter/kernel_manager_test.go +++ b/pkg/agent/jupyter/kernel_manager_test.go @@ -17,7 +17,7 @@ func TestKernelManagerStartRestartStop(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) defer cancel() - started, err := manager.Start(ctx, "python3") + started, err := manager.Start(ctx, testPythonLaunchSpec(t)) if err != nil { t.Fatalf("Start() error = %v", err) } @@ -100,7 +100,7 @@ func TestKernelManagerUnexpectedExitBecomesDead(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) defer cancel() - model, err := manager.Start(ctx, "python3") + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) if err != nil { t.Fatal(err) } @@ -136,7 +136,7 @@ func TestKernelManagerCloseReapsAllKernels(t *testing.T) { var commands []*exec.Cmd var runtimeDirs []string for range 2 { - model, err := manager.Start(ctx, "python3") + model, err := manager.Start(ctx, testPythonLaunchSpec(t)) if err != nil { t.Fatal(err) } @@ -158,27 +158,26 @@ func TestKernelManagerCloseReapsAllKernels(t *testing.T) { t.Errorf("runtime directory %d still exists or stat failed unexpectedly: %v", i, err) } } - if _, err := manager.Start(ctx, "python3"); err == nil { + if _, err := manager.Start(ctx, testPythonLaunchSpec(t)); err == nil { t.Fatal("Start() succeeded after Close()") } } -func TestKernelManagerRejectsArbitraryLaunchCommands(t *testing.T) { - python := findIPyKernelPython(t) - _, err := NewKernelManager(KernelManagerConfig{ - RuntimeDir: t.TempDir(), - Profiles: map[string]LaunchProfile{ - "python3": {Name: "python3", Command: python, Args: []string{"-c", "pass"}}, - }, +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("NewKernelManager() error = %v, want missing placeholder error", err) + t.Fatalf("Start() error = %v, want missing placeholder error", err) } - - manager := newTestKernelManager(t) - defer manager.Close(context.Background()) - if _, err := manager.Start(context.Background(), "browser-supplied-command"); err == nil { - t.Fatal("Start() accepted a non-allowlisted profile") + 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) } } @@ -197,11 +196,8 @@ func TestBoundedWriterKeepsOnlyTail(t *testing.T) { func newTestKernelManager(t *testing.T) *KernelManager { t.Helper() - python := findIPyKernelPython(t) - profile := PythonLaunchProfile(python) manager, err := NewKernelManager(KernelManagerConfig{ RuntimeDir: filepath.Join(t.TempDir(), "runtime"), - Profiles: map[string]LaunchProfile{profile.Name: profile}, StartupTimeout: 10 * time.Second, ReadinessTimeout: 5 * time.Second, GracefulTimeout: 2 * time.Second, @@ -214,6 +210,11 @@ func newTestKernelManager(t *testing.T) *KernelManager { return manager } +func testPythonLaunchSpec(t *testing.T) KernelLaunchSpec { + t.Helper() + return PythonKernelLaunchSpec(findIPyKernelPython(t)) +} + func testKernelProcess(manager *KernelManager, id string) (*exec.Cmd, string) { manager.mu.RLock() defer manager.mu.RUnlock() diff --git a/pkg/agent/server/jupyter_kernels.go b/pkg/agent/server/jupyter_kernels.go index d74b10615..f0f9645c6 100644 --- a/pkg/agent/server/jupyter_kernels.go +++ b/pkg/agent/server/jupyter_kernels.go @@ -36,13 +36,13 @@ type jupyterKernelsHandler struct { } type kernelCreateRequest struct { - Name string `json:"name"` - Path string `json:"path"` + 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") + return nil, errors.New("jupyter kernel manager is required") } bridge, err := jupyterbridge.NewKernelChannelsBridge(manager, nil, jupyterbridge.BridgeConfig{}) if err != nil { @@ -110,6 +110,7 @@ func (h *jupyterKernelsHandler) handleCollection(w http.ResponseWriter, r *http. 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 @@ -118,21 +119,17 @@ func (h *jupyterKernelsHandler) handleCollection(w http.ResponseWriter, r *http. writeHTTPError(w, http.StatusBadRequest, "invalid kernel request") return } - if strings.TrimSpace(request.Path) != "" { - writeHTTPError(w, http.StatusBadRequest, "custom kernel paths are not supported") - return - } - profile := strings.TrimSpace(request.Name) - if profile == "" { - profile = "python3" - } - kernel, err := h.manager.Start(r.Context(), profile) + 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", "profile", profile) + 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, "profile", profile) + 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") diff --git a/pkg/agent/server/jupyter_kernels_test.go b/pkg/agent/server/jupyter_kernels_test.go index 34ab24dc2..6b6005e4f 100644 --- a/pkg/agent/server/jupyter_kernels_test.go +++ b/pkg/agent/server/jupyter_kernels_test.go @@ -22,10 +22,8 @@ import ( func newTestJupyterKernelsHandler(t *testing.T) (*jupyterKernelsHandler, *jupyterbridge.KernelManager) { t.Helper() - profile := jupyterbridge.PythonLaunchProfile("python3") manager, err := jupyterbridge.NewKernelManager(jupyterbridge.KernelManagerConfig{ RuntimeDir: t.TempDir(), - Profiles: map[string]jupyterbridge.LaunchProfile{profile.Name: profile}, }) if err != nil { t.Fatal(err) @@ -50,7 +48,8 @@ func TestJupyterKernelsRouteContract(t *testing.T) { }{ {name: "list", method: http.MethodGet, path: jupyterKernelsRoute, status: http.StatusOK}, {name: "collection method", method: http.MethodDelete, path: jupyterKernelsRoute, status: http.StatusMethodNotAllowed}, - {name: "reject custom path", method: http.MethodPost, path: jupyterKernelsRoute, body: `{"name":"python3","path":"/tmp/connection.json"}`, status: http.StatusBadRequest}, + {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}, @@ -137,10 +136,8 @@ func TestDirectJupyterKernelAPIAndWebSocket(t *testing.T) { } else if err := exec.Command(python, "-c", "import ipykernel").Run(); err != nil { t.Fatalf("RUNME_TEST_PYTHON does not provide ipykernel: %v", err) } - profile := jupyterbridge.PythonLaunchProfile(python) manager, err := jupyterbridge.NewKernelManager(jupyterbridge.KernelManagerConfig{ RuntimeDir: t.TempDir(), - Profiles: map[string]jupyterbridge.LaunchProfile{profile.Name: profile}, StartupTimeout: 15 * time.Second, }) if err != nil { @@ -154,10 +151,17 @@ func TestDirectJupyterKernelAPIAndWebSocket(t *testing.T) { server := httptest.NewServer(handler) t.Cleanup(server.Close) + createBody, err := json.Marshal(kernelCreateRequest{ + Name: "python3", + Argv: jupyterbridge.PythonKernelLaunchSpec(python).Argv, + }) + if err != nil { + t.Fatal(err) + } createResponse, err := http.Post( server.URL+jupyterKernelsRoute, "application/json", - strings.NewReader(`{"name":"python3"}`), + strings.NewReader(string(createBody)), ) if err != nil { t.Fatal(err) diff --git a/pkg/agent/server/server.go b/pkg/agent/server/server.go index 03ac8f4de..280d3a0af 100644 --- a/pkg/agent/server/server.go +++ b/pkg/agent/server/server.go @@ -331,14 +331,8 @@ func (s *Server) registerServices() error { } if s.jupyterManager == nil { - pythonCommand := "" - if s.serverConfig.Jupyter != nil { - pythonCommand = s.serverConfig.Jupyter.PythonCommand - } - pythonProfile := jupyter.PythonLaunchProfile(pythonCommand) s.jupyterManager, err = jupyter.NewKernelManager(jupyter.KernelManagerConfig{ RuntimeDir: filepath.Join(s.configDir, "jupyter-kernels"), - Profiles: map[string]jupyter.LaunchProfile{pythonProfile.Name: pythonProfile}, }) if err != nil { return errors.Wrap(err, "failed to initialize direct Jupyter kernel manager") From 993edeb06514d7e78fe9d986f0b9256d84927050 Mon Sep 17 00:00:00 2001 From: Jeremy lewi Date: Fri, 14 Aug 2026 10:54:52 +0200 Subject: [PATCH 4/8] Use OpenTelemetry for Jupyter metrics Signed-off-by: Jeremy lewi --- pkg/agent/jupyter/bridge.go | 18 ++++----- pkg/agent/jupyter/kernel_manager.go | 12 +++--- pkg/agent/jupyter/metrics.go | 62 ++++++++++++++++++++++------- 3 files changed, 61 insertions(+), 31 deletions(-) diff --git a/pkg/agent/jupyter/bridge.go b/pkg/agent/jupyter/bridge.go index b964d6ce2..fe9455a0b 100644 --- a/pkg/agent/jupyter/bridge.go +++ b/pkg/agent/jupyter/bridge.go @@ -101,8 +101,8 @@ func (b *KernelChannelsBridge) Bridge(ctx context.Context, kernelID string, clie if client == nil { return errors.New("jupyter channels client is required") } - jupyterActiveBridges.Inc() - defer jupyterActiveBridges.Dec() + addActiveBridge(1) + defer addActiveBridge(-1) connection, release, err := b.manager.Connect(kernelID) if err != nil { return err @@ -211,7 +211,7 @@ func (b *KernelChannelsBridge) Bridge(ctx context.Context, kernelID string, clie case isClosed(generationChanged): closeMetricReason = "kernel_restart" } - jupyterBridgeCloses.WithLabelValues(closeMetricReason).Inc() + observeBridgeClose(closeMetricReason) _ = client.Close(closeCode, closeReason) return bridgeErr } @@ -424,8 +424,7 @@ func (b *KernelChannelsBridge) readClient( if err := socket.SendMultipart(frames); err != nil { return err } - jupyterChannelMessages.WithLabelValues("browser_to_kernel", string(channel)).Inc() - jupyterChannelBytes.WithLabelValues("browser_to_kernel", string(channel)).Add(float64(len(payload))) + observeChannelMessage("browser_to_kernel", channel, len(payload)) } } @@ -455,7 +454,7 @@ func (b *KernelChannelsBridge) readKernel( } if channel == ChannelIOPub { if err := limiter.Allow(len(payload)); err != nil { - jupyterRateLimits.Inc() + addCounter(jupyterRateLimits) return err } b.observeKernelState(kernelID, generation, message) @@ -463,8 +462,7 @@ func (b *KernelChannelsBridge) readKernel( if err := queue.Enqueue(ctx, payload); err != nil { return err } - jupyterChannelMessages.WithLabelValues("kernel_to_browser", string(channel)).Inc() - jupyterChannelBytes.WithLabelValues("kernel_to_browser", string(channel)).Add(float64(len(payload))) + observeChannelMessage("kernel_to_browser", channel, len(payload)) } } @@ -547,7 +545,7 @@ func (b *KernelChannelsBridge) monitorHeartbeat(ctx context.Context, initial Soc } if misses < b.config.HeartbeatMisses { if misses > 0 { - jupyterHeartbeatMisses.Inc() + addCounter(jupyterHeartbeatMisses) _ = current.Close() replacement, err := b.openSocket(ctx, SocketTypeReq, nil, connection, ChannelHeartbeat) if err != nil { @@ -558,7 +556,7 @@ func (b *KernelChannelsBridge) monitorHeartbeat(ctx context.Context, initial Soc } continue } - jupyterHeartbeatMisses.Inc() + addCounter(jupyterHeartbeatMisses) return fmt.Errorf("jupyter heartbeat missed %d times", misses) } } diff --git a/pkg/agent/jupyter/kernel_manager.go b/pkg/agent/jupyter/kernel_manager.go index 047cd9625..68c3ef731 100644 --- a/pkg/agent/jupyter/kernel_manager.go +++ b/pkg/agent/jupyter/kernel_manager.go @@ -215,13 +215,13 @@ func (m *KernelManager) Start(ctx context.Context, spec KernelLaunchSpec) (Kerne return KernelModel{}, errors.New("jupyter kernel manager is closed") } m.kernels[kernel.id] = kernel - jupyterKernelStates.WithLabelValues(string(kernel.state)).Inc() + 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() - jupyterKernelStates.WithLabelValues(string(kernel.state)).Dec() + addKernelState(kernel.state, -1) delete(m.kernels, kernel.id) m.mu.Unlock() return KernelModel{}, fmt.Errorf("start Jupyter kernel: %w", err) @@ -351,7 +351,7 @@ func (m *KernelManager) waitForProcess(kernel *managedKernel, command *exec.Cmd, close(exited) if kernel.state != KernelStateStopping && kernel.state != KernelStateRestarting { m.setKernelStateLocked(kernel, KernelStateDead) - jupyterUnexpectedExits.Inc() + addCounter(jupyterUnexpectedExits) kernel.lastActivity = time.Now().UTC() if kernel.connection != nil { kernel.connection.cancel() @@ -486,7 +486,7 @@ func (m *KernelManager) Stop(ctx context.Context, id string) error { return err } m.mu.Lock() - jupyterKernelStates.WithLabelValues(string(kernel.state)).Dec() + addKernelState(kernel.state, -1) delete(m.kernels, id) m.mu.Unlock() return nil @@ -597,9 +597,9 @@ func (m *KernelManager) setKernelStateLocked(kernel *managedKernel, state Kernel if kernel.state == state { return } - jupyterKernelStates.WithLabelValues(string(kernel.state)).Dec() + addKernelState(kernel.state, -1) kernel.state = state - jupyterKernelStates.WithLabelValues(string(kernel.state)).Inc() + addKernelState(kernel.state, 1) } // Diagnostics returns bounded process output for local troubleshooting. diff --git a/pkg/agent/jupyter/metrics.go b/pkg/agent/jupyter/metrics.go index f00a8beb6..d0bb5c770 100644 --- a/pkg/agent/jupyter/metrics.go +++ b/pkg/agent/jupyter/metrics.go @@ -1,32 +1,64 @@ package jupyter import ( + "context" "errors" "strings" "time" - "github.com/prometheus/client_golang/prometheus" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + otelmetric "go.opentelemetry.io/otel/metric" ) var ( - jupyterKernelStates = prometheus.NewGaugeVec(prometheus.GaugeOpts{Name: "runme_jupyter_kernels", Help: "Current directly managed Jupyter kernels by lifecycle state."}, []string{"state"}) - jupyterLifecycleLatency = prometheus.NewHistogramVec(prometheus.HistogramOpts{Name: "runme_jupyter_lifecycle_duration_seconds", Help: "Duration of direct Jupyter kernel lifecycle operations.", Buckets: prometheus.DefBuckets}, []string{"operation"}) - jupyterActiveBridges = prometheus.NewGauge(prometheus.GaugeOpts{Name: "runme_jupyter_active_bridges", Help: "Current direct browser-to-kernel Jupyter bridges."}) - jupyterChannelMessages = prometheus.NewCounterVec(prometheus.CounterOpts{Name: "runme_jupyter_channel_messages_total", Help: "Direct Jupyter channel messages by direction and channel."}, []string{"direction", "channel"}) - jupyterChannelBytes = prometheus.NewCounterVec(prometheus.CounterOpts{Name: "runme_jupyter_channel_bytes_total", Help: "Direct Jupyter channel payload bytes by direction and channel."}, []string{"direction", "channel"}) - jupyterProtocolErrors = prometheus.NewCounterVec(prometheus.CounterOpts{Name: "runme_jupyter_protocol_errors_total", Help: "Rejected direct Jupyter protocol messages by bounded reason."}, []string{"reason"}) - jupyterRateLimits = prometheus.NewCounter(prometheus.CounterOpts{Name: "runme_jupyter_rate_limits_total", Help: "Direct Jupyter bridges closed by IOPub rate limits."}) - jupyterHeartbeatMisses = prometheus.NewCounter(prometheus.CounterOpts{Name: "runme_jupyter_heartbeat_misses_total", Help: "Missed direct Jupyter kernel heartbeat probes."}) - jupyterUnexpectedExits = prometheus.NewCounter(prometheus.CounterOpts{Name: "runme_jupyter_unexpected_exits_total", Help: "Unexpected exits of directly managed Jupyter kernels."}) - jupyterBridgeCloses = prometheus.NewCounterVec(prometheus.CounterOpts{Name: "runme_jupyter_bridge_closes_total", Help: "Direct Jupyter bridge closes by bounded reason."}, []string{"reason"}) + 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 init() { - prometheus.MustRegister(jupyterKernelStates, jupyterLifecycleLatency, jupyterActiveBridges, jupyterChannelMessages, jupyterChannelBytes, jupyterProtocolErrors, jupyterRateLimits, jupyterHeartbeatMisses, jupyterUnexpectedExits, jupyterBridgeCloses) +func mustMetric[T any](instrument T, err error) T { + if err != nil { + panic(err) + } + return instrument } func observeLifecycle(operation string, started time.Time) { - jupyterLifecycleLatency.WithLabelValues(operation).Observe(time.Since(started).Seconds()) + 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) { @@ -36,5 +68,5 @@ func observeProtocolError(err error) { } else if err != nil && strings.Contains(strings.ToLower(err.Error()), "signature") { reason = "bad_signature" } - jupyterProtocolErrors.WithLabelValues(reason).Inc() + addCounter(jupyterProtocolErrors, attribute.String("reason", reason)) } From c6d5dc7854426f1e6e524bb22cc52db860ce321f Mon Sep 17 00:00:00 2001 From: Jeremy lewi Date: Fri, 14 Aug 2026 14:52:41 +0200 Subject: [PATCH 5/8] Remove Python kernel launch helper Signed-off-by: Jeremy lewi --- pkg/agent/jupyter/kernel_manager.go | 11 ----------- pkg/agent/jupyter/kernel_manager_test.go | 5 ++++- pkg/agent/server/jupyter_kernels_test.go | 2 +- 3 files changed, 5 insertions(+), 13 deletions(-) diff --git a/pkg/agent/jupyter/kernel_manager.go b/pkg/agent/jupyter/kernel_manager.go index 68c3ef731..df8178540 100644 --- a/pkg/agent/jupyter/kernel_manager.go +++ b/pkg/agent/jupyter/kernel_manager.go @@ -54,17 +54,6 @@ type KernelLaunchSpec struct { Argv []string } -// PythonKernelLaunchSpec returns the standard ipykernel launch specification. -func PythonKernelLaunchSpec(command string) KernelLaunchSpec { - if strings.TrimSpace(command) == "" { - command = "python3" - } - return KernelLaunchSpec{ - Name: "python3", - Argv: []string{command, "-m", "ipykernel_launcher", "-f", connectionFilePlaceholder}, - } -} - // KernelManagerConfig configures local kernel lifecycle boundaries. type KernelManagerConfig struct { RuntimeDir string diff --git a/pkg/agent/jupyter/kernel_manager_test.go b/pkg/agent/jupyter/kernel_manager_test.go index 27f10ea9d..4b99cecf5 100644 --- a/pkg/agent/jupyter/kernel_manager_test.go +++ b/pkg/agent/jupyter/kernel_manager_test.go @@ -212,7 +212,10 @@ func newTestKernelManager(t *testing.T) *KernelManager { func testPythonLaunchSpec(t *testing.T) KernelLaunchSpec { t.Helper() - return PythonKernelLaunchSpec(findIPyKernelPython(t)) + return KernelLaunchSpec{ + Name: "python3", + Argv: []string{findIPyKernelPython(t), "-m", "ipykernel_launcher", "-f", connectionFilePlaceholder}, + } } func testKernelProcess(manager *KernelManager, id string) (*exec.Cmd, string) { diff --git a/pkg/agent/server/jupyter_kernels_test.go b/pkg/agent/server/jupyter_kernels_test.go index 6b6005e4f..441c45cf4 100644 --- a/pkg/agent/server/jupyter_kernels_test.go +++ b/pkg/agent/server/jupyter_kernels_test.go @@ -153,7 +153,7 @@ func TestDirectJupyterKernelAPIAndWebSocket(t *testing.T) { createBody, err := json.Marshal(kernelCreateRequest{ Name: "python3", - Argv: jupyterbridge.PythonKernelLaunchSpec(python).Argv, + Argv: []string{python, "-m", "ipykernel_launcher", "-f", "{connection_file}"}, }) if err != nil { t.Fatal(err) From adfb1ae6301cc6cb1a1422eb7e1f73ef021174f6 Mon Sep 17 00:00:00 2001 From: Jeremy lewi Date: Sat, 29 Aug 2026 16:42:19 -0700 Subject: [PATCH 6/8] Fix Jupyter bridge reconnect readiness Signed-off-by: Jeremy lewi --- pkg/agent/jupyter/bridge.go | 37 +++++++++++++++++++++++-------------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/pkg/agent/jupyter/bridge.go b/pkg/agent/jupyter/bridge.go index fe9455a0b..40b72e2e7 100644 --- a/pkg/agent/jupyter/bridge.go +++ b/pkg/agent/jupyter/bridge.go @@ -303,15 +303,7 @@ func (b *KernelChannelsBridge) readinessNudge( } defer control.Close() - 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 - } - infoReply := make(chan struct{}, 1) + infoReply := make(chan string, 2) iopubMessage := make(chan struct{}, 1) errorsCh := make(chan error, 3) readInfo := func(socket Socket) { @@ -333,9 +325,9 @@ func (b *KernelChannelsBridge) readinessNudge( errorsCh <- err return } - if messageType == "kernel_info_reply" && parentID == requestID { + if messageType == "kernel_info_reply" { select { - case infoReply <- struct{}{}: + case infoReply <- parentID: default: } return @@ -362,7 +354,17 @@ func (b *KernelChannelsBridge) readinessNudge( 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 } @@ -374,11 +376,18 @@ func (b *KernelChannelsBridge) readinessNudge( for !gotInfo || !gotIOPub { select { case <-ctx.Done(): - return ctx.Err() + return fmt.Errorf( + "readiness incomplete (kernel_info=%t iopub=%t): %w", + gotInfo, + gotIOPub, + ctx.Err(), + ) case err := <-errorsCh: return err - case <-infoReply: - gotInfo = true + case parentID := <-infoReply: + if _, ok := sentRequestIDs[parentID]; ok { + gotInfo = true + } case <-iopubMessage: gotIOPub = true case <-ticker.C: From 5920278f1e3f7e79efc2c5a076aefd6bab4f945a Mon Sep 17 00:00:00 2001 From: Jeremy lewi Date: Thu, 22 Oct 2026 07:43:41 -0700 Subject: [PATCH 7/8] Declare OpenTelemetry metric dependency directly Signed-off-by: Jeremy lewi --- go.mod | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/go.mod b/go.mod index f201ae2bb..80159164c 100644 --- a/go.mod +++ b/go.mod @@ -71,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 @@ -230,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 From 194b88cb0bc38f12ec13495197334a25dc526c5c Mon Sep 17 00:00:00 2001 From: Jeremy lewi Date: Sat, 31 Oct 2026 14:30:57 -0700 Subject: [PATCH 8/8] Gate Jupyter kernels behind service config Signed-off-by: Jeremy lewi --- pkg/agent/README.md | 6 +- pkg/agent/config/config.go | 6 ++ pkg/agent/config/config_test.go | 19 +++++++ pkg/agent/server/jupyter_kernels_test.go | 71 +++++++++++++++++++++++- pkg/agent/server/server.go | 50 ++++++++++------- 5 files changed, 129 insertions(+), 23 deletions(-) diff --git a/pkg/agent/README.md b/pkg/agent/README.md index 0a5bdb42e..dee3610f4 100644 --- a/pkg/agent/README.md +++ b/pkg/agent/README.md @@ -20,7 +20,9 @@ The server provides the runtime services required by Runme Web: `assistantServer.parserService` is enabled. - The Jupyter kernel API starts client-specified kernels, manages their lifecycle, and bridges browser channel WebSockets directly to kernel ZeroMQ - sockets. Callers need `RunnerUserRole`; no separate Jupyter Server is used. + 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 @@ -45,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/server/jupyter_kernels_test.go b/pkg/agent/server/jupyter_kernels_test.go index 441c45cf4..80d7bcf15 100644 --- a/pkg/agent/server/jupyter_kernels_test.go +++ b/pkg/agent/server/jupyter_kernels_test.go @@ -84,7 +84,7 @@ func TestNewJupyterKernelsHandlerRequiresManager(t *testing.T) { } } -func TestServerRegistersOnlyKernelScopedJupyterRoutes(t *testing.T) { +func TestServerDoesNotRegisterJupyterRoutesByDefault(t *testing.T) { runmeServer, err := NewServer(Options{ Server: &config.AssistantServerConfig{ ParserService: true, @@ -106,6 +106,60 @@ func TestServerRegistersOnlyKernelScopedJupyterRoutes(t *testing.T) { 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) @@ -125,6 +179,21 @@ func TestServerRegistersOnlyKernelScopedJupyterRoutes(t *testing.T) { } } +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 == "" { diff --git a/pkg/agent/server/server.go b/pkg/agent/server/server.go index 280d3a0af..f3ac49eed 100644 --- a/pkg/agent/server/server.go +++ b/pkg/agent/server/server.go @@ -83,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 { @@ -330,30 +333,35 @@ func (s *Server) registerServices() error { mux.HandleProtected(runnerSvcPath, runnerSvcHandler, s.checker, api.RunnerUserRole) } - if s.jupyterManager == nil { - s.jupyterManager, err = jupyter.NewKernelManager(jupyter.KernelManagerConfig{ - RuntimeDir: filepath.Join(s.configDir, "jupyter-kernels"), - }) + 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 direct Jupyter kernel manager") + 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") } - 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, - ) // Health check should be public checker := grpchealth.NewStaticChecker()