From 108e82a0c425a09c3145a9d485d979f717cc15d4 Mon Sep 17 00:00:00 2001 From: rgarcia <72655+rgarcia@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:33:41 +0000 Subject: [PATCH] Add an ACP WebSocket proxy using acpremote --- .github/workflows/server-test.yaml | 11 + server/README.md | 8 + server/cmd/api/main.go | 15 ++ server/cmd/config/config.go | 4 + server/lib/agentproxy/README.md | 142 ++++++++++++ server/lib/agentproxy/bridge.go | 128 +++++++++++ server/lib/agentproxy/config.go | 120 +++++++++++ server/lib/agentproxy/config_test.go | 121 +++++++++++ server/lib/agentproxy/disconnect_test.go | 28 +++ server/lib/agentproxy/handler.go | 93 ++++++++ server/lib/agentproxy/proxy_test.go | 262 +++++++++++++++++++++++ server/lib/agentproxy/testdata/agent.py | 81 +++++++ server/lib/agentproxy/testdata/client.py | 65 ++++++ server/lib/wsproxy/proxy_options_test.go | 73 +++++++ server/lib/wsproxy/wsproxy.go | 38 +++- server/runtime/acp/requirements.txt | 3 + 16 files changed, 1190 insertions(+), 2 deletions(-) create mode 100644 server/lib/agentproxy/README.md create mode 100644 server/lib/agentproxy/bridge.go create mode 100644 server/lib/agentproxy/config.go create mode 100644 server/lib/agentproxy/config_test.go create mode 100644 server/lib/agentproxy/disconnect_test.go create mode 100644 server/lib/agentproxy/handler.go create mode 100644 server/lib/agentproxy/proxy_test.go create mode 100644 server/lib/agentproxy/testdata/agent.py create mode 100644 server/lib/agentproxy/testdata/client.py create mode 100644 server/lib/wsproxy/proxy_options_test.go create mode 100644 server/runtime/acp/requirements.txt diff --git a/.github/workflows/server-test.yaml b/.github/workflows/server-test.yaml index b3bc43e60..a24af5ae5 100644 --- a/.github/workflows/server-test.yaml +++ b/.github/workflows/server-test.yaml @@ -41,6 +41,17 @@ jobs: with: node-version: 22 + - name: Set up Python for ACP interoperability tests + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install pinned ACP bridge and client + run: | + python -m venv "$RUNNER_TEMP/acp" + "$RUNNER_TEMP/acp/bin/python" -m pip install -r server/runtime/acp/requirements.txt + echo "AGENT_PROXY_TEST_ACPREMOTE=$RUNNER_TEMP/acp/bin/acpremote" >> "$GITHUB_ENV" + # categorygen's checks (unclassified route, category that isn't control or # platform, classified route with no handler) only run when the generator # does, and its only other caller is `make oapi-generate`, which needs the diff --git a/server/README.md b/server/README.md index cc80cb414..81ab20485 100644 --- a/server/README.md +++ b/server/README.md @@ -57,6 +57,7 @@ Configure the server using environment variables: | `MAX_SIZE_MB` | `500` | Default maximum file size (MB) | | `OUTPUT_DIR` | `.` | Directory to save recordings | | `FFMPEG_PATH` | `ffmpeg` | Path to the ffmpeg binary | +| `AGENT_CONFIG_PATH` | empty | Trusted launch catalog enabling the optional ACP WebSocket proxy | #### Example Configuration @@ -68,6 +69,13 @@ export OUTPUT_DIR=/tmp/recordings ./bin/api ``` +### Optional ACP agents + +The [ACP proxy](lib/agentproxy/README.md) exposes one WebSocket endpoint using +`acpremote expose` and its per-connection process lifecycle. It is disabled by +default and requires separately provisioned bridge and harness executables. +The public declarative installer/configuration API is not implemented yet. + ### API Documentation - **YAML Spec**: `GET /spec.yaml` diff --git a/server/cmd/api/main.go b/server/cmd/api/main.go index b3384c6f5..08d770a64 100644 --- a/server/cmd/api/main.go +++ b/server/cmd/api/main.go @@ -27,6 +27,7 @@ import ( serverpkg "github.com/kernel/kernel-images/server" "github.com/kernel/kernel-images/server/cmd/api/api" "github.com/kernel/kernel-images/server/cmd/config" + "github.com/kernel/kernel-images/server/lib/agentproxy" "github.com/kernel/kernel-images/server/lib/chromedriverproxy" "github.com/kernel/kernel-images/server/lib/devtoolsproxy" "github.com/kernel/kernel-images/server/lib/events" @@ -285,6 +286,20 @@ func main() { apiService.HandleProcessAttachWS(w, r, id, wsRegistry) }) + if config.AgentConfigPath != "" { + agentConfig, err := agentproxy.Load(config.AgentConfigPath) + if err != nil { + slogger.Error("agent endpoints disabled: invalid launch catalog", "err", err) + } else { + agents, err := agentproxy.New(ctx, agentConfig, slogger, wsRegistry) + if err != nil { + slogger.Error("agent endpoints disabled", "err", err) + } else { + r.Handle("/agent/v1/*", agents) + } + } + } + // Serve extension files for Chrome policy-installed extensions // This allows Chrome to download .crx and update.xml files via HTTP extensionsDir := "/home/kernel/extensions" diff --git a/server/cmd/config/config.go b/server/cmd/config/config.go index 70cce9142..bf7ea152e 100644 --- a/server/cmd/config/config.go +++ b/server/cmd/config/config.go @@ -15,6 +15,9 @@ type Config struct { // Server configuration Port int `envconfig:"PORT" default:"10001"` + // Optional trusted launch catalog. Empty leaves the ACP endpoints disabled. + AgentConfigPath string `envconfig:"AGENT_CONFIG_PATH" default:""` + // Port for the Prometheus metrics endpoint. Served on a separate // listener so scrapes bypass the scale-to-zero middleware and the // external API surface. @@ -89,6 +92,7 @@ func (c *Config) LogValue() slog.Value { } return slog.GroupValue( slog.Int("port", c.Port), + slog.String("agent_config_path", c.AgentConfigPath), slog.Int("metrics_port", c.MetricsPort), slog.Int("frame_rate", c.FrameRate), slog.Int("display_num", c.DisplayNum), diff --git a/server/lib/agentproxy/README.md b/server/lib/agentproxy/README.md new file mode 100644 index 000000000..6066d528f --- /dev/null +++ b/server/lib/agentproxy/README.md @@ -0,0 +1,142 @@ +# ACP WebSocket proxy + +This is the first integration milestone for an ACP-first agent API: a real +kernel-images route and an existing-client compatibility gate. It deliberately +contains no conversation REST protocol, runtime resource, prompt journal, +idempotency layer, session-ID translation, or delivery replay. + +## Implemented surface + +- `GET /agent/v1/harnesses` returns configured harness names. This is configuration + discovery, not a claim that a harness/model or every ACP capability is validated. +- WebSocket `GET /agent/v1/acp?harness=pi` starts `acpremote expose` with the trusted + launch definition and proxies ACP messages unchanged. +- Names are `pi`, `codex`, `claude`, and `gemini`. Only provisioned names are enabled. +- `AGENT_CONFIG_PATH` enables the routes. Empty or invalid configuration leaves + them unavailable without breaking the browser's other APIs. + +These routes use the existing browser API authentication/routing boundary, like +`/process/{process_id}/attach`; they do not introduce another public token scheme. +Do not expose the image server directly to an untrusted network. Each internal +bridge binds only to loopback and requires a generated bearer token; browser +credentials and client headers are not forwarded to that bridge. + +## Lifecycle + +For each client connection the server starts a dedicated `acpremote expose` +listener on an OS-assigned loopback port. Opening its upstream WebSocket starts +one harness/adapter subprocess. The listener and subprocess are scoped to that +attachment and are cleaned up when the connection ends. Multiple concurrent +attachments are allowed within `maxConnections`; each may contain multiple ACP +sessions if the selected harness supports them. + +A disconnect is termination, not detach. On reconnect, a fresh subprocess may load +an exact persisted native conversation with ACP `session/load` or `session/resume`. +Native storage must survive connection cleanup. No output is retained by the proxy, +and an interrupted tool's side effect is not automatically reconciled or repeated. + +The proxy retains no ACP state. It does not intercept initialize, permissions, +model selection, MCP definitions, session IDs, or results/errors. Subprotocol +`acp.v1` is accepted when offered; no subprotocol is required. Frame size is limited +to 1 MiB to match acpremote 1.7.0's default. Downstream ping/pong uses a 20-second +interval and timeout so a dead network connection cannot leave the agent running +indefinitely while the proxy answers upstream pings. + +Server shutdown cancels startup and active connections. Bridge cleanup first asks +the CLI to shut down gracefully, then kills the owned process group after at most +six seconds so adapters cannot leave ordinary descendant processes behind. This +is process cleanup, not a security sandbox for programs that deliberately daemonize +outside their process group. + +## Trusted launch catalog (bootstrap, not the declarative public API) + +Install Python 3.11+ and the pinned transport dependencies in a private environment: + +```sh +python3 -m venv /opt/kernel-agent/venv +/opt/kernel-agent/venv/bin/python -m pip install -r server/runtime/acp/requirements.txt +``` + +The image does not install these dependencies or any harness automatically in this +milestone. An operator must provision the environment and pinned harness/adapter +binaries before enabling the feature. The bridge executable is expected to be +acpremote 1.7.0; startup reads that CLI's loopback readiness banner, not agent output. + +Example catalog, assuming the executable and working directory are provisioned: + +```json +{ + "acpremote": "/opt/kernel-agent/venv/bin/acpremote", + "maxConnections": 8, + "harnesses": { + "gemini": { + "command": "/opt/kernel-agent/gemini/node_modules/.bin/gemini", + "args": ["--acp"], + "cwd": "/workspace", + "env": {"HOME": "/home/kernel"}, + "inheritEnv": ["GEMINI_API_KEY"] + } + } +} +``` + +Set `AGENT_CONFIG_PATH` to this private file and restart the image server. Launch +configuration is read once; editing the file does not hot-reload running processes. +The catalog is limited to 64 KiB. Commands and working directories must be absolute. +It is trusted operator input, not a mechanism for a WebSocket caller to supply a +command, arguments, or environment. Agent stderr is discarded by this first +integration; no transcript or secret-bearing log stream is exposed. + +Only PATH, HOME, USER, LANG, TMPDIR, TERM and certificate paths are inherited by +default. Provider credentials must be explicitly named by `inheritEnv`; a missing +or empty required value rejects connection startup. `env` overrides inherited +values and should contain nonsecret settings only. Neither is returned by discovery. + +All connections to a harness use the same configured environment and working +directory. Shared configuration, extensions, MCP files and native history are not +copied or deleted per connection. For Pi, explicitly setting `PI_CODING_AGENT_DIR` +selects shared Pi state; project-local discovery remains Pi/adapter behavior. +Clients must supply a working directory meaningful on the remote host. Local IDE +files, terminal callbacks and local stdio MCP servers do not become remotely +available automatically. + +## Acceptance checks + +The integration suite uses real acpremote processes and TCP WebSockets. Its +stdlib-only ACP peer is deterministic: no provider credentials or model calls. + +```sh +python3 -m venv /tmp/acp-test +/tmp/acp-test/bin/python -m pip install -r runtime/acp/requirements.txt +AGENT_PROXY_TEST_ACPREMOTE=/tmp/acp-test/bin/acpremote \ + go test -race -v ./lib/agentproxy ./lib/wsproxy +``` + +Run from `server/`. CI installs the pinned bridge and sets the environment variable; +without it, the external-process cases explicitly skip rather than claim a pass. + +The gates cover independent concurrent processes, multiple native sessions per +connection, connection admission, termination during an in-flight prompt, exact-ID +history load after reconnect, permissions, opaque content/metadata, server shutdown, +message limits, and unresponsive clients. A separate case drives the unmodified +`acpremote mirror` CLI using the official ACP Python SDK. No Kernel-specific client +messages or headers are needed. + +ACP UI 0.1.16 at commit `cd9c3cb464a4b321bff652101953a64c07473e31` was also tested +locally without source changes: initialization, session creation, streamed output, +permission approval, disconnect, and exact-session restoration in a fresh process. +This was a local WebSocket test with the deterministic peer, not a TLS/gateway or +real-model test. The UI check is not yet part of CI. + +## Remaining work + +The public declarative configuration GET/PUT API, transactional preparation, +harness/extension installers, and model-provider secret bindings remain a separate +milestone. The bootstrap catalog is not their replacement. They should produce +trusted launch definitions while making shared filesystem effects explicit. + +Each harness still needs pinned native installation, capability/MCP/media tests, +permission and native restoration checks, and shared-configuration verification. +Successful fixture/client interoperability is not a blanket harness support claim. +Platform-level routing/authentication and packaged-image tests with the feature +enabled are also required before public availability. diff --git a/server/lib/agentproxy/bridge.go b/server/lib/agentproxy/bridge.go new file mode 100644 index 000000000..4fad2c1ad --- /dev/null +++ b/server/lib/agentproxy/bridge.go @@ -0,0 +1,128 @@ +package agentproxy + +import ( + "bufio" + "context" + "crypto/rand" + "encoding/hex" + "errors" + "io" + "net/http" + "net/url" + "os" + "os/exec" + "strconv" + "strings" + "sync" + "syscall" + "time" +) + +const bridgeTokenEnv = "KERNEL_ACP_BRIDGE_TOKEN" + +// Each attachment gets its own expose listener/process group. No connection or +// child process is retained for reattachment after the downstream closes. +type bridge struct { + url string + headers http.Header + command *exec.Cmd + done chan struct{} + once sync.Once +} + +func startBridge(ctx context.Context, executable string, harness Harness) (*bridge, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + secret := make([]byte, 32) + if _, err := rand.Read(secret); err != nil { + return nil, err + } + token := hex.EncodeToString(secret) + env, err := harness.environment(token) + if err != nil { + return nil, err + } + args := []string{"expose", "--host", "127.0.0.1", "--port", "0", "--cwd", harness.Cwd, "--stderr-mode", "discard", "--token-env", bridgeTokenEnv, "--", harness.Command} + args = append(args, harness.Args...) + command := exec.Command(executable, args...) + command.Env = env + command.Dir = harness.Cwd + command.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + // Keep our reader separate from exec.Wait's pipe cleanup. + reader, writer, err := os.Pipe() + if err != nil { + return nil, err + } + command.Stderr = writer + if err := command.Start(); err != nil { + reader.Close() + writer.Close() + return nil, errors.New("could not start ACP bridge") + } + writer.Close() + b := &bridge{command: command, done: make(chan struct{}), headers: http.Header{"Authorization": {"Bearer " + token}}} + go func() { + _ = command.Wait() + close(b.done) + }() + ready := make(chan string, 1) + go func() { + defer reader.Close() + defer close(ready) + scanner := bufio.NewScanner(reader) + for scanner.Scan() { + if address := bridgeAddress(scanner.Text()); address != "" { + ready <- address + // Never forward agent/bridge stderr into request logs or responses. + _, _ = io.Copy(io.Discard, reader) + return + } + } + }() + startup, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + select { + case address := <-ready: + if address != "" { + b.url = address + return b, nil + } + case <-b.done: + case <-startup.Done(): + } + b.close() + return nil, errors.New("ACP bridge did not become ready") +} + +// acpremote 1.7.0 reports its OS-selected port on stderr. Restrict the advertised +// address to the listener we requested; never proxy an arbitrary URL from output. +func bridgeAddress(line string) string { + address, ok := strings.CutPrefix(line, "Serving ACP WebSocket at ") + if !ok { + return "" + } + u, err := url.Parse(address) + if err != nil || u.Scheme != "ws" || u.Hostname() != "127.0.0.1" || u.Path != "/acp/ws" || u.User != nil || u.RawQuery != "" || u.Fragment != "" { + return "" + } + port, err := strconv.Atoi(u.Port()) + if err != nil || port < 1 || port > 65535 { + return "" + } + return address +} + +func (b *bridge) close() { + b.once.Do(func() { + // The CLI handles SIGINT by closing its server and awaiting child cleanup. + _ = b.command.Process.Signal(os.Interrupt) + select { + case <-b.done: + case <-time.After(6 * time.Second): + } + // Terminate remaining descendants if an adapter outlives its bridge. + _ = syscall.Kill(-b.command.Process.Pid, syscall.SIGKILL) + <-b.done + }) +} diff --git a/server/lib/agentproxy/config.go b/server/lib/agentproxy/config.go new file mode 100644 index 000000000..45ed22ab2 --- /dev/null +++ b/server/lib/agentproxy/config.go @@ -0,0 +1,120 @@ +// Package agentproxy exposes ACP WebSockets without interpreting ACP messages. +package agentproxy + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strings" +) + +// Config is trusted, operator-provisioned launch configuration. It is not a +// public command-execution API. Harness-specific installers will produce it. +type Config struct { + ACPRemote string `json:"acpremote"` + MaxConnections int `json:"maxConnections"` + Harnesses map[string]Harness `json:"harnesses"` +} + +type Harness struct { + Command string `json:"command"` + Args []string `json:"args"` + Cwd string `json:"cwd"` + Env map[string]string `json:"env"` + InheritEnv []string `json:"inheritEnv"` +} + +func Load(path string) (Config, error) { + file, err := os.Open(path) + if err != nil { + return Config{}, err + } + defer file.Close() + data, err := io.ReadAll(io.LimitReader(file, (64<<10)+1)) + if err != nil { + return Config{}, err + } + if len(data) > 64<<10 { + return Config{}, errors.New("agent configuration exceeds 64 KiB") + } + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.DisallowUnknownFields() + var config Config + if err := decoder.Decode(&config); err != nil { + return Config{}, errors.New("invalid agent configuration") + } + var extra any + if err := decoder.Decode(&extra); err != io.EOF { + return Config{}, errors.New("expected one agent configuration") + } + if err := config.validate(); err != nil { + return Config{}, err + } + return config, nil +} + +func (c Config) validate() error { + if !filepath.IsAbs(c.ACPRemote) { + return errors.New("acpremote must be an absolute executable path") + } + if c.MaxConnections < 1 || c.MaxConnections > 64 { + return errors.New("maxConnections must be between 1 and 64") + } + if len(c.Harnesses) == 0 { + return errors.New("at least one harness is required") + } + for name, harness := range c.Harnesses { + switch name { + case "pi", "codex", "claude", "gemini": + default: + return fmt.Errorf("unknown harness %q", name) + } + if !filepath.IsAbs(harness.Command) || !filepath.IsAbs(harness.Cwd) { + return fmt.Errorf("%s requires absolute command and cwd paths", name) + } + for name, value := range harness.Env { + if !validEnvName(name) || strings.ContainsRune(value, '\x00') { + return errors.New("invalid harness environment") + } + } + for _, name := range harness.InheritEnv { + if !validEnvName(name) { + return errors.New("invalid inherited environment name") + } + } + } + return nil +} + +func validEnvName(name string) bool { + return name != "" && !strings.ContainsAny(name, "=\x00") && name != bridgeTokenEnv +} + +func (h Harness) environment(token string) ([]string, error) { + values := make(map[string]string) + for _, name := range []string{"PATH", "HOME", "USER", "LANG", "TMPDIR", "TERM", "SSL_CERT_FILE", "SSL_CERT_DIR"} { + if value, ok := os.LookupEnv(name); ok { + values[name] = value + } + } + for _, name := range h.InheritEnv { + value, ok := os.LookupEnv(name) + if !ok || value == "" { + return nil, errors.New("required harness environment is unavailable") + } + values[name] = value + } + for name, value := range h.Env { + values[name] = value + } + values[bridgeTokenEnv] = token + env := make([]string, 0, len(values)) + for name, value := range values { + env = append(env, name+"="+value) + } + return env, nil +} diff --git a/server/lib/agentproxy/config_test.go b/server/lib/agentproxy/config_test.go new file mode 100644 index 000000000..a7190346f --- /dev/null +++ b/server/lib/agentproxy/config_test.go @@ -0,0 +1,121 @@ +package agentproxy + +import ( + "context" + "encoding/json" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestLoadConfig(t *testing.T) { + valid := `{"acpremote":"/opt/acp/bin/acpremote","maxConnections":8,"harnesses":{"pi":{"command":"/opt/pi-acp","cwd":"/workspace","args":[],"env":{},"inheritEnv":[]}}}` + for _, test := range []struct { + name string + data string + ok bool + }{ + {"valid", valid, true}, + {"unknown field", strings.Replace(valid, `"maxConnections":8`, `"typo":8`, 1), false}, + {"extra JSON", valid + `{}`, false}, + {"oversized", valid + strings.Repeat(" ", 64<<10), false}, + {"relative executable", strings.Replace(valid, "/opt/pi-acp", "pi-acp", 1), false}, + {"relative cwd", strings.Replace(valid, "/workspace", "workspace", 1), false}, + {"unknown harness", strings.Replace(valid, `"pi":`, `"other":`, 1), false}, + {"no connection bound", strings.Replace(valid, `"maxConnections":8`, `"maxConnections":0`, 1), false}, + {"reserved environment", strings.Replace(valid, `"env":{}`, `"env":{"KERNEL_ACP_BRIDGE_TOKEN":"override"}`, 1), false}, + } { + t.Run(test.name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.json") + if err := os.WriteFile(path, []byte(test.data), 0600); err != nil { + t.Fatal(err) + } + _, err := Load(path) + if (err == nil) != test.ok { + t.Fatalf("load: %v", err) + } + }) + } +} + +func TestEnvironmentIsExplicit(t *testing.T) { + t.Setenv("AGENT_PROXY_TEST_SECRET", "do-not-inherit") + t.Setenv("AGENT_PROXY_TEST_ALLOWED", "inherited") + h := Harness{InheritEnv: []string{"AGENT_PROXY_TEST_ALLOWED"}, Env: map[string]string{"HOME": "/shared/pi"}} + env, err := h.environment("internal-token") + if err != nil { + t.Fatal(err) + } + joined := strings.Join(env, "\n") + if strings.Contains(joined, "do-not-inherit") || !strings.Contains(joined, "AGENT_PROXY_TEST_ALLOWED=inherited") || !strings.Contains(joined, "HOME=/shared/pi") { + t.Fatal("environment inheritance or override is incorrect") + } + t.Setenv("AGENT_PROXY_TEST_ALLOWED", "") + if _, err := h.environment("internal-token"); err == nil { + t.Fatal("empty required environment accepted") + } +} + +func TestBridgeAddress(t *testing.T) { + for _, address := range []string{"ws://127.0.0.1:1234/acp/ws", "ws://example.com:1234/acp/ws", "ws://127.0.0.1:0/acp/ws", "ws://127.0.0.1:70000/acp/ws", "ws://127.0.0.1:1234/acp/ws?secret=x", "ws://user@127.0.0.1:1234/acp/ws", "ws://127.0.0.1:1234/other"} { + got := bridgeAddress("Serving ACP WebSocket at " + address) + if (got != "") != (address == "ws://127.0.0.1:1234/acp/ws") { + t.Fatalf("unexpected address accepted: %q", got) + } + } +} + +func TestHTTPAdmissionAndDiscovery(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + config := Config{ACPRemote: "/nonexistent/acpremote", MaxConnections: 1, Harnesses: map[string]Harness{ + "pi": {Command: "/nonexistent/pi", Cwd: t.TempDir(), Env: map[string]string{"PRIVATE_KEY": "secret-value"}}, + }} + h, err := New(ctx, config, slog.New(slog.NewTextHandler(io.Discard, nil)), nil) + if err != nil { + t.Fatal(err) + } + response := httptest.NewRecorder() + h.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/agent/v1/harnesses", nil)) + var discovery struct { + Configured []string `json:"configured"` + } + if err := json.Unmarshal(response.Body.Bytes(), &discovery); err != nil || len(discovery.Configured) != 1 || discovery.Configured[0] != "pi" { + t.Fatal("invalid configured harness list") + } + if strings.Contains(response.Body.String(), "secret-value") { + t.Fatal("discovery leaked launch environment") + } + for _, test := range []struct { + path string + code int + }{ + {"/agent/v1/acp?harness=unknown", 404}, + {"/agent/v1/acp?harness=pi", 400}, + {"/agent/v1/commands", 404}, + } { + response := httptest.NewRecorder() + h.ServeHTTP(response, httptest.NewRequest(http.MethodGet, test.path, nil)) + if response.Code != test.code { + t.Fatalf("%s: %d", test.path, response.Code) + } + } + request := httptest.NewRequest(http.MethodGet, "/agent/v1/acp?harness=pi", nil) + request.Header.Set("Upgrade", "websocket") + response = httptest.NewRecorder() + h.ServeHTTP(response, request) + if response.Code != 503 || len(h.slots) != 0 { + t.Fatal("failed startup did not release admission slot") + } + cancel() + response = httptest.NewRecorder() + h.ServeHTTP(response, request) + if response.Code != 503 { + t.Fatal("shutdown accepted a new connection") + } +} diff --git a/server/lib/agentproxy/disconnect_test.go b/server/lib/agentproxy/disconnect_test.go new file mode 100644 index 000000000..fdf6432cf --- /dev/null +++ b/server/lib/agentproxy/disconnect_test.go @@ -0,0 +1,28 @@ +package agentproxy + +import ( + "syscall" + "testing" +) + +func TestAbruptDisconnect(t *testing.T) { + for _, prompt := range []string{"block", "permission"} { + t.Run(prompt, func(t *testing.T) { + handler, address, dir, _ := fixture(t, 1) + conn := connect(t, address) + pid := initialize(t, conn) + session := newSession(t, conn, 2, dir) + send(t, conn, map[string]any{"jsonrpc": "2.0", "id": 3, "method": "session/prompt", "params": textPrompt(session, prompt)}) + message := read(t, conn) + expected := "session/update" + if prompt == "permission" { + expected = "session/request_permission" + } + if message.Method != expected { + t.Fatalf("agent did not reach %s", prompt) + } + conn.CloseNow() // No WebSocket close handshake or ACP cancellation. + eventually(t, func() bool { return syscall.Kill(pid, 0) == syscall.ESRCH && len(handler.slots) == 0 }) + }) + } +} diff --git a/server/lib/agentproxy/handler.go b/server/lib/agentproxy/handler.go new file mode 100644 index 000000000..5fb376712 --- /dev/null +++ b/server/lib/agentproxy/handler.go @@ -0,0 +1,93 @@ +package agentproxy + +import ( + "context" + "encoding/json" + "log/slog" + "net/http" + "sort" + "strings" + "time" + + "github.com/coder/websocket" + "github.com/kernel/kernel-images/server/lib/wsdrain" + "github.com/kernel/kernel-images/server/lib/wsproxy" +) + +type Handler struct { + ctx context.Context + config Config + logger *slog.Logger + registry *wsdrain.Registry + slots chan struct{} +} + +func New(ctx context.Context, config Config, logger *slog.Logger, registry *wsdrain.Registry) (*Handler, error) { + if err := config.validate(); err != nil { + return nil, err + } + return &Handler{ctx: ctx, config: config, logger: logger, registry: registry, slots: make(chan struct{}, config.MaxConnections)}, nil +} + +func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodGet && r.URL.Path == "/agent/v1/harnesses": + names := make([]string, 0, len(h.config.Harnesses)) + for name := range h.config.Harnesses { + names = append(names, name) + } + sort.Strings(names) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(struct { + Configured []string `json:"configured"` + }{names}) + case r.Method == http.MethodGet && r.URL.Path == "/agent/v1/acp": + h.connect(w, r) + default: + http.NotFound(w, r) + } +} + +func (h *Handler) connect(w http.ResponseWriter, r *http.Request) { + name := r.URL.Query().Get("harness") + harness, ok := h.config.Harnesses[name] + if !ok { + http.Error(w, "harness is not configured", http.StatusNotFound) + return + } + if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") { + http.Error(w, "WebSocket upgrade required", http.StatusBadRequest) + return + } + ctx, cancel := context.WithCancel(r.Context()) + stop := context.AfterFunc(h.ctx, cancel) + defer stop() + defer cancel() + if h.ctx.Err() != nil { + http.Error(w, "agent service is shutting down", http.StatusServiceUnavailable) + return + } + select { + case h.slots <- struct{}{}: + defer func() { <-h.slots }() + default: + http.Error(w, "agent connection limit reached", http.StatusTooManyRequests) + return + } + b, err := startBridge(ctx, h.config.ACPRemote, harness) + if err != nil { + h.logger.Warn("ACP bridge startup failed", "harness", name, "err", err) + http.Error(w, "ACP bridge unavailable", http.StatusServiceUnavailable) + return + } + defer b.close() + wsproxy.Proxy(w, r.WithContext(ctx), b.url, wsproxy.ProxyOptions{ + // Same browser-scoped authentication boundary as the process attach API. + AcceptOptions: &websocket.AcceptOptions{OriginPatterns: []string{"*"}, Subprotocols: []string{"acp.v1"}}, + DialOptions: &websocket.DialOptions{HTTPHeader: b.headers}, + Logger: h.logger, + Registry: h.registry, + ReadLimit: 1 << 20, // acpremote 1.7.0's default message limit + PingInterval: 20 * time.Second, + }) +} diff --git a/server/lib/agentproxy/proxy_test.go b/server/lib/agentproxy/proxy_test.go new file mode 100644 index 000000000..0306a52fc --- /dev/null +++ b/server/lib/agentproxy/proxy_test.go @@ -0,0 +1,262 @@ +package agentproxy + +import ( + "bytes" + "context" + "encoding/json" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "syscall" + "testing" + "time" + + "github.com/coder/websocket" +) + +type rpcMessage struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + Params json.RawMessage `json:"params"` + Result json.RawMessage `json:"result"` + Error json.RawMessage `json:"error"` +} + +func fixture(t *testing.T, limit int) (*Handler, string, string, context.CancelFunc) { + t.Helper() + remote := os.Getenv("AGENT_PROXY_TEST_ACPREMOTE") + if remote == "" { + t.Skip("set AGENT_PROXY_TEST_ACPREMOTE to the pinned acpremote venv executable") + } + peer, err := filepath.Abs("testdata/agent.py") + if err != nil { + t.Fatal(err) + } + dir := t.TempDir() + config := Config{ACPRemote: remote, MaxConnections: limit, Harnesses: map[string]Harness{ + "pi": {Command: filepath.Join(filepath.Dir(remote), "python"), Args: []string{peer}, Cwd: dir, Env: map[string]string{"AGENT_PROXY_TEST_DIR": dir}}, + }} + ctx, cancel := context.WithCancel(context.Background()) + handler, err := New(ctx, config, slog.New(slog.NewTextHandler(io.Discard, nil)), nil) + if err != nil { + cancel() + t.Fatal(err) + } + server := httptest.NewServer(handler) + t.Cleanup(func() { cancel(); server.Close(); eventually(t, func() bool { return len(handler.slots) == 0 }) }) + return handler, "ws" + strings.TrimPrefix(server.URL, "http") + "/agent/v1/acp?harness=pi", dir, cancel +} + +func connect(t *testing.T, address string) *websocket.Conn { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + conn, _, err := websocket.Dial(ctx, address, &websocket.DialOptions{Subprotocols: []string{"acp.v1"}}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { conn.CloseNow() }) + if conn.Subprotocol() != "acp.v1" { + t.Fatal("client subprotocol not negotiated") + } + return conn +} + +func send(t *testing.T, conn *websocket.Conn, value any) { + t.Helper() + data, err := json.Marshal(value) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + // ACP UI sends newline-terminated JSON in its WebSocket text frames. + if err := conn.Write(ctx, websocket.MessageText, append(data, '\n')); err != nil { + t.Fatal(err) + } +} + +func read(t *testing.T, conn *websocket.Conn) rpcMessage { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _, data, err := conn.Read(ctx) + if err != nil { + t.Fatal(err) + } + var message rpcMessage + if err := json.Unmarshal(data, &message); err != nil { + t.Fatal(err) + } + if len(message.Error) != 0 { + t.Fatalf("RPC error: %s", message.Error) + } + return message +} + +func call(t *testing.T, conn *websocket.Conn, id int, method string, params any) (json.RawMessage, []json.RawMessage) { + t.Helper() + send(t, conn, map[string]any{"jsonrpc": "2.0", "id": id, "method": method, "params": params}) + return result(t, conn, id) +} + +func result(t *testing.T, conn *websocket.Conn, id int) (json.RawMessage, []json.RawMessage) { + t.Helper() + var updates []json.RawMessage + for { + message := read(t, conn) + if string(message.ID) == strconv.Itoa(id) { + return message.Result, updates + } + if message.Method != "session/update" { + t.Fatalf("unexpected message: %+v", message) + } + updates = append(updates, message.Params) + } +} + +func initialize(t *testing.T, conn *websocket.Conn) int { + t.Helper() + data, _ := call(t, conn, 1, "initialize", map[string]any{"protocolVersion": 1, "clientCapabilities": map[string]any{}}) + var initialized struct { + Meta struct { + PID int `json:"pid"` + } `json:"_meta"` + } + if err := json.Unmarshal(data, &initialized); err != nil || initialized.Meta.PID == 0 { + t.Fatalf("initialize failed: %s", data) + } + return initialized.Meta.PID +} + +func newSession(t *testing.T, conn *websocket.Conn, id int, dir string) string { + t.Helper() + data, _ := call(t, conn, id, "session/new", map[string]any{"cwd": dir, "mcpServers": []any{}}) + var session struct { + ID string `json:"sessionId"` + } + if err := json.Unmarshal(data, &session); err != nil || session.ID == "" { + t.Fatalf("new session failed: %s", data) + } + return session.ID +} + +func textPrompt(sessionID, text string) map[string]any { + return map[string]any{"sessionId": sessionID, "prompt": []map[string]string{{"type": "text", "text": text}}} +} + +func eventually(t *testing.T, condition func() bool) { + t.Helper() + deadline := time.Now().Add(10 * time.Second) + for !condition() { + if time.Now().After(deadline) { + t.Fatal("condition timed out") + } + time.Sleep(10 * time.Millisecond) + } +} + +func TestConnectionLifecycle(t *testing.T) { + handler, address, dir, _ := fixture(t, 2) + first := connect(t, address) + firstPID := initialize(t, first) + firstSession := newSession(t, first, 2, dir) + if other := newSession(t, first, 3, dir); other == firstSession { + t.Fatal("ACP sessions were restricted to one per connection") + } + _, updates := call(t, first, 4, "session/prompt", textPrompt(firstSession, "remember this")) + if len(updates) != 1 || !bytes.Contains(updates[0], []byte("remember this")) { + t.Fatal("missing output") + } + second := connect(t, address) + secondPID := initialize(t, second) + if firstPID == secondPID { + t.Fatal("connections shared an agent process") + } + secondSession := newSession(t, second, 2, dir) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + if conn, response, err := websocket.Dial(ctx, address, nil); err == nil { + conn.CloseNow() + t.Fatal("connection limit not enforced") + } else if response == nil || response.StatusCode != http.StatusTooManyRequests { + t.Fatalf("expected 429, got %+v: %v", response, err) + } + // Disconnect an in-flight turn: the agent must terminate, not detach. + send(t, first, map[string]any{"jsonrpc": "2.0", "id": 5, "method": "session/prompt", "params": textPrompt(firstSession, "block")}) + if message := read(t, first); message.Method != "session/update" { + t.Fatal("blocking turn did not start") + } + first.Close(websocket.StatusNormalClosure, "") + eventually(t, func() bool { return syscall.Kill(firstPID, 0) == syscall.ESRCH && len(handler.slots) == 1 }) + call(t, second, 3, "session/prompt", textPrompt(secondSession, "still alive")) + reconnected := connect(t, address) + if pid := initialize(t, reconnected); pid == firstPID || pid == secondPID { + t.Fatal("reconnect did not start a fresh agent") + } + _, restored := call(t, reconnected, 2, "session/load", map[string]any{"sessionId": firstSession, "cwd": dir, "mcpServers": []any{}}) + if len(restored) != 1 || !bytes.Contains(restored[0], []byte("remember this")) { + t.Fatal("exact native session history was not restored") + } + second.Close(websocket.StatusNormalClosure, "") + reconnected.Close(websocket.StatusNormalClosure, "") +} + +func TestACPContentAndPermissions(t *testing.T) { + _, address, dir, _ := fixture(t, 1) + conn := connect(t, address) + initialize(t, conn) + session := newSession(t, conn, 2, dir) + content := json.RawMessage(`[{"type":"image","data":"AAEC","mimeType":"image/png","_meta":{"n":9007199254740993}},{"type":"resource_link","uri":"file:///source","name":"source"}]`) + completion, updates := call(t, conn, 3, "session/prompt", map[string]any{"sessionId": session, "prompt": content}) + if len(updates) != 2 || !bytes.Contains(updates[0], []byte("9007199254740993")) || !bytes.Contains(completion, []byte("preserved")) { + t.Fatal("opaque ACP content or completion metadata changed") + } + send(t, conn, map[string]any{"jsonrpc": "2.0", "id": 4, "method": "session/prompt", "params": textPrompt(session, "permission")}) + permission := read(t, conn) + if permission.Method != "session/request_permission" || string(permission.ID) != `"approval"` { + t.Fatalf("permission request changed: %+v", permission) + } + send(t, conn, map[string]any{"jsonrpc": "2.0", "id": permission.ID, "result": map[string]any{"outcome": map[string]string{"outcome": "selected", "optionId": "allow"}}}) + _, updates = result(t, conn, 4) + if len(updates) != 1 || !bytes.Contains(updates[0], []byte("selected")) { + t.Fatal("permission response did not reach agent") + } + conn.Close(websocket.StatusNormalClosure, "") +} + +func TestShutdownClosesAgent(t *testing.T) { + handler, address, _, stop := fixture(t, 1) + conn := connect(t, address) + pid := initialize(t, conn) + stop() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + if _, _, err := conn.Read(ctx); err == nil { + t.Fatal("shutdown did not close connection") + } + eventually(t, func() bool { return syscall.Kill(pid, 0) == syscall.ESRCH && len(handler.slots) == 0 }) +} + +func TestExistingACPMirrorClient(t *testing.T) { + _, address, dir, _ := fixture(t, 1) + remote := os.Getenv("AGENT_PROXY_TEST_ACPREMOTE") + client, err := filepath.Abs("testdata/client.py") + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + command := exec.CommandContext(ctx, filepath.Join(filepath.Dir(remote), "python"), client, remote, address, dir) + output, err := command.CombinedOutput() + if err != nil || !bytes.Contains(output, []byte(`"ok": true`)) { + t.Fatalf("existing client failed: %v\n%s", err, output) + } + t.Log(strings.TrimSpace(string(output))) +} diff --git a/server/lib/agentproxy/testdata/agent.py b/server/lib/agentproxy/testdata/agent.py new file mode 100644 index 000000000..fcebe1472 --- /dev/null +++ b/server/lib/agentproxy/testdata/agent.py @@ -0,0 +1,81 @@ +"""Deterministic ACP stdio peer; no model calls or provider credentials.""" +import json +import os +from pathlib import Path +import sys +import uuid + +root = Path(os.environ["AGENT_PROXY_TEST_DIR"]) +pending = None + + +def send(value): + print(json.dumps(value), flush=True) + + +def reply(request_id, result): + send({"jsonrpc": "2.0", "id": request_id, "result": result}) + + +def update(session_id, content): + send({"jsonrpc": "2.0", "method": "session/update", "params": { + "sessionId": session_id, + "update": {"sessionUpdate": "agent_message_chunk", "content": content}, + }}) + + +for line in sys.stdin: + if not line.strip(): + continue + message = json.loads(line) + method = message.get("method") + params = message.get("params", {}) + request_id = message.get("id") + if method == "initialize": + reply(request_id, {"protocolVersion": 1, "agentInfo": {"name": "test-peer", "version": "1"}, + "agentCapabilities": {"loadSession": True, + "promptCapabilities": {"image": True, "audio": True, "embeddedContext": True}, + "sessionCapabilities": {"list": {}, "resume": {}}}, + "authMethods": [], "_meta": {"pid": os.getpid()}}) + elif method == "session/new": + session_id = str(uuid.uuid4()) + (root / (session_id + ".json")).write_text(json.dumps([])) + reply(request_id, {"sessionId": session_id}) + elif method in ("session/load", "session/resume"): + session_id = params["sessionId"] + if str(uuid.UUID(session_id)) != session_id: + raise ValueError("invalid fixture session ID") + history = json.loads((root / (session_id + ".json")).read_text()) + if method == "session/load": + for content in history: + update(session_id, content) + reply(request_id, {}) + elif method == "session/list": + reply(request_id, {"sessions": [{"sessionId": path.stem, "cwd": str(root)} for path in root.glob("*.json")]}) + elif method == "session/prompt": + session_id = params["sessionId"] + content = params["prompt"] + text = content[0].get("text", "") + if text == "permission": + pending = (request_id, session_id) + send({"jsonrpc": "2.0", "id": "approval", "method": "session/request_permission", "params": { + "sessionId": session_id, "toolCall": {"toolCallId": "tool-1", "title": "checkpoint", "status": "pending"}, + "options": [{"optionId": "allow", "kind": "allow_once", "name": "Allow"}], + }}) + elif text == "block": + pending = (request_id, session_id) + update(session_id, {"type": "text", "text": "started"}) + else: + (root / (session_id + ".json")).write_text(json.dumps(content)) + for block in content: + update(session_id, block) + reply(request_id, {"stopReason": "end_turn", "_meta": {"preserved": True}}) + elif method == "session/cancel" and pending: + reply(pending[0], {"stopReason": "cancelled"}) + pending = None + elif method is None and request_id == "approval" and pending: + update(pending[1], {"type": "text", "text": message["result"]["outcome"]["outcome"]}) + reply(pending[0], {"stopReason": "end_turn"}) + pending = None + elif method and request_id is not None: + send({"jsonrpc": "2.0", "id": request_id, "error": {"code": -32601, "message": "not implemented"}}) diff --git a/server/lib/agentproxy/testdata/client.py b/server/lib/agentproxy/testdata/client.py new file mode 100644 index 000000000..e79b25014 --- /dev/null +++ b/server/lib/agentproxy/testdata/client.py @@ -0,0 +1,65 @@ +"""Exercise the unmodified acpremote mirror with the official ACP Python client.""" +import asyncio +import json +import sys + +from acp import connect_to_agent, text_block +from acp.schema import ClientCapabilities, RequestPermissionResponse + + +class Client: + def __init__(self): + self.texts = asyncio.Queue() + self.permissions = 0 + + def on_connect(self, connection): + pass + + async def session_update(self, session_id, update, **kwargs): + content = getattr(update, "content", None) + if getattr(content, "type", None) == "text": + self.texts.put_nowait(content.text) + + async def request_permission(self, options, session_id, tool_call, **kwargs): + self.permissions += 1 + return RequestPermissionResponse.model_validate({ + "outcome": {"outcome": "selected", "optionId": options[0].option_id} + }) + + async def ext_notification(self, method, params): + pass + + +async def main(): + process = await asyncio.create_subprocess_exec( + sys.argv[1], "mirror", sys.argv[2], + stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, + ) + client = Client() + connection = connect_to_agent(client, process.stdin, process.stdout) + try: + initialized = await connection.initialize(protocol_version=1, client_capabilities=ClientCapabilities()) + assert initialized.protocol_version == 1 + first = await connection.new_session(cwd=sys.argv[3], mcp_servers=[]) + second = await connection.new_session(cwd=sys.argv[3], mcp_servers=[]) + assert first.session_id != second.session_id + result = await connection.prompt(session_id=first.session_id, prompt=[text_block("mirror works")]) + assert result.stop_reason == "end_turn" + assert await client.texts.get() == "mirror works" + await connection.prompt(session_id=second.session_id, prompt=[text_block("permission")]) + assert client.permissions == 1 + assert await client.texts.get() == "selected" + await connection.load_session(session_id=first.session_id, cwd=sys.argv[3], mcp_servers=[]) + assert await client.texts.get() == "mirror works" + print(json.dumps({"ok": True, "client": "ACP Python SDK via acpremote mirror", "sessions": 2, "permissions": 1})) + finally: + await connection.close() + process.stdin.close() + try: + await asyncio.wait_for(process.wait(), 5) + except TimeoutError: + process.kill() + await process.wait() + + +asyncio.run(asyncio.wait_for(main(), 20)) diff --git a/server/lib/wsproxy/proxy_options_test.go b/server/lib/wsproxy/proxy_options_test.go new file mode 100644 index 000000000..9bc895d10 --- /dev/null +++ b/server/lib/wsproxy/proxy_options_test.go @@ -0,0 +1,73 @@ +package wsproxy + +import ( + "context" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/coder/websocket" +) + +func proxyWithOptions(t *testing.T, opts ProxyOptions) (*websocket.Conn, <-chan struct{}) { + t.Helper() + closed := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := websocket.Accept(w, r, nil) + if err != nil { + return + } + defer close(closed) + defer conn.CloseNow() + for { + kind, data, err := conn.Read(r.Context()) + if err != nil { + return + } + if err := conn.Write(r.Context(), kind, data); err != nil { + return + } + } + })) + t.Cleanup(upstream.Close) + opts.Logger = slog.New(slog.NewTextHandler(io.Discard, nil)) + proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + Proxy(w, r, "ws"+strings.TrimPrefix(upstream.URL, "http"), opts) + })) + t.Cleanup(proxy.Close) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + conn, _, err := websocket.Dial(ctx, "ws"+strings.TrimPrefix(proxy.URL, "http"), nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { conn.CloseNow() }) + return conn, closed +} + +func TestProxyReadLimit(t *testing.T) { + conn, _ := proxyWithOptions(t, ProxyOptions{ReadLimit: 64}) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + if err := conn.Write(ctx, websocket.MessageText, []byte(strings.Repeat("x", 128))); err != nil { + t.Fatal(err) + } + if _, _, err := conn.Read(ctx); err == nil { + t.Fatal("oversized frame was forwarded") + } +} + +func TestProxyPingDetectsUnresponsiveClient(t *testing.T) { + _, closed := proxyWithOptions(t, ProxyOptions{PingInterval: 20 * time.Millisecond}) + // Without a client read loop, pings are not answered. The proxy must close + // the upstream too instead of keeping the remote agent alive indefinitely. + select { + case <-closed: + case <-time.After(3 * time.Second): + t.Fatal("unresponsive client did not close upstream") + } +} diff --git a/server/lib/wsproxy/wsproxy.go b/server/lib/wsproxy/wsproxy.go index 65ad81f46..aa867412d 100644 --- a/server/lib/wsproxy/wsproxy.go +++ b/server/lib/wsproxy/wsproxy.go @@ -38,6 +38,11 @@ type ProxyOptions struct { Logger *slog.Logger Transform MessageTransform Observe Observer + // ReadLimit overrides the default 100 MiB limit on both connections. + ReadLimit int64 + // PingInterval enables downstream liveness checks with the same pong timeout. + // Zero leaves keepalive behavior unchanged. + PingInterval time.Duration // Registry, when set, tracks the accepted client connection so it is // closed with a Going Away frame on server shutdown. Registry *wsdrain.Registry @@ -140,7 +145,11 @@ func Proxy(w http.ResponseWriter, r *http.Request, upstreamURL string, opts Prox logger.Error("websocket accept failed", slog.String("err", err.Error())) return } - clientConn.SetReadLimit(100 * 1024 * 1024) + readLimit := opts.ReadLimit + if readLimit <= 0 { + readLimit = 100 * 1024 * 1024 + } + clientConn.SetReadLimit(readLimit) untrack := opts.Registry.Track(clientConn) defer untrack() @@ -151,10 +160,16 @@ func Proxy(w http.ResponseWriter, r *http.Request, upstreamURL string, opts Prox clientConn.Close(websocket.StatusInternalError, "failed to connect to upstream") return } - upstreamConn.SetReadLimit(100 * 1024 * 1024) + upstreamConn.SetReadLimit(readLimit) logger.Debug("proxying websocket", slog.String("url", upstreamURL)) + if opts.PingInterval > 0 { + pingCtx, cancel := context.WithCancel(r.Context()) + defer cancel() + go pingClient(pingCtx, clientConn, opts.PingInterval) + } + var once sync.Once cleanup := func(_ PumpExitCause) { once.Do(func() { @@ -165,3 +180,22 @@ func Proxy(w http.ResponseWriter, r *http.Request, upstreamURL string, opts Prox Pump(r.Context(), clientConn, upstreamConn, cleanup, logger, opts.Transform, opts.Observe) } + +func pingClient(ctx context.Context, client *websocket.Conn, interval time.Duration) { + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + pingCtx, cancel := context.WithTimeout(ctx, interval) + err := client.Ping(pingCtx) + cancel() + if err != nil { + client.CloseNow() + return + } + } + } +} diff --git a/server/runtime/acp/requirements.txt b/server/runtime/acp/requirements.txt new file mode 100644 index 000000000..ea3ef8a56 --- /dev/null +++ b/server/runtime/acp/requirements.txt @@ -0,0 +1,3 @@ +acpremote==1.7.0 +agent-client-protocol==0.11.0 +websockets==15.0.1