diff --git a/.agents/workflows/verify-pr.md b/.agents/workflows/verify-pr.md index df43fc8..cee8315 100644 --- a/.agents/workflows/verify-pr.md +++ b/.agents/workflows/verify-pr.md @@ -1,6 +1,6 @@ --- description: Verify a PR before submission --- -1. Read the `skills/contribute/SKILL.md` file using the `view_file` tool to understand the 5-step gate function. -2. Execute all 5 steps of the gate function exactly as described. +1. Read the `skills/contribute/SKILL.md` file using the `view_file` tool to understand the 6-step gate function. +2. Execute all 6 steps of the gate function exactly as described. 3. Report the result of each step to the user. diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json index b77da27..f1826b6 100644 --- a/.claude-plugin/marketplace.json +++ b/.claude-plugin/marketplace.json @@ -7,7 +7,7 @@ "plugins": [ { "name": "smg", - "description": "4 process-enforcing skills for the Shepherd Model Gateway — orientation, implementation (with 19 subsystem recipes), PR review, and pre-ship quality gates", + "description": "4 process-enforcing skills for the Shepherd Model Gateway — orientation, implementation (with 20 subsystem recipes), PR review, and pre-ship quality gates", "version": "0.2.0", "source": "./", "author": { diff --git a/.claude-plugin/plugin.json b/.claude-plugin/plugin.json index a0e98c3..0718c79 100644 --- a/.claude-plugin/plugin.json +++ b/.claude-plugin/plugin.json @@ -1,12 +1,12 @@ { "name": "smg", - "description": "4 process-enforcing skills for the Shepherd Model Gateway — orientation, implementation (with 19 subsystem recipes), PR review, and pre-ship quality gates", + "description": "4 process-enforcing skills for the Shepherd Model Gateway — orientation, implementation (with 20 subsystem recipes), PR review, and pre-ship quality gates", "version": "0.2.0", "author": { "name": "LightSeek" }, - "homepage": "https://github.com/lightseekorg/smg-dev-guide", - "repository": "https://github.com/lightseekorg/smg-dev-guide", + "homepage": "https://github.com/smg-project/smg-dev-guide", + "repository": "https://github.com/smg-project/smg-dev-guide", "license": "Apache-2.0", "keywords": ["smg", "llm-gateway", "rust", "kubernetes", "routing", "grpc"], "skills": "./skills/", diff --git a/.cursor-plugin/plugin.json b/.cursor-plugin/plugin.json index a2ca0e5..bb8270f 100644 --- a/.cursor-plugin/plugin.json +++ b/.cursor-plugin/plugin.json @@ -1,13 +1,13 @@ { "name": "smg", "displayName": "SMG Development Assistant", - "description": "4 process-enforcing skills for the Shepherd Model Gateway — orientation, implementation (with 19 subsystem recipes), PR review, and pre-ship quality gates", + "description": "4 process-enforcing skills for the Shepherd Model Gateway — orientation, implementation (with 20 subsystem recipes), PR review, and pre-ship quality gates", "version": "0.2.0", "author": { "name": "LightSeek" }, - "homepage": "https://github.com/lightseekorg/smg-dev-guide", - "repository": "https://github.com/lightseekorg/smg-dev-guide", + "homepage": "https://github.com/smg-project/smg-dev-guide", + "repository": "https://github.com/smg-project/smg-dev-guide", "license": "Apache-2.0", "keywords": ["smg", "llm-gateway", "rust", "kubernetes", "routing", "grpc"], "skills": "./skills/", diff --git a/README.md b/README.md index 0edc72d..d6452d4 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # SMG Dev Guide -AI-powered development guide for the [Shepherd Model Gateway](https://github.com/lightseekorg/smg) — 4 process-enforcing skills that change what your AI coding agent **does**, not just what it **knows**. +AI-powered development guide for the [Shepherd Model Gateway](https://github.com/smg-project/smg) — 4 process-enforcing skills that change what your AI coding agent **does**, not just what it **knows**. Works with **Google Antigravity**, **Gemini CLI**, **Claude Code**, **Codex**, and **Cursor**. @@ -13,7 +13,7 @@ Works with **Google Antigravity**, **Gemini CLI**, **Claude Code**, **Codex**, a **Global**: To make these skills and workflows available across all projects, copy or symlink them into your global Antigravity `~/.gemini/antigravity/` directory: ```bash -git clone https://github.com/lightseekorg/smg-dev-guide.git ~/.gemini/antigravity/repos/smg-dev-guide +git clone https://github.com/smg-project/smg-dev-guide.git ~/.gemini/antigravity/repos/smg-dev-guide mkdir -p ~/.gemini/antigravity/workflows ~/.gemini/antigravity/skills ln -s ~/.gemini/antigravity/repos/smg-dev-guide/.agents/workflows/* ~/.gemini/antigravity/workflows/ ln -s ~/.gemini/antigravity/repos/smg-dev-guide/skills/* ~/.gemini/antigravity/skills/ @@ -25,12 +25,12 @@ The Gemini CLI natively supports Agent Skills. You can install these skills dire **Global Installation (Available in all projects)**: ```bash -gemini skills install https://github.com/lightseekorg/smg-dev-guide.git +gemini skills install https://github.com/smg-project/smg-dev-guide.git ``` **Workspace Installation (Only in current project)**: ```bash -gemini skills install https://github.com/lightseekorg/smg-dev-guide.git --scope workspace +gemini skills install https://github.com/smg-project/smg-dev-guide.git --scope workspace ``` ### Claude Code @@ -38,14 +38,14 @@ gemini skills install https://github.com/lightseekorg/smg-dev-guide.git --scope From the shell: ```bash -claude plugin marketplace add lightseekorg/smg-dev-guide +claude plugin marketplace add smg-project/smg-dev-guide claude plugin install smg@smg-dev-guide ``` Or inside Claude Code: ``` -/plugin marketplace add lightseekorg/smg-dev-guide +/plugin marketplace add smg-project/smg-dev-guide /plugin install smg@smg-dev-guide ``` @@ -56,7 +56,7 @@ Run `/reload-plugins` (or restart Claude Code) to load the skills. Copy or symlink the skills into your user skills directory: ```bash -git clone https://github.com/lightseekorg/smg-dev-guide.git ~/.agents/repos/smg-dev-guide +git clone https://github.com/smg-project/smg-dev-guide.git ~/.agents/repos/smg-dev-guide ln -s ~/.agents/repos/smg-dev-guide/.agents/skills/* ~/.agents/skills/ ``` @@ -75,14 +75,14 @@ Install as a Cursor plugin via `.cursor-plugin/plugin.json`. | `map` | Orient | Crate map, layering rules, config propagation, request flow, label pipeline | | `implement` | Build | Detects subsystem, loads recipe, creates tasks, enforces step-by-step execution with verification | | `review-pr` | Review | Maps changed files to checklist sections, creates review tasks per subsystem, cites file:line | -| `contribute` | Ship | 5-step quality gate (fmt → clippy → test → bindings → commit format) with enforcement | +| `contribute` | Ship | 6-step quality gate (fmt → clippy → test → hooks → python/go → names and commits) with enforcement | **Google Antigravity** — invoke workflows using slash commands in the chat: ``` /map → discover codebase structure and ownership /implement-feature → guides you through building a feature /review-pr → checks your work against anti-patterns -/verify-pr → runs the full 5-step quality gate +/verify-pr → runs the full 6-step quality gate ``` **Gemini CLI** — skills trigger automatically based on your prompt: @@ -110,7 +110,7 @@ Unlike passive reference docs, these skills **enforce workflows**: - **Hard Gates** block progression without prerequisites (e.g. must identify touched subsystems before reviewing) - **Rationalization Tables** counter common excuses for cutting corners - **Skill Chaining** ensures `implement` → `contribute` → `review-pr` flow -- **19 Implementation Recipes** provide step-by-step guidance with exact file paths, code patterns, and verification commands for every subsystem +- **20 Implementation Recipes** provide step-by-step guidance with exact file paths, code patterns, and verification commands for every subsystem ## Implementation Recipes @@ -120,11 +120,12 @@ Unlike passive reference docs, these skills **enforce workflows**: |--------|-----------| | config-plumbing | CLI flags, config fields, the critical two-path rule | | routing-policy | Load balancing, dual HTTP/gRPC mode | -| tool-parser | Tool/function call formats (14) | -| reasoning-parser | Reasoning extraction (8 model families) | +| tool-parser | Tool/function call formats (18 parsers) | +| reasoning-parser | Reasoning extraction (10 model families) | | bindings-update | Python PyO3 + Go FFI | | discovery-feature | K8s discovery, label pipeline | | grpc-backend | gRPC client, trace injection | +| zmq-backend | Direct `ipc://` ZMQ engines (vLLM EngineCore, TokenSpeed) | | storage-backend | Data connectors, hooks | | wasm-plugin | WASM middleware, WIT interface | | auth-feature | API keys, JWT/OIDC | @@ -135,7 +136,7 @@ Unlike passive reference docs, these skills **enforce workflows**: | multimodal-feature | Vision processors, media pipeline | | scheduler-feature | Priority admission: classes, queues, reservations, preemption | | tenancy-feature | Multi-tenant identity, tenant resolution, per-tenant state | -| rate-limit-feature | Token-bucket rate limiting + concurrency caps | +| rate-limit-feature | Token-bucket rate limiting, concurrency caps, per-tenant quotas | | provider-api | Provider-compatible API routers (Anthropic, Gemini) | ## Directory Structure @@ -149,7 +150,7 @@ skills/ # Skill source files map/SKILL.md contribute/SKILL.md review-pr/SKILL.md - implement/SKILL.md # + 19 recipe files + implement/SKILL.md # + 20 recipe files commands/ # Claude Code /smg command router smg.md ``` diff --git a/skills/contribute/SKILL.md b/skills/contribute/SKILL.md index e899777..4b0cdb8 100644 --- a/skills/contribute/SKILL.md +++ b/skills/contribute/SKILL.md @@ -23,29 +23,65 @@ BEFORE claiming code is ready, opening a PR, or committing: 1. FORMAT Run: cargo +nightly fmt --all Expected: No output (silent success) - If fails: Fix formatting, re-run + If fails: Fix formatting, re-run 2. LINT Run: cargo clippy --workspace --all-targets --all-features -- -D warnings Expected: No warnings, no errors - If fails: Fix all warnings, re-run + Note: --all-features enables multimodal's opencv-video, which links + system OpenCV. One-time: bash scripts/install_opencv.sh (or + make opencv-deps). Without it, drop --all-features. + If fails: Fix all warnings, re-run 3. TEST Run: cargo test Expected: "test result: ok" with 0 failures - If fails: Fix failures, re-run - -4. BINDINGS (if config/types.rs, protocols/, or bindings/ changed) + Note: Touched any Cargo.toml? Stage the updated Cargo.lock — it is + tracked, and every Rust CI job starts with cargo metadata --locked. + Note: Touched crates/wasm/ or examples/wasm/? Run + bash crates/wasm/tests/fixtures/build_fixtures.sh first, or the + guest integration tests skip silently. + If fails: Fix failures, re-run + +4. HOOKS + Run: pre-commit run --all-files + Expected: every hook Passed — codespell, trailing-whitespace, + end-of-file-fixer, check-yaml, check-toml, mixed-line-ending, + ruff, ruff-format + Setup: pip install pre-commit && pre-commit install && + pre-commit install --hook-type commit-msg + If fails: Take the hook's fixes, re-stage, re-run + +5. PYTHON / GO (if config/types.rs, protocols/, bindings/, e2e_test/, + scripts/ or grpc_servicer/ changed) Run: make python-dev - Expected: Successful compilation - If fails: Update struct literals, re-run - -5. COMMIT FORMAT - Verify: Conventional commit (feat/fix/docs/refactor/test/perf/chore/ci) - Verify: DCO sign-off present (git commit -s) - Verify: No AI attribution - -Skip any step = not verified. Run ALL five. + Run: ruff check e2e_test/ bindings/python/ scripts/ + ruff format --check e2e_test/ bindings/python/ scripts/ + Run: mypy e2e_test/ --config-file mypy.ini + mypy bindings/python/ --config-file mypy.ini + Run: cd bindings/python && pytest -q tests --cov=smg + --cov-config=.coveragerc --cov-fail-under=80 + Run: PYTHONPATH=e2e_test pytest -q --noconftest e2e_test/infra e2e_test/fixtures + Run: cd bindings/golang && make test (if bindings/golang/ changed) + Expected: wheel builds; ruff/mypy clean; tests pass at >= 80% coverage + If a new config field never reaches the bindings: add it to the Router + pyclass and its to_router_config() builder chain in + bindings/python/src/lib.rs, to RouterArgs in + bindings/python/src/smg/router_args.py, and to Go ClientConfig. + +6. NAMES AND COMMITS + Verify: Branch is / or / (e.g. feat/add-auth) — + CI comments on and closes internal PRs that do not match + Verify: PR title is "type(scope): summary", scope lowercase — types + feat/fix/docs/style/refactor/perf/test/chore/ci (PR titles also + allow lint/revert/build); CI rejects any other title + Verify: Every commit DCO-signed (git commit -s), matching your git + user.name/user.email exactly + Verify: No AI attribution — no Co-authored-by/Signed-off-by line naming + Claude or noreply@anthropic.com, in commits, PR body, or review + replies + +Skip any step = not verified. Run ALL six. ``` ## Clippy Rules You'll Hit @@ -53,28 +89,31 @@ Skip any step = not verified. Run ALL five. | Rule | Meaning | |------|---------| | `unwrap_used = "deny"` | No `.unwrap()`. Use `?`, `.ok_or()`, `.unwrap_or_else()` | -| `expect_used = "warn"` | Prefer `?` over `.expect()` | +| `expect_used = "warn"`, `panic = "warn"` | Prefer `?` over `.expect()`/`panic!` in production code | | `print_stdout/print_stderr = "warn"` | Use `tracing` crate, not `println!`/`eprintln!` | -| `todo/unimplemented/unreachable = "deny"` | No placeholder code | +| `todo/unimplemented/unreachable = "deny"`, `dbg_macro = "deny"` | No placeholder or debug macros | +| `unsafe_code = "deny"` (rust lints) | No `unsafe` outside crates that opt out | | `allow_attributes = "warn"` | Use `#[expect(lint, reason = "...")]` not `#[allow]` | -| Disallowed: `tokio::task::spawn` | Use project's task spawning utilities | +| `uninlined_format_args = "warn"` | `format!("{x}")`, not `format!("{}", x)` | +| `absolute_paths = "warn"` (max 3 segments) | `use` the item; no inline `std::collections::HashMap::new()` | +| `large_futures = "warn"` (> 20 KB) | `Box::pin` an oversized future | +| Disallowed: `tokio::task::spawn` / `tokio::spawn` | Only for tasks that may die at shutdown — justify at the call site with `#[expect(clippy::disallowed_methods, reason = "...")]` (see `model_gateway/src/health.rs`); otherwise keep and join the `JoinHandle` | | Disallowed: `uuid::Uuid::new_v4` | Use `now_v7` | | Disallowed: `std::process::exit` | Use normal shutdown logic, not a hard process exit | -Disallowed macros (`clippy.toml`): `unimplemented!` and `todo!` — remove before merging to main. +`clippy.toml` also caps cognitive complexity at 25 and type complexity at 250, and disallows the `unimplemented!` / `todo!` macros outright. Every lint above is fatal under `-D warnings`. ## Import Organization +`rustfmt.toml` sets `group_imports = "StdExternalCrate"` and `imports_granularity = "Crate"`, so step 1 produces this shape for you: + ```rust -// 1. Standard library -use std::collections::HashMap; +use std::collections::HashMap; // 1. Standard library -// 2. External crates (alphabetically) -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Serialize}; // 2. External crates use tokio::sync::Mutex; -// 3. Internal crates -use crate::config::types::RouterConfig; +use crate::config::types::RouterConfig; // 3. This crate ``` ## Error Handling Pattern @@ -99,39 +138,50 @@ let value = parse(input).context("parsing model config")?; |-----------------|-------| | Add CLI flag | `model_gateway/src/main.rs` (CliArgs) | | Change config | `model_gateway/src/config/types.rs` | -| Change worker creation | `model_gateway/src/worker/` (lifecycle steps: `model_gateway/src/workflow/steps/local/`) | +| Change worker creation / registration | `model_gateway/src/worker/` (steps: `model_gateway/src/workflow/steps/{local,shared,external}/`, DAG in `steps/mod.rs`) | | Change service discovery | `model_gateway/src/service_discovery.rs` | | Change API types | `crates/protocols/src/` (careful — shared by all crates) | | Add routing policy | `model_gateway/src/policies/` | | Change request scheduling / admission | `model_gateway/src/middleware/scheduler/` | | Change multi-tenancy | `model_gateway/src/tenant.rs`, `model_gateway/src/middleware/tenant_resolution.rs` | -| Change rate limiting | `model_gateway/src/middleware/token_bucket.rs`, `model_gateway/src/middleware/concurrency.rs` | +| Change per-request concurrency / global token bucket | `model_gateway/src/middleware/concurrency.rs`, `model_gateway/src/middleware/token_bucket.rs` | +| Change per-tenant token/request rate limiting | `model_gateway/src/rate_limit/` (`RateLimitManager`, YAML policy via `--tenant-rate-limit-config`) | | Add/modify a provider API | `model_gateway/src/routers/{openai,anthropic,gemini}/` | +| Change the ZMQ direct-backend path | wire structs `crates/engine_zmq_client/src/protocol/{vllm,tokenspeed}/` (append-only field order); translation `model_gateway/src/routers/grpc/zmq_client.rs` + `zmq_multimodal.rs`; dispatch `model_gateway/src/routers/grpc/backend_client.rs` (`BackendClient`, `ZmqDialect`) | | Add tool parser | `crates/tool_parser/src/parsers/` | | Add reasoning parser | `crates/reasoning_parser/src/parsers/` | -| Update Python bindings | `bindings/python/src/lib.rs` | -| Update Go SDK | `bindings/golang/` | +| Add tokenizer backend | `crates/tokenizer/src/factory.rs` (`TokenizerType`), implementing `crates/tokenizer/src/traits.rs` (`Tokenizer`); registration in `crates/tokenizer/src/registry.rs` | +| Add chat-template renderer | `crates/tokenizer/src/encoders/`, wired by `detect_renderer_from_config` in `crates/tokenizer/src/huggingface.rs` (DeepSeek) or `crates/tokenizer/src/tiktoken.rs` (Kimi); gateway autoload in `model_gateway/src/workflow/tokenizer_registration.rs` | +| Update Python bindings | `bindings/python/src/lib.rs` (`Router`) + `bindings/python/src/smg/router_args.py` (`RouterArgs`) | +| Update Go SDK | `bindings/golang/` (`ClientConfig` in `client.go`) | | Regenerate client SDKs (after protocol changes) | `make generate-clients` (`clients/openapi-gen/`) — regenerates Python + Java types only; the Go SDK is hand-maintained | -| Add storage backend | `crates/data_connector/src/` | -| Add E2E tests | `e2e_test/` | -| Add WASM middleware | `examples/wasm/` (guests) + `crates/wasm/src/` (host) | +| Add storage backend | `crates/data_connector/src/` (CI also runs `cargo test -p data-connector --test postgres_integration -- --ignored` against a live Postgres) | +| Launch engines / pick a connection mode | `bindings/python/src/smg/serve.py` (`--connection-mode grpc/http/zmq`; `_zmq_handshake_port` mirrors `zmq_client.rs::derive_handshake_port`) | +| Test routing without GPUs | `cargo run -p mock-worker -- --http-count N` / `--grpc-count N` / `--zmq-handshake --zmq-count N` (`scripts/scale_test.sh`, `scripts/sim_ab.sh`) | +| Add E2E tests | `e2e_test/` (markers in `e2e_test/fixtures/hooks.py`; CPU-only harness tests in `e2e_test/infra/` and `e2e_test/fixtures/`) | +| Add WASM middleware | `examples/wasm/` (guests) + `crates/wasm/src/` (host) — build guests with `crates/wasm/tests/fixtures/build_fixtures.sh` (needs `rustup target add wasm32-wasip2`) | | Add MCP tool support | `crates/mcp/src/` | +| Update user/contributor docs | the separate `smg-project/smg-docs` repo (`src/lib/content/`) — `docs/` and mkdocs were removed from smg; only the root Markdown (README, CONTRIBUTING, REVIEW, GOVERNANCE, CODE_OF_CONDUCT) lives here | ## Rationalization Prevention | Excuse | Reality | |--------|---------| | "Clippy is clean enough with a few warnings" | `-D warnings` means zero. One warning = not clean. | -| "I didn't change bindings, skip step 4" | If you touched `config/types.rs` or `crates/protocols/`, the struct literal in `bindings/python/src/lib.rs` may need a default. Check. | -| "Only touched one file, don't need full gate" | The two-path config rule means a one-file change can silently break propagation. Run all five. | +| "I didn't change bindings, skip step 5" | A new config field compiles fine but silently never reaches Python/Go. Wire it through `Router`/`to_router_config` in `bindings/python/src/lib.rs`, `router_args.py`, and Go `ClientConfig`. Check. | +| "Cargo.lock is generated, don't commit it" | It is tracked, and CI runs `cargo metadata --locked` before every Rust job. Commit the diff. | +| "Only touched one file, don't need full gate" | The two-path config rule means a one-file change can silently break propagation. Run all six. | | "Tests are slow, I'll run them later" | "Later" means shipping untested code. Run them now. | -| "It's just a docs change" | Even docs PRs need clean formatting and conventional commits. Steps 1 and 5 still apply. | +| "It's just a docs change" | Step 4 still runs codespell and whitespace hooks over every file, and step 6 has no exemption for the PR title, branch name, or sign-off. | +| "I'll rename the branch after review starts" | CI comments and runs `gh pr close` on a badly named branch. Rename before you push. | ## Red Flags — STOP - About to open a PR without running the gate function in this session - Using "should pass" or "looks clean" without command output -- Skipping the bindings check after config changes +- Skipping the bindings/Python check after config or protocol changes +- Changing a `Cargo.toml` without staging the updated `Cargo.lock` +- Pushing a branch that isn't `/` or `/` - Committing without `-s` (DCO sign-off) - Thinking "just this once" about any gate step - Trusting a subagent's claim without verifying independently diff --git a/skills/implement/SKILL.md b/skills/implement/SKILL.md index 1765177..9ade535 100644 --- a/skills/implement/SKILL.md +++ b/skills/implement/SKILL.md @@ -22,7 +22,7 @@ Do NOT write implementation code until you have: 3. Created a task for each step in the recipe -**Escape hatch:** Single-file changes under 20 lines that don't touch `config/types.rs`, `crates/protocols/`, `main.rs` (CliArgs or conversion functions), or `bindings/` may skip the full recipe. You MUST still chain to `smg:contribute` before PR. +**Escape hatch:** Single-file changes under 20 lines that don't touch `config/types.rs`, `crates/protocols/`, `main.rs` (CliArgs or conversion functions), `bindings/`, or `crates/engine_zmq_client/src/protocol/` (positional msgpack — field order is the wire contract, append only) may skip the full recipe. You MUST still chain to `smg:contribute` before PR. ## Detection Table @@ -34,7 +34,8 @@ Do NOT write implementation code until you have: | Reasoning, thinking, chain-of-thought parser | @reasoning-parser.md | | Python binding, Go SDK, FFI, PyO3, maturin | @bindings-update.md | | K8s, service discovery, pod, worker lifecycle, label | @discovery-feature.md | -| gRPC, backend client, tonic, streaming | @grpc-backend.md | +| gRPC, backend client, tonic, streaming (an `ipc://` worker is ZMQ, next row) | @grpc-backend.md | +| ZMQ, `ipc://` worker, direct backend, EngineCore, headless vLLM/TokenSpeed, handshake, msgpack wire, `zmq_engine_count`, grouped engine worker | @zmq-backend.md | | Storage, database, PostgreSQL, Oracle, Redis, data connector | @storage-backend.md | | WASM, WebAssembly, plugin, middleware hook | @wasm-plugin.md | | Auth, API key, JWT, OIDC, role, permission | @auth-feature.md | @@ -46,7 +47,7 @@ Do NOT write implementation code until you have: | Provider-compatible API router (Anthropic /v1/messages, Gemini /v1/interactions), RoutingMode variant, new API surface | @provider-api.md | | Priority scheduler, admission, preemption, queue, reservation, autoscaling | @scheduler-feature.md | | Multi-tenancy, tenant, tenant policy, tenant resolution | @tenancy-feature.md | -| Rate limit, token bucket, concurrency cap | @rate-limit-feature.md | +| Rate limit, token bucket, concurrency cap, per-tenant token/request quota, reserve/settle | @rate-limit-feature.md | ### Subsystems without a dedicated recipe yet @@ -55,13 +56,16 @@ These are real, actively-developed subsystems that don't yet have a step-by-step | Signal in User Request | Where it lives | |------------------------|----------------| | Provider API: Responses, Conversations, Realtime (for Anthropic/Gemini-style API routers see @provider-api.md) | `model_gateway/src/routers/{responses,conversations,common/realtime}/` | +| Tokenizer, chat template, prompt renderer/encoder, stop sequence, EOS | `crates/tokenizer/src/` (`encoders/`, `huggingface.rs::detect_renderer_from_config`, `stop.rs`, `eos.rs`, `registry.rs`); gateway side `model_gateway/src/workflow/tokenizer_registration.rs`, `routers/tokenize/`; `cargo test -p llm-tokenizer` | +| Workflow engine, registration step, DAG, job queue | `crates/workflow/` (the `wfaas` crate: `StepExecutor`, `StepDefinition`, `StepResult`) plus `model_gateway/src/workflow/steps/{local,shared,external}/` wired in `steps/mod.rs`, and `workflow/job_queue.rs` | +| E2E test, pytest harness, CI lane, kind discovery | `e2e_test/` — markers and filtering in `fixtures/hooks.py`, backends in `fixtures/setup_backend.py`, launchers in `infra/worker.py`; harness units `PYTHONPATH=e2e_test pytest -q --noconftest e2e_test/infra e2e_test/fixtures`; `SMG_KIND_E2E=1 pytest e2e_test/kind_discovery -m kind` | | Client SDK generation | `clients/openapi-gen/` (`make generate-clients`) | **If multiple match:** Load all matching recipes. `config-plumbing.md` almost always co-triggers with other recipes (most features need a config field). **If none match:** This is a novel change. Read the codebase architecture with `smg:map` first, then follow the general pattern: implement → test → verify → chain to `smg:contribute`. -All recipes were rewritten and verified against the codebase (2026-06; re-verified against HEAD 2026-07-15). They reflect current trait/type names and paths; still run each recipe's `cargo check`/`cargo test` step as you go. +All recipes were rewritten and verified against the codebase (2026-06; re-verified against HEAD 732a42bf, post-v1.9.0, 2026-08-22). They reflect current trait/type names and paths; still run each recipe's `cargo check`/`cargo test` step as you go. ## Process @@ -95,6 +99,6 @@ All recipes were rewritten and verified against the codebase (2026-06; re-verifi ## Skill Chaining -When implementation is complete: **invoke `smg:contribute`** to run the 5-step quality gate before PR. +When implementation is complete: **invoke `smg:contribute`** to run the 6-step quality gate before PR. Before submitting: self-review with `smg:review-pr`. diff --git a/skills/implement/auth-feature.md b/skills/implement/auth-feature.md index f649b90..e2910db 100644 --- a/skills/implement/auth-feature.md +++ b/skills/implement/auth-feature.md @@ -1,35 +1,40 @@ # Adding Auth Features to SMG -Two-factor auth: API key (SHA-256) + JWT/OIDC. Roles: Admin (control plane) and User (data plane). +Two independent auth layers — know which one you are touching. + +**Serving path** (`/v1/*`) — `model_gateway/src/middleware/auth.rs`. `AuthConfig::with_tenant_keys(api_key, &tenant_api_keys)` SHA-256-hashes every key into a map; `auth_middleware` resolves the Bearer token and inserts a `DataPlaneCaller` extension carrying a `TenantKey` (`auth:` for `--tenant-api-key tenant_id:key`, hash-derived for the shared `--api-key`). Config: `RouterConfig.tenant_api_keys: Vec`, checked by `ConfigValidator::validate_tenant_api_keys` (non-empty tenant_id/key, no surrounding whitespace, no duplicate credential incl. collision with `api_key`). Tenant keys never reach the control plane (`admin_auth_config = AuthConfig::new(api_key)`; tenant-keys-only → `deny_all_middleware`), and `/v1/models` short-circuits on `AuthConfig::contains_token` so a tenant key is never forwarded as a BYOK credential. + +**Control plane** (`admin_routes` + `worker_routes`) — crate `smg-auth`. Two methods, either one suffices: JWT/OIDC is tried first, then API key (`--control-plane-api-keys id:name:role:key`, SHA-256). A 3-segment token that fails JWT validation is rejected outright, not retried as an API key. Roles: `Admin` (control plane) and `User` (data plane). **The rest of this skill covers this layer.** ## Steps ### Adding a New Role -1. Extend `Role` enum in `crates/auth/src/config.rs` -2. Update permission checks in middleware -3. Update role mapping for JWT claims -4. Add audit logging for new role actions +1. Extend `Role` in `crates/auth/src/config.rs` — including its `Display`, `FromStr` and `is_admin()` impls +2. Update the hard-coded `"admin" | "user"` matches that bypass `FromStr`: `parse_control_plane_api_key` and `parse_role_mapping` in `model_gateway/src/main.rs`, and `PyRole` + the `role_mapping` match in `bindings/python/src/lib.rs`. A variant added only to the enum is silently dropped there (warn + skip, or default to `User`) +3. Update `check_admin_role` in `crates/auth/src/middleware.rs` +4. Update `JwtValidator::extract_role` in `crates/auth/src/jwt.rs` ### Adding a New OIDC Provider -1. Configure `JwtConfig` with provider's JWKS URI -2. Set `role_claim` to match provider's claim name (default: `"roles"`) -3. Map IDP roles to gateway roles via `role_mapping`: +1. Configure `JwtConfig`: `--jwt-issuer` and `--jwt-audience` (env `JWT_ISSUER` / `JWT_AUDIENCE`) are BOTH required — one without the other prints a warning and disables JWT entirely. `--jwt-jwks-uri` is optional; without it `JwksProvider::from_issuer` discovers via `/.well-known/openid-configuration`. +2. Every JWKS/discovery URL passes `validate_url` (`crates/auth/src/jwks.rs`): HTTPS only (plain http for `localhost` / `127.0.0.1` / `::1` only), literal private/link-local/CGNAT/metadata IPs and `metadata` / `*.internal` / `*.local` hostnames rejected (no DNS resolution — the check is on the literal host), response capped at 1 MiB. A JWKS on an internal network fails at startup. +3. Set `--jwt-role-claim` to the provider's claim name (default `roles`; if absent, `extract_role` falls back to `role` / `roles` / `groups` / `group`) +4. Map IDP roles to gateway roles via repeatable `--jwt-role-mapping idp_role=gateway_role`, or `JwtConfig::with_role_mapping`: ```rust role_mapping: HashMap::from([ ("admin".to_string(), Role::Admin), ("reader".to_string(), Role::User), ]) ``` -4. Test: Token validation, claim extraction, role mapping +5. Test: Token validation, claim extraction, role mapping ### Adding a Custom Auth Method 1. Implement validation logic in `crates/auth/src/` 2. Extract `Principal` from request (via `PrincipalExt`, `crates/auth/src/middleware.rs`) -3. Integrate via `control_plane_auth_middleware` on the merged `admin_routes` router (control-plane endpoints like `/wasm`, `/parse/*` — there is no literal `/admin` URL prefix) -4. Add audit event for the new method +3. Integrate via the `apply_control_plane_auth` closure in `model_gateway/src/server.rs:build_app`, which layers `control_plane_auth_middleware` on BOTH `admin_routes` (`/flush_cache`, `/get_loads`, `/parse/*`, `/wasm*`, `/v1/tokenizers*` — no literal `/admin` prefix) and `worker_routes` (`/workers*`). With no control-plane auth configured the closure falls back to `auth_middleware` with the shared `--api-key`, or to `deny_all_middleware` when only tenant keys exist. +4. Emit the matching `AuditLogger` call inside `crates/auth` (see below) ### JWT Flow (7 Steps) @@ -37,21 +42,21 @@ Two-factor auth: API key (SHA-256) + JWT/OIDC. Roles: Admin (control plane) and 2. Fetch signing key from JWKS (cached, TTL 1hr) 3. Verify algorithm matches key 4. Validate: expiry, issuer, audience -5. Extract role from `role_claim` -6. Map via `role_mapping` -7. *(Opt-in, off by default)* Check JTI cache for replay (10k token cache, `DEFAULT_JTI_CACHE_SIZE`). `JwtValidator::from_config` sets `enable_jti_check = false`, so the default control-plane path skips this — it runs only when JTI replay protection is explicitly enabled (`from_config_with_options`). +5. *(Opt-in, off by default)* Check JTI cache for replay (10k-entry LRU, `DEFAULT_JTI_CACHE_SIZE`). `JwtValidator::from_config` passes `enable_jti_check = false`; only `from_config_with_options(config, true)` enables it. Runs BEFORE subject/role extraction, not last. +6. Extract role from `role_claim` (falls back to `role` / `roles` / `groups` / `group`) +7. Map via `role_mapping`; with an empty mapping the value is parsed directly as `admin` / `user`. Anything unmatched → `Role::User` with a warn — never an error. ## Audit Events -All control plane mutations must be logged. Emit via `AuditLogger` (`log_success` / `log_denied` / `log_auth_failure`) — don't construct `AuditEvent` by hand: +Audit logging is per-request and lives in `control_plane_auth_middleware` (`crates/auth/src/middleware.rs`): `log_success` on authenticated admin access, `log_denied` for a non-admin principal, `log_auth_failure` for a missing/invalid token. Handlers in `model_gateway` emit nothing. Add an auth path inside `crates/auth` → build an `AuditContext` and call the matching `AuditLogger` method. `AuditContext` is NOT re-exported from `lib.rs`, so from outside the crate only `AuditLogger::log(&AuditEvent)` and `log_auth_failure` are callable: ```rust -// AuditEvent fields (audit.rs): timestamp, principal, auth_method, role, +// AuditEvent fields (audit.rs), all pub: timestamp, principal, auth_method, role, // method, path, resource, outcome (Success | Denied), request_id, details ``` ## Key Rules -- API keys: SHA-256 hash at load time, constant-time comparison +- API keys: SHA-256 hash at load time, constant-time comparison; `find_api_key` scans every entry with no early return - JWT: Always validate expiry, issuer, audience -- Audit: Log all admin operations (add/remove modules, config changes) -- Never store plaintext keys +- Audit: emitted once per request by `control_plane_auth_middleware`, never by a handler +- Never store plaintext keys — `TenantApiKeyEntry`'s `Debug` impl redacts `key` diff --git a/skills/implement/bindings-update.md b/skills/implement/bindings-update.md index 704e751..c4fbad2 100644 --- a/skills/implement/bindings-update.md +++ b/skills/implement/bindings-update.md @@ -4,15 +4,15 @@ Required whenever config types or public APIs change. The #2 contributor mistake ## Python Bindings (PyO3 + maturin) -`bindings/python/src/lib.rs` exposes ONE pyclass `Router`. There is no `__init__`: the constructor is a `#[new] fn new(...)` with a flat `#[pyo3(signature = (...))]` of 110+ params (the `fn new` inside `impl Router`'s `#[pymethods]` block). `RouterConfig` is NOT a struct literal — it is assembled in `to_router_config()` via `RouterConfig::builder()....build()`. A new field is threaded through four spots. +`bindings/python/src/lib.rs` exposes the router as a single pyclass `Router`, alongside enum/config helper pyclasses (`PolicyType`, `BackendType`, `HistoryBackendType`, `PyRole`, `Py*Config`, `PyApiKeyEntry`) — all registered in `#[pymodule] fn smg_rs`. There is no `__init__`: the constructor is a `#[new] fn new(...)` with a flat `#[pyo3(signature = (...))]` of 140+ params (the `fn new` inside `impl Router`'s `#[pymethods]` block). `RouterConfig` is NOT a struct literal — it is assembled in `to_router_config()` via `RouterConfig::builder()....build()`. A new field is threaded through four spots. -> **Keep `PolicyType` in sync with `PolicyConfig`.** The Python `enum PolicyType` (`bindings/python/src/lib.rs`) must mirror `PolicyConfig` (`model_gateway/src/config/types.rs`) — both currently list the same 10 policies (`Random`, `RoundRobin`, `Passthrough`, `CacheAware`, `PowerOfTwo`, `LeastLoad`, `Bucket`, `Manual`, `ConsistentHashing`, `PrefixHash`). Adding a policy means updating both plus the `convert_policy` mapping below. +> **Keep `PolicyType` in sync with `PolicyConfig`.** The Python `enum PolicyType` (`bindings/python/src/lib.rs`) must mirror `PolicyConfig` (`model_gateway/src/config/types.rs`) — both currently list the same 10 policies (`Random`, `RoundRobin`, `Passthrough`, `CacheAware`, `PowerOfTwo`, `LeastLoad`, `Bucket`, `Manual`, `ConsistentHashing`, `PrefixHash`). Adding a policy means updating both, the `convert_policy` mapping below, the `policy_map` dict in `router.py`'s `policy_from_str` (an unknown name raises `KeyError`), and the `COMMON_POLICY_CHOICES` / `PREFILL_POLICY_CHOICES` / `ENCODE_POLICY_CHOICES` lists at the top of `router_args.py` (argparse `choices=`). ### Step 1: Add to the `#[new]` signature + store on the pyclass **File:** `bindings/python/src/lib.rs` -1. Add a default to the `#[pyo3(signature = (...))]` list (the `#[new]` block in `impl Router`). APPEND at the end — new params MUST go last so positional `_Router(...)` callers don't break (see the "Appended last to match the `#[pyo3(signature)]` order" comment in `fn new`; the current tail of the list is `multimodal_shm_min_bytes`). +1. Add a default to the `#[pyo3(signature = (...))]` list (the `#[new]` block in `impl Router`). APPEND after the last entry before the closing `))]` — new params MUST go last so positional `_Router(...)` callers don't break (see the "Appended last (not inserted mid-list)" comment there). Read the actual tail before editing; the list grows every release. 2. Add the typed arg to `fn new(...)` (the `Router` constructor) and set it in the returned `Router { ... }`. 3. Add the field to the `struct Router` pyclass. @@ -27,13 +27,14 @@ Required whenever config types or public APIs change. The #2 contributor mistake **File:** `bindings/python/src/lib.rs` (the `to_router_config` method) -Add a builder call in the `RouterConfig::builder()...` chain (the `config::RouterConfig::builder()` call near the end of `to_router_config`). Use `maybe_*` for `Option` fields: +Add a builder call in the `RouterConfig::builder()...` chain (the `config::RouterConfig::builder()` call near the end of `to_router_config`). `RouterConfigBuilder` (`model_gateway/src/config/builder.rs`) is hand-written — use whichever setter it actually exposes, and add one there first if the field has none: ```rust -.maybe_my_field(self.my_field.as_ref()) +.maybe_my_field(self.my_field.as_ref()) // maybe_* setter (Option / sub-config fields) +.zmq_engine_count(self.zmq_engine_count) // plain setter that takes the Option directly ``` -If it belongs to discovery, add it to the `DiscoveryConfig { .. }` literal instead (the `Some(DiscoveryConfig { .. })` block in `to_router_config`). For a new enum, mirror the `convert_policy` closure (the `let convert_policy = |policy: &PolicyType| -> ...` closure in `to_router_config`): `match` on the variant, return `config::ConfigError::InvalidValue` on a bad string. +If it belongs to discovery, add it to BOTH the `Some(DiscoveryConfig { .. })` literal in `to_router_config` AND the `Some(service_discovery::ServiceDiscoveryConfig { .. })` literal in `Router::start` (passed as `server::ServerConfig.service_discovery_config`) — different types, and a field added to only the first never reaches the k8s watcher. For a new enum, mirror the `convert_policy` closure (the `let convert_policy = |policy: &PolicyType| -> ...` closure in `to_router_config`): `match` on the variant, return `config::ConfigError::InvalidValue` on a bad string. **Anti-pattern:** grepping for `RouterConfig {` — there is no such literal. The only struct literals here are `DiscoveryConfig`, `MetricsConfig`, `RetryConfig`, etc. @@ -41,26 +42,29 @@ If it belongs to discovery, add it to the `DiscoveryConfig { .. }` literal inste **File:** `bindings/python/src/smg/router_args.py` -`RouterArgs` (the `@dataclasses.dataclass class RouterArgs` in `router_args.py`) is a dataclass; `router.py`'s `Router.from_args` calls `_Router(**args_dict)`. A param missing here is never passed. +`RouterArgs` (the `@dataclasses.dataclass class RouterArgs` in `router_args.py`) is a dataclass; `router.py`'s `Router.from_args` calls `_Router(**args_dict)`. A param missing here is never passed. Conversely, a `RouterArgs` field that is NOT a `#[new]` param must be popped in `from_args`'s `fields_to_remove` list (as the `oracle_*`/`postgres_*`/`redis_*`/`jwt_*` fields are) or `_Router(**args_dict)` raises `TypeError`. -1. Add the dataclass field (with default) to `RouterArgs`. -2. Add the matching `--my-field` in `RouterArgs.add_cli_args(...)` under the right `add_argument_group`. +1. APPEND the dataclass field (with default) at the tail of `RouterArgs` — field order is a public contract (positional `RouterArgs(...)`) frozen by `bindings/python/tests/test_arg_parser.py::TestRouterArgsFieldOrder`; add the name to the end of its `EXPECTED_FIELD_SEQUENCE` too. +2. Add the matching `--my-field` in `RouterArgs.add_cli_args(...)` under the right `add_argument_group`. List-valued flags use `action="extend"` (or `"append"`) so repeated occurrences accumulate like the Rust CLI. -**Verify:** field name matches the `#[new]` param exactly (underscores). +**Verify:** field name matches the `#[new]` param exactly (underscores) — `from_cli_args` copies values by dataclass field name, so the argparse dest must equal the field name. ### Step 4: Build and smoke-test ```bash make python-dev # Makefile -> maturin develop in bindings/python -python -c "from smg import Router; Router(worker_urls=['http://127.0.0.1:8000'], my_field='x')" +python -c "from smg.smg_rs import Router; Router(worker_urls=['http://127.0.0.1:8000'], my_field='x')" +cd bindings/python && pytest -q tests # CI gate (field-order + argparse tests) ``` -**Verify:** build succeeds AND import + construction work (package is `smg`; native ext is `smg.smg_rs`). +**Verify:** build succeeds AND import + construction work. The flat-kwarg class is the native `smg.smg_rs.Router` — `smg/__init__.py` exports only `__version__`, and `smg.router.Router` merely wraps a prebuilt handle (`from_args`). ## Go SDK (cgo FFI + gRPC) Two layers. Rust FFI exports live in `bindings/golang/src/*.rs` (client.rs, policy.rs, tokenizer.rs, preprocessor.rs, postprocessor.rs, stream.rs, grpc_converter.rs, tool_parser.rs, memory.rs) as `#[no_mangle] pub unsafe extern "C" fn sgl_*`. `src/lib.rs` only re-exports + wires modules (its `pub use` / `mod` block). The Go cgo bridge lives in `internal/ffi/*.go` (those `import "C"`); top-level wrappers (`multi_client.go`, `client.go`) do NOT import C — they delegate to `internal/ffi` and `internal/grpc`. Module path: `github.com/lightseek/smg/go-grpc-sdk`. +`policy.rs` also carries its own `impl Worker for GrpcWorker` plus a `WorkerMetadata { .. }` literal: any change to the `Worker` trait or to `WorkerMetadata` in `model_gateway/src/worker/` must be mirrored there. `bindings/golang` is a root workspace member, so `cargo check --manifest-path bindings/golang/Cargo.toml` (a CI gate) fails until it is. + ### Step 1: Add the Rust FFI export **File:** the relevant `bindings/golang/src/*.rs` (e.g. `client.rs` for client ops) @@ -78,8 +82,6 @@ pub unsafe extern "C" fn sgl_my_function( Re-export it from `src/lib.rs` (in the `pub use` block): `pub use client::sgl_my_function;`. Follow `sgl_client_create` (in `src/client.rs`) for the null-check + error-out contract. -**Anti-pattern:** adding the `extern "C" fn` to `src/lib.rs` — it only does `pub use` + `mod`. - ### Step 2: Add the cgo bridge + high-level wrapper **File:** `bindings/golang/internal/ffi/.go` (e.g. `client.go`) @@ -97,21 +99,26 @@ if ErrorCode(result) != ErrorSuccess { /* GoString(errorPtr); C.sgl_free_string( func (c *MultiClient) MyFunction(arg string) error { return c.ffiClient.MyFunction(arg) } ``` -**Anti-pattern:** calling `C.sgl_*` from top-level `client.go`/`multi_client.go` — they have no `import "C"`. - ### Step 3: Test ```bash -cd bindings/golang && go test ./... +cargo check --manifest-path bindings/golang/Cargo.toml # CI gate for the Rust side (from repo root) +cd bindings/golang && make test # cargo build --release, then go test ./... ``` +Bare `go test ./...` cannot link: `internal/ffi` declares `#cgo LDFLAGS: -lsmg_go -ldl` with no `-L`, and the `make test` target builds the `libsmg_go` cdylib and exports `CGO_LDFLAGS` + `LD_LIBRARY_PATH`/`DYLD_LIBRARY_PATH` first. + ## Common Mistakes | Mistake | Consequence | |---------|-------------| | Looking for a `RouterConfig {` literal | None exists — it's `builder()...build()` in `to_router_config()` | | Inserting a `#[new]` param mid-list | Breaks positional `_Router(...)` callers | +| Inserting a `RouterArgs` field mid-list | `TestRouterArgsFieldOrder` fails in CI | | Skipping `router_args.py` | CLI/serve path never passes the new param | +| `from smg import Router` | `smg/__init__.py` exports only `__version__` — the flat-kwarg class is `smg.smg_rs.Router` | +| Discovery field added to only one literal | k8s watcher never sees it — `Router::start` builds a second config | | Putting `extern "C" fn` in golang `src/lib.rs` | Wrong file — exports live in `src/*.rs`, lib.rs only re-exports | | `C.sgl_*` in top-level Go wrapper | No `import "C"` there — bridge belongs in `internal/ffi` | +| Bare `go test ./...` | Link failure — use `make test`, which builds the cdylib and sets the paths | | Go/Rust FFI type mismatch | Segfault or data corruption | diff --git a/skills/implement/config-plumbing.md b/skills/implement/config-plumbing.md index 44fb4fd..6e442b0 100644 --- a/skills/implement/config-plumbing.md +++ b/skills/implement/config-plumbing.md @@ -14,26 +14,45 @@ pub my_field: Option, ``` **Verify:** `cargo build` -**Anti-pattern:** Missing `#[serde(default)]` — existing YAML configs will fail to deserialize. +**Anti-pattern:** Missing `#[serde(default)]` — already-serialized configs that lack the field stop deserializing (the `serde_json` backward-compat tests in `types.rs` `mod tests` catch it). A non-`Option` field with a non-zero default takes `#[serde(default = "default_my_field")]` plus a `fn default_my_field()` that the `Default` impl reuses (cf. `default_job_queue_capacity`, `default_stream_body_stall_timeout_secs`). ### Step 2: Update Default impl **File:** `model_gateway/src/config/types.rs` -Add `my_field: None` (or appropriate default) to the `Default` impl for the struct. +Add `my_field: None` (or `my_field: default_my_field()`) to the `Default` impl for the struct. + +**Verify:** `cargo build` + +### Step 2b: Add a builder setter + +**File:** `model_gateway/src/config/builder.rs` + +`RouterConfigBuilder` is hand-written, not derived — `.my_field(..)` does not exist until you write it. Add a setter beside the existing ones: + +```rust +pub fn my_field(mut self, my_field: Option) -> Self { + self.config.my_field = my_field; + self +} +``` + +cf. `RouterConfigBuilder::health_check_port`; for an `Option`-of-string use the `maybe_api_key(Option>)` shape. `RouterConfigBuilder::new()` starts from `RouterConfig::default()`, so Step 2's default is the builder's base. **Verify:** `cargo build` ### Step 3: Add CLI flag -**File:** `model_gateway/src/main.rs` (CliArgs struct) +**File:** `model_gateway/src/main.rs` (`CliArgs` struct) ```rust -#[arg(long, value_parser = parse_my_type)] -pub my_field: Option, +#[arg(long, value_parser = parse_my_type, help_heading = "Worker Configuration")] +my_field: Option, ``` -Add validation function: +`CliArgs` fields are private (no `pub`), and nearly every flag carries `help_heading = "..."` to group `--help` output (cf. `health_check_port`). Renaming an existing flag keeps the old spelling alive via `visible_alias = "old-name"` (cf. `--cache-match-threshold`, test `alias_flags_parse_identically_to_canonical`). + +Add a validation function — return the typed value when the field is typed (cf. `parse_transport_mode -> Result`): ```rust fn parse_my_type(s: &str) -> Result { MyType::parse(s).map_err(|e| e.to_string())?; @@ -48,24 +67,40 @@ fn parse_my_type(s: &str) -> Result { **File:** `model_gateway/src/main.rs` -Both are `CliArgs` methods. `to_router_config(&self, ...)` sources fields from `self` (the parsed CLI args). `to_server_config(&self, router_config: RouterConfig)` *also* takes a `RouterConfig` and, for a few discovery-derived fields (e.g. `model_id_source`, `router_selector`), merges CLI-over-config with `.or_else(...)` — but for a plain new scalar flag you still just set it from `self` in both. Find BOTH and add the field, but note they build differently: +Both are `CliArgs` methods sourcing from `self` — there is no config-file loader for `RouterConfig`. + +- `to_router_config(&self, prefill_urls, encode_urls) -> ConfigResult` assembles via `RouterConfig::builder()....build()`. Add the Step 2b setter to the chain. +- `to_server_config(&self, router_config: RouterConfig) -> ConfigResult` returns a `ServerConfig { .. }` literal whose `router_config` field carries the whole `RouterConfig`. A RouterConfig-only field needs **nothing** here — it rides along nested. ```rust -// to_router_config: assembles via RouterConfig::builder()....build() +// to_router_config let builder = RouterConfig::builder() // ... - .my_field(self.my_field); // or .maybe_my_field(self.my_field.as_ref()) for an Option -// cf. .health_check_port(self.health_check_port), .enable_wasm(self.enable_wasm) + .my_field(self.my_field); // cf. .health_check_port(self.health_check_port) +``` + +Only a listener/runtime-level setting that `server::startup` (or `main` around it) reads directly off `ServerConfig` (e.g. `config.host`, `config.port`, `config.health_check_port`, `config.runtime_worker_threads`, `config.max_payload_size`, `config.request_timeout_secs`, `config.shutdown_grace_period_secs`, `config.webrtc_bind_addr`, `config.control_plane_auth`) also needs a top-level copy. For those, add the field to `pub struct ServerConfig` in `model_gateway/src/server.rs` and set it in ALL THREE literals — `main.rs:to_server_config`, `bindings/python/src/lib.rs` (`Router::start`), and the `server.rs` test helper `minimal_server_config`: -// to_server_config: returns a ServerConfig { ... } struct literal -Ok(ServerConfig { +```rust +Ok(ServerConfig { // main.rs:to_server_config // ... my_field: self.my_field, // cf. health_check_port: self.health_check_port, }) ``` -**Verify:** `grep -n "my_field" model_gateway/src/main.rs` — must show assignments in BOTH functions. -**Anti-pattern:** Only wiring one path. CLI flag works but config file doesn't (or vice versa). This is silent — no error, just ignored. +`to_server_config` also derives `ServiceDiscoveryConfig` from `router_config.discovery`: only `model_id_source` merges CLI (`--model-id-from`) over config with `.or_else(...)`; `router_selector` and `router_mesh_port_annotation` are copied straight out with no CLI override. + +**Verify:** add `my_field_flows_into_both_configs` to `main.rs` `mod tests` — `cli_args_from(&["--my-field", ..])`, then `to_router_config(vec![], vec![])`, then `to_server_config(router_config)`. Assert `server_config.my_field` for a top-level field (cf. `health_check_port_flows_into_both_configs`) or `server_config.router_config.my_field` for a nested one (cf. `engine_metrics_flows_into_both_configs`). +**Anti-pattern:** Wiring only the builder for a top-level `ServerConfig` field. `server::startup` reads the `ServerConfig` copy and silently sees the default — no error, just ignored. + +### Step 4b: Add config validation (if the field has invariants) + +**File:** `model_gateway/src/config/validation.rs` + +Add the check to the matching `ConfigValidator::validate_*` (`validate_server_settings`, `validate_policy`, `validate_discovery`, `validate_compatibility`, ...), returning `ConfigError::InvalidValue { field, value, reason }`. This gate runs from `RouterConfigBuilder::build()` for BOTH the CLI and the Python bindings; the Step 3 `value_parser` only guards the CLI. + +**Verify:** `cargo test` +**Anti-pattern:** Range or cross-field invariants enforced only by `value_parser` — Python callers reach the builder directly and bypass clap entirely. ### Step 5: Update Python bindings @@ -73,20 +108,18 @@ Ok(ServerConfig { The `Router` constructor is `#[new]` with a flat `#[pyo3(signature = (...))]`, and `RouterConfig` is assembled via `RouterConfig::builder()` in `to_router_config()` — there is NO `RouterConfig` struct literal. Append `my_field` to the signature, store it on the `Router` pyclass, thread it through the builder chain, and mirror it in `src/smg/router_args.py`. Full procedure: @bindings-update.md. -**Verify:** `make python-dev` -**Anti-pattern:** Grepping for a `RouterConfig {` literal — none exists; add the field to the `builder()` chain instead. - -### Step 6: Update Go SDK (if exposed) +**Verify:** `make python-dev`, then `cd bindings/python && pytest -q tests` (what CI runs; `make python-test` runs `e2e_test/` instead). Append the new field to the END of the `RouterArgs` dataclass — after the `# Append new fields here` marker — and to the end of `EXPECTED_FIELD_SEQUENCE` in `bindings/python/tests/test_arg_parser.py`. Positional field order is a public contract; the snapshot test fails on any mid-list insertion. +**Anti-pattern:** Grepping for a `RouterConfig {` literal — none exists; add the field to the `builder()` chain instead. But `ServerConfig {` and `DiscoveryConfig {` literals DO exist in `lib.rs` (`Router::start`, `to_router_config`) as well as in `main.rs`; a field on either of those structs must be added to every literal or the bindings stop compiling. -**Files:** `bindings/golang/` — add to Go struct + FFI bridge if the field is user-facing. +### Step 6: Go SDK — no change -**Verify:** Go tests pass. +`bindings/golang/` is a gRPC client SDK plus a tokenizer/parser cdylib; it carries no `RouterConfig`/`ServerConfig` surface. Skip. ### Step 7: Add tests -- Config deserialization test: YAML with and without the new field -- Validation test: invalid value rejected at parse time -- Update existing test struct literals that construct the modified config type +- Serde backward-compat test in `config/types.rs` `mod tests`: `serde_json` round-trip with and without the field (cf. `test_health_check_port_serde_roundtrip_and_backward_compat`) — not YAML +- `main.rs` `mod tests`, the established trio per flag: `my_field_flows_into_both_configs`, `my_field_defaults_to_none_in_both_configs`, `my_field_..._rejected_at_parse_time`, all built on `cli_args_from(&[..])` +- Update full struct literals of the modified type. `RouterConfig` test literals are all `..Default::default()` spreads (nothing to do); `DiscoveryConfig` has a full literal in `types.rs` tests plus production literals in `main.rs:to_router_config` and `lib.rs:to_router_config`; `ServerConfig` has the three from Step 4 **Verify:** `cargo test` @@ -94,7 +127,10 @@ The `Router` constructor is `#[new]` with a flat `#[pyo3(signature = (...))]`, a | Mistake | Consequence | Prevention | |---------|-------------|------------| -| Only wiring `to_router_config()` | Config file value silently ignored | Always grep for BOTH functions | -| Missing `#[serde(default)]` | Existing configs break on deserialize | Every `Option` field needs it | +| Only wiring `to_router_config()` for a `ServerConfig`-level field | `server::startup` reads the default off `ServerConfig` | Add a `*_flows_into_both_configs` test in `main.rs` | +| Adding `my_field: self.my_field` to the `ServerConfig` literal for a RouterConfig-only field | Does not compile — it rides in `ServerConfig.router_config` | Only top-level `ServerConfig` fields get a literal entry | +| Calling `.my_field(..)` without writing the setter | Does not compile — no such method on `RouterConfigBuilder` | It is hand-written, not derived: add the setter in `config/builder.rs` first | +| Missing `#[serde(default)]` | Serialized configs without the field break on deserialize | Every `Option` field needs it | +| Invariants only in clap `value_parser` | Python bindings bypass clap and accept the bad value | Put the check in `config/validation.rs` | +| Inserting a `RouterArgs` field mid-list | `EXPECTED_FIELD_SEQUENCE` snapshot fails; positional callers rebind | Append at the tail and to the snapshot list | | String instead of typed enum at runtime | Re-parsing on every request | Parse at boundary, store typed | -| Looking for a `RouterConfig {` literal | None exists — it's `builder()...build()` | Add to the builder chain + `router_args.py` (see @bindings-update.md) | diff --git a/skills/implement/discovery-feature.md b/skills/implement/discovery-feature.md index 3765bbd..68dd967 100644 --- a/skills/implement/discovery-feature.md +++ b/skills/implement/discovery-feature.md @@ -9,41 +9,61 @@ DiscoverMetadataStep: Backend probe → discovered_labels HashMap ↓ CreateLocalWorkerStep: Merge with config.labels (config wins) ↓ Extract special keys (kv_connector, kv_role, kv_engine_id) - ↓ Resolve model_id (priority chain): + ↓ Resolve model_id (`resolve_model_id()` priority chain): ↓ 1. config.models.primary() ↓ 2. labels["served_model_name"] ↓ 3. labels["model_id"] ↓ 4. labels["model_path"] ↓ 5. UNKNOWN_MODEL_ID ↓ - build_model_card() → ModelCard with full metadata + build_model_card(model_id, config, &labels, &router_config.model_aliases) + → ModelCard with full metadata ``` **To inject new metadata:** Add it as a label. The pipeline handles the rest. -Note: gRPC backend detection (sglang/vllm/trtllm/tokenspeed/mlx — `tokenspeed` is new) lives in `workflow/steps/local/detect_backend.rs`. +Note: `--model-alias alias=canonical` attaches client-facing aliases to discovered cards. Aliases never change the canonical id — `resolve_model_id()` deliberately ignores them. + +Note: gRPC backend detection (sglang/vllm/trtllm/tokenspeed/mlx) lives in `workflow/steps/local/detect_backend.rs`; an OpenAI-compatible HTTP backend matching no fingerprint registers as `generic` rather than being rejected. ## Steps ### Step 1: Add config field -Follow @config-plumbing.md to add the field to `ServiceDiscoveryConfig`. +Follow @config-plumbing.md for `DiscoveryConfig` (`config/types.rs`), then mirror the field on the runtime `ServiceDiscoveryConfig` (`service_discovery.rs`, plus its `Default` impl) — these are two distinct structs, and `main.rs` `to_server_config` builds the second one. `worker_ports_annotation` shows all four sites. Use typed enum, not String: ```rust pub my_source: Option, // NOT Option ``` -### Step 2: Use in pod handling +### Step 2: Extract in discovery and stamp as a label **File:** `model_gateway/src/service_discovery.rs` -Extract from pod metadata in `PodInfo::from_pod()` or `handle_pod_event()`: +Discovery is a level-triggered reconcile, not per-event handling (`handle_pod_event()` no longer exists): + +``` +PodInfo::from_pod(pod, Some(config)) per-pod parse; also reads the + ↓ worker_ports_annotation multi-port list +compute_desired_state() one DesiredWorker per pod data port + ↓ +compute_actions() diff vs k8s_owned_workers() + ↓ +build_worker_spec() → Job::AddWorker { registration_mode: Upsert } / Job::RemoveWorker +``` + +Thread the value through that chain: + +1. `PodInfo::from_pod()` — read the pod label/annotation into a new `PodInfo` field. +2. `compute_desired_state()` — copy it onto `DesiredWorker`. +3. `build_worker_spec()` — inject as a label; the pipeline merges automatically: ```rust -// Inject as label — the pipeline merges automatically -worker_spec.labels.insert("my_key".to_string(), extracted_value); +spec.labels.insert("my_key".to_string(), desired.my_value.clone()); ``` +Use the existing `model_id_override` → `served_model_name` path as the template. `smg.ai/pod-name` / `smg.ai/pod-uid` (`POD_NAME_LABEL` / `POD_UID_LABEL`) are reserved — the reconciler identifies the workers it owns by pod uid. + ### Step 3: Consume in worker creation (if needed) **File:** `model_gateway/src/workflow/steps/local/create_worker.rs` @@ -61,11 +81,12 @@ Follow @bindings-update.md. ### Step 5: Write tests -- Pod with metadata → label injected → model card reflects it +- Unit: `build_worker_spec()` stamps the label from `DesiredWorker` — copy `test_build_worker_spec_stamps_ownership_labels` / `test_compute_desired_state_carries_model_id_override` in `service_discovery.rs` +- Integration: add a case to `model_gateway/tests/k8s_discovery_test.rs` — scripted fake K8s API server driving the real reconciler, JobQueue and registry via `start_service_discovery_with_client` (`test-util` feature) — assert the label reaches the model card - Config label override wins over discovered label - Missing metadata → graceful default -**Verify:** `cargo test` +**Verify:** `cargo test -p smg service_discovery && cargo test -p smg --test k8s_discovery_test`. Real cluster: `SMG_KIND_E2E=1 pytest e2e_test/kind_discovery -m kind` (manual workflow `.github/workflows/e2e-kind-discovery.yml`). ## Pod Types diff --git a/skills/implement/grpc-backend.md b/skills/implement/grpc-backend.md index 397f2ed..c39ab84 100644 --- a/skills/implement/grpc-backend.md +++ b/skills/implement/grpc-backend.md @@ -1,6 +1,8 @@ # Adding a gRPC Backend Client to SMG -gRPC clients live in the `smg-grpc-client` crate (one file per engine) and talk to an inference backend (SGLang, vLLM, TRT-LLM, MLX, TokenSpeed — TokenSpeed is the newest and a good template for a fresh engine). Each engine wraps its generated tonic client and shares connection/health/tokenizer/admin logic via four macros (see Step 1). The router layer (`model_gateway/src/routers/grpc/`) wraps all *generative* engines behind the `GrpcClient` enum and owns request building, streaming, and tool/reasoning parsing — the client crate does NOT parse output. +gRPC clients live in the `smg-grpc-client` crate (one file per engine) and talk to an inference backend (SGLang, vLLM, TRT-LLM, MLX, TokenSpeed — model a fresh engine on `vllm_engine.rs`; TokenSpeed is the odd one out, see the proto-naming note in Step 1). Each engine wraps its generated tonic client and shares connection/health/tokenizer/admin logic via four macros (see Step 1). The router layer (`model_gateway/src/routers/grpc/`) wraps all *generative* engines behind the `GrpcClient` enum and owns request building, streaming, and tool/reasoning parsing — the client crate does NOT parse output. + +**Not this recipe:** a worker URL of `ipc://` is the direct-ZMQ backend (no proto, no probe, no metadata RPC) — see @zmq-backend.md. `BackendClient` (`routers/grpc/backend_client.rs`) is the gRPC-vs-ZMQ seam; never add ZMQ arms to `GrpcClient`. > **Exception — `TokenSpeedEncoderClient` (`tokenspeed_encoder.rs`, proto `tokenspeed.grpc.encoder`).** A special-purpose EPD *encode* client that does **not** fit the engine pattern: it is not a `GrpcClient` enum variant, uses no `impl_engine_client_basics!`, and exposes a bespoke pooled `connect_cached` (round-robin over `ENCODE_CONNS_PER_ENDPOINT` channels) plus a single `encode()` RPC. It is called directly from the EPD encode stage (`routers/grpc/common/stages/encode.rs`), still injects trace context, and connects via `connect_channel`. If you are adding a generative engine, follow the enum pattern below; the encoder is its own thing. @@ -26,6 +28,7 @@ use tracing::warn; use crate::BoxedTraceInjector; +#[expect(clippy::allow_attributes)] pub mod proto { #![allow(clippy::all, clippy::absolute_paths, unused_qualifications)] tonic::include_proto!("myengine.grpc.engine"); @@ -64,9 +67,11 @@ impl MyengineEngineClient { } ``` +The `#[expect(clippy::allow_attributes)]` is load-bearing, not decoration: the workspace sets `allow_attributes = "warn"`, CI runs `cargo clippy --all-targets --all-features -- -D warnings`, and the *service* code tonic generates for your proto (register it in `crates/grpc_client/build.rs`, whose engine pass uses `build_server(true)`/`build_client(true)`) carries outer `#[allow(...)]` attributes that trip the lint. Every engine file carries it; `common_proto` — messages only, no service codegen — does not, and an unfulfilled `#[expect]` is itself a warning. + `impl_engine_client_basics!` requires the macro's `proto::HealthCheckRequest` / `GetModelInfoRequest` / `GetServerInfoRequest` / response types to exist in your `proto` module. `impl_get_tokenizer!`/`impl_subscribe_kv_events!` use `common_proto` types and need the matching RPCs on the generated client. -**Proto module naming — two conventions.** Older engines name the module `proto` and re-export it aliased: `pub mod proto { include_proto!(...) }` then `pub use {engine}_engine::{proto as {engine}_proto, ...}` (see `vllm_engine.rs` + `lib.rs`). The newer TokenSpeed pattern names the module `tokenspeed_proto` directly and re-exports it **without** an alias: `pub mod tokenspeed_proto { include_proto!("tokenspeed.grpc.scheduler") }` then `pub use tokenspeed_scheduler::{tokenspeed_proto, TokenSpeedSchedulerClient};` (see `tokenspeed_scheduler.rs` + `lib.rs`). Either works; pick one and stay consistent. +**Proto module naming — two conventions.** Older engines name the module `proto` and re-export it aliased: `pub mod proto { include_proto!(...) }` then `pub use {engine}_engine::{proto as {engine}_proto, ...}` (see `vllm_engine.rs` + `lib.rs`). The TokenSpeed pattern names the module `tokenspeed_proto` directly and re-exports it **without** an alias (see `tokenspeed_scheduler.rs` + `lib.rs`). Either naming works for the three `common_proto` macros, but `impl_engine_client_basics!` expands those `proto::*` paths *unqualified*, resolved at the call site — name the module `proto` if you want that macro. TokenSpeed did not, which is exactly why it hand-rolls connect/health/info. **Four shared macros — use whichever your engine's proto supports** (all defined in `crates/grpc_client/src/lib.rs`): - `impl_engine_client_basics!($proto_client, $display_name)` — the two `connect` constructors, `with_trace_injector`, and `health_check` / `get_model_info` / `get_server_info`. @@ -74,7 +79,7 @@ impl MyengineEngineClient { - `impl_subscribe_kv_events!()` — `subscribe_kv_events(start_sequence_number)` (uses `common_proto`). - `impl_admin_ops!()` — `flush_cache(timeout_s)` (FlushCache RPC), `start_profile(req)`, `stop_profile()`, with local deadlines from `FLUSH_RPC_DEADLINE_MARGIN` / `PROFILE_RPC_DEADLINE`. -Macro use is per-engine/partial — pick only the ones whose RPCs exist on your generated client: sglang uses all four; vllm and trtllm use basics + `impl_get_tokenizer!` + `impl_subscribe_kv_events!`; mlx uses basics + `impl_get_tokenizer!` only (no kv_events); TokenSpeed uses only `impl_admin_ops!` + `impl_subscribe_kv_events!` and **hand-rolls** connect/health/info (no `impl_engine_client_basics!`) with no tokenizer. +Macro use is per-engine/partial — pick only the ones whose RPCs exist on your generated client: sglang uses all four; vllm and trtllm use basics + `impl_get_tokenizer!` + `impl_subscribe_kv_events!`; mlx uses basics + `impl_get_tokenizer!` only (no kv_events); TokenSpeed uses `impl_get_tokenizer!` + `impl_admin_ops!` + `impl_subscribe_kv_events!` and **hand-rolls** connect/health/info (no `impl_engine_client_basics!` — see the proto-naming note above). If `generate` is streaming, prefer the auto-abort wrapper: return `crate::AbortOnDropStream`, `impl AbortOnDropClient for MyengineEngineClient` (its `abort_for_drop` calls your `abort_request`), and build it with `AbortOnDropStream::new(stream, request_id, self.clone())`. The router calls `mark_completed()` on success. See `abort_on_drop.rs`. @@ -82,29 +87,37 @@ If `generate` is streaming, prefer the auto-abort wrapper: return `crate::AbortO **File:** `crates/grpc_client/src/lib.rs` -Add `pub mod {ENGINE}_engine;` and re-export: `pub use {ENGINE}_engine::{proto as myengine_proto, MyengineEngineClient};` (mirror the existing vLLM/SGLang lines), or, if you named the proto module after the engine, the no-alias TokenSpeed form (e.g. `pub use tokenspeed_scheduler::{tokenspeed_proto, TokenSpeedSchedulerClient};`). The four shared macros (`impl_engine_client_basics!`, `impl_get_tokenizer!`, `impl_subscribe_kv_events!`, `impl_admin_ops!`) and the deadline constants `FLUSH_RPC_DEADLINE_MARGIN` / `PROFILE_RPC_DEADLINE` are defined here. Shared infra also lives here: `channel.rs` (`connect_channel`, `normalize_grpc_endpoint`), `abort_on_drop.rs`, `tokenizer_bundle.rs`, and the `TraceInjector` trait (with `NoopTraceInjector` default + `BoxedTraceInjector` alias). Do not add a free trace-injection function — injection goes through the trait. +Add `pub mod {ENGINE}_engine;` and re-export: `pub use {ENGINE}_engine::{proto as myengine_proto, MyengineEngineClient};` (mirror the existing vLLM/SGLang lines), or, if you named the proto module after the engine, the no-alias TokenSpeed form (e.g. `pub use tokenspeed_scheduler::{tokenspeed_proto, TokenSpeedSchedulerClient};`). The four shared macros (`impl_engine_client_basics!`, `impl_get_tokenizer!`, `impl_subscribe_kv_events!`, `impl_admin_ops!`) and the deadline constants `FLUSH_RPC_DEADLINE_MARGIN` / `PROFILE_RPC_DEADLINE` / `ABORT_RPC_DEADLINE` (10s bound on the detached drop-abort RPC) are defined here. Shared infra also lives here: `channel.rs` (`connect_channel` — applies `DEFAULT_CONNECT_TIMEOUT` = 10s — plus `connect_channel_with_timeout`, `normalize_grpc_endpoint`), `abort_on_drop.rs` (`AbortOnDropStream`, incl. `defer_abort_until_first_item()`), `tokenizer_bundle.rs`, and the `TraceInjector` trait (with `NoopTraceInjector` default + `BoxedTraceInjector` alias). Do not add a free trace-injection function — injection goes through the trait. ### Step 3: Wire into the router's GrpcClient enum **File:** `model_gateway/src/routers/grpc/client.rs` -Add a `Myengine(MyengineEngineClient)` variant, an arm in `connect()` (`"myengine" => Ok(Self::Myengine(MyengineEngineClient::connect(url).await?))`), and arms in `health_check()`, `get_model_info()`, `get_server_info()`, and the `ModelInfo`/`ServerInfo` `to_labels()` matches. Also add arms (or fall through to the existing error default) in: -- `get_loads()` — GetLoads RPC; currently supported for sglang/vllm/tokenspeed (those clients expose `get_loads`), and returns `Status::unimplemented` for the rest. -- `subscribe_kv_events(start_seq)` — supported for sglang/vllm/trtllm/tokenspeed; MLX is explicitly unsupported. The KV-events consumer that drives this is `model_gateway/src/worker/kv_event_monitor.rs`. +Add a `Myengine(MyengineEngineClient)` variant, an arm in `connect()` (`"myengine" => Ok(Self::Myengine(MyengineEngineClient::connect(url).await?))`), and arms in every exhaustive match: `health_check()`, `get_model_info()`, `get_server_info()`, `runtime_type()`, `get_tokenizer()`, `build_{chat,messages,completion,generate}_request()`, and the `ModelInfo`/`ServerInfo` `to_labels()` matches. `generate()`/`embed()` are tuple matches on `(Self, ProtoGenerateRequest)` / `(Self, ProtoEmbedRequest)` ending in `_ => panic!(...)` — a missing arm compiles and panics on the first request, so add them by hand. Opt-in RPCs — add a supported arm, or join the existing unsupported branch: +- `get_loads()` — GetLoads RPC; sglang/vllm/tokenspeed. Falls through a `_ =>` unimplemented default, so it compiles without an arm. +- `flush_cache()` / `start_profile()` / `stop_profile()` — sglang/tokenspeed only; vllm/trtllm/mlx share one explicit `unimplemented` arm (no wildcard). +- `subscribe_kv_events(start_seq)` — sglang/vllm/trtllm/tokenspeed; MLX is an explicit unsupported arm (no wildcard). The consumer is `model_gateway/src/worker/kv_event_monitor.rs`. + +**More files the enum drags in:** +- `crates/protocols/src/worker.rs` — add a `Myengine` variant to `RuntimeType` plus arms in `as_str()` and its `FromStr` chain. Its lowercase string is the runtime key `GrpcClient::connect` and `detect_grpc_backend` pass around; `GrpcClient::runtime_type()`, `common/stages/request_execution.rs`, and `multimodal/capability.rs` match on the enum. +- `routers/grpc/proto_wrapper.rs` — `ProtoGenerateRequest`, `ProtoGenerateResponse`, `ProtoGenerateStreamChunk`, `ProtoGenerateComplete`, `ProtoStream`, and `ProtoEmbedRequest`/`ProtoEmbedComplete` are matched exhaustively; add a variant plus arms to each surface your engine supports (`MultimodalData` only if it is multimodal — MLX is absent there). +- `routers/grpc/harmony/stages/request_building.rs` — exhaustive `(BackendClient::Grpc(GrpcClient::…), HarmonyBody::{Chat,Responses})` match, no engine wildcard; add Chat and Responses arms, or one body-generic arm delegating to a helper (as vllm/tokenspeed do). +- `routers/grpc/regular/stages/embedding/request_building.rs` — exhaustive per-engine match; add a `build_embed_request` arm or an explicit `not_implemented` one. -Request building, streaming, and tool/reasoning parsing live under `routers/grpc/regular/` and `utils/parsers.rs` — extend those, not the client crate. See the router-architecture note below for where these fit. +The `build_*_request` arms delegate to proto-level builders on the engine client itself (`build_generate_request_from_{chat,messages,completion}`, `build_plain_generate_request`), while the harmony and embedding stages call `build_generate_request_from_responses` and `build_embed_request` — all of them belong in your Step 1 file (see `vllm_engine.rs`). Tool/reasoning parsing stays a router concern under `routers/grpc/regular/` + `utils/parsers.rs` — see the architecture note below. ## gRPC router architecture (Mode-parameterized pipeline) -Regular, PD (prefill/decode), and EPD (encode/prefill/decode) serving were **unified into one `Mode`-parameterized router** — there is no longer a separate `pd_router.rs`. Understand this shape before touching request handling: +Regular, PD (prefill/decode), and EPD (encode/prefill/decode) serving were **unified into one `Mode`-parameterized router** — there is no longer a separate gRPC `pd_router.rs` (the HTTP PD router, `routers/http/pd_router.rs:PDRouter`, `router_type() == "pd"`, still exists and is unrelated to this recipe). Understand this shape before touching request handling: -- **`GrpcRouter` (`router.rs`)** — a single `struct GrpcRouter { mode: Mode, .. }` with `GrpcRouter::new(ctx, mode)` and one `impl RouterTrait`. Regular-only components (harmony/embedding/classify) are `Some` only in `Mode::Regular`. `RouterFactory::create_router` derives `mode = grpc_mode(cfg)` and calls `create_grpc_router(ctx, mode)`. -- **`Mode` (`mode.rs`)** — `enum Mode { Regular, PrefillDecode, EncodePrefillDecode }` + `grpc_mode(&RouterConfig) -> Option`. Each `Mode` maps to a `WorkerSelectionMode`, an `ExecutionPlanKind`, an `inject_pd_metadata()` flag, and a `router_type()` label (`"grpc"` / `"grpc_pd"` / `"grpc_epd"`). -- **`RequestPipeline` (`pipeline.rs`)** — `RequestPipeline::build(endpoint, mode, deps)` composes a `Vec>` per endpoint, inserting `EncodeStage` only for EPD mode. A route method (e.g. `route_chat_impl`) builds the pipeline and runs the ordered stages. -- **Shared stages (`common/stages/`)** — the `PipelineStage` trait plus `WorkerSelectionStage`, `ClientAcquisitionStage`, `DispatchMetadataStage`, `RequestExecutionStage`, and `EncodeStage` (the relocated EPD encode, formerly `epd_encode.rs`). +- **`GrpcRouter` (`router.rs`)** — a single `struct GrpcRouter { mode: Mode, .. }` with `GrpcRouter::new(ctx, mode)` and one `impl RouterTrait`. `new()` builds one `RequestPipeline` per endpoint up front and stores it: `embedding_pipeline`/`classify_pipeline` are `Some` only in `Mode::Regular`; `harmony_pipeline`, `responses_context`, and `harmony_responses_context` are `Some` in Regular **and** PD and `None` only in EPD (those endpoints 501). `RouterFactory::create_router` derives `mode = grpc_mode(cfg)` and calls `create_grpc_router(ctx, mode)`. +- **`Mode` (`mode.rs`)** — `enum Mode { Regular, PrefillDecode, EncodePrefillDecode }` + `grpc_mode(&RouterConfig) -> Option`, which returns `Some` for `ConnectionMode::Grpc | ConnectionMode::Zmq` (ZMQ reuses this pipeline). Each `Mode` maps to a `WorkerSelectionMode`, an `ExecutionPlanKind`, an `inject_pd_metadata()` flag, and a `router_type()` label (`"grpc"` / `"grpc_pd"` / `"grpc_epd"`). +- **`RequestPipeline` (`pipeline.rs`)** — `RequestPipeline::build(endpoint, mode, deps)` composes a `Vec>` per endpoint, inserting `EncodeStage` only for EPD mode. A route method (e.g. `route_chat_impl`) picks the regular or harmony pipeline and calls `execute_chat`, which runs the ordered stages. +- **Shared stages (`common/stages/`)** — the `PipelineStage` trait plus `RateLimitReserveStage` (`rate_limit.rs`; tenant token reserve, inserted right after the preparation stage for chat/messages/completion/harmony only), `WorkerSelectionStage`, `ClientAcquisitionStage`, `DispatchMetadataStage`, `RequestExecutionStage`, and `EncodeStage` (the relocated EPD encode, formerly `epd_encode.rs`). +- **`BackendClient` (`backend_client.rs`)** — `enum BackendClient { Grpc(GrpcClient), Zmq(ZmqEngineClient) }`: what `ClientAcquisitionStage` yields and what a worker holds (`BasicWorkerBuilder::backend_client`). `GrpcClient` stays pure gRPC and `BackendClient::Grpc` delegates to it, so a new gRPC engine adds no arms here. - **Endpoint stages (`regular/`)** — per-endpoint request-building stages (`stages/chat|completion|generate|messages|embedding|classify/`), `processor.rs`, and `streaming.rs`. Despite the module name, `regular/` now backs **all** modes (Regular/PD/EPD). -**When adding an engine (this recipe), you rarely touch the router internals** — the `Mode`/pipeline machinery is engine-agnostic. Add token-aware request building under `regular/stages/*` only if your engine needs endpoint-specific shaping. +**When adding an engine (this recipe) you do not touch the `Mode`/pipeline machinery** — it is engine-agnostic — but you do add every per-engine arm listed in Step 3. ### Step 4: Register in runtime detection diff --git a/skills/implement/kv-index-feature.md b/skills/implement/kv-index-feature.md index eed8ea3..a8094f2 100644 --- a/skills/implement/kv-index-feature.md +++ b/skills/implement/kv-index-feature.md @@ -6,16 +6,18 @@ Radix trees for cache-aware routing. Tracks which workers have which prompt pref | Tree | Key Type | Used By | |------|----------|---------| -| `StringTree` | `&str` (characters) | HTTP routing — prompt text prefix matching | -| `TokenTree` | `&[u32]` (token IDs) | gRPC routing — token sequence prefix matching | +| `StringTree` | `&str` (characters) | HTTP requests carrying no token ids — prompt text prefix matching | +| `TokenTree` | `&[u32]` (token IDs) | gRPC **and** pre-tokenized HTTP — token sequence prefix matching | + +Token ids win over text. `GenerationRequest::routing_tokens()` (`crates/protocols/src/common.rs`, implemented in `generate.rs`: `input_ids` Single, or a Batch's first sequence; empty falls back to text) and the `x-smg-routing-tokens` header hint (`model_gateway/src/routers/common/header_utils.rs:parse_routing_tokens_hint`, which wins over the body) both feed `SelectWorkerInfo { tokens, .. }`; `extract_text_for_routing` is only materialized when there are no tokens. `StringTree` is a re-export alias of `string_tree::Tree` (`crates/kv_index/src/lib.rs`). -Both implement the `RadixTree` trait (`crates/kv_index/src/lib.rs`): prefix insertion, longest-prefix-match (`prefix_match`/`prefix_match_with_counts`), per-tenant eviction (`evict(tenant, max_units)`), concurrent access (`DashMap` **and** `parking_lot::RwLock`). +Both implement the `RadixTree` trait (`crates/kv_index/src/lib.rs`): prefix insertion, longest-prefix-match (`prefix_match`/`prefix_match_with_counts`), per-tenant eviction (`evict(tenant, max_units)`), concurrent access (`DashMap` **and** `parking_lot::RwLock`). Match results carry `tenant` **plus** `matched_tenants` — every holder of the deepest matched node, capped at `common.rs:MATCHED_TENANTS_CAP` (8), empty on a miss. `CacheAware` selects from `matched_tenants`, so a new match path that leaves it empty silently disables cache hits. ### Match+insert (the routing hot path) -`cache_aware.rs` does not call `insert_text`/`insert_tokens` then `match_prefix_with_counts` separately (two tree descents). It calls the FUSED `match_and_insert(key, tenant)` / `match_and_insert_with(key, select)` on both `Tree` (`string_tree.rs`) and `TokenTree` (`token_tree.rs`): one descent does longest-prefix-match, then inserts the unmatched remainder. The legacy two-descent `insert_*` + `match_prefix_with_counts` pair still exists. Prefer `match_and_insert*` for any read-then-populate path (see `match_and_insert` in `model_gateway/src/policies/cache_aware.rs`). +`cache_aware.rs` does not call `insert_text`/`insert_tokens` then `match_prefix_with_counts` separately (two tree descents). It calls the FUSED `match_and_insert(key, tenant)` / `match_and_insert_with(key, select)` on both `Tree` (`string_tree.rs`) and `TokenTree` (`token_tree.rs`): one descent does longest-prefix-match, then inserts the unmatched remainder. The legacy two-descent `insert_*` + `match_prefix_with_counts` pair still exists. Prefer `match_and_insert*` for any read-then-populate path (call sites: `select_worker_with_tokens` / `select_worker_with_text` in `model_gateway/src/policies/cache_aware.rs`; `select_worker_min_load` uses `match_and_insert` only when the mesh hash index is populated, plain `insert_*` otherwise). ## Steps @@ -23,9 +25,10 @@ Both implement the `RadixTree` trait (`crates/kv_index/src/lib.rs`): prefix inse 1. Implement in `crates/kv_index/src/` 2. Ensure `Send + Sync` (accessed from routing hot path) -3. Support both String and Token variants if applicable +3. Support both String and Token variants if applicable. Token-tree code must read the tree's runtime `page_size` (`TokenTree::with_config(page_size, policy)` / `page_size()`), never the `PAGE_SIZE` const — `PAGE_SIZE = 16` is only the `new()`/`with_policy()` default, and `CacheAwarePolicy::new_token_tree` sizes each tree to the backend KV page (`block_size`, CLI `--block-size`). Children are keyed by `TokenPageKey`, a `u64` digest of the edge's first page 4. Add eviction/cleanup mechanism (prevent unbounded memory) -5. Consider mesh sync if state should be cluster-wide +5. Purge per-tenant state on worker removal — extend `remove_tenant_all` (on `Tree`/`TokenTree`, deliberately not on the `RadixTree` trait), which `CacheAwarePolicy::remove_worker_by_url` calls for every model's trees. Size-based eviction never fires for a tenant whose count stops growing, so anything not purged there leaks forever +6. Consider mesh sync if state should be cluster-wide ### PositionalIndexer @@ -33,23 +36,30 @@ Event-driven cache-aware routing at block level — finer granularity than tree- There is no longer a 2048-worker cap (removed): `PositionalIndexer`/`TreeSizes` (`crates/kv_index/src/event_tree.rs`) use a segmented growable array supporting unbounded worker counts (2048 is just the first segment). `intern_worker` returns `Result` — erroring only on u32 id-space exhaustion — and must NEVER panic, because it runs inside per-worker subscription tasks (a panic would silently stop KV event indexing for that worker). KV subscription failures are now surfaced rather than swallowed. +Bounding: the indexer is unbounded by default. `PositionalIndexer::prune(ttl_secs, max_entries) -> PruneStats { scanned, evicted_ttl, evicted_capacity, remaining }` runs a last-touch TTL pass plus a capacity pass that evicts oldest-touched entries down to 90% of the ceiling; `None`/`Some(0)` disables each pass. Driven by `KvEventMonitor::start_prune_task` (`model_gateway/src/worker/kv_event_monitor.rs`) from `kv_indexer_ttl_secs` / `kv_indexer_max_entries` (CLI `--kv-indexer-ttl-secs`, `--kv-indexer-max-entries`; unset or 0 = off). New indexer state must be covered by `prune`, `remove_worker` and `apply_cleared`. + ## CacheAware Integration The `CacheAware` routing policy maintains per-model trees: ``` -DashMap> // HTTP -DashMap> // gRPC +DashMap> // text-keyed (no token ids) +DashMap> // token-keyed (gRPC + pre-tokenized HTTP) ``` -Selection: longest-prefix-match → prefer cached worker (weighted by prefix length) → fall back to least-loaded. +Only in the default `cache_index = tree` mode. Under `--cache-index hash` (`CacheIndexKind::Hash`, `--cache-ttl-secs` default 180, `--cache-boundaries`) `select_worker_hash` uses a TTL'd exact-match placement map keyed on request heads and the radix trees are **neither consulted nor populated** — tree features are invisible there, and untokenized requests stay load-balanced. + +Token requests also bypass the trees when the model has a populated `PositionalIndexer`: `select_worker` checks `has_event_indexer` (indexer present, `current_size() > 0`) before the tree paths and routes them to `select_worker_event_driven`, which scores block overlap and neither consults nor populates the token tree. + +Tree selection (`select_worker` → `select_worker_with_tokens` / `select_worker_with_text`): one `match_and_insert_with` descent → if `matched/input > cache_threshold` (CLI `--cache-threshold` default 0.3; `CacheAwareConfig::default()` is 0.5), `select_matched_candidate` pressure-selects among `result.matched_tenants` (least-loaded holder by default; `overlap_decay` / `selection_temperature` tune it), then `gate_selected_candidate` spills that pick to the min-load worker when its load exceeds **both** `avg_load * balance_rel_threshold` **and** `avg_load + balance_abs_threshold`, inserting the spill target as a new tenant → on a miss, route to the least-loaded worker and insert for it. Backend KV pressure (`is_kv_imbalanced`, from `token_usage`) abandons affinity fleet-wide for `select_worker_min_load`. There is no prefix-length weighting. ## Eviction -Configurable interval (`cache_aware.eviction_interval_secs`). Entries not accessed within window are pruned via LRU. +Size-based, not time-based. Every `eviction_interval_secs` (`CacheAwareConfig`, CLI `--eviction-interval`, default 120) a `PeriodicTask` calls `evict_tenant_by_size(max_tree_size)` on every model's string and token tree. `max_tree_size` (CLI `--max-tree-size`) is a budget **shared by all tenants of one tree**, not per-tenant: eviction returns immediately unless the tree-wide total (`total_char_size()` / `total_token_size()`) exceeds it, then removes leaves in LRU order across all tenants until it is back under budget. The trees have no access-time window — TTL exists only for the `PositionalIndexer` (`kv_indexer_ttl_secs`) and the hash placement index (`cache_ttl_secs`). ## Key Rules - All state must be `Send + Sync` — the trees use both `DashMap` and `parking_lot::RwLock` - Support both String and Token variants for HTTP/gRPC dual mode -- Always add eviction — unbounded trees cause OOM +- Always add eviction (unbounded trees cause OOM) **and** a purge path in `remove_tenant_all` — size eviction alone never reclaims a departed worker - Test concurrent access from multiple routing tasks +- Any change under `crates/kv_index/` (tests excluded) triggers `.github/workflows/benchmark-radix-tree.yml` on the PR; sanity-check locally with `cargo bench --bench radix_tree_benchmark -- benchmark_summary` (trees) or `cargo bench -p kv-index --bench throughput_bench` (PositionalIndexer) diff --git a/skills/implement/mcp-feature.md b/skills/implement/mcp-feature.md index fcaedff..011e6b4 100644 --- a/skills/implement/mcp-feature.md +++ b/skills/implement/mcp-feature.md @@ -1,12 +1,14 @@ # Adding MCP Features to SMG -`smg-mcp` (`crates/mcp/`) is the Model Context Protocol client: it discovers tools on external servers, gates execution behind an approval policy, and proxies calls. The crate is **OpenAI-protocol-free** — response-format adapter logic lives in `model_gateway::routers::common::openai_bridge`, not here (`lib.rs` crate doc). Built on `rmcp` 1.7 (`Cargo.toml`, `[dependencies] rmcp`). Note: rmcp 1.7 **dropped the standalone SSE client transport** — the `Sse` variant now returns an error (see Task B). +`smg-mcp` (`crates/mcp/`) is the Model Context Protocol client: it discovers tools on external servers, gates execution behind an approval policy, and proxies calls. The crate is **OpenAI-protocol-free** — response-format adapter logic lives in `model_gateway::routers::common::openai_bridge`, not here (`lib.rs` crate doc). Built on `rmcp` ^1.7 (`Cargo.toml`, `[dependencies] rmcp`; the committed `Cargo.lock` resolves 1.8.0). Note: rmcp 1.7 **dropped the standalone SSE client transport** — the `Sse` variant now returns an error (see Task B). Central type: `McpOrchestrator` (`core/orchestrator.rs`) owns the inventory, connection pool, and an `Arc` built from `McpConfig` (`core/config.rs`). ``` crates/mcp/src/ annotations.rs → ToolAnnotations, AnnotationType + error.rs → McpError, McpResult, ApprovalError + tenant.rs → TenantContext, TenantId, SessionId core/ orchestrator.rs (central) · session.rs · config.rs · pool.rs handler.rs · proxy.rs · reconnect.rs · metrics.rs approval/ manager.rs · policy.rs · audit.rs @@ -21,7 +23,7 @@ The two realistic tasks are **customizing the approval policy** and **adding a t `PolicyEngine` (`approval/policy.rs`, `struct PolicyEngine`) decides per call. `evaluate()` checks, in order: 1. Explicit tool policy (`tool_policies`, keyed `QualifiedToolName`) -2. Server policy + `TrustLevel` (`evaluate_with_trust`) +2. Server policy + `TrustLevel` (`evaluate_with_trust`) — short-circuits only when the server is `Trusted` or the result is not `Allow`; an `Allow` from a non-`Trusted` server falls through to steps 3-4 3. Pattern `rules` (`Vec`, evaluated in insertion order) 4. Annotation default (read_only→Allow, destructive→deny, else `default_policy`) @@ -50,8 +52,7 @@ Config enums are the `*Config` mirrors (`TrustLevelConfig`, `PolicyDecisionConfi `PolicyRule { name, pattern: RulePattern, condition: RuleCondition, decision: PolicyDecision }` (`policy.rs`, `struct PolicyRule`). `RulePattern` (`enum RulePattern`) = `Server(Regex) | Tool(Regex) | Qualified(Regex) | Any`; `RuleCondition` (`enum RuleCondition`) = `Always | HasAnnotation(AnnotationType) | LacksAnnotation(AnnotationType)`. Add via the builder `with_rule`, modeled on the `test_pattern_rule` test and the `Default` impl (`impl Default for PolicyEngine`): ```rust -use smg_mcp::approval::{PolicyRule, RulePattern, RuleCondition, PolicyDecision}; -use smg_mcp::AnnotationType; +use smg_mcp::approval::{PolicyDecision, PolicyEngine, PolicyRule, RuleCondition, RulePattern}; use regex::Regex; let engine = PolicyEngine::new(audit_log) @@ -79,6 +80,7 @@ Transports are the enum `McpTransport` (`core/config.rs`, `enum McpTransport`), 2. `core/orchestrator.rs` `connect_server_impl` — build the rmcp transport and call `handler.serve(transport)`. Model on the **`Streamable` arm** (the working reference): proxy via `super::proxy::resolve_proxy_config`, client via `build_http_client`, transport from `rmcp::transport::*` (`StreamableHttpClientTransport::with_client`). The `Sse` arm is **not** a template — it now returns `Err(sse_unsupported(..))` since rmcp 1.7 dropped the SSE client transport. 3. `core/orchestrator.rs` `connect_dynamic_server_with_tenant` — dynamic path (Stdio is rejected here, "Stdio not supported for dynamic connections"). 4. `core/orchestrator.rs` `server_key` and `core/pool.rs` `PoolKey::from_config` — derive the pool key (url + `hash_auth(token, headers)`). +5. (Dynamic/request-supplied servers only) `model_gateway/src/routers/common/mcp_utils.rs` `connect_mcp_servers` — the gateway builds an `McpTransport` from the request's `server_url` here (`/sse` → `Sse`, else `Streamable`). This is a **construction, not a match**, so the compiler will not flag it. **Anti-pattern:** implementing an `rmcp::Transport` trait. SMG does not define transports as trait impls — it matches the `McpTransport` enum and hands an rmcp-provided transport (`StreamableHttpClientTransport`, `TokioChildProcess`) to `handler.serve()`. Also update the `Debug` impl (`config.rs`, `impl fmt::Debug for McpTransport`) so secrets stay redacted. diff --git a/skills/implement/mesh-feature.md b/skills/implement/mesh-feature.md index c50cd8c..e3314ca 100644 --- a/skills/implement/mesh-feature.md +++ b/skills/implement/mesh-feature.md @@ -26,7 +26,7 @@ Mesh state is namespaced by key prefix (`worker:`, `policy:`, `rl:`, `config:`). 1. Choose a `MergeStrategy` (`crdt_kv/merge_strategy.rs`): - `LastWriterWins` — higher `(timestamp, replica_id)` wins. Default for `worker:`, `policy:`, `config:`. - - `EpochMaxWins` — rate-limit counters; raw put payload MUST be exactly 16 bytes `(epoch u64-be, count i64-be)` via `epoch_max_wins::encode`. + - `EpochMaxWins` — rate-limit counters; raw put payload MUST be exactly 16 bytes `(epoch u64-be, count i64-be)` via `smg_mesh::encode_epoch_count(epoch, count)` (decode with `decode_epoch_count`; see `adapters/rate_limit_sync.rs`). 2. Configure the prefix once on the shared `MeshKV` (from `handle.mesh_kv()`): ```rust let ns = mesh_kv.configure_crdt_prefix("worker:", MergeStrategy::LastWriterWins); @@ -39,7 +39,7 @@ Mesh state is namespaced by key prefix (`worker:`, `policy:`, `rl:`, `config:`). ns.delete("worker:7"); // tombstone ``` 4. React to remote changes with `ns.subscribe(sub_prefix)` → `Subscription { receiver }`. Events are `(key, Option>)`; `None` = delete. Delivered for both local writes and remote merges, carrying the canonical post-merge value. -5. Replication is automatic — no gossip wiring needed. Each round `collect_round_batch()` snapshots the op-log; `transport/crdt_batch.rs` (`build_crdt_batches`) frames it under `MAX_MESSAGE_SIZE`; the peer's `dispatch_crdt_batch` merges via `MeshKV::merge_crdt_ops`. Merge is idempotent by op-id. +5. Replication is automatic — no gossip wiring needed. Each round `collect_round_batch()` snapshots the op-log; `transport/crdt_batch.rs` (`build_crdt_batches`) frames it under `MAX_STREAM_CHUNK_BYTES` (`transport/limits.rs`; `MAX_MESSAGE_SIZE` minus a 64 KiB envelope margin); the peer's `dispatch_crdt_batch` merges via `MeshKV::merge_crdt_ops`. Merge is idempotent by op-id. Model on the integration tests in `crates/mesh/src/tests/crdt_integration.rs` (LWW `worker:` and EpochMaxWins `rl:`). @@ -66,7 +66,7 @@ Only if neither LWW nor EpochMaxWins fits. Strategy logic lives entirely inside - Op-id `(replica_id, timestamp)` must be globally unique per node — all engines share one `LamportClock`. - `apply_remote_ops` must emit a `CrdtChange` only when `get` actually changes (suppress idempotent re-applies), or subscribers get spurious events. - Ephemeral/lossy traffic (tenant deltas, tree repair) belongs in `StreamNamespace`, not CRDT — it is dropped under backpressure, not retried. -- Gateway-side wiring: there are three adapters in `model_gateway/src/mesh/adapters/` — `worker_sync.rs` (`worker:` CRDT, LWW), `rate_limit_sync.rs` (`rl:` CRDT, EpochMaxWins), and `tree_sync.rs` (the `td:` *stream* adapter for the distributed prefix tree, owned/wired by the cache-aware policy, not a CRDT). `MeshAdapters::start` (`model_gateway/src/mesh/wiring.rs`) registers the two CRDT prefixes/engines (`worker:` = LWW, `rl:` = EpochMaxWins) and starts each adapter's **inbound** sync **before** gossip starts (order matters — an unregistered engine drops to the default merge). A new replicated value that needs its own sync adapter is registered there. -- Outbound worker mesh sync: `WorkerSyncAdapter::start` (`adapters/worker_sync.rs`) also spawns an **outbound** publish loop over the registry's `WorkerEvent` stream — it publishes each locally-owned worker as `worker:{id}` under single-writer/owner ownership (mesh-imported workers are filtered out by registration origin so a peer's state is never re-published), tombstones on removal, and on broadcast-lag re-publishes all local workers + tombstones any it published that no longer exist (recovery). +- Gateway-side wiring: three adapters in `model_gateway/src/mesh/adapters/` — `worker_sync.rs` (`worker:` CRDT, LWW), `rate_limit_sync.rs` (`rl:` CRDT, EpochMaxWins), and `tree_sync.rs` (the distributed prefix tree's *stream* adapter: `td:` Broadcast tenant deltas + `tree:req:` / `tree:page:` Targeted repair — not a CRDT). `MeshAdapters::start(mesh_kv, node_name, worker_registry, cluster_state, policy_registry)` (`model_gateway/src/mesh/wiring.rs`) is the single composition root: it registers the two CRDT prefixes/engines and the three stream prefixes (`configure_stream_prefix` with `StreamConfig { max_buffer_bytes, routing }`), builds `TreeSyncAdapter` over a `ClusterState`-backed `PeerList` and a `PolicyRegistry`-backed `TreeHandle`, and starts every adapter's **inbound** loop **before** gossip starts (order matters — an unregistered engine drops to the default merge; a `td:` delta arriving before the subscription is dropped). Only THEN does it attach the tree adapter via `policy_registry.set_mesh_tree_sync(Some(tree))`, which propagates to every existing and future `CacheAwarePolicy` and flips `populate_hash_index` (`None` detaches) — the policy never constructs or wires the adapter itself. A new replicated value that needs its own sync adapter is registered there. +- Outbound worker mesh sync: `WorkerSyncAdapter::start` (`adapters/worker_sync.rs`) also spawns an **outbound** publish loop over the registry's `WorkerEvent` stream — it publishes each locally-owned worker as `worker:{id}` under single-writer/owner ownership, tombstones on removal, and on broadcast-lag re-publishes all local workers + tombstones any it published that no longer exist (recovery). One `is_publishable` gate covers every outbound path: mesh-imported workers are filtered out by registration origin (a peer's state is never re-published) and `ConnectionMode::Zmq` workers are never published at all (their `ipc://` endpoint is host-local). A `Replaced` event that turns an already-published worker non-publishable tombstones it. **Verify:** `cargo test -p smg-mesh` diff --git a/skills/implement/multimodal-feature.md b/skills/implement/multimodal-feature.md index ad4dca7..7c759b6 100644 --- a/skills/implement/multimodal-feature.md +++ b/skills/implement/multimodal-feature.md @@ -6,14 +6,14 @@ Image preprocessing for vision LLMs lives in the `llm-multimodal` crate (lib `ll ``` MediaContentPart (Text | ImageUrl | ImageData | ImageEmbeds | AudioUrl | AudioData | VideoUrl | VideoData) // types.rs - -> MediaConnector::fetch_image(MediaSource::{Url,DataUrl,InlineBytes,File}) // media.rs, Blake3 hash + -> MediaConnector::fetch_image(MediaSource::{Url,DataUrl,InlineBytes,File}, ImageFetchConfig { detail }) // media.rs, Blake3 hash -> Arc { image: DynamicImage, raw_bytes, detail, source, hash } -> VisionPreProcessor::preprocess(&[DynamicImage], &PreProcessorConfig) -> PreprocessedEncoderInputs -> ModelProcessorSpec::prompt_replacements(...) -> Vec // expands placeholder tokens -> tracker emits TrackerOutput { data: MultiModalData, uuids: MultiModalUUIDs } ``` -Video is a first-class modality: `MediaConnector::fetch_video(MediaSource::...)` returns an `Arc`, `VisionPreProcessor` has default `preprocess_video` / `preprocess_video_rgb` hooks, and `Modality::Video` flows through the same `PreprocessedEncoderInputs` contract (see `media.rs`, `vision/processor.rs`, `types.rs`). +Video is a first-class modality: `MediaConnector::fetch_video(MediaSource::..., VideoFetchConfig { min_frames, max_frames, sample_fps })` returns an `Arc`, `VisionPreProcessor` has default `preprocess_video` / `preprocess_video_rgb` hooks, and `Modality::Video` flows through the same `PreprocessedEncoderInputs` contract (see `media.rs`, `vision/processor.rs`, `types.rs`). The worked example below mirrors the existing **Phi3-Vision** pair: `vision/processors/phi3_vision.rs` (`Phi3VisionProcessor`) and `registry/phi3_v.rs` (`Phi3VisionSpec`). @@ -79,7 +79,7 @@ pub mod mymodel; pub use mymodel::MyModelProcessor; ``` -Then add it to `VisionProcessorRegistry::with_defaults()` in `vision/processor.rs`, registering every lowercase id substring the model uses (matching is case-insensitive `contains`): +Then add it to `VisionProcessorRegistry::with_defaults()` in `vision/processor.rs`, registering every lowercase id substring the model uses **plus the HF `config.json` `model_type` string** (e.g. `"phi3_v"`, `"kimi_k3"`). `VisionProcessorRegistry::find(model_id, model_type)` matches case-insensitive `contains` on the id and falls back to `model_type`, so a processor registered only under marketing ids is never found for renamed or custom checkpoints: ```rust registry.register("mymodel", Box::new(super::processors::MyModelProcessor::new())); @@ -91,7 +91,7 @@ registry.register("mymodel", Box::new(super::processors::MyModelProcessor::new() ### Step 3: Implement the spec -Implement `ModelProcessorSpec` (`registry/traits.rs`). `matches` is keyed by `ModelMetadata { model_id, tokenizer, config }`; pull token ids from `metadata.token_id(...)` or `metadata.config_u32(&["image_token_id"])`. `prompt_replacements` builds one `PromptReplacement` per image, expanding the placeholder to `feature_token_counts` copies. +Implement `ModelProcessorSpec` (`registry/traits.rs`). `matches` is keyed by `ModelMetadata { model_id, tokenizer, config }`; pull token ids from `metadata.token_id(...)` or `metadata.config_u32(&["image_token_id"])`. `prompt_replacements` builds one `PromptReplacement` per image, expanding the placeholder to `feature_token_counts` copies. `modality_limits` is a default, not a hard cap: `validate_media_request` (`registry/traits.rs`) lets a deployment raise or lower it via `SMG_IMAGE_MAX_COUNT` / `SMG_VIDEO_MAX_COUNT` / `SMG_AUDIO_MAX_COUNT`, but an override never enables a modality the spec did not declare. **File:** `crates/multimodal/src/registry/mymodel.rs` @@ -99,9 +99,9 @@ Implement `ModelProcessorSpec` (`registry/traits.rs`). `matches` is keyed by `Mo use std::collections::HashMap; use serde_json::{json, Value}; use crate::{ + encoder_inputs::PreprocessedEncoderInputs, registry::{ModelMetadata, ModelProcessorSpec, RegistryResult}, types::{Modality, PromptReplacement, TokenId}, - vision::processor::PreprocessedEncoderInputs, }; pub(super) struct MyModelSpec; @@ -167,14 +167,14 @@ LazySpec::new(|| Box::new(MyModelSpec)), ### Step 5: Quality gate -Invoke `smg:contribute` to run fmt -> clippy -> test -> bindings. +Invoke `smg:contribute` to run fmt -> clippy -> test -> bindings. Clippy runs `--all-features`, which turns on this crate's `opencv-video` feature and needs system OpenCV — run `make opencv-deps` (`scripts/install_opencv.sh`) once first. ## Critical Rules - Verify against `crates/multimodal/src/lib.rs` exports: content type is `MediaContentPart` (not `ChatContentPart`); the tracker yields `TrackerOutput`, not any `MultiModalInputs`. -- `field_layouts` defaults to `pixel_values: Batched`. Override it (like `qwen3_vl.rs`) only for patchified/flat tensors, declaring the sizes tensor via `FieldLayout::flat("patches_per_image")`. The layout key stays the logical `"pixel_values"` HF/vLLM kwarg even though the struct field is `encoder_input`. +- `field_layouts` defaults to `{"pixel_values": Batched}` and **any key not listed is treated as shared — replicated across every media item** (`registry/traits.rs:ModelProcessorSpec::field_layouts`; `model_gateway/src/routers/grpc/zmq_multimodal.rs`: batched keys index row `i`, flat keys slice by the cumulative sizes tensor, everything else is shared). So declare **every** per-item `model_specific` side tensor as `Batched` (Phi3-Vision: `image_sizes`; Qwen3-VL: `image_grid_thw`, `patches_per_image`), and reserve `FieldLayout::flat("patches_per_image")` for a patchified/flat `pixel_values` (see `qwen3_vl.rs`). The layout key stays the logical `"pixel_values"` HF/vLLM kwarg even though the struct field is `encoder_input`. That map is the legacy shape: multi-modality specs override `encoder_field_layouts_for(modality) -> EncoderFieldLayouts` instead (`registry/qwen3_omni.rs`, `registry/qwen3_asr.rs`, `registry/inkling.rs`), and the trait default bridges via `EncoderFieldLayouts::from_legacy_fields(self.field_layouts())`. - `PreprocessedEncoderInputs::new` is generic over dimensionality (calls `.into_dyn()`), so the same constructor takes 4D and 5D `encoder_input` arrays; the `channels`/`height`/`width` accessors error on non-4D/5D shapes. - Registry lookups are substring `contains` (processors) / `matches` (specs). Ordering is load-bearing only for the **spec** registry (`ModelRegistry.specs` is a `Vec`, first `matches` wins — so a specific spec like `Qwen3VLVisionSpec` must precede a general one like `QwenVLVisionSpec`). The **processor** registry is a `HashMap` with unspecified iteration order — there, correctness comes from specific, non-overlapping patterns, not registration order. - Processor and spec sets are NOT 1:1: `Phi4VisionProcessor` and `PixtralProcessor` exist (`vision/processors/`) with no registered `ModelProcessorSpec`. A processor without a spec preprocesses tensors but has no placeholder/prompt-expansion contract. -- For video, override the `VisionPreProcessor::preprocess_video` / `preprocess_video_rgb` defaults (they error by default) and emit `Modality::Video` replacements via `prompt_replacements_for` (see `qwen3_vl.rs`). -- All media flows through `MediaConnector`: honor `MediaConnectorConfig` (allowed_domains, `fetch_timeout` default 10s); images are Blake3-hashed (`hasher.rs`) for dedup. +- For video, all four pieces are required (see `qwen3_vl.rs`): add `Modality::Video` to `modality_limits`; override `placeholder_token_for` / `placeholder_token_id_for`, whose defaults return `UnsupportedModality` for anything but `Image` and which `model_gateway/src/routers/grpc/multimodal/plan.rs` + `process.rs` call *before* expansion; override the erroring `VisionPreProcessor::preprocess_video` / `preprocess_video_rgb` defaults; emit `Modality::Video` replacements from `prompt_replacements_for`. Miss any one and video requests fail at planning. +- All media flows through `MediaConnector`: honor `MediaConnectorConfig` (`allowed_domains`; `allowed_local_media_path`, which gates `MediaSource::File` — otherwise `MediaConnectorError::DisallowedLocalPath`; `fetch_timeout` default 10s); image, video and audio bytes are Blake3-hashed (`hasher.rs`) for dedup. diff --git a/skills/implement/observability-feature.md b/skills/implement/observability-feature.md index 4bfcf4a..e7c39e0 100644 --- a/skills/implement/observability-feature.md +++ b/skills/implement/observability-feature.md @@ -1,30 +1,37 @@ # Adding Observability Features to SMG -Three pillars: Prometheus metrics (40+), OpenTelemetry tracing, structured logging via `tracing` crate. +Three pillars: Prometheus metrics (90+ described, `smg_`-prefixed except the legacy `router_*` names in `crates/mesh/src/metrics.rs`), OpenTelemetry tracing, structured logging via `tracing` crate. ## Adding Metrics -### Step 1: Describe at startup +### Step 1: Describe in `init_metrics()` + +**File:** `model_gateway/src/observability/metrics.rs::init_metrics` (called by `start_prometheus`; it also calls `runtime_metrics::describe`, `smg_mesh::init_mesh_metrics`, `middleware/scheduler/metrics.rs::describe` and `metrics.rs::allocator_stats::describe`). ```rust -describe_counter!("my_metric_total", "Description of what this counts"); -describe_histogram!("my_latency_ms", "Description of what this measures"); +describe_counter!("smg_my_metric_total", "Description of what this counts"); +describe_histogram!("smg_my_stage_duration_seconds", "Description of what this measures"); ``` -### Step 2: Record on hot path with string interning +`start_prometheus` attaches histogram buckets **by name**: `Matcher::Suffix("duration_seconds")`, `Matcher::Suffix("ttft_seconds")`, `Matcher::Suffix("tpot_seconds")`, plus a `Matcher::Full` for the event-loop canary. Any other histogram name renders as a summary (quantile lines only) unless you add a `Matcher` there — so name duration histograms `_duration_seconds` and record `as_secs_f64()`, never milliseconds. + +### Step 2: Record on hot path with the right interner ```rust use crate::observability::metrics::intern_string; -// Dynamic labels MUST use intern_string to avoid allocation per request -let model = intern_string(&model_id); -counter!("my_metric_total", "model" => model).increment(1); -histogram!("my_latency_ms").record(elapsed_ms as f64); +// intern_string is for SERVER-controlled labels only (worker URLs, matched +// route templates, gateway-set error codes). Interned strings are never freed. +let worker = intern_string(worker_url); +counter!("smg_my_metric_total", "worker" => worker).increment(1); +histogram!("smg_my_stage_duration_seconds").record(elapsed.as_secs_f64()); ``` -**Anti-pattern:** Using raw strings for labels on hot paths — unbounded allocations, label cardinality explosion. +**Client-controlled labels** (request `model` IDs, MCP tool names) must NOT go through `intern_string`. They use the bounded interners in `observability/metrics.rs`: `intern_model_label` / `intern_tool_label`, both over `intern_bounded_label(map, cap, s)` — 1024 distinct values per kind, everything past the cap collapses to the `"other"` sentinel. All three are **private**, so a metric with a client-controlled label is added as a `Metrics::` method inside `metrics.rs` (pattern: `Metrics::record_router_request`); a new label kind needs its own `DashMap` + cap constant passed to `intern_bounded_label`. -**Anti-pattern:** interning dynamic per-request path segments. The HTTP metrics layer labels by the matched axum route *template* (`matched_path_label` / `MatchedPath` in `model_gateway/src/middleware/metrics.rs`), with an `"other"` fallback when no route matched, to bound cardinality — the interner never evicts. +**Anti-pattern:** `intern_string(&model_id)` or any other raw client input — the interner never evicts, so it is an unbounded memory and Prometheus series leak. + +**Anti-pattern:** interning dynamic per-request path segments. The HTTP metrics layer labels by the matched axum route *template* (`matched_path_label` / `MatchedPath` in `model_gateway/src/middleware/metrics.rs`), with an `"other"` fallback when no route matched. Static values for common cases (zero allocation): ```rust @@ -36,7 +43,7 @@ bool_to_static_str(true) // → "true" Metrics exposed via Prometheus `/metrics` endpoint (served on port **29000** by default, set via `--prometheus-port`). -**Verify:** `curl localhost:29000/metrics | grep my_metric` +**Verify:** `curl localhost:29000/metrics | grep smg_my_metric` ## Adding Tracing @@ -52,7 +59,7 @@ span.record("worker_url", url.as_str()); ``` - OTel context is bridged in `model_gateway/src/observability/otel_trace.rs`; trace context propagates through gRPC metadata (see @grpc-backend.md) -- Runtime self-observability lives in `model_gateway/src/observability/runtime_metrics.rs`: a background observer task (`spawn_observer`) runs an event-loop canary (sleeps 10ms in a loop, records `smg_tokio_event_loop_delay_seconds`, increments `smg_tokio_event_loop_stalls_total` when wake drift exceeds the threshold) plus a ~1s `RuntimeMetrics` sampler (queue depth, alive tasks, worker count, per-worker busy ratio, parks) — all exposed on the Prometheus `/metrics` endpoint. +- Runtime self-observability: `observability/runtime_metrics.rs::spawn_observer` runs an event-loop canary (10ms sleep loop, records `smg_tokio_event_loop_delay_seconds`, increments `smg_tokio_event_loop_stalls_total` past the drift threshold) plus a ~1s `RuntimeMetrics` sampler (queue depth, alive tasks, worker count, per-worker busy ratio, parks). ## Logging Rules @@ -64,14 +71,17 @@ info!(worker_url = %url, model_id = %id, "Worker registered"); warn!(error = %e, "Health check failed"); ``` -Module-specific levels: +Probe routes (`/health`, `/readiness`, `/liveness`) attach the `ProbeResponse` marker via `health.rs::mark_probe`; `ResponseLogger` (`middleware/logging.rs`) then logs them at DEBUG. Any new probe route must do the same or its 503 ERROR-floods once per poll. + +Module-specific levels — targets are module paths under the `smg` crate, and a set `RUST_LOG` replaces the default workspace filter entirely (`observability/logging.rs::init_logging` / `build_workspace_filter`), so always include a base level: ```bash -RUST_LOG=smg::routing=trace cargo run +RUST_LOG=info,smg::policies=debug cargo run -p smg --bin smg -- ``` +`--bin smg` is required: the package declares two bin targets (`smg`, `amg`) and no `default-run`. Routing decisions log at DEBUG from `policies/registry.rs` ("Sticky routing decision") and `policies/cache_aware.rs::log_tree_decision` ("Cache-aware selection"). ## Key Rules -- `intern_string()` for all dynamic label values +- `intern_string()` for server-controlled label values only; client-controlled ones go through the bounded `intern_model_label`/`intern_tool_label` (cap 1024, then `"other"`) - `tracing` crate only, no `println!` - Structured fields on spans, not string interpolation - GaugeHistogram for in-flight tracking (see `gauge_histogram.rs`) diff --git a/skills/implement/provider-api.md b/skills/implement/provider-api.md index 6fd3472..b9d1e32 100644 --- a/skills/implement/provider-api.md +++ b/skills/implement/provider-api.md @@ -45,7 +45,7 @@ Adding a variant breaks every exhaustive match — the compiler lists them. Patc } ``` -Then add a `RoutingMode::MyProvider { worker_urls } =>` arm to the `all_urls` match (also in `to_router_config`). `determine_connection_mode` returns `Http` unless a URL is `grpc://` — provider routers need HTTP (see Step 4). YAML config users select via `type: myprovider` (the serde tag); see @config-plumbing.md. +Then add a `RoutingMode::MyProvider { worker_urls } =>` arm to the `all_urls` match (also in `to_router_config`). `determine_connection_mode` returns `Http` unless a URL is `grpc://`/`grpcs://` (→ `Grpc`) or `ipc://` (→ `Zmq`), per `ConnectionMode::from_url` in `crates/protocols/src/worker.rs` — provider routers need HTTP (see Step 4). YAML config users select via `type: myprovider` (the serde tag); see @config-plumbing.md. ## Step 3: Implement the router @@ -64,7 +64,7 @@ impl RouterTrait for MyProviderRouter { &self, headers: Option<&HeaderMap>, tenant_meta: &TenantRequestMeta, - body: &CreateMessageRequest, // from openai_protocol::messages (struct CreateMessageRequest in crates/protocols/src/messages.rs) + body: CreateMessageRequest, // by value — from openai_protocol::messages (struct CreateMessageRequest in crates/protocols/src/messages.rs) model_id: &str, ) -> Response { // 1. (optional) MCP: if header_utils::is_smg_mcp_enabled(headers) && body.has_mcp_toolset() @@ -77,12 +77,14 @@ impl RouterTrait for MyProviderRouter { } ``` +Every typed-JSON route method takes the body **by value** — `route_generate`, `route_chat`, `route_completion`, `route_responses`, `route_embeddings`, `route_classify`, `route_rerank`, `route_messages`, `route_interactions` — so the dispatching router owns it and can free it once the upstream bytes exist. Writing `&CreateMessageRequest` will not match the trait. Only `route_audio_transcriptions` and the realtime session-create methods still take `&body` (`route_realtime_ws`/`route_realtime_webrtc` take a raw `Request`); `route_interactions` takes `model_id: Option<&str>`. + Worker selection (in `AnthropicRouter::route_messages`): `WorkerSelector::new(®istry, &client).select_worker(&SelectWorkerRequest { model_id, headers, provider: Some(ProviderType::Anthropic), ..Default::default() })`. `ProviderType` lives in `crates/protocols/src/worker.rs` (`enum ProviderType`: `OpenAI | XAI | Anthropic | Gemini | Custom(String)`) — add a dedicated variant only if your provider needs distinct request shaping; otherwise `ProviderType::Custom("myprovider".into())` covers passthrough providers without touching the enum. -Streaming split (also in `route_messages`): `let is_streaming = body.stream.unwrap_or(false);` then call into separate modules. Anthropic's `streaming::execute` / `non_streaming::execute` take `(&RouterContext, RequestContext)` (`execute` in `anthropic/streaming.rs` and `anthropic/non_streaming.rs`); `RequestContext` (`struct RequestContext` in `anthropic/context.rs`) carries the cloned request, headers, `model_id`, `tenant_request_meta`, connected `mcp_servers`, and the pre-selected `worker`. +Streaming split (also in `route_messages`): take ownership once (`let request = body;`), then `let is_streaming = request.stream.unwrap_or(false);` and call into separate modules. Anthropic's `streaming::execute` / `non_streaming::execute` take `(&RouterContext, RequestContext)` (`execute` in `anthropic/streaming.rs` and `anthropic/non_streaming.rs`); `RequestContext` (`struct RequestContext` in `anthropic/context.rs`) carries that owned request, headers, `model_id`, `tenant_request_meta`, connected `mcp_servers`, and the pre-selected `worker`. **Provider concerns** (split into sibling modules, like Anthropic): -- **SSE:** use the **shared SSE codec** in `routers/common/sse.rs` rather than hand-rolling framing. Encode downstream via `SseEncoder` (`encode_data(value)`, `encode_event(event_type, value)`, `SseEncoder::done()` for `data: [DONE]`); decode upstream via `SseDecoder` (`push(chunk)` then drain `next_frame()` to `None`, then `compact()`; `flush()` at end-of-stream). `SseFrame` is the parsed `{ event_type, data }` (with `decode_data::()` / `is_done()`), and `parse_block(&str)` parses a single already-extracted block. Both the Anthropic router (`anthropic/sse.rs`, `anthropic/streaming.rs`) and the OpenAI Responses router (`openai/responses/streaming.rs`) now stream **through** this codec. Event **names** stay vendor-specific (match the vendor exactly), but framing/parsing is shared and DoS-bounded (≈1 MB default decode buffer, `DEFAULT_MAX_BUFFER_SIZE` in `sse.rs`). `build_sse_response` still sets the `text/event-stream` headers (the shared `build_sse_response` in `common/sse.rs` builds from an `mpsc` receiver; Anthropic keeps its own `build_sse_response` in `anthropic/sse.rs` to forward upstream headers). +- **SSE:** use the **shared SSE codec** in `routers/common/sse.rs` rather than hand-rolling framing. Encode downstream via `SseEncoder` (`encode_data(value)`, `encode_event(event_type, value)`, `SseEncoder::done()` for `data: [DONE]`); decode upstream via `SseDecoder` (`push(chunk)` then drain `next_frame()` to `None`, then `compact()`; `flush()` at end-of-stream). `SseFrame` is the parsed `{ event_type, data }` (with `decode_data::()` / `is_done()`), and `parse_block(&str)` parses a single already-extracted block. Both the Anthropic router (`anthropic/sse.rs`, `anthropic/streaming.rs`) and the OpenAI Responses router (`openai/responses/streaming.rs`) now stream **through** this codec. Event **names** stay vendor-specific (match the vendor exactly), but framing/parsing is shared and DoS-bounded (≈1 MB default decode buffer, `DEFAULT_MAX_BUFFER_SIZE` in `sse.rs`). `build_sse_response` still sets the `text/event-stream` headers (the shared `build_sse_response` in `common/sse.rs` builds from the **bounded** receiver of `sse_channel()` — `SSE_CHANNEL_BUFFER` = 32 frames, for backpressure; never hand-roll an unbounded channel. Anthropic keeps its own `build_sse_response` in `anthropic/sse.rs` to forward upstream headers). - **MCP tool interception:** when the request carries MCP toolsets, run an agentic tool loop instead of a plain proxy — Anthropic branches to `execute_mcp_streaming` (in `anthropic/streaming.rs`) / the MCP path in `anthropic/non_streaming.rs`, capping at `mcp_utils::DEFAULT_MAX_ITERATIONS`. Per-server allowed tools come from `mcp::collect_allowed_tools_per_server` (`collect_allowed_tools_per_server` in `anthropic/mcp.rs`). See @mcp-feature.md. ## Step 4: Wire the factory @@ -90,7 +92,7 @@ Streaming split (also in `route_messages`): `let is_streaming = body.stream.unwr **File:** `model_gateway/src/routers/factory.rs`. 1. Add a `RouterId` const in the `router_ids` module, e.g. `pub const HTTP_MYPROVIDER: RouterId = RouterId::new("http-myprovider");`. 2. Add a `create_myprovider_router` mirroring `create_anthropic_router`: `Ok(Box::new(MyProviderRouter::new(ctx.clone())?))`. -3. In `create_router`, add a `RoutingMode::MyProvider { .. } =>` arm under **both** connection modes: the `ConnectionMode::Grpc` block returns `Err("MyProvider mode requires HTTP connection_mode")`; the `ConnectionMode::Http` block calls your factory. +3. In `create_router`, add a `RoutingMode::MyProvider { .. } =>` arm under **both** connection-mode arms: the `ConnectionMode::Grpc | ConnectionMode::Zmq` arm returns `Err("MyProvider mode requires HTTP connection_mode")`; the `ConnectionMode::Http` arm calls your factory. 4. Add the router to `create_igw_routers` so multi-router IGW mode serves it. Import your router in the `super::{...}` block at the top of `factory.rs`. @@ -98,12 +100,13 @@ Import your router in the `super::{...}` block at the top of `factory.rs`. ## Step 5: Mount the route (only if it is a new endpoint) `/v1/messages` and `/v1/interactions` are **already** mounted (`server.rs` route builder) and dispatch via existing trait methods. For a genuinely new path: -1. Add an axum handler beside `v1_messages` (the `v1_messages` fn in `server.rs`): `State>`, `HeaderMap`, `Extension`, a `PreemptionGuard`, and `ValidatedJson(body)`; call `state.router.route_*(...)` inside `cancel.guard(...)`. +1. Add an axum handler beside `v1_messages` (the `v1_messages` fn in `server.rs`): `State>`, `HeaderMap`, `Extension`, a `PreemptionGuard`, and `ValidatedJson(body)`. Clone the model id out first (`let model = body.model.clone();`), then call `state.router.route_*(Some(&headers), &tenant_meta, body, &model)` inside `cancel.guard(...)` — `body` is moved into the call. 2. Register it in the router builder next to `.route("/v1/messages", post(v1_messages))`: `.route("/v1/mypath", post(v1_mypath))`. +3. Leave `--stream-request-bodies-over` (`RouterConfig::stream_request_bodies_over`, default 0 = off — streamed pass-through of large bodies) alone: `stream_large_request_bodies` (`model_gateway/src/routers/http/router.rs`) matches the path, then downcasts to the **regular** HTTP `Router` and declines otherwise (it also declines when IGW has more than one router), so a provider router never receives a streamed body. Add your path to its `match req.uri().path()` list only if the regular router serves it too — that is why `/v1/messages` is listed. The body type must be a request struct from the protocols crate (e.g. `CreateMessageRequest`, `InteractionsRequest`) with `Deserialize` so `ValidatedJson` can parse it. That crate lives at `crates/protocols/` but its package is `openai-protocol` (Rust crate path `openai_protocol`). -**Protocol fidelity:** the protocols crate uses `#[serde(flatten)]` to preserve unknown vendor fields and keep `type` discriminators on content/system blocks, so proxied requests round-trip losslessly (e.g. Anthropic `/v1/messages` preserves system "text" + unknown fields — see `messages.rs`). Tool-call routers must keep the `function_call` / `function_call_output` rename pairing (the `#[serde(rename = "function_call")]` / `function_call_output` items in `responses.rs`). When extending a request struct, mirror these patterns rather than dropping unknown fields. +**Protocol fidelity:** the protocols crate uses `#[serde(flatten)]` to preserve unknown vendor fields and keep `type` discriminators on content/system blocks, so proxied requests round-trip losslessly (`messages.rs`). Tool-call routers must keep the `function_call` / `function_call_output` rename pairing (`responses.rs`). Mirror these patterns when extending a request struct; never drop unknown fields. --- @@ -116,4 +119,4 @@ cargo test -p smg routers # router-scoped tests Then invoke `smg:contribute` for fmt → clippy → test → commit. -**Anti-patterns:** implementing many `RouterTrait` methods (override only what you serve — the rest 501 by design); allowing gRPC for an HTTP-only provider (return the `Err` in the `ConnectionMode::Grpc` block of `create_router`); inlining streaming + SSE + MCP in `route_messages` (keep them in sibling modules as Anthropic does); inventing a request struct in the router instead of using/extending the protocols crate; hand-rolling SSE framing instead of going through `SseEncoder`/`SseDecoder`. +**Anti-patterns:** implementing many `RouterTrait` methods (override only what you serve — the rest 501 by design); allowing gRPC/ZMQ for an HTTP-only provider (return the `Err` in the `ConnectionMode::Grpc | ConnectionMode::Zmq` arm of `create_router`); inlining streaming + SSE + MCP in `route_messages` (keep them in sibling modules as Anthropic does); inventing a request struct in the router instead of using/extending the protocols crate; hand-rolling SSE framing instead of going through `SseEncoder`/`SseDecoder`. diff --git a/skills/implement/rate-limit-feature.md b/skills/implement/rate-limit-feature.md index 6a2fcc7..e8af52e 100644 --- a/skills/implement/rate-limit-feature.md +++ b/skills/implement/rate-limit-feature.md @@ -1,27 +1,26 @@ # Rate Limiting & Concurrency in SMG -One mechanism, two knobs. A single **global** (per-process) `TokenBucket` (`model_gateway/src/middleware/token_bucket.rs`) backs both the in-flight concurrency cap and the request rate limit. It is **not** per-tenant or per-route: `app_state.context.rate_limiter: Option>` is one bucket, and `concurrency_limit_middleware` acquires exactly `1.0` token per request regardless of tenant (`concurrency.rs::concurrency_limit_middleware`). Per-tenant admission is a *separate* subsystem — the priority scheduler (`AdmissionMode::Priority`, `server.rs::with_admission_layer`); don't conflate them. +One bucket, two knobs. A single **global** (per-process) `TokenBucket` (`model_gateway/src/middleware/token_bucket.rs`) backs both the in-flight concurrency cap and the optional request rate limit. It is **not** per-tenant or per-route: `app_state.context.rate_limiter: Option>` is one bucket, and `concurrency_limit_middleware` acquires exactly `1.0` token per request regardless of tenant (`concurrency.rs::concurrency_limit_middleware`). Two *other* admission layers exist — the priority scheduler (`AdmissionMode::Priority`, `middleware/scheduler/`) and per-tenant token/request limits (`rate_limit/`, below). Don't conflate the three. The bucket's two parameters map straight to config (`app_context.rs::maybe_rate_limiter`): | `TokenBucket::new` arg | RouterConfig field | Meaning | |---|---|---| -| `capacity` | `max_concurrent_requests: i32` | Burst size / max in-flight. `<= 0` disables rate limiting entirely (`rate_limiter = None`). | -| `refill_rate` | `rate_limit_tokens_per_second: Option`, coerced via `.filter(|&t| t > 0).unwrap_or(max_concurrent_requests)` | Tokens/sec. **`None`, `0`, or negative all fall back to `max_concurrent_requests`** — a configured `0` behaves identically to leaving the field unset, so semaphore mode is *not* reachable through this config path. | +| `capacity` | `max_concurrent_requests: i32` | Burst size / max in-flight. `<= 0` disables the bucket entirely (`rate_limiter = None`). | +| `refill_rate` | `rate_limit_tokens_per_second: Option`, coerced via `.filter(\|&t\| t > 0).unwrap_or(0)` | Tokens/sec. **Unset or `0` ⇒ refill `0`, a pure semaphore**: tokens only come back when a response body drops. Negative is rejected earlier by `config/validation.rs`. Set `> 0` to add a steady admission rate on top of the burst cap. | -A token is **not** returned when the handler returns — `TokenGuardBody` wraps the response body and calls `return_tokens_sync(1.0)` on `Drop`, i.e. only after the whole (possibly streamed) body is delivered (`concurrency.rs`, `TokenGuardBody::drop`). So `capacity` genuinely bounds concurrent *streams*, not just handler invocations. +A token is **not** returned when the handler returns — the acquired `TokenPermit` is moved into `TokenGuardBody`, which wraps the response body, and `TokenPermit::drop` calls `return_tokens_sync(tokens)`. That fires only after the whole (possibly streamed) body is delivered, or the request is cancelled (`concurrency.rs::TokenPermit`). So `capacity` genuinely bounds concurrent *streams*, not just handler invocations. -When the bucket is empty, the request is enqueued on a bounded `mpsc` channel (`queue_size`); `QueueProcessor::run` drains it, granting a token or returning `408 REQUEST_TIMEOUT` after `queue_timeout_secs`. A full queue → `429 TOO_MANY_REQUESTS`; queue disabled (`queue_size == 0`) → `429` immediately (`concurrency.rs::concurrency_limit_middleware`). +When the bucket is empty, the request takes a slot on the bounded `AdmissionQueue` (a `Semaphore` of `queue_size` slots, `concurrency.rs::AdmissionQueue`) and parks on the bucket via `TokenPermit::acquire_timeout(bucket, 1.0, queue_timeout)` — no background processor. With `refill_rate == 0` (the default shape) it joins the bucket's FIFO waiter list and `try_acquire` cannot barge past it (`token_bucket.rs::acquire_fifo`); with `refill_rate > 0` there is no waiter list — `acquire` polls under its own `tokens_needed / refill_rate` timeout, so a fresh `try_acquire` can barge. Queue disabled (`queue_size == 0` ⇒ `AppState.admission_queue == None`) or queue full → `429 TOO_MANY_REQUESTS` (`admission_queue_full`); timing out (`queue_timeout_secs`, or `acquire`'s own wait, whichever fires first) → `503 SERVICE_UNAVAILABLE` (`admission_queue_timeout`). Every shed carries `Retry-After: 2` (`middleware/mod.rs::SHED_RETRY_AFTER_SECS`). ## Primary task: tune the limits -These are plain `RouterConfig` fields with CLI flags already wired in `main.rs` (`--max-concurrent-requests`, `--queue-size`, `--queue-timeout-secs`, `--rate-limit-tokens-per-second`, all under `help_heading = "Rate Limiting"` on `CliArgs`). No code change needed to *set* them — set via CLI or YAML. Defaults: `max_concurrent_requests: -1` (disabled), `queue_size: 100`, `queue_timeout_secs: 60` (`config/types.rs`, `RouterConfig`'s `Default`). +These are plain `RouterConfig` fields with CLI flags already wired in `main.rs` (`--max-concurrent-requests`, `--queue-size`, `--queue-timeout-secs`, `--rate-limit-tokens-per-second`, all under `help_heading = "Rate Limiting"` on `CliArgs`). No code change needed to *set* them — set via CLI or YAML. Defaults: `max_concurrent_requests: -1` (disabled), `queue_size: 100`, `queue_timeout_secs: 60`, `rate_limit_tokens_per_second: None` (`config/types.rs`, `RouterConfig`'s `Default`). Common configurations: -- **Burst == steady rate (the closest to a pure concurrency cap):** `max_concurrent_requests=64`, leave `rate_limit_tokens_per_second` unset (or any `<= 0`, which coerces to the same). 64-burst + 64 req/s refill. Note: there is **no** config value that yields a true semaphore (refill_rate 0) via `maybe_rate_limiter` — a configured `0` is filtered out and falls back to `max_concurrent_requests`. -- **Rate limit + burst:** `max_concurrent_requests=200` (burst), `rate_limit_tokens_per_second=50` (steady 50 req/s). -- **Default (field unset):** refill rate falls back to `max_concurrent_requests`, so it acts as both a 64-burst and a 64 req/s limit. +- **Pure concurrency cap (the default shape):** `max_concurrent_requests=64`, leave `rate_limit_tokens_per_second` unset (or `0`). Refill is `0`: at most 64 in-flight responses, and a slot frees only when a response body drops. +- **Rate limit + burst:** `max_concurrent_requests=200` (burst), `rate_limit_tokens_per_second=50` (steady 50 req/s on top of the 200 cap). ## Adding a new knob @@ -35,49 +34,58 @@ Mirror `rate_limit_tokens_per_second` end to end per @config-plumbing.md. ### Step 2: Consume it where the bucket is built or acquired -Construction lives in `maybe_rate_limiter` (`app_context.rs`). Per-request behavior (token cost, branching) lives in `concurrency_limit_middleware` (`concurrency.rs`). The `TokenBucket` API you have: `try_acquire(f64) -> Result<(),()>`, `acquire_timeout(f64, Duration)`, `return_tokens_sync(f64)`, `available_tokens() -> f64`. Acquire and return amounts **must match** or the bucket drifts. +Construction lives in `maybe_rate_limiter` (`app_context.rs`). Per-request behavior (token cost, branching) lives in `concurrency_limit_middleware` (`concurrency.rs`). The `TokenBucket` API you have: `try_acquire(f64) -> Result<(),()>`, `acquire_timeout(f64, Duration)`, `return_tokens_sync(f64)`, `available_tokens() -> f64`. In the middleware go through `TokenPermit`, not the bucket directly — it pairs acquire with return and moves the admission gauges. ```rust // in concurrency_limit_middleware, if cost is configurable: let cost = app_state.context.router_config.rate_limit_cost as f64; -if token_bucket.try_acquire(cost).is_ok() { - let response = next.run(request).await; - let (parts, body) = response.into_parts(); - // TokenGuardBody MUST return the SAME amount on drop: - let guarded = TokenGuardBody::new(body, token_bucket, cost); - Response::from_parts(parts, Body::new(guarded)) +if let Ok(permit) = TokenPermit::try_acquire(token_bucket.clone(), cost) { + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED); + // the permit carries `cost`; its Drop returns exactly that when the body drops + return run_with_permit(next, request, permit).await; } +// ...and pass the same `cost` to TokenPermit::acquire_timeout on the queued path. ``` **Verify:** `cargo check -p smg` -**Anti-pattern:** `try_acquire(cost)` but `TokenGuardBody::new(.., 1.0)` — tokens leak/accumulate and the limiter silently stops working. The `1.0` literal appears in **three** places in `concurrency.rs` (immediate path, queued path, guard); keep them consistent. +**Anti-pattern:** acquiring a different amount on the immediate vs. queued path. The `1.0` literal appears in **two** non-test places in `concurrency.rs` (`TokenPermit::try_acquire`, `TokenPermit::acquire_timeout`); the permit returns whatever it acquired, so the body guard cannot drift — only those two sites can. ### Step 3: Record the outcome metric -Both admit/reject paths already call `Metrics::record_http_rate_limit(...)` with `metrics_labels::RATE_LIMIT_ALLOWED` / `RATE_LIMIT_REJECTED` (`observability/metrics.rs::record_http_rate_limit`). Reuse it; don't add a parallel counter. See @observability-feature.md for hot-path label rules. +Both admit/reject paths already call `Metrics::record_http_rate_limit(...)` with `metrics_labels::RATE_LIMIT_ALLOWED` / `RATE_LIMIT_REJECTED` (`observability/metrics.rs::record_http_rate_limit`). Queue depth, in-flight and rejection reason are covered too — `record_admission_queue_entered/exited` (`smg_admission_queue_depth`), `record_admission_inflight_acquired/released` (`smg_admission_inflight`), `record_admission_rejected(ADMISSION_REJECTED_FULL|ADMISSION_REJECTED_TIMEOUT)` (`smg_admission_queue_rejected_total{reason}`). The gauge pairs are owned by `QueueDepthGuard` and `TokenPermit`; `record_admission_rejected` fires inline in `concurrency_limit_middleware`, on the queue-full and timeout arms only (not the queue-disabled 429). Reuse them and keep the pairs balanced; don't add a parallel counter. See @observability-feature.md for hot-path label rules. **Verify:** `cargo check -p smg` ### Step 4: Test -`token_bucket.rs` has a `#[cfg(test)]` module covering refill, `refill_rate=0`, and notify-on-return — extend it. For middleware behavior, assert the 429-on-full-queue / 408-on-timeout paths. +`token_bucket.rs`'s `#[cfg(test)]` module covers refill, `refill_rate=0` FIFO order, no-barge, cancelled and timed-out waiters — extend it. For middleware behavior extend `concurrency.rs`'s own test module (helpers `test_app_state`, `echo_app`, `stream_app`) and assert the 429-on-full-queue / 503-on-timeout paths plus `Retry-After`. **Verify:** `cargo test -p smg` +## Per-tenant token/request limits (`rate_limit/`) + +A separate subsystem with its own config file, wired only into the gRPC pipeline. Don't extend the global bucket to do per-tenant work. + +- Enable: `--tenant-rate-limit-enabled --tenant-rate-limit-config ` (`RouterConfig.tenant_rate_limit_enabled` / `tenant_rate_limit_config`, `help_heading = "Tenant Rate Limit"`). Disabled is `Ok(None)`; enabled with missing or invalid YAML **fails startup** (`app_context.rs::maybe_rate_limit_manager` -> `manager.rs::RateLimitManager::from_config`). No Python/Go binding fields yet. +- YAML (`rate_limit/config.rs::RateLimitYaml`, `deny_unknown_fields`): `default_policy: TenantPolicySpec` plus `tenants: [TenantPolicySpec]`, each `{tenant_key, tokens_per_minute, requests_per_minute, model_rules: [ModelRuleSpec]}` with `matcher: {type: exact|prefix, value}`. `tenant_key` must be canonical (`auth:`, `header:`, `ip:`, `anonymous`) and `header:` keys additionally require `tenant_resolution.trust_tenant_header`. `CompiledPolicySet::compile` (`policy.rs`): exact beats prefix, longest prefix wins, rules never stack. +- Engine: `RateLimitBackend` (`backend.rs` — `reserve`/`settle_success`/`close_reserved_only`/`abandon`, keyed by `request_charge_id`) implemented by `LocalRateLimitBackend` (`local_backend.rs`, per-instance only, `ScopeBucket` in `bucket.rs`). `RateLimitManager::reserve(ReserveRequest) -> Reservation::{Admitted(Arc), Denied { retry_after_secs }}`; denial → `rejection_response(retry_after_secs)` = `429` + `tenant_rate_limit_exceeded` (`rejection.rs`). +- Lifecycle: reserve once per *logical* request, before any retry loop; resolve the handle exactly once (CAS-guarded, first resolution wins). Streaming defers resolution by attaching `ReservationAttachment` to the body via `AttachedBody::wrap_response` (`routers/grpc/common/stages/helpers.rs`), so a disconnect abandons instead of leaking a reservation. +- Wiring: `RateLimitReserveStage` sits right after the preparation stage in `routers/grpc/pipeline.rs` (Chat/Messages/Completion/Harmony) and no-ops unless the caller threads in a `RateLimitCell` (`router.rs`: `route_chat_impl`, `route_generate_impl`, `route_messages_impl`, `route_completion_impl`). Responses/embeddings/classify and the whole HTTP router are **not** covered. + ## Wiring reference (cite, don't reinvent) - Bucket built once at startup: `app_context.rs::maybe_rate_limiter`, stored on `AppContext.rate_limiter`. -- Limiter + queue processor created and spawned in `server.rs::startup` — `ConcurrencyLimiter::new(rate_limiter.clone(), queue_size, Duration::from_secs(queue_timeout_secs))` returns `(ConcurrencyLimiter, Option)`; `proc.run()` is spawned for the server's lifetime. `limiter.queue_tx` lands on `AppState.concurrency_queue_tx`. -- Middleware installed as a `route_layer` on protected/realtime/multipart routes via `with_admission_layer` (`server.rs::with_admission_layer`). That function is **always** called for those routes and switches internally on `AdmissionMode`: it installs `concurrency_limit_middleware` in `Legacy` mode and `priority_admission_middleware` in `Priority` mode. So when `priority_scheduler_enabled`, the token bucket is **bypassed** entirely in favor of the priority scheduler — verify which path your deployment uses before assuming the token bucket runs. -- All types re-exported from `middleware`: `ConcurrencyLimiter, QueueProcessor, QueuedRequest, TokenGuardBody, concurrency_limit_middleware, TokenBucket` (`middleware/mod.rs`). +- Queue built in `server.rs::startup`: `Arc::new(AdmissionQueue::new(queue_size, Duration::from_secs(queue_timeout_secs)))`, only when `rate_limiter.is_some() && queue_size > 0`; lands on `AppState.admission_queue: Option>`. No background task — waiters park on the bucket itself. +- Middleware installed as a `route_layer` on protected/realtime/multipart routes via `server.rs::with_admission_layer`. That function is **always** called for those routes and switches internally on `AdmissionMode`: `concurrency_limit_middleware` in `Legacy`, `priority_admission_middleware` in `Priority`. Under `priority_scheduler_enabled` the scheduler owns concurrency and the bucket's cap is ignored — but when `rate_limit_tokens_per_second > 0` the same bucket survives on `SchedulerState.rate_limiter` as a pre-admission RPS check (`scheduler/admission.rs`: `try_acquire(1.0)`, never returned). Verify which path your deployment uses. +- Re-exported from `middleware` (`middleware/mod.rs`): `concurrency_limit_middleware, AdmissionQueue, TokenGuardBody, TokenBucket`. `TokenPermit`, `QueueDepthGuard` and `run_with_permit` are private to `concurrency.rs`. ## Critical rules -- One global bucket. Per-request cost is `1.0`. No per-tenant/per-route limiting here — that's the priority scheduler (`middleware/scheduler/`). -- `max_concurrent_requests <= 0` ⇒ rate limiting fully off (`None`, pass-through). The `-1` sentinel is the disable signal, not a value. -- `refill_rate == 0.0` ⇒ `acquire` waits **indefinitely** for a returned token (`token_bucket.rs`); always pair with `acquire_timeout` (the queue processor does, `concurrency.rs::QueueProcessor::run`). -- Acquire amount must equal the `TokenGuardBody` return amount. The guard returns on body-drop, not handler-return — required for streaming correctness. +- One global bucket, cost `1.0` per request. Per-tenant limiting is `rate_limit/`; priority admission is `middleware/scheduler/`. +- `max_concurrent_requests <= 0` ⇒ the bucket is `None` and the middleware passes through. The `-1` sentinel is the disable signal, not a value. +- Unset or `0` `rate_limit_tokens_per_second` ⇒ `refill_rate == 0.0`, and `acquire` then waits **indefinitely** in FIFO order for a returned token (`token_bucket.rs::acquire_fifo`); always use `acquire_timeout` (the middleware does). +- Acquire and return amounts must match — hold the `TokenPermit` for as long as the work lasts instead of returning tokens by hand. The permit rides in the response body, so streaming and client disconnects are both covered. - Bucket math uses `f64`; `try_acquire`/`return_tokens` `debug_assert` finite, non-negative inputs. ## Quality gate diff --git a/skills/implement/reasoning-parser.md b/skills/implement/reasoning-parser.md index 4fde59a..93d228f 100644 --- a/skills/implement/reasoning-parser.md +++ b/skills/implement/reasoning-parser.md @@ -6,9 +6,9 @@ Most parsers are thin wrappers around `BaseReasoningParser`, configured with two | Input | Example | Notes | |-------|---------|-------| -| `MODEL_NAME` | `mymodel` | Snake case. Used as file name, struct prefix (`MymodelParser`), factory key, `model_type` | +| `MODEL_NAME` | `mymodel` | Snake case. Used as file name, struct prefix (`MymodelParser`), factory key, `model_type`. The factory key is user-facing: `--reasoning-parser `, the per-model `reasoning_parser` worker label / `ModelCard` override (validated via `registry().has_parser`), and Python `get_available_reasoning_parsers()` | | Tokens | `` / `` | Start/end delimiters. Confirm against vLLM/SGLang, HF `tokenizer_config.json` (`added_tokens_decoder` / `chat_template`), or model card | -| `always_in_reasoning` | `true` / `false` | `true` if the template injects the start token in the prefill so output begins as reasoning (DeepSeek-R1, MiniMax M2). `false` if output is normal until an explicit start token (Qwen3, GLM-4.5) | +| `always_in_reasoning` | `true` / `false` | `true` only when the template has NO thinking toggle and output always begins inside reasoning (DeepSeek-R1, MiniMax M2, Step3, Qwen3-Thinking). `false` whenever the template exposes an `enable_thinking` / `thinking` toggle (Qwen3, GLM-4.5, Kimi-K2.5, Nemotron, DeepSeek-V4) even if it prefills `` — the gateway arms the parser per request via `should_mark_reasoning_started` (`routers/grpc/utils/parsers.rs`), keyed on the `ThinkingToggle` detected from the chat template | | Patterns | `["my-model", "mymodel-v2"]` | Case-insensitive substrings matched against the model ID. List every variant/alias | ## Steps @@ -17,7 +17,7 @@ Most parsers are thin wrappers around `BaseReasoningParser`, configured with two **File:** `crates/reasoning_parser/src/parsers/{MODEL_NAME}.rs` -Model this on `parsers/deepseek_r1.rs`. The `ReasoningParser` trait has SEVEN methods; delegate all to `self.base`. Use `DEFAULT_MAX_BUFFER_SIZE` (4 MB), not a literal. +Model this on `parsers/deepseek_r1.rs`. The `ReasoningParser` trait has seven required methods — delegate all to `self.base` — plus one defaulted method, `requires_special_tokens()` (default `false`). Override it to `true` only when your delimiters are tokenizer special tokens rather than plain text (`<|content_thinking|>`, `<|open|>think<|sep|>`); otherwise detokenization strips them before the parser sees them (see `parsers/inkling.rs`, `parsers/kimi_k3.rs`). Use `DEFAULT_MAX_BUFFER_SIZE` (4 MB), not a literal. ```rust use crate::{ @@ -107,7 +107,7 @@ mod tests { **Verify:** `cargo check -p reasoning-parser` -**Anti-pattern:** Omitting `mark_reasoning_started` / `mark_think_start_stripped`, or using a non-existent field like `initial_in_reasoning`. Both fail to compile — the trait requires all seven methods and the field is `always_in_reasoning`. +**Anti-pattern:** Omitting `mark_reasoning_started` / `mark_think_start_stripped`, or using a non-existent field like `initial_in_reasoning`. Both fail to compile — the trait requires all seven of these methods and the field is `always_in_reasoning`. ### Step 2: Export the parser @@ -121,7 +121,7 @@ pub use mymodel::MymodelParser; **Verify:** `cargo check -p reasoning-parser` -**Anti-pattern:** Updating only `mod.rs`. The factory imports from the crate root, so a missing `lib.rs` re-export breaks the factory import in Step 3. +**Anti-pattern:** Updating only `mod.rs`. The factory imports from `crate::parsers`, so it still compiles — but every parser is also re-exported from the crate root for external callers (`reasoning_parser::XParser`); keep the `lib.rs` list in sync or the public API drifts. ### Step 3: Register in the factory @@ -146,8 +146,8 @@ Invoke `smg:contribute` to run fmt -> clippy -> test -> bindings -> commit. ## Critical Rules -- Compose `BaseReasoningParser`; delegate all seven trait methods to `self.base`. It already solves partial tokens, buffer overflow, and start/end stripping. -- For a model whose template injects the start token in the prefill (output starts mid-reasoning), set `always_in_reasoning: true` — that is the entire customization (file-based examples: `parsers/minimax.rs` `MiniMaxParser`, and `parsers/qwen3.rs` `QwenThinkingParser`, factory key `qwen3_thinking`). No `is_first_chunk` field or method overrides. +- Compose `BaseReasoningParser`; delegate all seven required trait methods to `self.base`. It already solves partial tokens, buffer overflow, and start/end stripping. Exception: multi-channel / typed-block formats that are not a single start/end pair implement `ReasoningParser` directly and override `requires_special_tokens()` — see `parsers/kimi_k3.rs` (`KimiK3Parser`) and `parsers/inkling.rs` (`InklingParser`); still enforce `DEFAULT_MAX_BUFFER_SIZE` and make `reset()` clear all state. +- For a model with no thinking toggle whose output always starts mid-reasoning, set `always_in_reasoning: true` — that is the entire customization (file-based examples: `parsers/minimax.rs` `MiniMaxParser`, and `parsers/qwen3.rs` `QwenThinkingParser`, factory key `qwen3_thinking`). No `is_first_chunk` field or method overrides. - `BaseReasoningParser::new` seeds `in_reasoning` from `always_in_reasoning`, and `reset()` restores it — so `is_in_reasoning()` equals the flag on a fresh or reset parser. - A new model only needs its own file when it has distinct tokens or behavior. Models reusing standard ``/`` can be a config-only `BaseReasoningParser::new(config).with_model_type(...)` registered inline in the factory (see `deepseek_v31`, `kimi_k25`, `kimi_thinking`) — no file in `parsers/`. - `ParserResult` has exactly two fields: `reasoning_text` and `normal_text`. diff --git a/skills/implement/routing-policy.md b/skills/implement/routing-policy.md index bfc9ad0..8f1f585 100644 --- a/skills/implement/routing-policy.md +++ b/skills/implement/routing-policy.md @@ -11,7 +11,7 @@ Note: `dp_min_token.rs` implements a *different* trait, `DPRankLoadPolicy` (sele | `POLICY_NAME` | `my_policy` | Snake case. File name, `name()` return, factory key | | Struct name | `MyPolicy` | PascalCase + `Policy` suffix | | State | `AtomicUsize`, `RwLock>`, none | Must be `Send + Sync`. Stateless policies are unit structs (see `random.rs`) | -| Routing input | none / `info.tokens` / `info.request_text` / `info.headers` | What drives selection (see Critical Rules) | +| Routing input | none / `info.tokens` / `info.routing_key` / `info.request_text` / `info.headers` | What drives selection (see Critical Rules) | | Config params | `load_factor: f64` | If tunable, needs a `PolicyConfig` variant — see @config-plumbing.md | ## Steps @@ -20,7 +20,7 @@ Note: `dp_min_token.rs` implements a *different* trait, `DPRankLoadPolicy` (sele **File:** `model_gateway/src/policies/{POLICY_NAME}.rs` -Model this on `random.rs` (stateless) or `power_of_two.rs` (cached load state). Required methods: `select_worker`, `name`, `as_any`. Default methods you may override: `on_request_complete`, `needs_request_text`, `update_loads`, `remove_worker`, `reset`. Always filter via the `get_healthy_worker_indices` helper — it applies `is_healthy() && circuit_breaker_can_execute()` per worker. +Model this on `random.rs` (stateless) or `least_load.rs` (cached load state; `power_of_two.rs` just wraps a `LeastLoadPolicy` scorer). Required methods: `select_worker`, `name`, `as_any`. Default methods you may override: `on_request_complete`, `needs_request_text`, `update_loads`, `needs_backend_loads` (must return `true` if you override `update_loads`), `remove_worker`, `reset`. Always filter via the `get_healthy_worker_indices` helper — it applies `Worker::is_available()` per worker (`is_healthy() && circuit_breaker_can_execute() && !is_overloaded()`). ```rust use std::sync::Arc; @@ -66,7 +66,7 @@ impl LoadBalancingPolicy for MyPolicy { **Verify:** `cargo check -p smg` -**Anti-pattern:** Adding `#[async_trait]` / `async fn`, or returning `Option>`. The trait is sync and returns `Option` — both fail to compile. Also: hand-rolling `w.circuit_breaker().can_execute()`. The method is `w.circuit_breaker_can_execute()`; just use `get_healthy_worker_indices`. +**Anti-pattern:** Adding `#[async_trait]` / `async fn`, or returning `Option>`. The trait is sync and returns `Option` — both fail to compile. Also: hand-rolling `is_healthy() && circuit_breaker_can_execute()` — that skips the overload veto the load monitor latches. Just use `get_healthy_worker_indices` (or `w.is_available()`). ### Step 2: Export the policy @@ -85,6 +85,8 @@ MyPolicy, Stateless policies are fieldless like `Random`/`RoundRobin`. For tunables, give the variant fields (`#[serde(default = "...")]`) and plumb them per @config-plumbing.md. +Also add an arm to the exhaustive `match` in `ConfigValidator::validate_policy` (`model_gateway/src/config/validation.rs`) — fieldless variants join the `PolicyConfig::Random | ... => {}` arm; tunables get their range checks there. Without it `cargo check` fails with a non-exhaustive-patterns error. + **Verify:** `cargo check -p smg` ### Step 4: Register in the factory @@ -99,12 +101,20 @@ PolicyConfig::MyPolicy => Arc::new(MyPolicy::new()), "my_policy" | "mypolicy" => Some(Arc::new(MyPolicy::new())), ``` -`PolicyRegistry` (`policies/registry.rs`) calls these to build the default/prefill/decode/per-model policies, so both arms are required: `create_from_config` for the configured policy, `create_by_name` for per-worker policy hints. +`PolicyRegistry` (`policies/registry.rs`) calls these to build the default/prefill/decode/encode/per-model policies, so both arms are required: `create_from_config` for the configured policy (and for a worker `policy_hint` naming the default policy's type, which then inherits the operator's tunables), `create_by_name` for any other per-worker hint. **Verify:** `cargo test -p smg` **Anti-pattern:** Updating only `create_from_config`. A worker that sends a `policy_hint` routes through `create_by_name`, which would silently fall back to the default policy on a miss. +### Step 4b: Expose the name on the CLI and bindings + +**File:** `model_gateway/src/main.rs` — add `"my_policy"` to the `--policy` (and `--prefill-policy` / `--decode-policy`) `value_parser = [...]` lists and an arm to `CliArgs::parse_policy`; its `_ => PolicyConfig::RoundRobin` fallback silently turns an unmapped name into round_robin. `--encode-policy` is deliberately restricted — leave it alone. + +**Python:** add the `PolicyType` variant and its `convert_policy` arm in `bindings/python/src/lib.rs`, the `policy_from_str` entry in `bindings/python/src/smg/router.py`, and the name in `COMMON_POLICY_CHOICES` (`bindings/python/src/smg/router_args.py`). See @bindings-update.md. + +**Verify:** `cargo check -p smg && make python-dev` + ### Step 5: Quality gate Invoke `smg:contribute` to run fmt -> clippy -> test -> bindings -> commit. @@ -112,10 +122,11 @@ Invoke `smg:contribute` to run fmt -> clippy -> test -> bindings -> commit. ## Critical Rules - The trait is `LoadBalancingPolicy`, sync, `Send + Sync + Debug`. No `#[async_trait]`. `select_worker` returns `Option` (index into the `workers` slice). -- Always filter with `get_healthy_worker_indices(workers)` and return `None` when it is empty. Index back into `workers` to read load/url or call `increment_processed()`. -- `SelectWorkerInfo` is a **struct**, not an enum — there are no `Http`/`Grpc` variants. Branch on its fields: `info.tokens: Option<&[u32]>` (token path) vs `info.request_text: Option<&str>` (text path); `cache_aware.rs::select_worker` does exactly this. If your policy reads request text/tokens, override `needs_request_text()` to return `true`. -- Header-based routing reads `info.headers: Option<&http::HeaderMap>` (e.g. `X-SMG-Target-Worker`, `X-SMG-Routing-Key`); consistent-hash policies use the prebuilt `info.hash_ring: Option>` rather than rebuilding per request. -- State must be `Send + Sync`. Prefer `AtomicUsize` (round_robin) or `DashMap`; `power_of_two.rs` uses `RwLock>` for cached loads updated via `update_loads`. Never `.unwrap()` on worker access or hold a lock across `select_worker`. -- Load-aware policies (those overriding `update_loads` — both `power_of_two` and `least_load`) are fed by the registry only if discoverable. `PolicyRegistry::get_all_load_aware_policies()` (`policies/registry.rs`) gathers them by name via an internal `is_load_aware(name)` check (currently `name == "power_of_two" || name == "least_load"`); a new load-aware policy must be added to that `is_load_aware()` name check to receive periodic `update_loads`. +- Always filter with `get_healthy_worker_indices(workers)` — it is `Worker::is_available()` (`is_healthy() && circuit_breaker_can_execute() && !is_overloaded()`; the overload veto is latched by the load monitor) — and return `None` when it is empty. Index back into `workers` to read load/url or call `increment_processed()`. (Only the ring policies differ: `consistent_hashing.rs` / `prefix_hash.rs` filter on `is_healthy_and_eligible()` — health plus the overload veto, no circuit breaker — to keep ring membership stable.) +- `SelectWorkerInfo` is a **struct**, not an enum — there are no `Http`/`Grpc` variants. Branch on its fields: `info.tokens: Option<&[u32]>` (token path) vs `info.request_text: Option<&str>` (text path); `cache_aware.rs::select_worker` does exactly this. If your policy reads request text/tokens, override `needs_request_text()` to return `true`. It derives `Default`, so tests write `SelectWorkerInfo { tokens: Some(&toks), ..Default::default() }`. +- Key-based routing reads the pre-validated `info.routing_key: Option<&str>` (non-empty UTF-8, <= 128 bytes, taken from the configured header names `routing_key_override.headers` / `--routing-key-headers`, default `x-smg-routing-key`) — never re-parse the header yourself. `consistent_hashing.rs` does `info.routing_key.or_else(|| extract_routing_key_hint(info.headers))`; `prefix_hash.rs` checks `info.routing_key` first. `info.rid_key` is the lineage-stripped body `rid`, `info.leg: WorkerLeg` (`Single`/`Prefill`/`Decode`) says which PD leg is being selected, and `info.tokens` may be an `x-smg-routing-tokens` hint that wins over body tokens. `X-SMG-Target-Worker` still comes from `info.headers`; consistent-hash policies use the prebuilt `info.hash_ring: Option>` rather than rebuilding per request. +- Routers call `PolicyRegistry::select_worker(&policy, workers, info)`, not your policy directly. With `routing_key_override.enabled`, keyed requests (`info.rid_key`, else the routing key) go to a shared sticky `ManualPolicy` for every policy except `manual`/`consistent_hashing` (`routing_key_override_applies`); only in the default `delegate` assignment mode is yours consulted at all — to place a first-seen/stale key, or to respill past the per-key in-flight cap (`STICKY_INFLIGHT_CAP = 2`). Keep `name()` identical to the serde rename; the registry matches it literally. +- State must be `Send + Sync`. Prefer `AtomicUsize` (round_robin) or `DashMap`; `least_load.rs` uses `RwLock>` for cached loads and since-poll in-flight tallies updated via `update_loads` (`power_of_two` delegates to it). Never `.unwrap()` on worker access or hold a lock across `select_worker`. +- Load-aware policies (those overriding `update_loads`) must also override `needs_backend_loads()` to return `true` — the trait default is `false`. `PolicyRegistry::get_all_load_aware_policies()` (`policies/registry.rs`) gathers policies by that method, not by a name list, and `WorkerMonitor` pushes `update_loads` only to what it returns (under `--disable-load-monitoring` an empty result also skips load polling entirely). `power_of_two`/`least_load` always return `true`; `cache_aware` only when a pressure knob is configured. - Load-aware policies should also override `remove_worker(&self, url)` to prune cached per-worker load when a worker is removed; the registry's `remove_worker_from_load_aware` calls it on each load-aware policy under worker churn so caches don't grow unbounded (see `least_load.rs` / `power_of_two.rs`). -- There are 10 `LoadBalancingPolicy` impls (random, round_robin, passthrough, power_of_two, cache_aware, least_load, bucket, manual, consistent_hashing, prefix_hash). Match an existing one rather than inventing API. (`passthrough` is a single-backend routing policy; `dp_min_token.rs` implements the separate `DPRankLoadPolicy` trait and is not counted.) +- There are 10 `LoadBalancingPolicy` impls (random, round_robin, passthrough, power_of_two, cache_aware, least_load, bucket, manual, consistent_hashing, prefix_hash). Match an existing one rather than inventing API. (`passthrough` is a single-backend routing policy; `dp_min_token.rs` is not counted — different trait.) diff --git a/skills/implement/scheduler-feature.md b/skills/implement/scheduler-feature.md index f423ab0..47c46d2 100644 --- a/skills/implement/scheduler-feature.md +++ b/skills/implement/scheduler-feature.md @@ -29,7 +29,7 @@ It is **off by default**: `priority_scheduler_enabled = false` keeps the legacy These four are the only scheduler knobs not in the YAML (`enabled`, `default_max_class`, `tenant_metric_top_n` are CLI-only; see `SchedulerSettings::from_cli_and_yaml`). `--priority-scheduler-config` is optional: an empty/absent file means every class uses `ClassConfig::default_for`. -**Verify:** `cargo run -p smg -- --help | rg priority-scheduler` +**Verify:** `cargo run -p smg --bin smg -- --help | rg priority-scheduler` (`--bin smg` is required — the package declares two bins, `smg` and `amg`, and no `default-run`) **Anti-pattern:** Setting `--max-concurrent-requests 0` and expecting the scheduler to use it as capacity. It only feeds the `WorkerCapacity` tier-4 fallback (`state.rs::try_build_priority`); real capacity is read live from `WorkerCapacity`. @@ -44,7 +44,7 @@ classes: reserved_floor: 128 # u16, absolute min slots; never drops below this reserved_per_slot: 0.25 # f64 >= 0, finite; effective = max(floor, ceil(share * capacity)) queue_size: 256 # u32, per-class queue depth -> 429 when full - queue_timeout_secs: 30 # u64 > 0; waiter past this -> 408 + queue_timeout_secs: 30 # u64 > 0; waiter past this -> 503 + Retry-After: 2 starvation_threshold_secs: 5 # u64 > 0; head-of-queue age the dispatcher promotes past can_preempt: true # may cancel a lower-class pre-TTFT request bulk: @@ -62,11 +62,11 @@ tenant_policies: max_class: system ``` -Every field of `ClassConfig` is required when a class appears (only `reserved_per_slot` has `#[serde(default)]`). Tenant keys must match `RouteRequestMeta::tenant_key()` (`model_gateway/src/tenant.rs`), e.g. `auth:`. +Every field of `ClassConfig` is required when a class appears (only `reserved_per_slot` has `#[serde(default)]`). Tenant keys must match `RouteRequestMeta::tenant_key()` (`model_gateway/src/tenant.rs`). A readable `auth:` only exists for per-tenant keys (`--tenant-api-key :`, `middleware/auth.rs::AuthConfig::with_tenant_keys`); the shared `--api-key` yields `auth:`, a trusted tenant header yields `header:`, otherwise `ip:` or `anonymous` (`middleware/tenant_resolution.rs::resolve_raw_tenant_key`). Scheduler keys are never validated at startup — a typo silently never matches. **Verify:** `cargo test -p smg scheduler::config` (the `PrioritySchedulerYaml` serde + `SchedulerSettings::from_cli_and_yaml` validation tests cover this exact shape). -**Anti-pattern:** Capitalized class keys (`System:`) or an unknown `max_class` value — both are hard serde errors (`test_yaml_unknown_class_value_is_serde_error`). Also: trying to encode capacity-vs-reservation limits in YAML; reservations are clamped to live capacity priority-ordered, so there is nothing to reject there (`SettingsValidationError` only flags zero timeouts / non-finite shares). +**Anti-pattern:** Capitalized class keys (`System:`) or an unknown `max_class` value — both are hard serde errors (`test_yaml_unknown_class_value_is_serde_error`). Also: trying to encode capacity-vs-reservation limits in YAML; reservations are clamped to live capacity priority-ordered, so there is nothing to reject there (`SettingsValidationError` only flags a zero `queue_timeout_secs`/`starvation_threshold_secs` and a non-finite or negative `reserved_per_slot`). ### Step 3 (extension): Custom `ClassQueue` discipline @@ -144,7 +144,7 @@ Invoke `smg:contribute` to run fmt -> clippy -> test -> bindings -> commit. - **Four fixed classes**, never add one — `Class` is `#[repr(u8)]` with numeric values packed into an `AtomicU64` slot count and indexed `class as usize` across `slots`, `engine.class_queues[4]`, and `SchedulerSettings.classes[4]`. Adding a variant breaks all four-element arrays. - The tenant clamp is `min(header_class, max_class)` (`admission.rs::resolve_priority`); a low-tier tenant cannot self-promote via `X-SMG-Priority`. An unknown header value silently degrades to `Class::Default` (`class.rs::parse_header`). - `ClassQueue` and `TenantPolicyResolver` are both **sync** `Send + Sync` traits. No `#[async_trait]`, no `async fn` — they run on the admission hot path. Queue state uses `parking_lot::Mutex`, not `tokio::sync::Mutex`. -- `admit(class, request_id, cancel) -> AdmitOutcome` (`engine.rs`) is the entry point: `Admitted(SchedulerPermit)` or `Rejected(RejectionReason)`. The four `RejectionReason` variants map 1:1 to HTTP in `error.rs`: `QueueFull`->429, `QueueTimeout`->408, `Preempted`->503 (`X-SMG-Preempted` + `Retry-After`), `ClientCancelled`->499. Don't add a reason without an `error.rs` arm. +- `admit(class, request_id, cancel) -> AdmitOutcome` (`engine.rs`) is the entry point: `Admitted(SchedulerPermit)` or `Rejected(RejectionReason)`. The four `RejectionReason` variants map 1:1 to HTTP in `error.rs`: `QueueFull`->429 + `Retry-After: 2`, `QueueTimeout`->503 + `Retry-After: 2` (both `middleware::SHED_RETRY_AFTER_SECS`; 408 was dropped because proxies retry it instantly), `Preempted`->503 (`X-SMG-Preempted: true` + `Retry-After: 1`), `ClientCancelled`->499. Don't add a reason without an `error.rs` arm. - `SchedulerPermit` is RAII — dropping it releases the slot, deregisters the inflight handle, and notifies the dispatcher. `admission.rs` wraps the response body in `SchedulerGuardBody` so the permit lives exactly as long as the response stream. - Reservation math is `effective = max(reserved_floor, ceil(reserved_per_slot * capacity))`, recomputed on every capacity change; capacity is read live from `WorkerCapacity`, never stored in `SchedulerSettings`. Don't hardcode capacity assumptions in YAML. - Settings are built once at startup and read-only. Editing the YAML requires a restart. diff --git a/skills/implement/storage-backend.md b/skills/implement/storage-backend.md index 7863b63..4dfc6ca 100644 --- a/skills/implement/storage-backend.md +++ b/skills/implement/storage-backend.md @@ -16,7 +16,7 @@ Storage backends live in `crates/data_connector/src/` as FLAT files (`memory.rs` **File:** `crates/data_connector/src/mybackend.rs` -Define a `Clone` store handle, then three structs each taking `store: MybackendStore`. Implement the three async traits. Honor schema overrides with `schema.col("field")` / `schema.is_skipped("field")` and append hook-provided columns from `current_extra_columns()` + `resolve_extra_column_values(...)` on every write (see `redis.rs` create paths). +Define a `Clone` store handle, then three structs each taking `store: MybackendStore`. Implement the three async traits. Honor schema overrides PER TABLE — `col` / `is_skipped` are methods on `TableConfig`, not on `SchemaConfig`: `let s = &self.store.schema.conversations;` then `s.col("field")` / `s.is_skipped("field")`. Append hook-provided columns from `current_extra_columns().unwrap_or_default()` + `resolve_extra_column_values(s, &hook_extra)` on every write (see `redis.rs` create paths). ```rust use std::sync::Arc; @@ -83,7 +83,7 @@ Add `MybackendConfig` (mirror `RedisConfig`: connection field, `#[serde(default) ### Step 3: Migrations (SQL backends only) -Key-value backends (redis) need none. SQL backends add a `mybackend_migrations.rs` + use `versioning.rs`, expose `init_schema` per struct, and a `run_migrations()` on the store that `create_storage` awaits (see `postgres.rs` `run_migrations` and `create_postgres_storage` in `factory.rs`). +Key-value backends (redis) need none. SQL backends add a `mybackend_migrations.rs` (`static MYBACKEND_HISTORY_MIGRATIONS: [versioning::Migration; N]`) and create tables at startup — either in each struct's `async fn new(store) -> Result` (Postgres style) or via per-struct `init_schema(conn, schema)` fns passed to the store constructor (Oracle style) — then a `run_migrations()` on the store that the factory awaits once tables exist (see `create_postgres_storage` in `factory.rs`, which also re-runs `ensure_response_indexes()` when migrations applied). The runners in `versioning.rs` are dialect-specific (`run_postgres_migrations` / `run_oracle_migrations`); a new dialect adds its own. ### Step 4: Wire into the factory @@ -103,15 +103,20 @@ HistoryBackend::Mybackend => { } ``` -Also add a `mybackend: Option<&'a MybackendConfig>` field to `StorageFactoryConfig` and update every existing `StorageFactoryConfig { .. }` literal (tests included) with `mybackend: None`. +Also add a `mybackend: Option<&'a MybackendConfig>` field to `StorageFactoryConfig` and update every existing `StorageFactoryConfig { .. }` literal: the `factory.rs` tests, `crates/data_connector/tests/postgres_integration.rs` (compiled by `cargo test -p data-connector` even though its tests are `#[ignore]`), and `model_gateway/src/app_context.rs` — that last one is the only production literal and must pass `config.mybackend.as_ref()`, not `None`. -### Step 5: Register the module + bindings +### Step 5: Register the module + gateway/bindings wiring -Add `mod mybackend;` to `lib.rs`. Expose the backend to Python by adding a `Mybackend` variant to `HistoryBackendType` in `bindings/python/src/lib.rs` and its `=> config::HistoryBackend::Mybackend` mapping arm (follow @bindings-update.md). +Add `mod mybackend;` to `lib.rs`. The crate enum alone is unreachable — the gateway selects the backend outside `data_connector`, so wire it too: + +- `model_gateway/src/config/types.rs`: `mybackend: Option` on `RouterConfig`, beside `oracle` / `postgres` / `redis`; `config/builder.rs`: a `maybe_mybackend` setter mirroring `maybe_redis`. +- `model_gateway/src/main.rs`: add `"mybackend"` to the `--history-backend` `value_parser = [...]` (clap rejects anything else), a `"mybackend" => HistoryBackend::Mybackend` arm in the string `match` (its `_` arm silently means `Memory`), a `build_mybackend_config` arm in the `(oracle, postgres, redis)` tuple match, and `.maybe_mybackend(mybackend)` in the builder chain. +- `model_gateway/src/app_context.rs`: `mybackend: config.mybackend.as_ref()` in its `StorageFactoryConfig` literal. +- Python: a `Mybackend` variant on `HistoryBackendType` in `bindings/python/src/lib.rs` with its `=> config::HistoryBackend::Mybackend` arm, plus a `PyMybackendConfig` / `to_config_mybackend` threaded like `redis_config: Option`; add `"mybackend"` to the `--history-backend` `choices=[...]` in `bindings/python/src/smg/router_args.py` and a branch in `router.py:history_backend_from_str` (follow @bindings-update.md). ### Step 6: Tests -Add `#[cfg(test)] mod tests` to `mybackend.rs` (round-trip each trait; cursor pagination via `ListParams`; not-found returns `Ok(None)`). Add a `create_storage` smoke test arm in `factory.rs` and a missing-config error test, mirroring `test_create_storage_redis_missing_config`. +Add `#[cfg(test)] mod tests` to `mybackend.rs` (round-trip each trait; cursor pagination via `ListParams`; not-found returns `Ok(None)`). Add a `create_storage` smoke test arm in `factory.rs` and a missing-config error test, mirroring `test_create_storage_redis_missing_config`. A backend needing a live server gets an opt-in `#[ignore]` suite in `crates/data_connector/tests/mybackend_integration.rs` gated on an env URL — model on `tests/postgres_integration.rs` (`DATA_CONNECTOR_TEST_POSTGRES_URL`; `cargo test -p data-connector --test postgres_integration -- --ignored --test-threads=1`, which CI runs in `pr-test-rust.yml`). **Verify:** `cargo test -p data-connector` @@ -119,5 +124,6 @@ Add `#[cfg(test)] mod tests` to `mybackend.rs` (round-trip each trait; cursor pa - Three structs / three traits, never a single `StorageBackend`. Hooks are wrapped by `create_storage`; the old `hooks.on_write/on_delete` calls do not exist. - Per-domain error enums — there is no unified `StorageError` type. -- Respect `SchemaConfig` (`col`/`is_skipped`) and append `resolve_extra_column_values` on writes so tenancy/hook columns persist. -- Package `data-connector` (v2.3.2). +- Respect the per-table `TableConfig` (`col`/`is_skipped`) and append `resolve_extra_column_values` on writes so tenancy/hook columns persist. +- The crate enum is not enough — a backend is unreachable until `model_gateway` (`main.rs` value_parser + match, `RouterConfig`, `builder.rs`, `app_context.rs`) and the Python launcher know its name. +- Package `data-connector` (v2.3.3; depended on from the workspace as `smg-data-connector`). diff --git a/skills/implement/tenancy-feature.md b/skills/implement/tenancy-feature.md index 1a6b26c..b907841 100644 --- a/skills/implement/tenancy-feature.md +++ b/skills/implement/tenancy-feature.md @@ -8,15 +8,18 @@ A `TenantKey` is the `Display` of a `TenantIdentity` variant with a prefix: `aut ``` auth_middleware (auth.rs) - ↓ inserts DataPlaneCaller (authenticated_from_sha256) into request extensions + ↓ hashes the Bearer token, looks it up in AuthConfig.keys (AuthConfig::with_tenant_keys: + ↓ shared --api-key → auth:, each tenant_api_keys entry → auth:) + ↓ inserts DataPlaneCaller::new(tenant_key) into request extensions route_request_meta_middleware (tenant_resolution.rs) - ↓ resolve_route_request_meta = RouteRequestMeta::new(resolve_raw_tenant_key(...)) + ↓ resolve_route_request_meta = RouteRequestMeta::new(resolve_raw_tenant_key(...)), + ↓ plus .with_extension(RequestId) when RequestIdLayer already set one ↓ resolve_raw_tenant_key, in priority order (tenant_resolution.rs::resolve_raw_tenant_key): ↓ 1. DataPlaneCaller extension → its TenantKey (auth wins) ↓ 2. trusted header (only if config.trust_tenant_header) → header: ↓ 3. ConnectInfo → ip: ↓ 4. else → anonymous - ↓ inserts RouteRequestMeta into extensions (resolution is infallible — no alias store, no fallible step) + ↓ inserts RouteRequestMeta into extensions (infallible) handler: Extension → passes &TenantRequestMeta into RouterTrait methods ``` @@ -39,11 +42,12 @@ let tenant = req .map(|m| m.tenant_key().clone()) .unwrap_or_else(|| TenantKey::new("anonymous")); ``` -In a **handler**, take it as an extractor (`server.rs`, the `generate` handler) and forward it: +In a **handler**, take it as an extractor (`server.rs`, the `generate` handler) and forward it. Typed-JSON bodies are passed **by value** so the router can free them after upstream dispatch, so clone the model id out first (`route_audio_transcriptions` and the realtime methods still take `&body`): ```rust Extension(tenant_meta): Extension, // ... -state.router.route_generate(Some(&headers), &tenant_meta, &body, &body.model) +let model = body.model.clone(); +state.router.route_generate(Some(&headers), &tenant_meta, body, &model) ``` In a **RouterTrait method**, you already receive `tenant_meta: &TenantRequestMeta`; read `tenant_meta.tenant_key().as_str()` (`router_manager.rs` route methods, e.g. `route_generate`, forward it to the underlying router). @@ -59,7 +63,7 @@ fn limit_for(&self, tenant: &TenantKey) -> Quota { } ``` -If you need to attach **derived per-request data** (not just read the key), append it to the meta with `RouteRequestMeta::with_extension` and read it back with `.extension::()` (tenant.rs). Prefer this typemap over new public fields: `RouteRequestMeta` already exposes two public fields — `tenant_key: TenantKey` and `request_charge_id: Uuid` (set via `Uuid::now_v7()` in `new()`, read via `request_charge_id()`) — and you should not add more. +If you need to attach **derived per-request data** (not just read the key), append it to the meta with `RouteRequestMeta::with_extension` and read it back with `.extension::()` (tenant.rs). `RequestId` is the in-tree example: `resolve_route_request_meta` attaches it and `routers/grpc/common/stages/helpers.rs::middleware_request_id` reads it back to derive backend request ids. Prefer this typemap over new public fields: `RouteRequestMeta` already exposes two public fields — `tenant_key: TenantKey` and `request_charge_id: Uuid` (set via `Uuid::now_v7()` in `new()`, read via `request_charge_id()`) — and you should not add more. **Verify:** `cargo check -p smg` @@ -71,21 +75,25 @@ Build a `RouteRequestMeta::new(TenantKey::from("auth:acme"))` (`router_manager.r ## Task B: add a new `TenantIdentity` source -Only when a new principal type must produce keys. Two coordinated edits: +Only when a new principal type must produce keys. **Not** for per-tenant API keys — those already resolve to `auth:` with no new variant: `TenantApiKeyEntry { tenant_id, key }` in `RouterConfig.tenant_api_keys` (`config/types.rs`), CLI `--tenant-api-key tenant_id:key` (repeatable, `main.rs::parse_tenant_api_key`), builder `.tenant_api_keys(..)`, validated by `config/validation.rs::validate_tenant_api_keys`. `server.rs` passes `AuthConfig::with_tenant_keys(api_key, &tenant_api_keys)` to `build_app` as the `serving_auth_config` and a shared-key-only `admin_auth_config`, so tenant keys never reach admin/worker routes. + +For a genuinely new source, two coordinated edits: ### Step 1: Add the variant + its key mapping -**File:** `model_gateway/src/tenant.rs` — add a variant to `TenantIdentity` and a matching arm in `into_key` with a **new unique prefix** (e.g. `apikey:`): +**File:** `model_gateway/src/tenant.rs` — add a variant to `TenantIdentity` and a matching arm in `into_key` with a **new unique prefix** (e.g. `mtls:`): ```rust -ApiKey(Arc) => format!("apikey:{id}"), +Mtls(Arc) => format!("mtls:{id}"), ``` Keep prefixes disjoint so two sources can never collide on one key string. `canonical_tenant_key(identity)` (tenant.rs) is just `identity.into_key()`. +Then teach `is_canonical_serving_tenant_key` (tenant.rs) the new prefix and add a case to its `canonical_forms_accepted` test. `rate_limit/config.rs::RateLimitYaml::validate` rejects any per-tenant `tenant_key` that fails it (`RateLimitConfigError::NonCanonicalTenantKey`), so a prefix added only to `into_key` compiles but can never be targeted by a tenant rate-limit policy. + ### Step 2: Emit it during resolution -**File:** `model_gateway/src/middleware/tenant_resolution.rs` — add a branch in `resolve_raw_tenant_key`, placed at the right **priority** (the function returns on the first match; `DataPlaneCaller` must stay first). Construct via `canonical_tenant_key(TenantIdentity::ApiKey(Arc::from(value)))`. +**File:** `model_gateway/src/middleware/tenant_resolution.rs` — add a branch in `resolve_raw_tenant_key`, placed at the right **priority** (the function returns on the first match; `DataPlaneCaller` must stay first). Construct via `canonical_tenant_key(TenantIdentity::Mtls(Arc::from(value)))`. -Header-sourced identities must stay behind the `state.trust_tenant_header` gate (`resolve_raw_tenant_key`) — an untrusted client can spoof a header. Auth-derived identities should instead be inserted upstream as a `DataPlaneCaller` in `auth_middleware` (`auth.rs`, via `DataPlaneCaller::authenticated_from_sha256` / `DataPlaneCaller::new(TenantKey)`). +Header-sourced identities must stay behind the `state.trust_tenant_header` gate (`resolve_raw_tenant_key`) — an untrusted client can spoof a header. Auth-derived identities belong upstream instead: `auth_middleware` (`auth.rs`) matches the token hash in `AuthConfig.keys` and inserts `DataPlaneCaller::new(TenantKey)`. ### Step 3: Config (if gated/named) @@ -103,6 +111,6 @@ Invoke `smg:contribute` to run fmt -> clippy -> test -> bindings -> commit. - Precedence in `resolve_raw_tenant_key` is fixed: `DataPlaneCaller` (auth) > trusted header > client IP > anonymous. Adding a source means choosing where it sits, not appending blindly. - Consumers must tolerate **any** key, including `anonymous` — always have a default-tenant fallback (`policy.rs` `unwrap_or(self.default)`; `admission.rs` `unwrap_or_else(|| TenantKey::new("anonymous"))`). Resolution is **infallible** — `resolve_route_request_meta` always produces a key and never errors (no alias store, no fallible step). - `route_request_meta_middleware` == `ordinary_tenant_resolution_middleware` (alias). Don't apply both. -- New `TenantIdentity` variant ⇒ new `into_key` arm with a **disjoint prefix**. Forgetting the arm fails to compile (non-exhaustive match); a duplicate prefix silently merges tenants. +- New `TenantIdentity` variant ⇒ new `into_key` arm with a **disjoint prefix**, plus a branch in `is_canonical_serving_tenant_key`. Forgetting the arm fails to compile (non-exhaustive match); a duplicate prefix silently merges tenants; skipping `is_canonical_serving_tenant_key` compiles but leaves the new tenants un-targetable by tenant rate-limit policy. - Extend per-request data via `RouteRequestMeta::with_extension`/`.extension::()`, not new public fields (the existing public fields are `tenant_key` and `request_charge_id`). The struct's `PartialEq` compares `tenant_key` and `request_charge_id`; it ignores `extensions` (tenant.rs). - Admin/control-plane paths build the meta directly from an explicit id (`resolve_admin_target_tenant_id` → `TenantIdentity::Explicit`, tenant.rs); they do not flow through `resolve_raw_tenant_key`. diff --git a/skills/implement/tool-parser.md b/skills/implement/tool-parser.md index c6145a3..9a6e844 100644 --- a/skills/implement/tool-parser.md +++ b/skills/implement/tool-parser.md @@ -153,10 +153,9 @@ registry.register_parser("acme", || Box::new(AcmeParser::new())); ``` And in `register_default_mappings`: ```rust -registry.map_model("acme-*", "acme"); -registry.map_model("Acme/Acme-*", "acme"); +registry.map_model("acme-*", "acme"); // one spelling is enough — matching is case-insensitive ``` -Model resolution (`resolve_model_to_parser`): exact match first, then the **longest** trailing-`*` prefix wins. Put more specific patterns alongside general ones; length, not order, decides. Unmatched models fall back to the `passthrough` default. +Model resolution (`resolve_model_to_parser`): exact match first, then the pattern's stem (trailing `*` stripped) is matched as a **case-insensitive substring** of the model ID; the **longest** matching stem wins — length, not order, decides. `acme-*` already covers `Acme/Acme-7B` and `org/acme-7b`, so don't add case or namespace variants. Pick a stem that cannot occur inside another family's ID (`nemotron-3*`, not `nemotron*`). Unmatched models fall back to the `passthrough` default. **Verify:** `cargo check -p tool-parser` @@ -224,7 +223,10 @@ Invoke `smg:contribute` to run fmt → clippy → test → bindings → commit. | Pure JSON, no tags (OpenAI/Claude/Gemini) | `parsers/json.rs` (`json`) | Reuse `JsonParser`; just `map_model(... , "json")`, no new file | | Pythonic `[func(arg=val)]` (Llama 4) | `parsers/pythonic.rs` (`pythonic`) | Reuse `PythonicParser` | | Unicode-delimited blocks (DeepSeek) | `parsers/deepseek.rs`, `deepseek31.rs`, `deepseek_dsml.rs` | Custom token matching | -| XML with `` children | `parsers/qwen_xml.rs` (`QwenXmlParser`, keys `qwen_xml`/`qwen_coder`), `glm4_moe.rs` | Custom XML extraction, not `handle_json_tool_streaming` | +| XML with `` children (Qwen3-Coder/3.5+, Nemotron-3) | `parsers/qwen_xml.rs` (`QwenXmlParser`, keys `qwen_xml`/`qwen_coder`/`nemotron`) | Custom XML extraction, not `handle_json_tool_streaming` — for a new model with this exact format just `map_model` or alias-register `QwenXmlParser`, no new file | +| XML ``/`` pairs (GLM-4.5/4.6 → `glm45_moe`, GLM-4.7/5 → `glm47_moe`) | `parsers/glm4_moe.rs` (`Glm4MoeParser::glm45()`/`::glm47()`) | Own parser, not the `` format above; the two constructors share a struct and differ only in the func-detail regex | | Namespaced/pipe-delimited XML | `minimax_m2.rs` (``), `step3.rs` (`steptml:`), `kimik2.rs` (`<\|tool_calls_section_begin\|>…`) | Read the parser for exact framing tokens — they are fragile; don't retype from memory | +| Channel/marker token streams (Kimi-K3 XTML, Inkling TML) | `kimi_k3.rs` (`<\|open\|>tools<\|sep\|>…<\|open\|>call tool="NAME" index="N"<\|sep\|>`), `inkling.rs` (`<\|message_model\|>NAME<\|content_invoke_tool_json\|>{…}<\|end_message\|>`) | Compiled-`Regex` marker set (Kimi-K3) or a `StreamingState` machine (Inkling); both expose `build_structural_tag` | +| Python-literal list of dicts `[{'name': …, 'arguments': {…}}]`, optional `<\|tool_calls\|>` prefix | `sarashina.rs` (`sarashina`) | Reuse `pythonic::{parse_python_expression, expression_to_json}`, not `serde_json` — a different grammar from `pythonic.rs`'s `[func(arg=val)]` | -Current parsers (factory keys): `passthrough`, `json`, `mistral`, `qwen`, `qwen_xml`, `qwen_coder`, `pythonic`, `llama`, `deepseek`, `deepseek31`, `deepseek32`, `deepseek_v4`, `glm45_moe`, `glm47_moe`, `step3`, `kimik2`, `minimax_m2`, `cohere`. +Current parsers (factory keys): `passthrough`, `json`, `mistral`, `qwen`, `qwen_xml`, `qwen_coder`, `nemotron`, `pythonic`, `llama`, `deepseek`, `deepseek31`, `deepseek32`, `deepseek_v4`, `glm45_moe`, `glm47_moe`, `step3`, `sarashina`, `kimik2`, `kimi_k3`, `inkling`, `minimax_m2`, `cohere`. The four that expose a structural tag (`mistral`, `kimik2`, `kimi_k3`, `inkling`) register via `register_parser_with_structural_tag(name, ctor, Parser::build_structural_tag)` instead of `register_parser`. diff --git a/skills/implement/wasm-plugin.md b/skills/implement/wasm-plugin.md index 45d22ad..7820746 100644 --- a/skills/implement/wasm-plugin.md +++ b/skills/implement/wasm-plugin.md @@ -9,22 +9,24 @@ Most "plugins" are new **guests** in `examples/wasm/` against the existing world | Variant | When | Typical use | |---------|------|-------------| | `OnRequest` | Before routing | Auth, rate limiting, request rewrite | -| `OnResponse` | After backend response | Logging, response rewrite | +| `OnResponse` | After backend response — **non-streaming only**: SSE / ndjson / chunked responses skip the whole phase (`is_streaming` in `model_gateway/src/middleware/wasm.rs`) | Logging, response rewrite | | `OnError` | On error path | **Not yet implemented** — the enum variant exists but dispatching it returns a runtime error, and the WIT world has no `middleware-on-error` interface. Don't attach modules here. | +`/v1/realtime`, `/v1/realtime/calls` and `/v1/audio/transcriptions` are mounted without the WASM middleware at all (`model_gateway/src/server.rs`) — no guest of any kind runs on them. + ## Actions (`action` variant in `spec.wit`) | Action | Effect | |--------|--------| | `Continue` | Pass through unchanged | | `Reject(u16)` | Block with that HTTP status | -| `Modify(ModifyAction)` | Set/add/remove headers, replace body, override status | +| `Modify(ModifyAction)` | Set/add/remove headers, replace body; `status` override applies on OnResponse only (OnRequest ignores `modify.status`) | ## Steps (new middleware guest — the common case) ### Step 1: Scaffold the guest crate -**Dir:** `examples/wasm/wasm-guest-{name}/`. Model this on `examples/wasm/wasm-guest-auth/` (on-request only) or `wasm-guest-logging/` (on-request + on-response with `Modify`). +**Dir:** `examples/wasm/wasm-guest-{name}/`. Model the guest code on `examples/wasm/wasm-guest-auth/` (on-request only) or `wasm-guest-logging/` (on-request + on-response with `Modify`) — but neither of those manifests has the `[workspace]` table below, so neither builds in-tree as-is. Only the `wasm-guest-storage-hook*` manifests get this right. `Cargo.toml`: @@ -34,6 +36,11 @@ name = "wasm-guest-myplugin" version = "0.1.0" edition = "2021" +# REQUIRED: examples/wasm/* are neither members nor `exclude`d in the root +# workspace, so without this cargo errors "current package believes it's in a +# workspace when it's not". It also keeps the target dir under the guest. +[workspace] + [lib] crate-type = ["cdylib"] @@ -76,13 +83,18 @@ You MUST implement both traits even if a hook is a no-op; the world exports both ### Step 3: Build the component -Target `wasm32-wasip2`, then wrap into a component. Copy `build.sh` from `wasm-guest-auth/` (handles target install + `wasm-tools component new`): +Target `wasm32-wasip2`. Simplest path: copy `build.sh` from `wasm-guest-auth/` and run it (installs the target, then does exactly this): ```bash cargo build --target wasm32-wasip2 --release -wasm-tools component new \ - target/wasm32-wasip2/release/wasm_guest_myplugin.wasm \ - -o target/wasm32-wasip2/release/wasm_guest_myplugin.component.wasm +# wit-bindgen 0.21+ on wasm32-wasip2 already emits a component, and +# `wasm-tools component new` rejects component input — wrap only a core module. +WASM=target/wasm32-wasip2/release/wasm_guest_myplugin.wasm +if wasm-tools print "$WASM" | grep -q '^(\s*component'; then + cp "$WASM" "${WASM%.wasm}.component.wasm" +else + wasm-tools component new "$WASM" -o "${WASM%.wasm}.component.wasm" +fi ``` ### Step 4: Upload via the admin API @@ -100,7 +112,7 @@ curl -X POST http://localhost:3000/wasm -H 'Content-Type: application/json' -d ' }' ``` -Modules execute in upload order; a `Reject` short-circuits the rest. +Execution order across modules at the same attach point is NOT guaranteed — `WasmModuleManager` holds `HashMap` and `get_modules_by_attach_point` iterates it unsorted. Never write a guest that depends on another guest running first. A `Reject` short-circuits the remaining modules for that phase. ## Adding a NEW hook or changing the WIT (host crate) @@ -108,12 +120,14 @@ Only when the existing world is insufficient: 1. Edit `crates/wasm/src/interface/spec.wit` (types/interfaces/world). 2. Host bindings regenerate via `wasmtime::component::bindgen!` in `crates/wasm/src/spec.rs` (`path: "src/interface"`, `world: "smg"`). -3. Extend `MiddlewareAttachPoint` in `crates/wasm/src/module.rs` and its handling in `module_manager.rs`. +3. Extend `MiddlewareAttachPoint` in `crates/wasm/src/module.rs`, add input/output variants in `crates/wasm/src/types.rs` (`WasmComponentInput` / `WasmComponentOutput`), dispatch the new export in `crates/wasm/src/runtime.rs` (`WasmThreadPool::execute_component_in_worker`, `match attach_point`), and add the phase loop in `model_gateway/src/middleware/wasm.rs`. `module_manager.rs` only forwards `attach_point` — touch it only if you add a `WasmComponentOutput` variant (`execute_module_for_attach_point` matches on it). 4. Tune limits in `crates/wasm/src/config.rs` (`WasmRuntimeConfig`). ## Storage hooks are a SEPARATE world -DB-layer interception is a different WIT world: `crates/wasm/src/interface/storage/storage-hooks.wit` (`package smg:storage; world storage-hook`, exports `storage-hook-before` + `storage-hook-after`), bridged by `WasmStorageHook` (`crates/wasm/src/storage_hook.rs`) and gated by the `storage-hooks` cargo feature. Model a storage guest on `examples/wasm/wasm-guest-storage-hook/` — note its `generate!` uses `path: ".../interface/storage", world: "storage-hook"`, not the gateway world. +DB-layer interception is a different WIT world: `crates/wasm/src/interface/storage/storage-hooks.wit` (`package smg:storage; world storage-hook`, exports `storage-hook-before` + `storage-hook-after`), bridged by `WasmStorageHook` (`crates/wasm/src/storage_hook.rs`). Model a storage guest on `examples/wasm/wasm-guest-storage-hook/` — note its `generate!` uses `path: ".../interface/storage", world: "storage-hook"`, not the gateway world. + +The gateway binary always enables the `storage-hooks` feature (`model_gateway/Cargo.toml`); the feature flag only matters for `cargo test -p smg-wasm --features storage-hooks` (fixtures via `crates/wasm/tests/fixtures/build_fixtures.sh`). A storage guest is NOT uploaded through `POST /wasm` and does not need `--enable-wasm` — pass it at startup with `--storage-hook-wasm-path /abs/guest.wasm` (Rust CLI, `bindings/python/src/smg/router_args.py`, Helm `router.wasm.path`); `app_context.rs` `with_storage` compiles it once. Uploading one via `POST /wasm` is silently accepted (validation only runs `Component::new`) and then skipped on every request. ## Step: Quality gate @@ -123,7 +137,8 @@ Invoke `smg:contribute` to run fmt -> clippy -> test -> bindings -> commit. - Two traits, exact method names: `on_request(req)` and `on_response(resp)`. `on_response` takes ONLY `resp`. There is no single `Guest` trait. - `generate!` and `bindgen!` both need the `path` arg; omitting it on the guest fails to find the WIT. -- Default limits: `max_memory_pages: 1024` (= 64MB), `max_execution_time_ms: 1000`, plus `max_body_size: 10MB` (`config.rs` `Default`). A guest exceeding them is killed. +- Default limits (`crates/wasm/src/config.rs:WasmRuntimeConfig`): `max_memory_pages: 1024` (= 64MB) and `max_execution_time_ms: 1000` trap the guest (store limiter / epoch deadline). `max_body_size: 10MB` is a HOST cap, not a guest trap — a larger body fails `axum::body::to_bytes` in `model_gateway/src/middleware/wasm.rs` before any guest runs and is rejected 400 (response phase: 500). +- Guest failures fail OPEN: a trap, timeout, or failed instantiation makes `execute_module_for_attach_point` return `None`, which `middleware/wasm.rs` handles as `continue` — the module is skipped and the request proceeds as if it returned `Continue`. An auth or rate-limit guest must never rely on a trap to block a request. - Upload URL is `/wasm`, not `/admin/wasm`, despite living on the admin router. - `--enable-wasm` is exposed in the Rust binary, the Python launcher (`bindings/python/src/smg/router_args.py`), and the Helm chart (`deploy/helm/`). With it off, the WASM middleware is skipped entirely (`model_gateway/src/middleware/wasm.rs`). - WIT identifiers are kebab-case; generated Rust is snake_case (`now-epoch-ms` -> `now_epoch_ms`) and PascalCase types (`modify-action` -> `ModifyAction`). diff --git a/skills/implement/zmq-backend.md b/skills/implement/zmq-backend.md new file mode 100644 index 0000000..0e62ea7 --- /dev/null +++ b/skills/implement/zmq-backend.md @@ -0,0 +1,92 @@ +# Extending SMG's ZMQ Direct Backend + +When SMG and the engine share a host, the gRPC path (gateway -> gRPC -> Python servicer -> ZMQ -> scheduler) collapses to a direct ZMQ connection over `ipc://`. The `engine-zmq-client` crate (`crates/engine_zmq_client/`) owns that wire; `model_gateway/src/routers/grpc/zmq_client.rs` adapts it to the existing proto request-execution pipeline, so `Mode`/pipeline/stages are reused unchanged. Two engines speak it: vLLM EngineCore (`vllm serve --headless`) and TokenSpeed (`ts serve --headless`). + +**Topology: SMG binds, the engine dials.** SMG binds a `tcp://` handshake ROUTER plus `ipc://` input (ROUTER) and output (PULL) sockets, drives HELLO -> INIT -> READY per engine, then waits for each engine's `EngineCoreReadyResponse` registration on the input socket (`transport.rs::connect_handshake`). `ConnectionMode::from_url` (`crates/protocols/src/worker.rs`) classifies an `ipc://` worker URL as `ConnectionMode::Zmq`; `uses_grpc_pipeline()` is true for both gRPC and ZMQ. + +**Hard limits — do not try to lift these in passing.** `vllm`/`tokenspeed` runtimes only (`ZmqEngineClient::connect`); `WorkerType::Regular` only (no PD/EPD leg — the wire carries no KV-transfer rendezvous and there is no encode dispatch); no admin RPCs (flush_cache / profile return `zmq_admin_unsupported`); no KV-event subscription (`kv_event_monitor.rs` is gRPC-only); no embeddings; not published to the mesh (`mesh/adapters/worker_sync.rs`); no dp-aware expansion. Health checks cannot be disabled — `resolve_zmq_health_config` forces them back on because the probe is the only path that reconnects a restarted engine. + +## Architecture map + +| Layer | Where | Key symbols | +|-------|-------|-------------| +| Protocol seam | `crates/engine_zmq_client/src/protocol/mod.rs` | `EngineProtocol`, `EngineOutput`, `EngineBatch`, `EngineLoad`, `WaveEvent` | +| Wire structs | `protocol/vllm/`, `protocol/tokenspeed/` | `EngineCoreRequest`, `EngineCoreOutput`, `TokenizedGenerateReqInput`, `BatchTokenIDOutSlim` | +| Handshake | `protocol/handshake.rs` | `ReadyMessage`, `HandshakeInitMessage`, `HandshakeAddresses`, `EngineCoreReadyResponse` | +| Codec | `codec/` | `encode_msgpack`, `decode_msgpack`, `TrailingTolerant`, `ModelDtype`, `WireNdArray` | +| Transport | `transport.rs` | `connect_handshake`, `ConnectedTransport`, `EngineId`, `ENGINE_CORE_DEAD_SENTINEL` | +| Connector | `connector.rs` | `Client

`, `EngineCoreClient`, `TokenSpeedClient`, `RequestStream

` | +| Gateway adapter | `model_gateway/src/routers/grpc/zmq_client.rs` | `ZmqEngineClient`, `ZmqDialect`, `translate_request`, `VllmGenerateStream`, `EosTokenIds` | +| Transport seam | `routers/grpc/backend_client.rs` | `BackendClient::{Grpc,Zmq}`, `finalize_generate_request`, `build_zmq_request` | +| Worker lifecycle | `model_gateway/src/worker/worker.rs` | `spawn_zmq_connect_driver`, `zmq_health_check`, `connect_zmq_backend` | +| Registration | `workflow/steps/local/create_worker.rs` | `validate_zmq_worker_type`, `validate_zmq_dp`, `validate_zmq_handshake_address` | +| Mock engine | `mock_engine.rs` (feature `mock-engine`), `crates/mock_worker/src/zmq.rs` | `connect_to_frontend`, `MockEngine`, `default_ready_response` | + +## Recipe A: Extend a wire field + +**Field order IS the wire contract.** Every data-plane message is a msgspec `array_like` struct: a positional msgpack array, not a map. Reordering or inserting silently misreads every later field. + +1. Read the upstream shape first: vLLM `rust/src/engine-core-client`, TokenSpeed `runtime/engine/io_struct.py`. `engine_zmq_client` is a clean-room port, not a dependency. +2. **Append** the field at the end of the struct, never insert. vLLM structs derive `Serialize_tuple`/`Deserialize_tuple` — mark the new field `#[serde(default)]` so an older engine's shorter array still decodes. +3. TokenSpeed structs hand-roll `visit_seq`: add a `next_field(&mut seq, "name")?` (required) or `seq.next_element::()?.unwrap_or_default()` (appended tail, zero from older senders) **before** the closing `drain_trailing(&mut seq)?`, and mirror it in `serialize_element` order. `expect_tag` validates element 0 (the `_tag` class name). +4. A newer engine always sends a longer array than this client models: keep the tolerance (`TrailingTolerant` / `deserialize_tolerant_seq` on the vLLM side, `drain_trailing` on TokenSpeed). +5. Add a msgpack roundtrip test beside the struct, plus a short-array case proving an older sender still decodes. + +**Verify:** `cargo test -p engine-zmq-client` + +## Recipe B: Surface a request/response feature on the ZMQ lane + +The ZMQ lane consumes the same `vllm_proto`/`tokenspeed_proto` `GenerateRequest` the gRPC builders produce, then translates it. + +1. Request building is already shared: `BackendClient::build_{chat,messages,completion,generate}_request` dispatch through `build_zmq_request`/`build_zmq_plain_request` on `ZmqDialect` and yield `ProtoGenerateRequest::{Vllm,TokenSpeed}`. +2. Map the proto field onto the wire in `zmq_client.rs::translate_request` (+ `translate_sampling`, `translate_constraint`) or `translate_request_tokenspeed` (+ `translate_sampling_tokenspeed`, `apply_tokenspeed_constraint`). +3. Shape the response in `VllmGenerateStream::map_output` / `TokenSpeedGenerateStream::map_output`. Both implement `MappedGenerateStream`; `poll_mapped` owns the `Stream` machinery and the parked terminal `Complete` — do not re-derive `Stream` by hand. +4. Never add stop-string handling here. This is a **token-only wire**: `BackendClient::finalize_generate_request` calls `helpers::resolve_string_stops(.., token_only_wire = true)` and `fold_tokenizer_eos_backstop`, returning the residual stop strings response processing must trim from output text. EOS ids come from `EosTokenIds` (local model dir at connect, else `adopt_tokenizer_eos`). +5. Dialect-specific behavior matches on `ZmqDialect` — a closed two-variant set, so write both arms and no wildcard. The existing sites are `backend_client.rs`, `harmony/stages/request_building.rs`, and `multimodal/assemble.rs`. Multimodal splits proto tensors into per-item `mm_features` in `zmq_multimodal.rs::build_mm_features` (inline bytes; no SHM/RDMA on this wire). + +**Verify:** `cargo test -p smg zmq` + +## Recipe C: Add a third dialect + +1. New `protocol//` module plus a zero-sized `EngineProtocol` impl (`add_frame`, `abort_frame`, `request_id`, `data_parallel_rank`, `validate`, `encode_add`, `encode_abort`, `encode_start_wave`, `decode_batch`), a `Client

` alias in `connector.rs`, and re-exports in `lib.rs`. Map the engine's native scheduler stats onto `EngineLoad` or the connector cannot load-balance its ranks. +2. `ZmqDialect` and `ZmqBackend` variants; a `RuntimeType` -> dialect arm in `ZmqEngineClient::connect`; arms in `dialect()`, `runtime()`, `generate()`, `is_alive()`, `engine_load()`. +3. A `ProtoGenerateRequest` variant and `ZmqBuilders` arms in `backend_client.rs`, then the other two `ZmqDialect` matches (harmony, multimodal). +4. Registration allowlist in `create_worker.rs` (the `RuntimeType::Vllm | RuntimeType::TokenSpeed` gate) and the unspecified-runtime default/warning in `detect_backend.rs`. +5. Launcher: a `WorkerLauncher` with `_build_zmq_command` in `bindings/python/src/smg/serve.py`, plus the e2e `infra/worker.py::_build__zmq_cmd`. + +**Verify:** `cargo test -p engine-zmq-client && cargo test -p smg zmq` + +## Recipe D: Lifecycle and config knobs + +- Per-worker: `WorkerSpec.zmq_handshake_address` — `tcp://` only, overrides the port `derive_handshake_port` folds out of the `ipc://` path (FNV-1a into 20000..=29999). `dp_size` on a ZMQ spec means a **grouped** worker (one socket set, N engines dialing in), not dp-aware expansion. +- Gateway-wide: `RouterConfig.zmq_engine_count` (CLI `--zmq-engine-count`, `parse_positive_usize`) and `startup_worker_runtime_type` (`--backend vllm|tokenspeed`, which pins a runtime **only** over ZMQ — the shared handshake carries no engine identity). Both are stamped onto startup workers by `job_queue.rs::apply_startup_worker_config`. Co-apply @config-plumbing.md and @bindings-update.md. +- New validation belongs in `create_worker.rs` beside `validate_zmq_handshake_override` / `validate_zmq_handshake_address` / `validate_zmq_worker_type` / `validate_zmq_dp` — reject at registration; a connect-time rejection strands the worker in `Pending` forever. +- Promotion is event-driven: `spawn_zmq_connect_driver` sends `WorkerConnected { url, revision }` to `manager.rs::apply_connect_signal` (wired via `registry.connect_signal_sender()`). The revision check discards a signal a same-URL replacement raced past. +- `update_worker_properties.rs` must preserve the ZMQ backend client when it rebuilds a worker. + +**Verify:** `cargo test -p smg` and `cd bindings/python && pytest -q tests/test_serve.py` + +## Local testing and CI + +- **No GPU:** `cargo run -p mock-worker -- --zmq-handshake tcp://127.0.0.1: --zmq-count 2` (`--zmq-start-index` sets the first rank's engine index). `` is `derive_handshake_port` of the worker's ipc path — mirror it with `python -c "from smg.serve import _zmq_handshake_port; print(_zmq_handshake_port('ipc:///tmp/...'))"`. The gateway needs `--model-path ` (model identity plus EOS ids), `--backend vllm`, and `--zmq-engine-count 2` for the grouped form. mock-worker speaks the vLLM dialect only. +- **In-process:** `mock_engine::{connect_to_frontend, MockEngine, default_ready_response}` under the `mock-engine` feature; see the tests in `zmq_client.rs`, `worker/worker.rs`, and `update_worker_properties.rs`. +- **Real engine:** `crates/engine_zmq_client/examples/live_probe.rs` against `vllm serve --headless --data-parallel-address 127.0.0.1 --data-parallel-rpc-port `. +- **E2E:** `E2E_CONNECTION_MODE=zmq` (plus `E2E_ZMQ_ENGINE_COUNT=2`) reruns the local chat suite over ZMQ; `fixtures/hooks.py::_filter_zmq_items` deselects PD/EPD/multi-worker cases and collapses the grpc/http twins. CI lanes `e2e-1gpu-chat-zmq` and `e2e-2gpu-chat-zmq-dp` are gated by a path filter that includes `crates/engine_zmq_client/**`. CPU harness gate: `PYTHONPATH=e2e_test pytest -q --noconftest e2e_test/infra e2e_test/fixtures`. + +## Key Rules + +- SMG binds every socket; the engine dials in. Never invert this — the operator launches the engine second. +- Positional field order is the wire contract. Append only; `#[serde(default)]` (vLLM) or a defaulting `next_element` before `drain_trailing` (TokenSpeed). +- Never await the handshake on a request, probe, or load-poll path. `get_backend_client` peeks the `OnceCell`, kicks `spawn_zmq_connect_driver`, and returns unavailable; a model load can outlast any caller's deadline. +- ZMQ workers stay `Pending` until `WorkerConnected` lands. Anything that can fail must fail at registration instead. +- Token-only wire: string stops and EOS are resolved frontend-side; the router keeps the residual trim obligation. +- Keep `derive_handshake_port` (`zmq_client.rs`) and `_zmq_handshake_port` (`serve.py`) byte-identical; `test_serve.py` pins the expected values. +- `n > 1` is fanned out into independent single-sample engine requests in `ZmqEngineClient::generate`, merged by `SelectAll` and tagged with the choice `index`. +- Liveness is local and latched (`is_alive`, set false on `ENGINE_CORE_DEAD_SENTINEL` or transport failure); there is no health RPC. Dropping a stream auto-aborts the engine-side request. + +## Anti-patterns + +- Adding a ZMQ arm to `GrpcClient` (`routers/grpc/client.rs`). The transport seam is `BackendClient`; `GrpcClient` stays pure gRPC. +- Following @grpc-backend.md for a ZMQ engine: there is no proto file, no `impl_engine_client_basics!`, no `detect_grpc_backend` probe entry, and no `fetch_grpc_metadata` — the handshake reports the metadata. +- A wildcard `_ =>` arm on `ZmqDialect`. It hides the site that a third dialect must update. +- Enabling `mock-engine` outside `dev-dependencies` / `mock-worker`. It is a test driver, not a shipped path. diff --git a/skills/map/SKILL.md b/skills/map/SKILL.md index 2b3a957..b61398f 100644 --- a/skills/map/SKILL.md +++ b/skills/map/SKILL.md @@ -7,34 +7,35 @@ description: Use when you need to understand the SMG codebase structure, find wh ## What Is SMG? -High-performance Rust gateway for LLM inference backends. Routes requests to workers running vLLM, SGLang, TensorRT-LLM, MLX (and more) with 10 routing policies, KV cache optimization, K8s service discovery, WASM plugins, MCP tool execution, and mesh HA. Exposes OpenAI-, Anthropic-, and Gemini-compatible APIs (plus Responses, Conversations, and Realtime/WebSocket), with a priority admission scheduler, multi-tenancy, and rate limiting. +High-performance Rust gateway for LLM inference backends. Routes requests to workers running vLLM, SGLang, TensorRT-LLM, MLX (and more) with 10 routing policies, KV cache optimization, K8s service discovery, WASM plugins, MCP tool execution, and mesh HA. Exposes OpenAI-, Anthropic-, and Gemini-compatible APIs (plus Responses, Conversations, and Realtime/WebSocket), with a priority admission scheduler, multi-tenancy, and rate limiting. Workers attach over HTTP, gRPC, or direct ZMQ (`ConnectionMode::{Http, Grpc, Zmq}`). ## Crate Map | Crate | Role | Key Types | |-------|------|-----------| | `model_gateway` | Main binary. HTTP/gRPC handlers, routing engine, service discovery, observability, CLI | `RouterConfig`, `ServerConfig`, `CliArgs` | -| `protocols` | OpenAI-compatible types shared by ALL consumers (config, bindings, API). Sacred — no impl-specific fields. | `WorkerSpec`, `ModelCard`, `WorkerModels`, `ChatCompletionRequest/Response` | +| `protocols` | OpenAI-compatible types shared by ALL consumers (config, bindings, API). Sacred — no impl-specific fields. `ConnectionMode::from_url` is the single scheme→mode classifier (`ipc://` = ZMQ) | `WorkerSpec`, `ModelCard`, `WorkerModels`, `ConnectionMode`, `ChatCompletionRequest/Response` | | `kv_index` | KV cache-aware routing. Radix trees (String for HTTP, Token for gRPC), positional indexer | `StringTree`, `TokenTree`, `RadixTree` trait, `PositionalIndexer` | | `auth` | API key (SHA-256 hashed), JWT/OIDC, role-based access (Admin/User), audit logging | `JwtConfig`, `ApiKeyEntry`, `Principal`, `Role` | | `mesh` | HA cluster via SWIM gossip. CRDT KV store, partition detection, consistent hashing | `ClusterState`, `WorkerState`, `NodeStatus` | | `wasm` | WebAssembly plugin system. WIT interface, middleware hooks (OnRequest/OnResponse), LRU cache | `WasmModule`, `Action` (Continue/Reject/Modify) | | `mcp` | MCP protocol client. Tool discovery, execution, approval workflows, response format translation | `McpConfig`, `McpOrchestrator`, `ToolAnnotations` | -| `grpc_client` | Per-engine gRPC clients for backends. Macros for shared logic; trace injection via `TraceInjector` | `SglangSchedulerClient`, `VllmEngineClient`, `TrtllmServiceClient` | +| `grpc_client` | Per-engine gRPC clients for backends. Macros for shared logic; trace injection via `TraceInjector` | `SglangSchedulerClient`, `VllmEngineClient`, `TrtllmServiceClient`, `MlxEngineClient`, `TokenSpeedSchedulerClient` | +| `engine_zmq_client` | Direct ZMQ transport to a same-host engine (crate `engine-zmq-client`), bypassing the Python gRPC servicer: `tcp://` handshake + `ipc://` data plane, gateway binds and engines dial in. Generic over `EngineProtocol`, with two dialects — vLLM EngineCore (clean-room port of vLLM's `vllm-engine-core-client`) and TokenSpeed. Feature `mock-engine` exposes a mock EngineCore | `EngineProtocol`, `Client

` (`EngineCoreClient`/`TokenSpeedClient`), `RequestStream`, `connect_handshake`, `EngineId` | | `data_connector` | Pluggable storage: PostgreSQL, Oracle, Redis, in-memory. Hook system for interception | `ConversationStorage`/`ConversationItemStorage`/`ResponseStorage` traits, `StorageHook` | -| `tool_parser` | 14 tool call parsers (JSON, Mistral, Qwen, DeepSeek, Pythonic, etc.). Streaming with incremental JSON | `ToolParser` trait, `ParserFactory`, `StreamingParseResult` | -| `reasoning_parser` | Reasoning extraction from 8 model families (DeepSeek-R1, Qwen3, Kimi, GLM, Step3, MiniMax, Cohere, Nano). Streaming | `ReasoningParser` trait, `ParserFactory`, `ParserResult` | -| `tokenizer` | LLM tokenization, chat templates | `Tokenizer` | -| `multimodal` | Image/audio processing (crate `llm-multimodal`). Per-model vision specs (LLaVA, Qwen-VL, Llama4, Phi3/4-V, Pixtral, Kimi-VL), media fetching | `ImageFrame`, `MediaContentPart`, `MediaConnector` | +| `tool_parser` | 17 tool call parsers (JSON, Mistral, Qwen, DeepSeek, Pythonic, Kimi-K2/K3, Inkling, Sarashina, etc.). Streaming with incremental JSON | `ToolParser` trait, `ParserFactory`, `StreamingParseResult` | +| `reasoning_parser` | Reasoning extraction from 9 model families (DeepSeek-R1, Qwen3, Kimi/Kimi-K3, GLM, Step3, MiniMax, Cohere, Nano, Inkling). Streaming | `ReasoningParser` trait, `ParserFactory`, `ParserResult` | +| `tokenizer` | LLM tokenization (crate `llm-tokenizer`): HF / tiktoken backends (Kimi-K2/K2.5/K2.6 is a tiktoken specialization, `kimi_k2_tokenizer.rs`), Jinja chat templates plus native renderers in `encoders/`, picked from `config.json::architectures` (DeepSeek-V3.2/V4 by `huggingface.rs:detect_renderer_from_config`, Kimi-K2.5 tools / Kimi-K3 XTML by `tiktoken.rs:detect_renderer_from_config`); EOS + stop-sequence decoding, L0/L1 encode cache, and the id+name registry the gateway loads into | `Tokenizer`, `TokenizerRegistry`, `StopSequenceDecoder`, `CachedTokenizer` | +| `multimodal` | Image/audio processing (crate `llm-multimodal`). Per-model specs in `registry/` (LLaVA/LLaVA-Next, Qwen-VL/Qwen3-VL/Qwen3-Omni, Llama4, Phi3-V, Kimi-K2.5/K3, Inkling, Qwen3-ASR audio), processors under `vision/processors/` + `audio/processors/`, media fetching | `ImageFrame`, `MediaContentPart`, `MediaConnector` | | `mm_rdma` | Multimodal pixel RDMA (NIXL) transport for the gateway (crate `smg-mm-rdma`) | | -| `workflow` | Step-based async workflow engine (wfaas) | `StepExecutor`, `WorkflowContext` | -| `bindings/python` | PyO3 bindings. `Router` class with ~110 constructor params, enum mapping | `Router`, `PolicyType` | +| `workflow` | Step-based async DAG workflow engine (crate `wfaas`); the gateway wraps it in `model_gateway/src/workflow/engines.rs` and feeds it from `workflow/job_queue.rs` (`Job::{AddWorker, RemoveWorker, UpdateWorker, AddTokenizer, ...}`). A step returning `StepResult::Skip` does NOT persist its context mutations — decide to skip before writing `context.data` | `WorkflowEngine`, `WorkflowDefinition`, `StepDefinition`, `StepExecutor`, `WorkflowContext`, `StepResult` | +| `bindings/python` | PyO3 bindings. `Router` class with ~140 constructor params, enum mapping | `Router`, `PolicyType` | | `bindings/golang` | Go SDK via FFI (cgo). OpenAI-style API, streaming, tool calling | `Client`, `ChatCompletionRequest` | | `clients/rust` | Rust client library | | | `clients/python`, `clients/java` | Client SDKs generated from the OpenAPI spec | | | `clients/openapi-gen` | Generates the OpenAPI spec + Python/Java client SDKs from protocol types (`make generate-clients`) | | -| `mock_worker` | Multi-port mock HTTP/gRPC inference-worker harness (package `mock-worker`, supports the TokenSpeed engine). Runs many fake workers in one process for routing/scale testing | (binary) | -| `grpc_servicer` | Python gRPC servicer wrapping vLLM/SGLang backends | | +| `mock_worker` | Multi-port mock HTTP/gRPC/ZMQ inference-worker harness (package `mock-worker`, lib `mock_worker` + binary, so in-process tests can drive it). gRPC side speaks the TokenSpeed scheduler; `--zmq-handshake --zmq-count N` runs mock vLLM EngineCore ranks that dial the gateway (via `engine-zmq-client`'s `mock-engine` feature); `--engine realistic` adds a continuous-batching simulator. Rigs: `scripts/scale_test.sh`, `scripts/sim_ab.sh` | `Config`, `http::serve`/`grpc::serve`/`zmq::serve` | +| `grpc_servicer` | Python gRPC servicer wrapping vLLM/SGLang/MLX/TokenSpeed backends | | ## Subsystems Inside `model_gateway` @@ -46,8 +47,10 @@ Beyond the crates, `model_gateway/src/` hosts several gateway-only subsystems. * | Provider routers | `routers/` | OpenAI, Anthropic, Gemini APIs + Responses, Conversations, Realtime/WebSocket, gRPC | `RouterManager` | | Priority scheduler | `middleware/scheduler/` | Priority-aware admission, per-class queues, slots, preemption, capacity reservations, autoscaling metrics | `PriorityScheduler`, `SchedulerPermit`, `Class`, `AdmitOutcome`, `TenantPolicy` | | Multi-tenancy | `tenant.rs` + `middleware/tenant_resolution.rs` | Canonical tenant identity + per-request resolution | `TenantIdentity`, `TenantKey`, `DataPlaneCaller`, `RouteRequestMeta` | -| Rate limiting | `middleware/token_bucket.rs`, `middleware/concurrency.rs` | Token-bucket rate limiting + concurrency caps | | +| Rate limiting | `rate_limit/` + `middleware/token_bucket.rs`, `middleware/concurrency.rs` | Per-tenant token/request budgets (`--tenant-rate-limit-enabled`, `--tenant-rate-limit-config`; reserve/settle runs in the gRPC pipeline stage `routers/grpc/common/stages/rate_limit.rs`) + the global token-bucket concurrency cap | `RateLimitManager`, `RateLimitYaml`, `CompiledPolicySet`, `ReserveRequest`, `TokenBucket` | | Worker lifecycle | `worker/` + `workflow/steps/local/` | Worker registry, health/circuit breaking, and the discovery→create DAG | `WorkerManager`, `CreateLocalWorkerStep` | +| ZMQ direct backend | `routers/grpc/{backend_client,zmq_client,zmq_multimodal}.rs` + `worker/worker.rs` (handshake driver) | Same-host `ipc://` workers ride the gRPC router pipeline behind `BackendClient::Zmq`; `ZmqEngineClient` translates the engine proto request to the `engine-zmq-client` wire and back, fans out `n`, folds tokenizer EOS ids into `stop_token_ids`, and drops string `stop` (the wire is token-only; the router's stop decoder trims them). Workers stay Pending until the detached handshake lands, then the `WorkerConnected` signal promotes them. No admin RPCs, KV events, PD/EPD legs, or mesh publication over ZMQ | `BackendClient`, `ZmqEngineClient`, `ZmqDialect`, `ZmqGenerateStream` | +| Tokenizer registration | `workflow/tokenizer_registration.rs`, `workflow/steps/local/submit_tokenizer_job.rs`, `routers/tokenize/` | gRPC/ZMQ workers are tokenized in the gateway: each non-HTTP local worker submits `Job::AddTokenizer` (path chain: label `tokenizer_path` → label `model_path` → `--tokenizer-path` → `--model-path`); `LoadTokenizerStep` loads locally, else streams a healthy gRPC worker's `GetTokenizer` bundle (the ZMQ backend serves none). `/readiness` holds at 503 until every healthy gRPC/ZMQ worker's tokenizer is registered (`health.rs`), unless `--disable-tokenizer-autoload` | `TokenizerRegistry`, `TokenizerConfigRequest`, `LoadTokenizerStep`, `SubmitTokenizerJobStep` | ## Layering Rule @@ -79,7 +82,9 @@ ServiceDiscoveryConfig / ServerConfig — typed, runtime ``` Client → HTTP/gRPC handler (OpenAI / Anthropic / Gemini router) - → Auth → Tenant resolution → Rate limit → Scheduler admission → WASM OnRequest + → WASM OnRequest → Auth → Tenant resolution + → Admission (priority scheduler OR legacy concurrency/token-bucket — never both) + → (gRPC pipeline only: per-tenant rate-limit reserve) → Routing policy selects worker → Proxy to backend → Stream response → Tool/reasoning parsing → WASM OnResponse → Client @@ -87,17 +92,33 @@ Realtime (WebSocket): Client → WS upgrade → Realtime session registry → Proxy to backend WS ``` +**Middleware order is the reverse of the `route_layer` calls** in `server.rs` (last added = outermost), which is why WASM runs before auth. `with_admission_layer` installs either the priority scheduler or the legacy concurrency limiter, not a chain of both. + +**gRPC and ZMQ workers share one pipeline** (`ConnectionMode::uses_grpc_pipeline()`), differing only in `BackendClient::Grpc` vs `BackendClient::Zmq` (metrics label `connection_mode="zmq"`); HTTP workers are proxied by their own router. Overload-vetoed workers are excluded from selection, and an all-overloaded fleet sheds 503 + `Retry-After` (`routers/common/overload.rs`). On the HTTP path, bodies over `--stream-request-bodies-over` bypass the typed path entirely — streamed verbatim, never parsed, never retried. + ## Worker Lifecycle (Discovery DAG) -Steps live under `model_gateway/src/workflow/steps/` (branches `local/`, `shared/`, `external/`, assembled in `steps/mod.rs`) — a DAG, not a fixed 5-step list: +Steps live under `model_gateway/src/workflow/steps/` (branches `local/`, `shared/`, `external/`, assembled in `steps/mod.rs`) — a DAG, not a fixed 5-step list. Discovery itself is level-triggered (a reconcile pass over a kube reflector store), not edge-triggered — there is no `handle_pod_event`: ``` -K8s Pod → PodInfo::from_pod() (service_discovery.rs) → handle_pod_event() → AddWorker - classify_worker_type → detect_connection_mode → detect_backend (sglang/vllm/trt/tokenspeed/mlx) +K8s pod store → reconcile_workers() (service_discovery.rs) + → compute_desired_state() [PodInfo::from_pod() per pod; one worker per port in + the `smg.ai/worker-ports` annotation] → compute_actions() + → Job::AddWorker / Job::RemoveWorker (workflow/job_queue.rs) + + classify_worker_type (waits worker_startup_delay_secs) + → detect_connection_mode (explicit http/grpc/ipc:// scheme honored, else probe + HTTP+gRPC every worker_startup_check_interval_secs) + → detect_backend (sglang/vllm/trtllm/tokenspeed/mlx; unidentified OpenAI-compatible + HTTP → `generic`; ZMQ defaults to vLLM unless runtime_type is set) → discover_metadata (flattens into labels HashMap) → discover_dp_info (rank/size) → create_local_worker (merge labels, resolve model_id, build ModelCard) + → ensure_harmony_encoding (gpt-oss gRPC/ZMQ only) → register_workers + → { update_policies | submit_tokenizer_job (non-HTTP) | activate_workers } ``` +**ZMQ workers short-circuit the probes**: `detect_connection_mode` is a no-op (SMG binds, the engine dials), `discover_metadata` yields no labels, `discover_dp_info` takes `dp_size` from the spec. `create_local_worker` requires `--model-path` (EngineCore reports no served model name), forces health checks on, and rejects Prefill/Decode/Encode worker types, dp-aware expansion, runtimes other than vllm/tokenspeed, and handshake-port collisions. `activate_workers` leaves them Pending until the `WorkerConnected` signal arrives. + ## The Label Pipeline Central integration pattern. All worker metadata flows as key-value labels: diff --git a/skills/review-pr/SKILL.md b/skills/review-pr/SKILL.md index 1f46fed..76e83a6 100644 --- a/skills/review-pr/SKILL.md +++ b/skills/review-pr/SKILL.md @@ -29,9 +29,9 @@ Do NOT write review comments, approve, or provide feedback until you have: 1. FETCH: Get the PR diff (gh pr diff ) 2. MAP: List changed files → match to sections using the file-to-section table 3. TASK: Create one review task per matched section -4. CHECK: Work through each task, flag issues as blocker/suggestion/nit +4. CHECK: Work through each task; prefix every inline comment with the repo's REVIEW.md severity marker — Important (bug, fix before merge) / Nit / Pre-existing (copy the exact marker glyphs from REVIEW.md) 5. ANTI-PATTERNS: Read @anti-patterns.md for the touched subsystems -6. SUMMARIZE: List all findings with severity and file:line citations +6. SUMMARIZE: List all findings with marker, file:line citations, and a count per severity ``` ### File-to-Section Mapping @@ -39,17 +39,19 @@ Do NOT write review comments, approve, or provide feedback until you have: | Files Changed | Review Sections | |---------------|-----------------| | `crates/protocols/src/` | 1 (Layering), 3 (Worker Lifecycle) | -| `model_gateway/src/config/` | 2 (Config Plumbing) | -| `model_gateway/src/main.rs` | 2 (Config Plumbing) | +| `model_gateway/src/config/`, `model_gateway/src/main.rs` | 2 (Config Plumbing) | | `model_gateway/src/service_discovery.rs` | 3 (Worker Lifecycle) | -| `model_gateway/src/worker/`, `model_gateway/src/workflow/steps/local/` | 3 (Worker Lifecycle) | +| `model_gateway/src/worker/`, `model_gateway/src/workflow/` (`steps/{local,shared,external}/`, `steps/mod.rs` DAG, `job_queue.rs`, `tokenizer_registration.rs`) | 3 (Worker Lifecycle) | +| `crates/workflow/src/` (wfaas engine) | 3 (Worker Lifecycle), 10 — a step `Skip` must not persist context; completion signals go out inside the tracker lock | | `model_gateway/src/policies/` | 4 (Routing Policy) | -| `model_gateway/src/routers/` (incl. `anthropic/`, `gemini/`, `responses/`, `conversations/`, `grpc/`) | 1 (Layering), 4 (Routing Policy) | -| `model_gateway/src/middleware/` (scheduler, tenant_resolution, rate limit) | 1 (Layering), 10 (Architecture) — no dedicated section yet | -| `crates/tool_parser/src/` | 5 (Parser Changes) | -| `crates/reasoning_parser/src/` | 5 (Parser Changes) | +| `model_gateway/src/routers/` (incl. `common/`, `http/`, `grpc/`, `anthropic/`, `gemini/`, `responses/`, `conversations/`) | 1 (Layering), 4 (Routing Policy) | +| `crates/engine_zmq_client/`, `model_gateway/src/routers/grpc/{backend_client,zmq_client,zmq_multimodal}.rs`, `crates/mock_worker/` | 3 (Worker Lifecycle), 4 (Routing Policy), 8 — ZMQ is a third transport; these paths trigger the merge-blocking `e2e-*-chat-zmq*` lanes | +| `model_gateway/src/middleware/` (scheduler, concurrency, tenant_resolution), `model_gateway/src/rate_limit/`, `model_gateway/src/tenant.rs`, `crates/auth/` | 1 (Layering), 10 (Architecture) — no dedicated section yet | +| `crates/tool_parser/src/`, `crates/reasoning_parser/src/` | 5 (Parser Changes) | +| `crates/tokenizer/src/` (encoders, stop/EOS, registry) | 5 (Parser Changes — streaming/stop semantics), 8 (Testing) | | `crates/data_connector/src/` | 6 (Storage) | | `bindings/` | 2 (Config Plumbing) | +| `e2e_test/infra/`, `e2e_test/fixtures/` | 8 (Testing) | | Any file | 7 (Error Handling), 8 (Testing), 9 (Code Quality) | Sections 7, 8, 9 always apply. Section 10 applies to PRs touching 3+ files or adding new types. @@ -62,6 +64,7 @@ Sections 7, 8, 9 always apply. Section 10 applies to PRs touching 3+ files or ad - [ ] Config types at correct layer: user-facing → `config/types.rs`, runtime → module-specific - [ ] No raw strings parsed at runtime — parse at boundary - [ ] WASM/MCP concerns stay in their crates, not leaking into core +- [ ] `serde(rename)` names match the OpenAI/Anthropic spec field names ### 2. Config Plumbing @@ -69,7 +72,8 @@ Sections 7, 8, 9 always apply. Section 10 applies to PRs touching 3+ files or ad - [ ] CLI flag has `value_parser` validation - [ ] `Default` impl includes new fields - [ ] `#[serde(default, skip_serializing_if)]` for backward compat -- [ ] Python bindings struct literals updated +- [ ] Python bindings threaded end to end: `config/builder.rs` setter → `router_args.py` field → `Router::new` pyo3 signature default → `.new_field(...)` on the `RouterConfig::builder()` in `to_router_config()` (`bindings/python/src/lib.rs`) — a missing builder call compiles and silently uses the Rust default +- [ ] ZMQ knobs reach `router_args.py` and `smg serve`, not just `lib.rs`; `derive_handshake_port` (`routers/grpc/zmq_client.rs`) and `_zmq_handshake_port` (`bindings/python/src/smg/serve.py`) stay in sync - [ ] Go SDK updated if new types exposed ### 3. Worker Lifecycle & Label Pipeline @@ -78,14 +82,19 @@ Sections 7, 8, 9 always apply. Section 10 applies to PRs touching 3+ files or ad - [ ] Model ID resolution chain not bypassed - [ ] WorkerSpec kept minimal at discovery time - [ ] No post-hoc ModelCard mutation — inject via labels before `build_model_card` +- [ ] ZMQ workers: the handshake is never awaited on a request or probe path (background connect driver + `WorkerConnected` signal); `activate` leaves them Pending; `update_worker_properties` re-applies `zmq_handshake_address` / `zmq_engine_group` / `connect_signal_tx`; misconfig is rejected in `create_worker`, not at connect time -### 4. Routing Policy +### 4. Routing Policy & Dispatch -- [ ] Works for both HTTP and gRPC paths (`SelectWorkerInfo`) +- [ ] Works for HTTP, gRPC and ZMQ backends (`SelectWorkerInfo`; `ConnectionMode` in `crates/protocols/src/worker.rs`, `BackendClient` in `routers/grpc/backend_client.rs`) - [ ] All state is `Send + Sync` (DashMap, Arc — no bare Mutex on hot paths) - [ ] No `.unwrap()` on worker slices — handle empty list -- [ ] Circuit breaker checked: `w.is_healthy() && w.circuit_breaker_can_execute()` +- [ ] Eligibility via `w.is_available()` (health + circuit breaker + overload veto) through `policies/mod.rs:get_healthy_worker_indices`; hash policies use `w.is_healthy_and_eligible()`. A hand-rolled `is_healthy() && circuit_breaker_can_execute()` drops the overload veto +- [ ] New selection/dispatch paths shed: `overload::shed_if_all_overloaded` on the candidate pool, `overload::shed_if_worker_overloaded` on the chosen worker (`routers/common/overload.rs`) - [ ] Policy implements `LoadBalancingPolicy` and is registered in `policies/factory.rs` with a `PolicyConfig` enum variant +- [ ] Dispatch memory: route methods take the parsed body by value and nothing outlives `RequestLease::release_dispatch()` (`routers/common/request_lease.rs`); upstream sends go through `attach_sized_body` / `serialize_json_sized` (`routers/common/mod.rs`) and upstream reads through `read_worker_body_capped` — no unbounded buffering +- [ ] Shed and other terminal responses call `retry::mark_non_retryable`; a declined streaming pass-through hands the request back with its body unconsumed +- [ ] No per-token `clone()` in the gRPC streaming path (`routers/grpc/regular/streaming.rs`, shared by Regular and PD; `routers/grpc/common/responses/streaming.rs` for the Responses SSE emitter) — `routers/grpc/common/stages/` is per-request, not per-token ### 5. Parser Changes (Tool / Reasoning) @@ -113,8 +122,11 @@ Sections 7, 8, 9 always apply. Section 10 applies to PRs touching 3+ files or ad - [ ] Unit tests for new types/parsing including error cases - [ ] Integration test for full flow - [ ] Existing test struct literals updated with new fields -- [ ] E2E tests if user-facing behavior changes (in `e2e_test/` — tests run sequentially with class-scoped backends) -- [ ] E2E test markers set: `@pytest.mark.model(...)`, `@pytest.mark.workers(count=, gpus=, extra_engine_args=)`, `@pytest.mark.gateway(policy=)` as needed (marker set is `model`/`workers`/`gateway`/`e2e`/`slow` — see `e2e_test/conftest.py`) +- [ ] E2E tests if user-facing behavior changes (in `e2e_test/`; `setup_backend` is class-scoped and items are ordered by (backend, model) so the session `WorkerPool` reuses workers) +- [ ] Every new e2e test carries `@pytest.mark.engine("sglang", "vllm", ...)` and `@pytest.mark.gpu(n)` — GPU lanes export `E2E_ENGINE`/`E2E_GPU_TIER` and `e2e_test/fixtures/hooks.py` deselects anything without a matching `engine` marker (a missing `gpu` marker defaults to `gpu(1)`, so it silently lands in the 1-GPU lanes instead; reported as an `e2e selection:` line; `E2E_MIN_SELECTED` fails the lane below its floor). Full marker set is registered in `hooks.py:pytest_configure` — `engine`, `vendor` (cloud), `gpu`, `model`, `skip_for_runtime`, `gateway(policy=, timeout=, extra_args=)`, `workers(count=, prefill=, decode=, gpus=, extra_engine_args=)`, `storage`, `external`, `e2e`, `slow`, `slowtest`, `nightly` — plus `kind` in `e2e_test/pyproject.toml` +- [ ] gRPC-pipeline changes also behave over ZMQ: the `e2e-*-chat-zmq*` lanes replay single-worker cases with `E2E_CONNECTION_MODE=zmq` (`hooks.py:_filter_zmq_items` drops PD/EPD/multi-worker there, so a ZMQ-only fix needs a single-worker chat case) +- [ ] Discovery changes (`service_discovery.rs`, multi-port) covered by the kind lane: `SMG_KIND_E2E=1 pytest e2e_test/kind_discovery --confcutdir e2e_test/kind_discovery -m kind` (Linux, needs kind/kubectl) +- [ ] Harness changes (`e2e_test/infra/`, `e2e_test/fixtures/`) pass the CPU gate: `PYTHONPATH=e2e_test pytest -q --noconftest e2e_test/infra e2e_test/fixtures` ### 9. Code Quality @@ -132,7 +144,8 @@ Sections 7, 8, 9 always apply. Section 10 applies to PRs touching 3+ files or ad - "Can existing config overrides or labels achieve this?" → may be unnecessary - "Does this compose with DP-aware mode, PD disagg, mesh HA?" → don't break existing - "Is this Send + Sync safe under concurrent load?" → all routing state thread-safe -- "Did I check both HTTP and gRPC paths?" → dual-mode is easy to forget +- "Did I check HTTP, gRPC and ZMQ paths?" → three transports now +- "Does it still hold for a ZMQ worker?" → ZMQ is host-local: excluded from mesh sync, PD/EPD legs, admin RPCs (flush/profile), KV events (cache_aware only warns), and dp-aware expansion (grouped engines use `dp_size` instead). A new cross-cutting feature must support the ZMQ lane or reject it loudly at registration See @anti-patterns.md for subsystem-specific anti-patterns. @@ -149,6 +162,6 @@ See @anti-patterns.md for subsystem-specific anti-patterns. - About to approve without mapping files to sections - Skipping a section because "it doesn't apply" without checking the mapping table -- No file:line citations in review comments +- No file:line citations in review comments, or comments missing the REVIEW.md severity marker - Approving a config change without verifying both conversion paths - Reviewing without fetching the actual diff first diff --git a/skills/review-pr/anti-patterns.md b/skills/review-pr/anti-patterns.md index 8570f62..7b5bfd7 100644 --- a/skills/review-pr/anti-patterns.md +++ b/skills/review-pr/anti-patterns.md @@ -8,24 +8,26 @@ Per-subsystem anti-patterns to check during PR review. |-------------|-------------|------------------| | Only wiring `to_router_config()`, missing `to_server_config()` | Field silently ignored in one code path | `grep -n "new_field" model_gateway/src/main.rs` — must appear in BOTH functions | | Missing `#[serde(default)]` on new optional field | Existing YAML configs fail to deserialize | Check all new `Option` fields in `config/types.rs` | -| No `value_parser` on CLI flag | Invalid values accepted, crash at runtime | Check `CliArgs` struct for new `#[clap]` fields | +| No `value_parser` on CLI flag | Invalid values accepted, crash at runtime | Check `CliArgs` in `main.rs` for new `#[arg(...)]` fields | ## Worker Lifecycle & Labels | Anti-Pattern | Consequence | What to Look For | |-------------|-------------|------------------| | Adding `_override` field to WorkerSpec | Bypasses label pipeline, creates parallel data path | New fields on `WorkerSpec` in `crates/protocols/src/worker.rs` | -| Post-hoc ModelCard mutation | Race conditions, stale data in routing | `model_card.model_id = ...` after `build_model_card()` | +| Post-hoc ModelCard mutation | Race conditions, stale data in routing | `card.id = ...` (or any field write) after `build_model_card()` in `workflow/steps/local/create_worker.rs` — inject via labels (`resolve_model_id`) instead | | Injecting K8s-specific data into `crates/protocols/` types | Tight coupling to K8s, breaks non-K8s deployments | New fields in `crates/protocols/` that reference namespaces, pods, labels | ## Routing | Anti-Pattern | Consequence | What to Look For | |-------------|-------------|------------------| -| Only testing HTTP path, missing gRPC | Feature breaks for gRPC backends | Test files that only use `RequestType::Http` | +| Only testing HTTP path, missing gRPC/ZMQ | Feature breaks for gRPC or ZMQ backends | Router tests whose workers are all `ConnectionMode::Http` (`crates/protocols/src/worker.rs:ConnectionMode`, `worker/builder.rs:connection_mode`); e2e tests parametrized on `setup_backend` with only one of `"grpc"`/`"http"` (ZMQ runs via the `E2E_CONNECTION_MODE=zmq` lane) | | Using `RwLock` on hot routing path | Contention under load | `RwLock` in routing policy structs (use `DashMap` instead) | | `.unwrap()` on empty worker slice | Panic when no workers available | `workers[0]` or `.unwrap()` on worker selection | -| Skipping circuit breaker check | Routing to unhealthy backends | Missing `w.is_healthy() && w.circuit_breaker_can_execute()` | +| Skipping the eligibility check, or re-deriving it | Routing to unhealthy or overload-vetoed backends | Missing `w.is_available()` (`w.is_healthy_and_eligible()` in hash policies); a hand-written `is_healthy() && circuit_breaker_can_execute()` drops the overload veto | +| Computing overload per request | Extra worker walk on the hot path | Overload is latched once per ingested load report (`worker/overload.rs`, `Worker::is_overloaded()`); selection must only read the flag | +| Holding request memory past the lease | Per-request memory doubles under load | Parsed request / `RoutingDerivatives` / body bytes cloned out of `RequestLease` (`routers/common/request_lease.rs`) into streaming tasks, or kept after `release_dispatch()` | ## Parsers (Tool / Reasoning) @@ -35,16 +37,26 @@ Per-subsystem anti-patterns to check during PR review. | Losing partial token prefix (e.g. ``) | Text silently dropped during streaming | Buffer handling when partial match fails | | Missing factory registration | Parser unreachable at runtime | New parser not in `ParserFactory::new()` | +## ZMQ Direct Backend + +| Anti-Pattern | Consequence | What to Look For | +|-------------|-------------|------------------| +| Reordering or inserting msgpack struct fields | Silent wire corruption against a real engine | `array_like` positional structs in `crates/engine_zmq_client/src/protocol/` (e.g. `vllm::EngineCoreRequest`) — field order is the wire contract, append only; cross-check the engine's Python msgspec struct | +| Open-ended transport dispatch | ZMQ silently takes the gRPC arm | `_ =>` arms over `BackendClient` / `ZmqDialect` (`routers/grpc/backend_client.rs`, `zmq_client.rs`); ZMQ arms bolted onto `GrpcClient` instead | +| Awaiting the handshake inline | Request or health path blocks for the whole connect timeout | `connect_handshake` reached from a request/probe path instead of `spawn_zmq_connect_driver` plus the `WorkerConnected` signal | +| Sending string stops over the ZMQ wire | EngineCore has no tokenizer — stops never fire, generation runs to the context window | Stop strings / EOS resolved anywhere but `BackendClient::finalize_generate_request` (`fold_tokenizer_eos_backstop`) — the wire is token-only | + ## Bindings | Anti-Pattern | Consequence | What to Look For | |-------------|-------------|------------------| -| Missing default in Python struct literal | `maturin develop` build failure | New config field without `new_field: None` in `bindings/python/src/lib.rs` | +| `RouterConfig` field not threaded through the Python chain | Compiles fine; Python launches silently run the Rust default | A new field needs a `model_gateway/src/config/builder.rs` setter, a `router_args.py` dataclass + argparse field, a `Router::new` pyo3 signature default, and a `.new_field(...)` call in `to_router_config()` (`bindings/python/src/lib.rs`) | +| Missing field in the `ServerConfig` struct literal | `maturin develop` build failure | `server::ServerConfig { .. }` in `bindings/python/src/lib.rs` has no `..Default::default()` | | Missing Go type mapping | Go SDK compile failure | New enum/type not mirrored in `bindings/golang/` | ## Storage | Anti-Pattern | Consequence | What to Look For | |-------------|-------------|------------------| -| Missing hook integration | Audit trail gaps | New backend without `on_write` / `on_delete` hook calls | +| Bypassing the hook wrappers | Audit trail gaps | New backend handed out raw instead of wrapped in `HookedConversationStorage` / `HookedConversationItemStorage` / `HookedResponseStorage` (`hooked.rs`); new `StorageOperation` without `before()`/`after()` coverage (`hooks.rs:StorageHook`) | | No schema migration | Data loss on upgrade | New fields without migration handling |