Synced from monorepo
Changes: - Classify clipboard delivery confidence - Add durable session update append - Scope the xAI session bearer to first-party memory embedding endpoints - Persist subagent outputs to disk and bound long-lived agent state - Add MiniSweAgent:bash for mini-swe-agent parity - Revert taking local sessions off the persistent shell - Contextual tip recommending grok wrap on SSH sessions - Voice STT bearer from model BYOK env_key/api_key - Define exact website policies for sandbox - Gate unsafe shell environments - Shared pin hoist; single require_sha gate for marketplace plugins - Server-signed is-managed claim (closes sidecar-removal downgrade) - Optional require_sha pin for remote plugin installs - Show session title and last exchange in the exit resume hint - Gate shell output redirects - Warn when fail_closed is present but not a boolean - Add canonical text editing core (ratatui-textarea) - Keep execution state out of goal scratch - Add acknowledged persistence primitives - Inherit child network restrictions in sandbox - Fail closed when hook matchers fail to recompile - Add MCP setup preferences for plugin MCPs - Gate sourced shell scripts - Gate file-typed project hooks - grok wrap: restore terminal modes on child death - Harden owner-only permissions on auth and MCP credentials - Create crash dump files with owner-only permissions - Write the agent_id cache owner-only (0600) - SessionMetrics mode skips Mixpanel profile sync - Dashboard: slim live-tail peek - Yank full queued prompt text, not (+N lines) - Defeat clock-rollback on the signed managed-config cache - Stop early session/cancel from overtaking the prompt and wedging the turn slot - Self-heal a diverged agent entrypoint on startup - Add matched inference expectations in test-support - Add AuthSingleFlight cancel/successor gap tests - Remove consumer from external OTEL allowlist and pin scrub coverage - Enable /copy in minimal mode - Surface capacity and API-key detail on 429 errors - Single-flight interactive auth - Fix PageUp/PageDown skipping lines behind sticky prompt header
This commit is contained in:
grokkybara[bot]
parent
8adf9013a0
commit
98c3b2438a
225 changed files with 18836 additions and 7156 deletions
|
|
@ -17,8 +17,9 @@ test-support surface.
|
|||
|
||||
| Module | What it provides |
|
||||
|--------|------------------|
|
||||
| `mock_server` | `MockInferenceServer` — `/v1/chat/completions`, `/v1/responses`, `/v1/messages`, `/v1/models`, `/v1/settings`, `/v1/user` on `127.0.0.1:0`. `/v1/models` entries are `MockModelEntry` (re-exported as `MockModel` for PTY tests): `new(id)` / `with_agent_type(id, ty)` plus chainable `with_api_backend`, `with_supports_backend_search(bool)` → `supportsBackendSearch`, `with_supports_reasoning_effort(bool)` → `supportsReasoningEffort`, `with_reasoning_effort(&str)` → `reasoningEffort`, `with_reasoning_efforts(Vec<Value>)` → `reasoningEfforts` (raw option tables/bare strings), all emitted top-level as `parse_remote_model_value` reads them. Three response modes for the inference endpoints, precedence **scripted > required-auth > mode**: (1) **echo** (default) streams `Echo: <last user message>`, whitespace-collapsing; (2) **fixed** via `set_response(text)`, byte-exact delta reconstruction (newlines preserved — fenced code blocks survive); (3) **scripted** via `enqueue_response`. Constructors (`start`, `start_with_models`, `start_with_required_auth`) return `anyhow::Result`. Settings are 404-until-set (`set_settings(impl Serialize)`, `preset_allow_access()` for the `{"allow_access": true}` gate); scripted `/v1/settings` one-shots (`enqueue_response`) take precedence over the steady-state value (stale-snapshot tests). `/v1/user` serves a minimal `UserInfo` whose `subscriptionTier` is controlled by `set_user_subscription_tier(Option<&str>)` (`None` = free); its log entries keep the query string (e.g. `/v1/user?include=subscription`) so subscription-check cadence is countable. Request log: `requests()` (`LogEntry` with body, `authorization`, full POST headers + `header(name)` accessor), `request_bodies()`, `request_count()`, `has_chat_completion_request()` / `has_responses_request()` (exact, per endpoint), `messages_request_count()`, `last_system_prompt()`, `request_log_summary()`. **Storage:** `POST /v1/storage` with flippable 401 (`set_storage_unauthorized`); accepted uploads via `storage_uploads()` → `StorageUpload { path, size, body, authorization }` (`body` retained up to 256 KiB, empty above; `authorization` is the raw header). Runtime knobs: `set_models`, `set_messages_stop_reason`. Shuts down on drop. |
|
||||
| `scripted` | Data-only scripting (no axum types in the surface): `SseEvent { event, data }` (`::data`, `::with_event`), `ScriptedBody::{Json, Sse, Raw}` (`Raw` = byte-controllable malformed SSE), `ScriptedResponse { status, headers, body }` (`::sse`, `::json`, `::text`). `enqueue_response(path, response)` queues FIFO **per path**; consumed by the three inference endpoints, falling back to the active mode when empty. Scripted SSE bodies honor the server's `set_chunk_delay` pacing, same as the echo/fixed modes. Validation is eager — bad status/header panics at the enqueue call site. Seed of the phase-2 scriptable-mock format. |
|
||||
| `inference_override` | Typed request matching and response precedence shared by all inference routes: endpoint + foreground/auxiliary classification, named expectation state, overlapping-duplicate fingerprint replay, per-expectation barriers, compatibility FIFO dispatch, auth rejection, and compatibility completion-gate policy. The module is crate-private; only `InferenceEndpoint`, `InferenceRequestMatcher`, and `InferenceExpectation` are re-exported. |
|
||||
| `mock_server` | `MockInferenceServer` — `/v1/chat/completions`, `/v1/responses`, `/v1/messages`, `/v1/models`, `/v1/settings`, `/v1/user` on `127.0.0.1:0`. `/v1/models` entries are `MockModelEntry` (re-exported as `MockModel` for PTY tests): `new(id)` / `with_agent_type(id, ty)` plus chainable `with_api_backend`, `with_supports_backend_search(bool)` → `supportsBackendSearch`, `with_supports_reasoning_effort(bool)` → `supportsReasoningEffort`, `with_reasoning_effort(&str)` → `reasoningEffort`, `with_reasoning_efforts(Vec<Value>)` → `reasoningEfforts` (raw option tables/bare strings), all emitted top-level as `parse_remote_model_value` reads them. Inference precedence is **matched expectation > compatibility FIFO > required-auth > echo/fixed mode**. Register a uniquely named response with `expect_response(name, InferenceRequestMatcher::{foreground,auxiliary}(InferenceEndpoint::{ChatCompletions,Responses,Messages}), ScriptedResponse)` or `expect_response_blocked`; duplicate names fail at registration and requests atomically claim one matching expectation. Overlapping duplicate requests replay by a deterministic fingerprint of endpoint, request kind, non-empty `x-grok-req-id`, and serialized request body; tool-result follow-ups reuse the turn id but change the body, so they claim the next expectation. Production exposes no explicit HTTP attempt/model-call identity, so completed sequential retries are intentionally not inferred from timing: after the active shared call settles, an identical request claims the next expectation. A foreground request normally carries a non-empty `x-grok-turn-idx`; a non-turn non-empty `x-grok-req-id` is auxiliary even if it uses tools, and empty headers fall through to the 2+-tool compatibility heuristic. The returned `InferenceExpectation` has watch-backed `wait_received`, `wait_blocked`, `release`, `wait_satisfied`, `is_satisfied`, and `assert_satisfied` lifecycle operations. `release` only opens the barrier; response-body/stream-owned RAII publishes `Satisfied` only when the primary crosses terminal and every active overlapping copy settles. Primary cancellation cleans up without satisfaction or replay retention, and dropping a handle safely releases blocked work. Echo (default) streams `Echo: <last user message>` and fixed mode via `set_response(text)` reconstructs bytes exactly. Constructors (`start`, `start_with_models`, `start_with_required_auth`) return `anyhow::Result`. Settings are 404-until-set (`set_settings(impl Serialize)`, `preset_allow_access()` for the `{"allow_access": true}` gate); scripted `/v1/settings` one-shots (`enqueue_response`) take precedence over the steady-state value (stale-snapshot tests). `/v1/user` serves a minimal `UserInfo` whose `subscriptionTier` is controlled by `set_user_subscription_tier(Option<&str>)` (`None` = free); its log entries keep the query string (e.g. `/v1/user?include=subscription`) so subscription-check cadence is countable. Request log: `requests()` (`LogEntry` with body, `authorization`, full POST headers + `header(name)` accessor), `request_bodies()`, `request_count()`, `has_chat_completion_request()` / `has_responses_request()` (exact, per endpoint), `messages_request_count()`, `last_system_prompt()`, `request_log_summary()`. **Storage:** `POST /v1/storage` with flippable 401 (`set_storage_unauthorized`); accepted uploads via `storage_uploads()` → `StorageUpload { path, size, body, authorization }` (`body` retained up to 256 KiB, empty above; `authorization` is the raw header). Runtime knobs: `set_models`, `set_messages_stop_reason`. Shuts down on drop. |
|
||||
| `scripted` | Data-only response bodies (no axum types in the public surface): `SseEvent { event, data }` (`::data`, `::with_event`), `ScriptedBody::{Json, Sse, Raw}` (`Raw` = byte-controllable malformed SSE), `ScriptedResponse { status, headers, body }` (`::sse`, `::json`, `::text`). Prefer request-matched expectations for inference calls; `enqueue_response(path, response)` remains a compatibility FIFO per path and is still used for non-inference one-shots such as `/v1/settings`. Scripted SSE honors `set_chunk_delay`; matched JSON, raw, SSE, and even empty SSE bodies all honor per-expectation completion barriers. The compatibility `hold_agent_completions` gate also covers foreground scripted SSE on all three inference endpoints. Validation is eager — bad status/header panics at registration. |
|
||||
| `sse` | The three wire formats as event-list builders: `chat_completion_events` / `responses_api_events` / `messages_api_events(text, model, stop_reason)` (echo-style, whitespace-collapsing) plus byte-exact variants `chat_completion_events_exact` / `responses_api_events_exact` (messages is single-delta, byte-exact by construction). The exact/echo split is load-bearing — see the in-module byte-exactness tests. Also the scripted-scenario builders returning `SseEvent`s (for `ScriptedResponse::sse`): `responses_api_reasoning_only_events(reasoning, model)` — reasoning summary deltas completing with a `reasoning` item but no message/output-text, so the shell collector classifies the turn `EmptyReason::ReasoningOnly` (the model-doomloop trigger); `responses_api_reasoning_and_text_events(reasoning, text, model)` — reasoning deltas then a normal text answer (the ordinary reasoning-model turn); `responses_api_reasoning_then_tool_call_events(reasoning, call_id, name, arguments, model)` + its Chat Completions twin `chat_completions_reasoning_then_tool_call_events(...)` — reasoning deltas then one tool call (the think-then-call turn whose tool call finishes the thought and keeps the turn non-empty); the doom-loop check trio: `responses_api_doom_loop_check_events(triggers, reasoning, model)` — a doomed reasoning-only turn with NAMED `response.doom_loop_check` frames re-sent per cumulative prefix of `triggers` plus the terminal `doom_loop_check.triggers` copy on `response.completed`, `responses_api_doom_loop_terminal_only_events(triggers, reasoning, text, model)` — a normal answer whose terminal response alone carries the field, and `responses_api_with_doom_loop_frame(check_frame_data, reasoning, text, model)` — splices one named check frame with a caller-supplied payload (byte-exact `xai_grok_sampling_types::doom_loop::SAMPLE_CHECK_EVENT_DATA{,_CUMULATIVE}` fixtures or malformed variants) into an ordinary turn. |
|
||||
| `acp_client` | `GrokStdioClient` — drives `grok agent stdio` over real pipes through `agent-client-protocol`: spawn variants (`spawn`, `spawn_with_home`, `spawn_with_home_and_env`, `spawn_with_home_env_and_args`), initialize/authenticate, session create/load, prompt, `*_with_timeout` wrappers, captured text + stderr. `RawStdioClient` — raw-wire sibling for bytes the typed `ClientSideConnection` can never produce (escaped-slash methods `"session\/prompt"`, string UUID ids — the Xcode/Foundation shape): `send_line` writes a line verbatim; `response_for_id` matches the response by exact string id (the match IS the id-echo assertion), skips notifications, auto-refuses agent→client requests with `-32601`, and panics on timeout with skipped-traffic diagnostics (count + last lines; `0 other messages` = true silence). Both spawn through one hermetic `spawn_agent_process` (sandbox env + debug-log kill-list exists once) atop `process::spawn_piped_with_stderr_capture` (crate-internal `process` module: pipes, `kill_on_drop`, stderr drain — also used by `leader::LeaderStdioClient`). |
|
||||
| `headless` | `run_headless(server, args, cwd)` / `run_headless_with_cmd(cmd)` → `HeadlessResult { status, stdout, stderr, timed_out }` (60s cap), `assert_headless_success`, `assert_no_crashes` (panic/SIGSEGV/linker patterns), `stderr_tail`. |
|
||||
|
|
@ -56,9 +57,16 @@ that models a specific completion the echo/fixed modes can't express, e.g.
|
|||
`ScriptedResponse::sse`, needs no `mock_server` mode wiring, and ships with an
|
||||
in-module shape test asserting its event shape.
|
||||
|
||||
**An expectation matcher** (`inference_override.rs`): keep the public matcher
|
||||
typed and narrow. Claim under the single expectation-state mutex before
|
||||
serving, replay only overlapping active duplicates by model-call fingerprint,
|
||||
and add focused tests for auxiliary non-consumption, concurrent one-claim
|
||||
behavior, lifecycle barriers, and useful unsatisfied diagnostics. Expectations
|
||||
and compatibility scripts must remain ahead of required auth and fallback modes.
|
||||
|
||||
**A scripted-body kind** (`scripted.rs`): new `ScriptedBody` variant + render
|
||||
arm in `into_response_paced` + eager checks in `validate` if the data can be
|
||||
invalid. Add an in-crate test asserting client-visible bytes (the `Raw`
|
||||
byte-exactness test is the template) and keep
|
||||
`scripted_response_takes_precedence_over_required_auth` green — precedence is
|
||||
part of the contract.
|
||||
byte-exactness test is the template), exercise terminal gating for the new
|
||||
body, and keep `scripted_response_takes_precedence_over_required_auth` green —
|
||||
precedence is part of the contract.
|
||||
|
|
|
|||
638
crates/codegen/xai-grok-test-support/src/inference_override.rs
Normal file
638
crates/codegen/xai-grok-test-support/src/inference_override.rs
Normal file
|
|
@ -0,0 +1,638 @@
|
|||
use std::collections::{HashMap, VecDeque};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use axum::Json;
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::scripted::{BoxWait, ScriptedResponse, TerminalWait};
|
||||
|
||||
/// Inference endpoint matched by a scripted expectation.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum InferenceEndpoint {
|
||||
ChatCompletions,
|
||||
Responses,
|
||||
Messages,
|
||||
}
|
||||
|
||||
impl InferenceEndpoint {
|
||||
pub(crate) fn path(self) -> &'static str {
|
||||
match self {
|
||||
Self::ChatCompletions => "/v1/chat/completions",
|
||||
Self::Responses => "/v1/responses",
|
||||
Self::Messages => "/v1/messages",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Coarse request kind used to keep auxiliary calls from stealing turn scripts.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
enum InferenceRequestKind {
|
||||
Foreground,
|
||||
Auxiliary,
|
||||
}
|
||||
|
||||
impl InferenceRequestKind {
|
||||
fn classify(headers: &HeaderMap, body: &Value) -> Self {
|
||||
if nonempty_header(headers, "x-grok-turn-idx").is_some() {
|
||||
return Self::Foreground;
|
||||
}
|
||||
if nonempty_header(headers, "x-grok-req-id").is_some() {
|
||||
return Self::Auxiliary;
|
||||
}
|
||||
if body
|
||||
.get("tools")
|
||||
.and_then(Value::as_array)
|
||||
.is_some_and(|tools| tools.len() >= 2)
|
||||
{
|
||||
Self::Foreground
|
||||
} else {
|
||||
Self::Auxiliary
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Typed match criteria for one named inference response.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct InferenceRequestMatcher {
|
||||
endpoint: InferenceEndpoint,
|
||||
kind: InferenceRequestKind,
|
||||
}
|
||||
|
||||
impl InferenceRequestMatcher {
|
||||
/// Match a user-facing agent turn on the selected endpoint.
|
||||
pub fn foreground(endpoint: InferenceEndpoint) -> Self {
|
||||
Self {
|
||||
endpoint,
|
||||
kind: InferenceRequestKind::Foreground,
|
||||
}
|
||||
}
|
||||
|
||||
/// Match title, classifier, prompt-suggestion, or other side-channel work.
|
||||
pub fn auxiliary(endpoint: InferenceEndpoint) -> Self {
|
||||
Self {
|
||||
endpoint,
|
||||
kind: InferenceRequestKind::Auxiliary,
|
||||
}
|
||||
}
|
||||
|
||||
fn matches(self, endpoint: InferenceEndpoint, kind: InferenceRequestKind) -> bool {
|
||||
self.endpoint == endpoint && self.kind == kind
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum ExpectationPhase {
|
||||
Pending,
|
||||
Received,
|
||||
Blocked,
|
||||
Satisfied,
|
||||
}
|
||||
|
||||
struct ExpectationControl {
|
||||
name: String,
|
||||
phase_tx: tokio::sync::watch::Sender<ExpectationPhase>,
|
||||
claims_tx: tokio::sync::watch::Sender<usize>,
|
||||
release_tx: tokio::sync::watch::Sender<bool>,
|
||||
}
|
||||
|
||||
impl ExpectationControl {
|
||||
fn set_phase(&self, phase: ExpectationPhase) {
|
||||
self.phase_tx.send_replace(phase);
|
||||
}
|
||||
|
||||
fn release(&self) {
|
||||
self.release_tx.send_replace(true);
|
||||
}
|
||||
|
||||
fn claim(&self) {
|
||||
self.claims_tx.send_modify(|claims| *claims += 1);
|
||||
}
|
||||
|
||||
async fn wait_for_release(&self) {
|
||||
let mut release_rx = self.release_tx.subscribe();
|
||||
if *release_rx.borrow_and_update() {
|
||||
return;
|
||||
}
|
||||
release_rx
|
||||
.wait_for(|released| *released)
|
||||
.await
|
||||
.expect("expectation release sender lives with the claimed response");
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
async fn wait_claims(&self, target: usize) {
|
||||
let mut claims_rx = self.claims_tx.subscribe();
|
||||
claims_rx
|
||||
.wait_for(|claims| *claims >= target)
|
||||
.await
|
||||
.expect("expectation claims sender lives with the control");
|
||||
}
|
||||
}
|
||||
|
||||
/// Deterministic lifecycle handle for one registered inference expectation.
|
||||
#[must_use = "expectation handles provide synchronization and satisfaction checks"]
|
||||
pub struct InferenceExpectation {
|
||||
control: Arc<ExpectationControl>,
|
||||
phase_rx: tokio::sync::watch::Receiver<ExpectationPhase>,
|
||||
}
|
||||
|
||||
impl InferenceExpectation {
|
||||
pub fn name(&self) -> &str {
|
||||
&self.control.name
|
||||
}
|
||||
|
||||
pub fn is_satisfied(&self) -> bool {
|
||||
*self.phase_rx.borrow() == ExpectationPhase::Satisfied
|
||||
}
|
||||
|
||||
/// Wait until one request atomically claims this expectation.
|
||||
pub async fn wait_received(&mut self) {
|
||||
self.wait_for(ExpectationPhase::Received).await;
|
||||
}
|
||||
|
||||
/// Wait until the response reaches its terminal-event barrier.
|
||||
pub async fn wait_blocked(&mut self) {
|
||||
self.wait_for(ExpectationPhase::Blocked).await;
|
||||
}
|
||||
|
||||
/// Wait until the primary response pipeline crosses its terminal boundary.
|
||||
pub async fn wait_satisfied(&mut self) {
|
||||
self.wait_for(ExpectationPhase::Satisfied).await;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) async fn wait_claims(&self, target: usize) {
|
||||
self.control.wait_claims(target).await;
|
||||
}
|
||||
|
||||
/// Release this expectation's terminal barrier.
|
||||
pub fn release(&self) {
|
||||
self.control.release();
|
||||
}
|
||||
|
||||
/// Panic with the expectation name and lifecycle state unless satisfied.
|
||||
pub fn assert_satisfied(&self) {
|
||||
assert!(
|
||||
self.is_satisfied(),
|
||||
"inference expectation `{}` was not satisfied (state: {:?})",
|
||||
self.name(),
|
||||
*self.phase_rx.borrow()
|
||||
);
|
||||
}
|
||||
|
||||
/// Describe the expectation for aggregation in test failure output.
|
||||
pub fn diagnostic(&self) -> String {
|
||||
format!(
|
||||
"inference expectation `{}` (state: {:?})",
|
||||
self.name(),
|
||||
*self.phase_rx.borrow()
|
||||
)
|
||||
}
|
||||
|
||||
async fn wait_for(&mut self, target: ExpectationPhase) {
|
||||
if self
|
||||
.phase_rx
|
||||
.wait_for(|phase| Self::phase_reached(*phase, target))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
panic!(
|
||||
"inference expectation `{}` closed before reaching {target:?} (state: {:?})",
|
||||
self.control.name,
|
||||
*self.phase_rx.borrow()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn phase_reached(current: ExpectationPhase, target: ExpectationPhase) -> bool {
|
||||
match target {
|
||||
ExpectationPhase::Pending => true,
|
||||
ExpectationPhase::Received => current != ExpectationPhase::Pending,
|
||||
ExpectationPhase::Blocked => matches!(
|
||||
current,
|
||||
ExpectationPhase::Blocked | ExpectationPhase::Satisfied
|
||||
),
|
||||
ExpectationPhase::Satisfied => current == ExpectationPhase::Satisfied,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for InferenceExpectation {
|
||||
fn drop(&mut self) {
|
||||
self.control.release();
|
||||
}
|
||||
}
|
||||
|
||||
struct PendingExpectation {
|
||||
matcher: InferenceRequestMatcher,
|
||||
response: ScriptedResponse,
|
||||
block_before_terminal: bool,
|
||||
control: Arc<ExpectationControl>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
struct ModelCallFingerprint {
|
||||
endpoint: InferenceEndpoint,
|
||||
kind: InferenceRequestKind,
|
||||
request_id: String,
|
||||
body: String,
|
||||
}
|
||||
|
||||
struct CallState {
|
||||
response: ScriptedResponse,
|
||||
block_before_terminal: bool,
|
||||
control: Arc<ExpectationControl>,
|
||||
active: usize,
|
||||
primary_crossed_terminal: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ExpectationState {
|
||||
pending: VecDeque<PendingExpectation>,
|
||||
in_flight: HashMap<ModelCallFingerprint, CallState>,
|
||||
}
|
||||
|
||||
type Expectations = Arc<std::sync::Mutex<ExpectationState>>;
|
||||
type ScriptQueues = Arc<std::sync::Mutex<HashMap<String, VecDeque<ScriptedResponse>>>>;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct InferenceOverrides {
|
||||
expectations: Expectations,
|
||||
scripted: ScriptQueues,
|
||||
completion_gate: Arc<CompletionGate>,
|
||||
required_token: Option<Arc<str>>,
|
||||
}
|
||||
|
||||
impl InferenceOverrides {
|
||||
pub(crate) fn new(required_token: Option<String>) -> Self {
|
||||
Self {
|
||||
expectations: Arc::new(std::sync::Mutex::new(ExpectationState::default())),
|
||||
scripted: Arc::new(std::sync::Mutex::new(HashMap::new())),
|
||||
completion_gate: Arc::new(CompletionGate::default()),
|
||||
required_token: required_token.map(Arc::from),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn classify(
|
||||
&self,
|
||||
endpoint: InferenceEndpoint,
|
||||
headers: &HeaderMap,
|
||||
body: &Value,
|
||||
) -> ClassifiedInferenceRequest {
|
||||
let kind = InferenceRequestKind::classify(headers, body);
|
||||
let fingerprint =
|
||||
nonempty_header(headers, "x-grok-req-id").map(|request_id| ModelCallFingerprint {
|
||||
endpoint,
|
||||
kind,
|
||||
request_id: request_id.to_owned(),
|
||||
body: serde_json::to_string(body).expect("serialize inference request fingerprint"),
|
||||
});
|
||||
ClassifiedInferenceRequest {
|
||||
endpoint,
|
||||
kind,
|
||||
fingerprint,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn response_override(
|
||||
&self,
|
||||
request: &ClassifiedInferenceRequest,
|
||||
headers: &HeaderMap,
|
||||
delay: Option<Duration>,
|
||||
) -> Option<Response> {
|
||||
if let Some(claimed) = self.claim_expectation(request) {
|
||||
let (response, wait) = claimed.into_parts();
|
||||
return Some(response.into_response_paced(delay, Some(wait)).await);
|
||||
}
|
||||
|
||||
if let Some(response) = self.pop_scripted(request.endpoint.path()) {
|
||||
let wait =
|
||||
(request.is_foreground() && response.is_sse()).then(|| self.global_terminal_wait());
|
||||
return Some(response.into_response_paced(delay, wait).await);
|
||||
}
|
||||
|
||||
self.auth_rejection(headers)
|
||||
}
|
||||
|
||||
pub(crate) fn register_expectation(
|
||||
&self,
|
||||
name: impl Into<String>,
|
||||
matcher: InferenceRequestMatcher,
|
||||
response: ScriptedResponse,
|
||||
block_before_terminal: bool,
|
||||
) -> InferenceExpectation {
|
||||
response.validate();
|
||||
let name = name.into();
|
||||
let mut expectations = self.expectations.lock().unwrap();
|
||||
assert!(
|
||||
expectations
|
||||
.pending
|
||||
.iter()
|
||||
.all(|expectation| expectation.control.name != name)
|
||||
&& expectations
|
||||
.in_flight
|
||||
.values()
|
||||
.all(|expectation| expectation.control.name != name),
|
||||
"duplicate inference expectation name `{name}`"
|
||||
);
|
||||
let (phase_tx, phase_rx) = tokio::sync::watch::channel(ExpectationPhase::Pending);
|
||||
let (claims_tx, _claims_rx) = tokio::sync::watch::channel(0);
|
||||
let (release_tx, _release_rx) = tokio::sync::watch::channel(!block_before_terminal);
|
||||
let control = Arc::new(ExpectationControl {
|
||||
name,
|
||||
phase_tx,
|
||||
claims_tx,
|
||||
release_tx,
|
||||
});
|
||||
expectations.pending.push_back(PendingExpectation {
|
||||
matcher,
|
||||
response,
|
||||
block_before_terminal,
|
||||
control: control.clone(),
|
||||
});
|
||||
InferenceExpectation { control, phase_rx }
|
||||
}
|
||||
|
||||
pub(crate) fn enqueue_response(&self, path: impl Into<String>, response: ScriptedResponse) {
|
||||
response.validate();
|
||||
self.scripted
|
||||
.lock()
|
||||
.unwrap()
|
||||
.entry(path.into())
|
||||
.or_default()
|
||||
.push_back(response);
|
||||
}
|
||||
|
||||
pub(crate) fn pop_scripted(&self, path: &str) -> Option<ScriptedResponse> {
|
||||
self.scripted
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get_mut(path)
|
||||
.and_then(VecDeque::pop_front)
|
||||
}
|
||||
|
||||
pub(crate) fn fallback_terminal_wait(
|
||||
&self,
|
||||
request: &ClassifiedInferenceRequest,
|
||||
) -> Option<TerminalWait> {
|
||||
request.is_foreground().then(|| self.global_terminal_wait())
|
||||
}
|
||||
|
||||
pub(crate) fn hold_completions(&self) {
|
||||
self.completion_gate.hold();
|
||||
}
|
||||
|
||||
pub(crate) fn release_completions(&self) {
|
||||
self.completion_gate.release();
|
||||
}
|
||||
|
||||
fn claim_expectation(
|
||||
&self,
|
||||
request: &ClassifiedInferenceRequest,
|
||||
) -> Option<ClaimedExpectation> {
|
||||
let mut expectations = self.expectations.lock().unwrap();
|
||||
if let Some(fingerprint) = request.fingerprint.as_ref()
|
||||
&& let Some(call) = expectations.in_flight.get_mut(fingerprint)
|
||||
&& call.active > 0
|
||||
{
|
||||
call.active += 1;
|
||||
call.control.claim();
|
||||
return Some(ClaimedExpectation {
|
||||
response: call.response.clone(),
|
||||
lease: ClaimLease::new(
|
||||
self.expectations.clone(),
|
||||
Some(fingerprint.clone()),
|
||||
call.control.clone(),
|
||||
call.block_before_terminal,
|
||||
ClaimRole::Replay,
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
let index = expectations
|
||||
.pending
|
||||
.iter()
|
||||
.position(|expectation| expectation.matcher.matches(request.endpoint, request.kind))?;
|
||||
let expectation = expectations
|
||||
.pending
|
||||
.remove(index)
|
||||
.expect("matched expectation index must remain valid");
|
||||
expectation.control.set_phase(ExpectationPhase::Received);
|
||||
expectation.control.claim();
|
||||
let lease = ClaimLease::new(
|
||||
self.expectations.clone(),
|
||||
request.fingerprint.clone(),
|
||||
expectation.control.clone(),
|
||||
expectation.block_before_terminal,
|
||||
ClaimRole::Primary,
|
||||
);
|
||||
if let Some(fingerprint) = request.fingerprint.clone() {
|
||||
let replaced = expectations.in_flight.insert(
|
||||
fingerprint,
|
||||
CallState {
|
||||
response: expectation.response.clone(),
|
||||
block_before_terminal: expectation.block_before_terminal,
|
||||
control: expectation.control,
|
||||
active: 1,
|
||||
primary_crossed_terminal: false,
|
||||
},
|
||||
);
|
||||
assert!(
|
||||
replaced.is_none(),
|
||||
"duplicate in-flight model-call fingerprint"
|
||||
);
|
||||
}
|
||||
Some(ClaimedExpectation {
|
||||
response: expectation.response,
|
||||
lease,
|
||||
})
|
||||
}
|
||||
|
||||
fn auth_rejection(&self, headers: &HeaderMap) -> Option<Response> {
|
||||
let expected = self.required_token.as_deref()?;
|
||||
let valid = headers
|
||||
.get("authorization")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.is_some_and(|value| {
|
||||
value
|
||||
.strip_prefix("Bearer ")
|
||||
.or_else(|| value.strip_prefix("bearer "))
|
||||
.is_some_and(|token| token == expected)
|
||||
});
|
||||
if valid {
|
||||
return None;
|
||||
}
|
||||
Some(
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
Json(json!({
|
||||
"error": "missing API key; set the x-api-key header or Authorization: Bearer header"
|
||||
})),
|
||||
)
|
||||
.into_response(),
|
||||
)
|
||||
}
|
||||
|
||||
fn global_terminal_wait(&self) -> TerminalWait {
|
||||
let completion_gate = self.completion_gate.clone();
|
||||
Box::new(move || Box::pin(async move { completion_gate.wait_if_held().await }))
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct ClassifiedInferenceRequest {
|
||||
endpoint: InferenceEndpoint,
|
||||
kind: InferenceRequestKind,
|
||||
fingerprint: Option<ModelCallFingerprint>,
|
||||
}
|
||||
|
||||
impl ClassifiedInferenceRequest {
|
||||
pub(crate) fn is_foreground(&self) -> bool {
|
||||
self.kind == InferenceRequestKind::Foreground
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
enum ClaimRole {
|
||||
Primary,
|
||||
Replay,
|
||||
}
|
||||
|
||||
struct ClaimedExpectation {
|
||||
response: ScriptedResponse,
|
||||
lease: ClaimLease,
|
||||
}
|
||||
|
||||
impl ClaimedExpectation {
|
||||
fn into_parts(self) -> (ScriptedResponse, TerminalWait) {
|
||||
let response = self.response;
|
||||
let mut lease = self.lease;
|
||||
let wait = Box::new(move || {
|
||||
Box::pin(async move {
|
||||
if lease.block_before_terminal {
|
||||
lease.mark_blocked();
|
||||
lease.control.wait_for_release().await;
|
||||
}
|
||||
lease.crossed_terminal = true;
|
||||
lease.finish();
|
||||
}) as BoxWait
|
||||
});
|
||||
(response, wait)
|
||||
}
|
||||
}
|
||||
|
||||
struct ClaimLease {
|
||||
expectations: Expectations,
|
||||
fingerprint: Option<ModelCallFingerprint>,
|
||||
control: Arc<ExpectationControl>,
|
||||
block_before_terminal: bool,
|
||||
role: ClaimRole,
|
||||
crossed_terminal: bool,
|
||||
finished: bool,
|
||||
}
|
||||
|
||||
impl ClaimLease {
|
||||
fn new(
|
||||
expectations: Expectations,
|
||||
fingerprint: Option<ModelCallFingerprint>,
|
||||
control: Arc<ExpectationControl>,
|
||||
block_before_terminal: bool,
|
||||
role: ClaimRole,
|
||||
) -> Self {
|
||||
Self {
|
||||
expectations,
|
||||
fingerprint,
|
||||
control,
|
||||
block_before_terminal,
|
||||
role,
|
||||
crossed_terminal: false,
|
||||
finished: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn mark_blocked(&self) {
|
||||
if matches!(&self.role, ClaimRole::Primary)
|
||||
&& *self.control.phase_tx.borrow() != ExpectationPhase::Satisfied
|
||||
{
|
||||
self.control.set_phase(ExpectationPhase::Blocked);
|
||||
}
|
||||
}
|
||||
|
||||
fn finish(&mut self) {
|
||||
if self.finished {
|
||||
return;
|
||||
}
|
||||
self.finished = true;
|
||||
self.update_shared_state();
|
||||
}
|
||||
|
||||
fn update_shared_state(&self) {
|
||||
let Some(fingerprint) = self.fingerprint.as_ref() else {
|
||||
if matches!(&self.role, ClaimRole::Primary) && self.crossed_terminal {
|
||||
self.control.set_phase(ExpectationPhase::Satisfied);
|
||||
}
|
||||
return;
|
||||
};
|
||||
let mut expectations = self.expectations.lock().unwrap();
|
||||
let Some(call) = expectations.in_flight.get_mut(fingerprint) else {
|
||||
return;
|
||||
};
|
||||
assert!(call.active > 0, "claim active count underflow");
|
||||
call.active -= 1;
|
||||
if matches!(&self.role, ClaimRole::Primary) && self.crossed_terminal {
|
||||
call.primary_crossed_terminal = true;
|
||||
}
|
||||
if call.active == 0 {
|
||||
let control = call.control.clone();
|
||||
let satisfied = call.primary_crossed_terminal;
|
||||
expectations.in_flight.remove(fingerprint);
|
||||
if satisfied {
|
||||
control.set_phase(ExpectationPhase::Satisfied);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ClaimLease {
|
||||
fn drop(&mut self) {
|
||||
self.finish();
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct CompletionGate {
|
||||
held: AtomicBool,
|
||||
notify: tokio::sync::Notify,
|
||||
}
|
||||
|
||||
impl CompletionGate {
|
||||
fn hold(&self) {
|
||||
self.held.store(true, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
fn release(&self) {
|
||||
self.held.store(false, Ordering::SeqCst);
|
||||
self.notify.notify_waiters();
|
||||
}
|
||||
|
||||
async fn wait_if_held(&self) {
|
||||
loop {
|
||||
let notified = self.notify.notified();
|
||||
if !self.held.load(Ordering::SeqCst) {
|
||||
return;
|
||||
}
|
||||
notified.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn nonempty_header<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> {
|
||||
headers
|
||||
.get(name)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
|
@ -23,6 +23,7 @@ pub mod acp_client;
|
|||
pub mod counting_server;
|
||||
pub mod env;
|
||||
pub mod headless;
|
||||
mod inference_override;
|
||||
#[cfg(unix)]
|
||||
pub mod leader;
|
||||
pub mod mock_server;
|
||||
|
|
@ -38,4 +39,5 @@ pub use headless::{
|
|||
HeadlessResult, assert_headless_success, assert_no_crashes, run_headless,
|
||||
run_headless_with_cmd, stderr_tail,
|
||||
};
|
||||
pub use inference_override::{InferenceEndpoint, InferenceExpectation, InferenceRequestMatcher};
|
||||
pub use mock_server::{MockInferenceServer, MockModelEntry, ScriptedResponse, SseEvent};
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -3,6 +3,8 @@
|
|||
//! time. Pure data — no router or handler types in the public surface.
|
||||
|
||||
use std::convert::Infallible;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use axum::Json;
|
||||
use axum::http::{HeaderName, HeaderValue, StatusCode};
|
||||
|
|
@ -11,6 +13,9 @@ use axum::response::{IntoResponse, Response};
|
|||
use futures_util::stream;
|
||||
use serde_json::Value;
|
||||
|
||||
pub(crate) type BoxWait = Pin<Box<dyn Future<Output = ()> + Send>>;
|
||||
pub(crate) type TerminalWait = Box<dyn FnOnce() -> BoxWait + Send>;
|
||||
|
||||
/// One SSE event as data: optional `event:` name plus the `data:` payload.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SseEvent {
|
||||
|
|
@ -45,9 +50,8 @@ pub enum ScriptedBody {
|
|||
Raw(String),
|
||||
}
|
||||
|
||||
/// A scripted reply for a single request on one path, consumed FIFO.
|
||||
/// Takes precedence over the response mode AND the required-auth check —
|
||||
/// a script is full control over the next reply.
|
||||
/// A scripted reply served by a matched expectation or compatibility FIFO.
|
||||
/// Scripted replies take precedence over required auth and fallback modes.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ScriptedResponse {
|
||||
pub status: u16,
|
||||
|
|
@ -83,6 +87,10 @@ impl ScriptedResponse {
|
|||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_sse(&self) -> bool {
|
||||
matches!(self.body, ScriptedBody::Sse(_))
|
||||
}
|
||||
|
||||
/// Validate status and headers eagerly so a bad script panics at the
|
||||
/// enqueue call site rather than far away at serve time.
|
||||
pub(crate) fn validate(&self) {
|
||||
|
|
@ -93,32 +101,60 @@ impl ScriptedResponse {
|
|||
}
|
||||
}
|
||||
|
||||
/// Render to HTTP with SSE events paced by `delay` (sleep before each
|
||||
/// event, mirroring the fixed/echo `paced_events` pacing) so
|
||||
/// `set_chunk_delay` also holds scripted turns open. `None` streams
|
||||
/// instantly. Non-SSE bodies ignore the delay.
|
||||
pub(crate) fn into_response_paced(self, delay: Option<std::time::Duration>) -> Response {
|
||||
use futures_util::StreamExt as _;
|
||||
/// Render to HTTP with SSE events paced by `delay` and optional terminal
|
||||
/// completion gating. Non-SSE bodies wait before returning so every body
|
||||
/// mode obeys the same release barrier.
|
||||
pub(crate) async fn into_response_paced(
|
||||
self,
|
||||
delay: Option<std::time::Duration>,
|
||||
before_terminal: Option<TerminalWait>,
|
||||
) -> Response {
|
||||
let mut resp = match self.body {
|
||||
ScriptedBody::Json(v) => Json(v).into_response(),
|
||||
ScriptedBody::Raw(s) => s.into_response(),
|
||||
ScriptedBody::Json(v) => {
|
||||
if let Some(wait) = before_terminal {
|
||||
wait().await;
|
||||
}
|
||||
Json(v).into_response()
|
||||
}
|
||||
ScriptedBody::Raw(s) => {
|
||||
if let Some(wait) = before_terminal {
|
||||
wait().await;
|
||||
}
|
||||
s.into_response()
|
||||
}
|
||||
ScriptedBody::Sse(events) => {
|
||||
let events: Vec<axum::response::sse::Event> = events
|
||||
.into_iter()
|
||||
.map(|e| {
|
||||
let ev = axum::response::sse::Event::default().data(e.data);
|
||||
match e.event {
|
||||
Some(name) => ev.event(name),
|
||||
None => ev,
|
||||
let last_idx = events.len().checked_sub(1);
|
||||
let mut events: Vec<_> = events.into_iter().enumerate().map(Some).collect();
|
||||
if events.is_empty() && before_terminal.is_some() {
|
||||
events.push(None);
|
||||
}
|
||||
let stream = stream::unfold(
|
||||
(events.into_iter(), before_terminal),
|
||||
move |(mut events, mut before_terminal)| async move {
|
||||
loop {
|
||||
let item = events.next()?;
|
||||
let Some((idx, scripted_event)) = item else {
|
||||
if let Some(wait) = before_terminal.take() {
|
||||
wait().await;
|
||||
}
|
||||
continue;
|
||||
};
|
||||
if let Some(d) = delay {
|
||||
tokio::time::sleep(d).await;
|
||||
}
|
||||
if Some(idx) == last_idx
|
||||
&& let Some(wait) = before_terminal.take()
|
||||
{
|
||||
wait().await;
|
||||
}
|
||||
let event =
|
||||
axum::response::sse::Event::default().data(scripted_event.data);
|
||||
let event = match scripted_event.event {
|
||||
Some(name) => event.event(name),
|
||||
None => event,
|
||||
};
|
||||
return Some((Ok::<_, Infallible>(event), (events, before_terminal)));
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let stream = stream::iter(events.into_iter().map(Ok::<_, Infallible>)).then(
|
||||
move |event| async move {
|
||||
if let Some(d) = delay {
|
||||
tokio::time::sleep(d).await;
|
||||
}
|
||||
event
|
||||
},
|
||||
);
|
||||
Sse::new(stream)
|
||||
|
|
|
|||
Loading…
Reference in a new issue