From 6a2f5aa5aac4d5ecee274a614b13748b271bd1bc Mon Sep 17 00:00:00 2001 From: BMAD CI Fix Agent Date: Sat, 5 Sep 2026 00:43:44 -0500 Subject: [PATCH] feat: add model profile ensure hooks --- CHANGELOG.md | 4 +- Makefile | 5 +- README.md | 16 +++++- configs/router.example.yaml | 12 +++++ docs/architecture.md | 31 ++++++++++++ internal/config/config.go | 73 ++++++++++++++++++++++++++++ internal/config/config_test.go | 70 ++++++++++++++++++++++++++ internal/server/server.go | 89 ++++++++++++++++++++++++++++++++-- internal/server/server_test.go | 63 ++++++++++++++++++++++++ 9 files changed, 354 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f3b874d..9ef786d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,9 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -### Changed +### Added -- Updated beta banner to v1 stable +- Opt-in command-backed model profile ensure hooks. ## [1.0.0] - 2026-03-01 diff --git a/Makefile b/Makefile index 2b38159..e76c593 100644 --- a/Makefile +++ b/Makefile @@ -27,6 +27,7 @@ VAGRANT_DESTROY ?= 1 DIST_DIR ?= dist BIN_DIR ?= bin PACKAGE_NAME := devrail-router_$(VERSION)_$(GOOS)_$(GOARCH) +BUILD_OUTPUT := $(BIN_DIR)/devrail-router_$(GOOS)_$(GOARCH) RELEASE_TARGETS ?= linux/amd64 linux/arm64 darwin/arm64 DOCKER_RUN := docker run --rm \ @@ -73,7 +74,7 @@ build: ## Build the devrail-router binary for GOOS/GOARCH CGO_ENABLED=0 GOOS="$(GOOS)" GOARCH="$(GOARCH)" go build \ -trimpath \ -ldflags "-s -w -X main.version=$(VERSION)" \ - -o "$(BIN_DIR)/devrail-router" \ + -o "$(BUILD_OUTPUT)" \ ./cmd/devrail-router changelog: ## Generate CHANGELOG.md from conventional commits @@ -131,7 +132,7 @@ lint: ## Run all linters package: build ## Build a Linux/macOS tarball package @rm -rf "$(DIST_DIR)/$(PACKAGE_NAME)" @mkdir -p "$(DIST_DIR)/$(PACKAGE_NAME)" - cp "$(BIN_DIR)/devrail-router" "$(DIST_DIR)/$(PACKAGE_NAME)/devrail-router" + cp "$(BUILD_OUTPUT)" "$(DIST_DIR)/$(PACKAGE_NAME)/devrail-router" cp LICENSE README.md CHANGELOG.md "$(DIST_DIR)/$(PACKAGE_NAME)/" mkdir -p "$(DIST_DIR)/$(PACKAGE_NAME)/configs" \ "$(DIST_DIR)/$(PACKAGE_NAME)/docs" \ diff --git a/README.md b/README.md index 502124f..2e6ab93 100644 --- a/README.md +++ b/README.md @@ -22,19 +22,25 @@ This repository is in early foundation work. The current service supports: - `/v1/models` - OpenAI-compatible `/v1/*` request proxying - model alias rewriting +- opt-in model profile ensure hooks - YAML configuration - Linux tarball packaging - Linux/systemd install script and unit - Docker image and Compose smoke testing with a mock OpenAI-compatible backend -Routing policy, auth, telemetry, LM Studio lifecycle integration, and Omarchy -integration are planned next. +Routing policy, auth, telemetry, native LM Studio lifecycle integration, and +Omarchy integration are planned next. Model aliases can also set basic concurrency guardrails with `max_concurrent_requests`, `max_queue_size`, and `queue_timeout`. This lets heavy local models wait or reject predictably instead of allowing multiple agents to dogpile the same backend. +Aliases can opt into a command-backed `ensure` hook before proxying. This is +intended for host adapters such as LM Studio profile loaders that need to +guarantee context length, parallelism, TTL, or model identifier before a client +request reaches the backend. + ## Quick Start Build and test locally: @@ -106,6 +112,12 @@ models: max_concurrent_requests: 2 max_queue_size: 4 queue_timeout: 30s + ensure: + mode: command + command: + - /usr/local/bin/lmstudio-load-profile + - local-coder + timeout: 30s backends: - id: lmstudio diff --git a/configs/router.example.yaml b/configs/router.example.yaml index 65fef21..487cb72 100644 --- a/configs/router.example.yaml +++ b/configs/router.example.yaml @@ -12,6 +12,12 @@ models: max_concurrent_requests: 2 max_queue_size: 4 queue_timeout: 30s + ensure: + mode: command + command: + - /usr/local/bin/lmstudio-load-profile + - local-coder + timeout: 30s - id: local-coder-large name: Local Coder Large backend: lmstudio @@ -22,6 +28,12 @@ models: max_concurrent_requests: 1 max_queue_size: 2 queue_timeout: 2m + ensure: + mode: command + command: + - /usr/local/bin/lmstudio-load-profile + - local-coder-large + timeout: 45s backends: - id: lmstudio diff --git a/docs/architecture.md b/docs/architecture.md index 2c12a0c..d9d192e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -32,6 +32,7 @@ parallelism, TTL, GPU offload, auth, or scheduling policy. The first router is deliberately simple: - Clients request a DevRail model alias. +- The router optionally runs an alias-specific readiness hook. - The router rewrites the request to the configured backend model. - The backend handles inference. @@ -62,3 +63,33 @@ full, DevRail Router returns an OpenAI-shaped `429` error. If the request waits longer than `queue_timeout`, it returns `503`. Successful queued requests get an `X-Devrail-Queue-Wait-Ms` response header, and queue decisions are logged with active count, queued count, and wait time. + +## Profile Ensure Hooks + +Backends such as LM Studio can Just-In-Time load models, but backend defaults may +not match the alias contract that clients see. For example, an alias may +advertise a 64k context window while the backend's default JIT load only creates +an 8k context. A model alias can opt into a command-backed ensure hook: + +```yaml +models: + - id: local-coder + backend: lmstudio + target_model: qwen3-coder-30b-a3b-instruct + context_window: 65536 + ensure: + mode: command + command: + - /usr/local/bin/lmstudio-load-profile + - local-coder + timeout: 30s +``` + +The ensure hook runs after queue slot acquisition and before request forwarding. +If it fails or times out, DevRail Router returns an OpenAI-shaped `503` error and +does not send the request to the backend. This keeps clients from hanging behind +a model that is unloaded, incorrectly loaded, or busy switching profiles. + +Command hooks are a conservative bridge for early local integrations. Native +backend adapters can later inspect runtime state directly and choose between +passive JIT loading, profile enforcement, or refusing unsafe model switches. diff --git a/internal/config/config.go b/internal/config/config.go index 89271ff..c186f6c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -33,6 +33,7 @@ type ModelConfig struct { MaxConcurrentRequests int `yaml:"max_concurrent_requests"` MaxQueueSize int `yaml:"max_queue_size"` QueueTimeout string `yaml:"queue_timeout"` + Ensure EnsureConfig } type BackendConfig struct { @@ -42,6 +43,39 @@ type BackendConfig struct { APIKeyEnv string `yaml:"api_key_env"` } +type EnsureConfig struct { + Mode string `yaml:"mode"` + Command CommandArgs `yaml:"command"` + Timeout string `yaml:"timeout"` +} + +type CommandArgs []string + +func (args *CommandArgs) UnmarshalYAML(value *yaml.Node) error { + switch value.Kind { + case yaml.ScalarNode: + var command string + if err := value.Decode(&command); err != nil { + return err + } + if command == "" { + *args = nil + return nil + } + *args = []string{"/bin/sh", "-c", command} + return nil + case yaml.SequenceNode: + var command []string + if err := value.Decode(&command); err != nil { + return err + } + *args = command + return nil + default: + return fmt.Errorf("command must be a string or list of strings") + } +} + func Load(path string) (Config, error) { raw, err := os.ReadFile(path) if err != nil { @@ -115,6 +149,9 @@ func (cfg Config) Validate() error { if _, err := model.QueueTimeoutDuration(); err != nil { return fmt.Errorf("model %q queue_timeout is invalid: %w", model.ID, err) } + if err := model.Ensure.Validate(); err != nil { + return fmt.Errorf("model %q ensure is invalid: %w", model.ID, err) + } if _, ok := models[model.ID]; ok { return fmt.Errorf("model %q is duplicated", model.ID) } @@ -159,3 +196,39 @@ func (model ModelConfig) QueueTimeoutDuration() (time.Duration, error) { return duration, nil } + +func (ensure EnsureConfig) Validate() error { + switch ensure.Mode { + case "", "disabled": + return nil + case "command": + if len(ensure.Command) == 0 { + return errors.New("command is required when mode is command") + } + for _, arg := range ensure.Command { + if arg == "" { + return errors.New("command arguments must not be empty") + } + } + _, err := ensure.TimeoutDuration() + return err + default: + return fmt.Errorf("unknown mode %q", ensure.Mode) + } +} + +func (ensure EnsureConfig) TimeoutDuration() (time.Duration, error) { + if ensure.Timeout == "" { + return 30 * time.Second, nil + } + + duration, err := time.ParseDuration(ensure.Timeout) + if err != nil { + return 0, err + } + if duration <= 0 { + return 0, errors.New("duration must be positive") + } + + return duration, nil +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 8fae1f8..0f1dfa0 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -70,6 +70,76 @@ func TestValidateQueueSettings(t *testing.T) { } } +func TestLoadEnsureCommandString(t *testing.T) { + t.Parallel() + + path := filepath.Join(t.TempDir(), "router.yaml") + raw := []byte(` +models: + - id: local-coder + backend: lmstudio + target_model: qwen3-coder-30b-a3b-instruct + ensure: + mode: command + command: /usr/local/bin/lmstudio-load-profile local-coder + timeout: 5s +backends: + - id: lmstudio + type: openai-compatible + base_url: http://127.0.0.1:1234/v1 +`) + if err := os.WriteFile(path, raw, 0o600); err != nil { + t.Fatalf("write config: %v", err) + } + + cfg, err := Load(path) + if err != nil { + t.Fatalf("load config: %v", err) + } + + model := cfg.Models[0] + if model.Ensure.Mode != "command" { + t.Fatalf("unexpected ensure mode: %q", model.Ensure.Mode) + } + if got := []string(model.Ensure.Command); len(got) != 3 || got[0] != "/bin/sh" || got[1] != "-c" { + t.Fatalf("unexpected command args: %#v", got) + } + duration, err := model.Ensure.TimeoutDuration() + if err != nil { + t.Fatalf("parse timeout: %v", err) + } + if duration != 5*time.Second { + t.Fatalf("unexpected timeout: %s", duration) + } +} + +func TestValidateEnsureCommandRequired(t *testing.T) { + t.Parallel() + + cfg := Config{ + Models: []ModelConfig{{ + ID: "local-coder", + Backend: "lmstudio", + TargetModel: "qwen/qwen3.6-35b-a3b", + Ensure: EnsureConfig{ + Mode: "command", + }, + }}, + Backends: []BackendConfig{{ + ID: "lmstudio", + BaseURL: "http://127.0.0.1:1234/v1", + }}, + } + + err := cfg.Validate() + if err == nil { + t.Fatal("expected validation error") + } + if !strings.Contains(err.Error(), "ensure") { + t.Fatalf("expected ensure error, got: %v", err) + } +} + func TestValidateInvalidQueueTimeout(t *testing.T) { t.Parallel() diff --git a/internal/server/server.go b/internal/server/server.go index 4375f16..aef9e0a 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -2,6 +2,7 @@ package server import ( "bytes" + "context" "encoding/json" "errors" "fmt" @@ -11,14 +12,17 @@ import ( "net/http/httputil" "net/url" "os" + "os/exec" "strings" + "sync" "github.com/devrail-dev/devrail-router/internal/config" ) type Server struct { - cfg config.Config - limiters map[string]*modelLimiter + cfg config.Config + limiters map[string]*modelLimiter + ensureLocks map[string]*sync.Mutex } func New(cfg config.Config) (*Server, error) { @@ -37,7 +41,14 @@ func New(cfg config.Config) (*Server, error) { } } - return &Server{cfg: cfg, limiters: limiters}, nil + ensureLocks := make(map[string]*sync.Mutex, len(cfg.Models)) + for _, model := range cfg.Models { + if model.Ensure.Mode != "" && model.Ensure.Mode != "disabled" { + ensureLocks[model.ID] = &sync.Mutex{} + } + } + + return &Server{cfg: cfg, limiters: limiters, ensureLocks: ensureLocks}, nil } func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { @@ -112,6 +123,10 @@ func (s *Server) proxyOpenAI(w http.ResponseWriter, r *http.Request) { defer release() } + if !s.ensureModelReady(w, r, model) { + return + } + body, err = rewriteModel(body, model.TargetModel) if err != nil { http.Error(w, err.Error(), http.StatusBadRequest) @@ -184,6 +199,74 @@ func (s *Server) acquireModelSlot(w http.ResponseWriter, r *http.Request, model return nil, false } +func (s *Server) ensureModelReady(w http.ResponseWriter, r *http.Request, model config.ModelConfig) bool { + if model.Ensure.Mode == "" || model.Ensure.Mode == "disabled" { + return true + } + + lock := s.ensureLocks[model.ID] + if lock == nil { + lock = &sync.Mutex{} + } + + lock.Lock() + defer lock.Unlock() + + switch model.Ensure.Mode { + case "command": + return s.ensureModelReadyWithCommand(w, r, model) + default: + writeOpenAIError(w, http.StatusInternalServerError, "model ensure mode is unsupported", "devrail_ensure_unsupported", "ensure_unsupported") + return false + } +} + +func (s *Server) ensureModelReadyWithCommand(w http.ResponseWriter, r *http.Request, model config.ModelConfig) bool { + timeout, err := model.Ensure.TimeoutDuration() + if err != nil { + writeOpenAIError(w, http.StatusInternalServerError, "model ensure timeout is invalid", "devrail_ensure_config_error", "ensure_config_error") + return false + } + + ctx, cancel := context.WithTimeout(r.Context(), timeout) + defer cancel() + + command := []string(model.Ensure.Command) + cmd := exec.CommandContext(ctx, command[0], command[1:]...) + output, err := cmd.CombinedOutput() + outputText := strings.TrimSpace(string(output)) + if len(outputText) > 2048 { + outputText = outputText[:2048] + "...[truncated]" + } + if err != nil { + status := http.StatusServiceUnavailable + message := "model profile is not ready" + code := "ensure_failed" + if errors.Is(ctx.Err(), context.DeadlineExceeded) { + message = "timed out ensuring model profile" + code = "ensure_timeout" + } + + slog.Warn( + "model ensure command failed", + "alias", model.ID, + "target_model", model.TargetModel, + "error", err, + "output", outputText, + ) + writeOpenAIError(w, status, message, "devrail_"+code, code) + return false + } + + slog.Info( + "model ensure command completed", + "alias", model.ID, + "target_model", model.TargetModel, + "output", outputText, + ) + return true +} + func requestModel(r *http.Request) (string, []byte, error) { if r.Body == nil { return "", nil, fmt.Errorf("request body is required") diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 99cbc8c..ed714aa 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -5,6 +5,8 @@ import ( "io" "net/http" "net/http/httptest" + "os" + "path/filepath" "strings" "sync/atomic" "testing" @@ -264,6 +266,67 @@ func TestModelLimiterTimesOutQueuedRequest(t *testing.T) { } } +func TestEnsureCommandRunsBeforeProxy(t *testing.T) { + t.Parallel() + + ensureFile := filepath.Join(t.TempDir(), "ensured") + ensureCommand := []string{"/bin/sh", "-c", "printf ready > " + ensureFile} + + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + if _, err := os.Stat(ensureFile); err != nil { + t.Errorf("ensure command did not run before backend request: %v", err) + http.Error(w, "not ready", http.StatusServiceUnavailable) + return + } + writeJSON(w, http.StatusOK, map[string]string{"status": "ok"}) + })) + t.Cleanup(backend.Close) + + srv := testServerWithBackend(t, backend.URL, config.ModelConfig{ + ID: "local-coder", + Backend: "lmstudio", + TargetModel: "target-model", + Ensure: config.EnsureConfig{ + Mode: "command", + Command: ensureCommand, + Timeout: "1s", + }, + }) + + if status := serveChat(t, srv, "local-coder"); status != http.StatusOK { + t.Fatalf("unexpected status: %d", status) + } +} + +func TestEnsureCommandFailureReturnsServiceUnavailable(t *testing.T) { + t.Parallel() + + var backendCalls atomic.Int32 + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + backendCalls.Add(1) + writeJSON(w, http.StatusOK, map[string]string{"status": "ok"}) + })) + t.Cleanup(backend.Close) + + srv := testServerWithBackend(t, backend.URL, config.ModelConfig{ + ID: "local-coder", + Backend: "lmstudio", + TargetModel: "target-model", + Ensure: config.EnsureConfig{ + Mode: "command", + Command: []string{"/bin/sh", "-c", "echo nope >&2; exit 7"}, + Timeout: "1s", + }, + }) + + if status := serveChat(t, srv, "local-coder"); status != http.StatusServiceUnavailable { + t.Fatalf("unexpected status: %d", status) + } + if backendCalls.Load() != 0 { + t.Fatalf("backend received %d calls, want 0", backendCalls.Load()) + } +} + func testServer(t *testing.T) *Server { t.Helper()