diff --git a/cmd/router/main.go b/cmd/router/main.go index 1690fcd8d..304a27816 100644 --- a/cmd/router/main.go +++ b/cmd/router/main.go @@ -54,6 +54,7 @@ import ( "workweave/router/internal/router/policy" "workweave/router/internal/router/rl" "workweave/router/internal/router/sessionpin" + "workweave/router/internal/router/sessionstrategy" "workweave/router/internal/server" "workweave/router/internal/websearch" "workweave/router/internal/wif" @@ -795,6 +796,7 @@ func main() { // hmm then routes through it. Unset fails closed with 503. var hmmRouter router.Router var hmmEmbeddingRouter router.Router + var hmmBetaRouter router.Router var hmmCapabilities policy.Capabilities var hmmReadinessChecker admin.HealthChecker var hmmRosterSource policy.RosterSource @@ -889,6 +891,73 @@ func main() { logger.Info("HMM policy routers disabled (ROUTER_HMM_SIDECAR_URL unset); HMM strategies will return 503") } + // Separate sidecar: /beta opts into an independently deployed beta policy; + // absent or unhealthy beta fails closed without affecting stable routing. + var hmmBetaCapabilities policy.Capabilities + if hmmBetaSidecarURL := config.GetOr("ROUTER_HMM_BETA_SIDECAR_URL", ""); hmmBetaSidecarURL != "" { + hmmBetaTimeout := parseEnvDurationMs("ROUTER_HMM_BETA_SIDECAR_TIMEOUT_MS", policyclient.DefaultTimeout) + hmmBetaAuthMode := config.GetOr("ROUTER_HMM_BETA_SIDECAR_AUTH", policySidecarAuthNone) + hmmBetaAttemptTimeout := parseEnvAttemptTimeoutMs( + "ROUTER_HMM_BETA_SIDECAR_ATTEMPT_TIMEOUT_MS", + policyclient.DeriveAttemptTimeout(hmmBetaTimeout), + ) + hmmBetaClient, clientErr := buildHMMBetaPolicyClient( + hmmBetaSidecarURL, + hmmBetaAuthMode, + hmmBetaTimeout, + policyclient.WithAttemptTimeout(hmmBetaAttemptTimeout), + ) + if clientErr != nil { + // Beta is an optional isolation ring. A malformed beta-only auth + // setting must not take the stable router down with it. + logger.Error("beta HMM policy sidecar client failed to build; beta disabled", "auth_mode", hmmBetaAuthMode, "err", clientErr) + } else { + capabilityCtx, cancelCapabilityDiscovery := context.WithTimeout(context.Background(), hmmBetaTimeout) + var capabilityErr error + hmmBetaCapabilities, capabilityErr = hmmBetaClient.Capabilities(capabilityCtx) + cancelCapabilityDiscovery() + if capabilityErr != nil { + logger.Warn("beta HMM policy sidecar capabilities unavailable at boot; optional behavior remains disabled", "sidecar_url", hmmBetaSidecarURL, "err", capabilityErr) + } + hmmBetaPolicyRouter := hmm.NewForStrategy( + router.StrategyHMMBeta, + hmmBetaClient, + availableProviders, + ) + hmmBetaPolicyRouter.WithCapabilities(hmmBetaCapabilities) + if capabilityErr != nil { + go func() { + retryErr := retryPolicyCapabilitiesUntilAvailable( + context.Background(), + hmmBetaClient, + hmmBetaTimeout, + hmmCapabilityRetryInterval, + func(capabilities policy.Capabilities) { + hmmBetaPolicyRouter.WithCapabilities(capabilities) + }, + ) + if retryErr != nil { + logger.Warn("beta HMM policy sidecar capability refresh stopped", "sidecar_url", hmmBetaSidecarURL, "err", retryErr) + return + } + logger.Info("beta HMM policy sidecar capabilities discovered after boot", "sidecar_url", hmmBetaSidecarURL) + }() + } + hmmBetaRouter = hmmBetaPolicyRouter + logger.Info( + "beta HMM policy router wired", + "sidecar_url", hmmBetaSidecarURL, + "auth_mode", hmmBetaAuthMode, + "timeout_ms", hmmBetaTimeout.Milliseconds(), + "attempt_timeout_ms", hmmBetaAttemptTimeout.Milliseconds(), + "candidate_models", len(routingTargets), + "strategy", router.StrategyHMMBeta, + ) + } + } else { + logger.Info("beta HMM policy router disabled (ROUTER_HMM_BETA_SIDECAR_URL unset); /beta will be unavailable") + } + // Wired only when ROUTER_BANDIT_POSTERIOR_FILE points at a ts_posterior.json; // x-weave-router-strategy: bandit then routes through it. Wraps the raw // cluster scorer, not the explore wrapper. Unset -> nil -> 503. @@ -952,7 +1021,12 @@ func main() { flags.KeyEmbedOnlyUserMessage: boolDefault(embedOnlyUser), }) + // Always wire even when beta is unavailable: existing beta sessions fail + // closed via nil policy registration rather than silently falling to stable. + var sessionStrategyStore sessionstrategy.Store = postgres.NewSessionStrategyRepo(pool) + proxySvc := proxy.NewService(routeEntry, providerMap, telemetryEmitter, embedOnlyUser, semanticCache, pinStore, hardPinExplore, hardPinProvider, hardPinModel, repo.Telemetry). + WithSessionStrategyStore(sessionStrategyStore). WithTranslationCompatibilityMode(proxy.TranslationCompatibilityMode(translationCompatibilityMode)). WithPolicyStrategy(policy.StrategySpec{Strategy: router.StrategyRL, Router: rlRouter, Unavailable: rl.ErrPolicyUnavailable}). WithPolicyStrategy(policy.StrategySpec{ @@ -963,6 +1037,10 @@ func main() { Strategy: router.StrategyHMMEmbedding, Router: hmmEmbeddingRouter, Unavailable: hmm.ErrHMMUnavailable, Capabilities: hmmCapabilities, }). + WithPolicyStrategy(policy.StrategySpec{ + Strategy: router.StrategyHMMBeta, Router: hmmBetaRouter, Unavailable: hmm.ErrHMMUnavailable, + Capabilities: hmmBetaCapabilities, + }). WithPolicyStrategy(policy.StrategySpec{Strategy: router.StrategyBandit, Router: banditRouter, Unavailable: bandit.ErrBanditUnavailable}). WithContentCapture(captureMode, captureMaxBytes, nil). WithFeedback(repo.Feedback, feedbackSigner, feedbackBaseURL). diff --git a/cmd/router/policy_sidecar_auth.go b/cmd/router/policy_sidecar_auth.go index c3e2cde55..14b0b8e5a 100644 --- a/cmd/router/policy_sidecar_auth.go +++ b/cmd/router/policy_sidecar_auth.go @@ -30,6 +30,22 @@ func buildHMMPolicyClient( ) } +func buildHMMBetaPolicyClient( + sidecarURL, authMode string, + timeout time.Duration, + opts ...policyclient.Option, +) (*policyclient.Client, error) { + return buildPolicyClientWithGoogleIDTokenFactory( + sidecarURL, + authMode, + timeout, + nil, + "ROUTER_HMM_BETA_SIDECAR_AUTH", + policyclient.NewGoogleIDToken, + opts..., + ) +} + func buildHMMPolicyClientWithGoogleIDTokenFactory( sidecarURL, authMode string, timeout time.Duration, diff --git a/cmd/router/policy_sidecar_auth_test.go b/cmd/router/policy_sidecar_auth_test.go index bc3e3378d..f08fbd376 100644 --- a/cmd/router/policy_sidecar_auth_test.go +++ b/cmd/router/policy_sidecar_auth_test.go @@ -27,6 +27,14 @@ func TestBuildHMMPolicyClientRejectsUnknownAuthMode(t *testing.T) { assert.Contains(t, err.Error(), "unsupported ROUTER_HMM_SIDECAR_AUTH") } +func TestBuildHMMBetaPolicyClientNamesBetaAuthSetting(t *testing.T) { + client, err := buildHMMBetaPolicyClient("https://sidecar.internal", "api-key", time.Second) + + require.Error(t, err) + assert.Nil(t, client) + assert.Contains(t, err.Error(), "unsupported ROUTER_HMM_BETA_SIDECAR_AUTH") +} + func TestBuildHMMPolicyClientFailsClosedWhenGoogleCredentialsCannotBuild(t *testing.T) { wantErr := errors.New("ADC unavailable") client, err := buildHMMPolicyClientWithGoogleIDTokenFactory( diff --git a/cmd/router/policy_sidecars.go b/cmd/router/policy_sidecars.go index dfbe368b3..cfc76a8a1 100644 --- a/cmd/router/policy_sidecars.go +++ b/cmd/router/policy_sidecars.go @@ -26,6 +26,7 @@ var reservedPolicyStrategies = map[router.Strategy]struct{}{ router.StrategyRL: {}, router.StrategyHMM: {}, router.StrategyHMMEmbedding: {}, + router.StrategyHMMBeta: {}, router.StrategyBandit: {}, } diff --git a/cmd/router/policy_sidecars_test.go b/cmd/router/policy_sidecars_test.go index d8e04adc2..c79235a26 100644 --- a/cmd/router/policy_sidecars_test.go +++ b/cmd/router/policy_sidecars_test.go @@ -74,6 +74,7 @@ func TestBuildConfiguredPolicySidecarsRejectsReservedAndInvalidConfiguration(t * for _, raw := range []string{ `{"hmm":"https://sidecar.internal"}`, `{"hmm_embedding":"https://sidecar.internal"}`, + `{"hmm_beta":"https://sidecar.internal"}`, `{"future":"not-a-url"}`, `{"future policy":"https://sidecar.internal"}`, `{"Future":"https://one.internal","future":"https://two.internal"}`, diff --git a/db/migrations/0070_session-strategy-preferences.down.sql b/db/migrations/0070_session-strategy-preferences.down.sql new file mode 100644 index 000000000..755d76cc0 --- /dev/null +++ b/db/migrations/0070_session-strategy-preferences.down.sql @@ -0,0 +1,8 @@ +BEGIN; + +DROP TABLE router.session_strategy_preferences; + +ALTER TABLE router.session_pins + DROP COLUMN routing_strategy; + +COMMIT; diff --git a/db/migrations/0070_session-strategy-preferences.up.sql b/db/migrations/0070_session-strategy-preferences.up.sql new file mode 100644 index 000000000..5a2a8186d --- /dev/null +++ b/db/migrations/0070_session-strategy-preferences.up.sql @@ -0,0 +1,19 @@ +BEGIN; + +ALTER TABLE router.session_pins + ADD COLUMN routing_strategy VARCHAR(32) NOT NULL DEFAULT ''; + +CREATE TABLE router.session_strategy_preferences ( + installation_id UUID NOT NULL, + session_key BYTEA NOT NULL CHECK (octet_length(session_key) = 16), + strategy VARCHAR(32) NOT NULL CHECK (strategy = 'hmm_beta'), + enabled BOOLEAN NOT NULL DEFAULT TRUE, + PRIMARY KEY (installation_id, session_key), + FOREIGN KEY (installation_id) + REFERENCES router.model_router_installations(id) ON DELETE CASCADE +); + +COMMENT ON TABLE router.session_strategy_preferences IS + 'Explicit per-session router strategy preferences'; + +COMMIT; diff --git a/db/queries/session_pins.sql b/db/queries/session_pins.sql index c579affbd..1d4f02679 100644 --- a/db/queries/session_pins.sql +++ b/db/queries/session_pins.sql @@ -10,12 +10,17 @@ FROM router.session_pins WHERE session_key = @session_key::bytea AND role = @role::varchar; --- Atomically consumes one active pin so a one-shot continuation cannot be --- reused by concurrent requests. Expired rows remain for the normal sweep. +-- Atomically consumes one active pin for the expected strategy so a stale +-- continuation cannot delete a replacement strategy's pin. Expired rows +-- remain for the normal sweep. -- name: DeleteSessionPin :one DELETE FROM router.session_pins WHERE session_key = @session_key::bytea AND role = @role::varchar + AND ( + routing_strategy = @expected_routing_strategy::varchar + OR (routing_strategy = '' AND @expected_routing_strategy::varchar <> 'hmm_beta') + ) AND pinned_until > CURRENT_TIMESTAMP RETURNING *; @@ -29,59 +34,63 @@ RETURNING *; -- them, so the at-start-of-turn refresh here cannot clobber the -- previous turn's usage with zeros before the planner reads it. -- --- consecutive_upstream_errors is preserved on a same-model refresh (so --- the two-strike eviction counter accumulates across turns of the same --- sticky pin) but reset to 0 on a switch (different model = clean --- slate). The reset on switch also covers the loop-break / force-model --- pin-expiry writes, which set pinned_model to the empty string. +-- consecutive_upstream_errors is preserved on a same-model, same-strategy +-- refresh (so the two-strike eviction counter accumulates across turns of the +-- same sticky pin) but reset to 0 on a model or strategy switch. The reset also +-- covers loop-break / force-model pin-expiry writes, which set pinned_model to +-- the empty string. -- -- paired_provider / paired_model hold the runner-up half of the band pair the -- scorer picks. On the conflict update they refresh to a fresh scorer runner-up --- (non-empty incoming pair), are preserved when the pinned model is unchanged --- (sticky refresh / same-model re-anchor carry an empty pair), and are cleared --- when the pinned model changes without a fresh pair (force-model, --- loop-escalation, eviction -- non-scorer writes). This keeps the stored pair --- consistent with the live decision: it tracks genuine re-routes, never --- inherits a stale runner-up across a non-scorer model change, and never --- collapses pinned_model and paired_model onto the same slug. A later per-turn --- swap policy reads the pair that matches the active decision. +-- (non-empty incoming pair), are preserved when both the pinned model and +-- strategy are unchanged (sticky refresh / same-model re-anchor carry an empty +-- pair), and are cleared when either changes without a fresh pair. This keeps +-- the stored pair consistent with the live decision: it tracks genuine +-- re-routes and never inherits a stale runner-up across a strategy change. -- -- policy_group follows the same three-way maintenance: a fresh policy decision --- supplies a non-empty group, a same-model refresh preserves the stored one, and --- a model change without a group (force-model, loop-break, eviction) clears it. +-- supplies a non-empty group, a same-model same-strategy refresh preserves the +-- stored one, and a model or strategy change without a group clears it. -- The pin-sticky arm-selector guard compares it against the fresh decision's -- group, so a stale group must never survive onto a different pinned model. -- name: UpsertSessionPin :exec INSERT INTO router.session_pins ( session_key, role, installation_id, pinned_provider, pinned_model, paired_provider, paired_model, - decision_reason, policy_group, turn_count, pinned_until + decision_reason, routing_strategy, policy_group, turn_count, pinned_until ) VALUES ( @session_key::bytea, @role::varchar, @installation_id::uuid, @pinned_provider::varchar, @pinned_model::varchar, @paired_provider::varchar, @paired_model::varchar, - @decision_reason::text, @policy_group::varchar, + @decision_reason::text, @routing_strategy::varchar, @policy_group::varchar, @turn_count::int, @pinned_until::timestamp ) ON CONFLICT (session_key, role) DO UPDATE SET pinned_provider = EXCLUDED.pinned_provider, pinned_model = EXCLUDED.pinned_model, decision_reason = EXCLUDED.decision_reason, - turn_count = router.session_pins.turn_count + 1, + routing_strategy = EXCLUDED.routing_strategy, + turn_count = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.turn_count + 1 + ELSE EXCLUDED.turn_count + END, pinned_until = EXCLUDED.pinned_until, + first_pinned_at = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.first_pinned_at + ELSE CURRENT_TIMESTAMP + END, last_seen_at = CURRENT_TIMESTAMP, -- Band pair maintenance, in priority order: -- 1. A fresh scorer decision supplies a non-empty pair -> take it. - -- 2. Empty incoming pair but the pinned model is unchanged (sticky refresh, - -- reconstructed re-anchor of the same model) -> preserve the stored pair. - -- 3. Empty incoming pair and the pinned model changed (force-model, - -- loop-escalation, eviction -- non-scorer writes) -> clear the pair, so a - -- model change never inherits a stale runner-up or collapses pinned_model - -- and paired_model onto the same slug. + -- 2. Empty incoming pair but model and strategy are unchanged -> preserve it. + -- 3. Empty incoming pair and either changed -> clear it. paired_provider = CASE WHEN EXCLUDED.paired_model <> '' THEN EXCLUDED.paired_provider WHEN EXCLUDED.pinned_model = router.session_pins.pinned_model + AND EXCLUDED.routing_strategy = router.session_pins.routing_strategy THEN router.session_pins.paired_provider ELSE '' END, @@ -89,6 +98,7 @@ ON CONFLICT (session_key, role) DO UPDATE SET WHEN EXCLUDED.paired_model <> '' THEN EXCLUDED.paired_model WHEN EXCLUDED.pinned_model = router.session_pins.pinned_model + AND EXCLUDED.routing_strategy = router.session_pins.routing_strategy THEN router.session_pins.paired_model ELSE '' END, @@ -96,11 +106,13 @@ ON CONFLICT (session_key, role) DO UPDATE SET WHEN EXCLUDED.policy_group <> '' THEN EXCLUDED.policy_group WHEN EXCLUDED.pinned_model = router.session_pins.pinned_model + AND EXCLUDED.routing_strategy = router.session_pins.routing_strategy THEN router.session_pins.policy_group ELSE '' END, consecutive_upstream_errors = CASE WHEN router.session_pins.pinned_model = EXCLUDED.pinned_model + AND router.session_pins.routing_strategy = EXCLUDED.routing_strategy THEN router.session_pins.consecutive_upstream_errors ELSE 0 END, @@ -111,8 +123,51 @@ ON CONFLICT (session_key, role) DO UPDATE SET -- requiring two genuine consecutive strikes on the SAME served provider. consecutive_overload_errors = CASE WHEN router.session_pins.pinned_model = EXCLUDED.pinned_model + AND router.session_pins.routing_strategy = EXCLUDED.routing_strategy THEN router.session_pins.consecutive_overload_errors ELSE 0 + END, + -- A strategy switch selects a different policy. Do not carry cache, + -- switch, or error evidence from the previous policy into it. + last_input_tokens = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_input_tokens + ELSE 0 + END, + last_cached_read_tokens = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_cached_read_tokens + ELSE 0 + END, + last_cached_write_tokens = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_cached_write_tokens + ELSE 0 + END, + last_output_tokens = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_output_tokens + ELSE 0 + END, + last_turn_ended_at = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_turn_ended_at + ELSE NULL + END, + last_served_model = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_served_model + ELSE '' + END, + has_ever_switched = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.has_ever_switched + ELSE FALSE + END, + disabled_providers = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.disabled_providers + ELSE '{}' END; -- Records the previous turn's upstream token usage on an existing pin @@ -121,7 +176,8 @@ ON CONFLICT (session_key, role) DO UPDATE SET -- turn to compute switch EV against eviction cost. The UPDATE matches -- by (session_key, role); if the pin has been evicted or never -- existed, zero rows are affected and the adapter wraps that as a --- successful no-op. last_served_model records the model that actually +-- successful no-op. A strategy mismatch is also a no-op, preventing a late +-- response from mutating a replacement strategy's pin. last_served_model records the model that actually -- served this turn; it lives here (not in UpsertSessionPin) so a -- /force-model upsert cannot overwrite the genuinely-last-served model -- before the next turn reads it to detect a mid-session model switch. @@ -146,7 +202,11 @@ SET last_input_tokens = @last_input_tokens::int, OR (@prior_served_model::varchar <> '' AND @prior_served_model::varchar <> @last_served_model::varchar), last_served_model = @last_served_model::varchar WHERE session_key = @session_key::bytea - AND role = @role::varchar; + AND role = @role::varchar + AND ( + routing_strategy = @expected_routing_strategy::varchar + OR (routing_strategy = '' AND @expected_routing_strategy::varchar <> 'hmm_beta') + ); -- Atomically increments consecutive_upstream_errors and returns the -- new value. The turn loop calls this after a non-retryable upstream @@ -159,6 +219,10 @@ UPDATE router.session_pins SET consecutive_upstream_errors = consecutive_upstream_errors + 1 WHERE session_key = @session_key::bytea AND role = @role::varchar + AND ( + routing_strategy = @expected_routing_strategy::varchar + OR (routing_strategy = '' AND @expected_routing_strategy::varchar <> 'hmm_beta') + ) RETURNING consecutive_upstream_errors; -- Clears the two-strike counter after a successful turn. UPDATE @@ -169,6 +233,10 @@ UPDATE router.session_pins SET consecutive_upstream_errors = 0 WHERE session_key = @session_key::bytea AND role = @role::varchar + AND ( + routing_strategy = @expected_routing_strategy::varchar + OR (routing_strategy = '' AND @expected_routing_strategy::varchar <> 'hmm_beta') + ) AND consecutive_upstream_errors > 0; -- Atomically increments consecutive_overload_errors and returns the new @@ -184,6 +252,10 @@ UPDATE router.session_pins SET consecutive_overload_errors = consecutive_overload_errors + 1 WHERE session_key = @session_key::bytea AND role = @role::varchar + AND ( + routing_strategy = @expected_routing_strategy::varchar + OR (routing_strategy = '' AND @expected_routing_strategy::varchar <> 'hmm_beta') + ) RETURNING consecutive_overload_errors; -- Clears the overload strike counter after a successful turn. UPDATE @@ -194,14 +266,17 @@ UPDATE router.session_pins SET consecutive_overload_errors = 0 WHERE session_key = @session_key::bytea AND role = @role::varchar + AND ( + routing_strategy = @expected_routing_strategy::varchar + OR (routing_strategy = '' AND @expected_routing_strategy::varchar <> 'hmm_beta') + ) AND consecutive_overload_errors > 0; -- Appends a provider to disabled_providers (deduped) and resets the -- overload strike counter in the same statement, fired once the --- two-strike threshold is reached. disabled_providers only grows for the --- life of this pin row -- UpsertSessionPin's ON CONFLICT update never --- touches it, so a struck-out provider stays disabled until the pin --- itself is evicted/expires, with no separate time-based cooldown. +-- two-strike threshold is reached. disabled_providers only grows within one +-- strategy's pin lifecycle; a strategy replacement resets it with the other +-- strategy-bound evidence. There is no separate time-based cooldown. -- name: DisableSessionPinProvider :exec UPDATE router.session_pins SET disabled_providers = CASE @@ -210,7 +285,11 @@ SET disabled_providers = CASE END, consecutive_overload_errors = 0 WHERE session_key = @session_key::bytea - AND role = @role::varchar; + AND role = @role::varchar + AND ( + routing_strategy = @expected_routing_strategy::varchar + OR (routing_strategy = '' AND @expected_routing_strategy::varchar <> 'hmm_beta') + ); -- Garbage-collects pins that have been expired for >24h. The 24h grace -- means a transient Postgres outage doesn't immediately prune live pins; diff --git a/db/queries/session_strategy_preferences.sql b/db/queries/session_strategy_preferences.sql new file mode 100644 index 000000000..a705dab5a --- /dev/null +++ b/db/queries/session_strategy_preferences.sql @@ -0,0 +1,34 @@ +-- Reads an explicit beta strategy for one installation-scoped session. A +-- missing or disabled row means the session uses stable routing. +-- name: GetSessionStrategyPreference :one +SELECT strategy +FROM router.session_strategy_preferences +WHERE installation_id = @installation_id::uuid + AND session_key = @session_key::bytea + AND enabled; + +-- Flips the session's explicit override and returns the state now persisted. +-- The row lock taken on conflict serializes overlapping toggles across router +-- instances, so each caller observes its own flip instead of a stale read. +-- The database constraint rejects any strategy other than hmm_beta. +-- name: UpsertToggledSessionStrategyPreference :one +INSERT INTO router.session_strategy_preferences ( + installation_id, session_key, strategy, enabled +) VALUES ( + @installation_id::uuid, @session_key::bytea, @strategy::varchar, TRUE +) +ON CONFLICT (installation_id, session_key) +DO UPDATE SET + strategy = EXCLUDED.strategy, + enabled = NOT router.session_strategy_preferences.enabled +RETURNING enabled; + +-- Turns the session's explicit override off and reports one affected row when +-- beta had been enabled. Callers use this instead of the toggle when the beta +-- policy is unavailable, so a concurrent command can never re-enable it. +-- name: UpdateSessionStrategyPreferenceDisabled :execrows +UPDATE router.session_strategy_preferences +SET enabled = FALSE +WHERE installation_id = @installation_id::uuid + AND session_key = @session_key::bytea + AND enabled; diff --git a/install/commands/beta.md b/install/commands/beta.md new file mode 100644 index 000000000..78ecb1396 --- /dev/null +++ b/install/commands/beta.md @@ -0,0 +1,5 @@ +--- +description: Toggle beta HMM routing for this session. +--- + +/beta $ARGUMENTS diff --git a/install/directives.tsv b/install/directives.tsv index 6d410861f..cf5fddae4 100644 --- a/install/directives.tsv +++ b/install/directives.tsv @@ -10,3 +10,4 @@ router-status||local-toggle|yes|yes|no|no|manual|command,skill router-session||prompt|yes|no|no|no|manual|command router-models|models|local-toggle|yes|yes|no|no|manual|command,skill disable-routing||local-toggle|no|yes|no|no|manual|skill +beta||prompt|yes|no|no|yes|manual|command diff --git a/install/install.sh b/install/install.sh index 6c95b64a6..d94382c64 100755 --- a/install/install.sh +++ b/install/install.sh @@ -1440,6 +1440,7 @@ router-status||local-toggle|yes|yes|no|no|manual|command,skill router-session||prompt|yes|no|no|no|manual|command router-models|models|local-toggle|yes|yes|no|no|manual|command,skill disable-routing||local-toggle|no|yes|no|no|manual|skill +beta||prompt|yes|no|no|yes|manual|command WEAVE_REGISTRY_EOF ) diff --git a/install/pi-router/README.md b/install/pi-router/README.md index dd05f742c..fcd2ac2d1 100644 --- a/install/pi-router/README.md +++ b/install/pi-router/README.md @@ -31,6 +31,9 @@ from npm on next start and loads this extension via its `pi.extensions` field. current router session; `/ufm` and `/unforce-model` resume automatic routing. The persistent status changes to `WEAVE ROUTER — [forced]` after the router validates and canonicalizes the requested model. +- **Beta routing toggle.** `/beta` toggles the router's beta HMM strategy for + the current session. The router confirms whether beta routing is enabled or + disabled; run `/beta` again to switch back. - **Per-process routing bias.** Static `x-weave-routing-*` knob headers bias the router: quality on the main loop and speed + cheap on subagents. - **Long tool-loop compaction.** Pi can cross its context threshold inside an diff --git a/install/pi-router/src/beta.ts b/install/pi-router/src/beta.ts new file mode 100644 index 000000000..fea045c1b --- /dev/null +++ b/install/pi-router/src/beta.ts @@ -0,0 +1,21 @@ +/** + * Forward Pi's local /beta command to the router as one canonical user turn. + * + * The router owns the session strategy state and the enabled/disabled reply; + * the client only validates the command shape and preserves busy-turn ordering. + */ + +import type { ExtensionAPI, ExtensionCommandContext } from "@mariozechner/pi-coding-agent"; + +export function registerBetaCommand(pi: ExtensionAPI): void { + pi.registerCommand("beta", { + description: "Toggle beta HMM routing for this Weave Router session", + handler: async (args: string, ctx: ExtensionCommandContext): Promise => { + if (args.trim()) { + ctx.ui.notify("Usage: /beta", "warning"); + return; + } + pi.sendUserMessage("/beta", ctx.isIdle() ? undefined : { deliverAs: "followUp" }); + }, + }); +} diff --git a/install/pi-router/src/index.ts b/install/pi-router/src/index.ts index 01f49277a..608a8efd6 100644 --- a/install/pi-router/src/index.ts +++ b/install/pi-router/src/index.ts @@ -22,6 +22,7 @@ import { fileURLToPath } from "node:url"; import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { registerBetaCommand } from "./beta.js"; import { isSubagent } from "./config.js"; import { registerCompaction } from "./compaction.js"; import { registerDispatch } from "./dispatch.js"; @@ -42,6 +43,7 @@ export default function (pi: ExtensionAPI): void { pi.on("session_start", () => registerWeave(pi)); registerMetadata(pi); + registerBetaCommand(pi); registerForceModelCommands(pi); registerRoutedModel(pi); registerCompaction(pi); diff --git a/install/pi-router/test/beta.test.ts b/install/pi-router/test/beta.test.ts new file mode 100644 index 000000000..404f32f9b --- /dev/null +++ b/install/pi-router/test/beta.test.ts @@ -0,0 +1,62 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import type { ExtensionAPI, ExtensionCommandContext } from "@mariozechner/pi-coding-agent"; +import { registerBetaCommand } from "../src/beta.js"; + +type CommandOptions = Parameters[1]; +type SendOptions = Parameters[1]; + +function extensionHarness() { + const commands = new Map(); + const sent: Array<{ content: string; options: SendOptions }> = []; + const pi = { + registerCommand(name: string, options: CommandOptions) { + commands.set(name, options); + }, + sendUserMessage(content: string, options?: SendOptions) { + sent.push({ content, options }); + }, + } as unknown as ExtensionAPI; + registerBetaCommand(pi); + return { commands, sent }; +} + +function commandContext(idle = true) { + const notifications: Array<{ message: string; level: string }> = []; + const ctx = { + isIdle: () => idle, + ui: { + notify(message: string, level: string) { + notifications.push({ message, level }); + }, + }, + } as unknown as ExtensionCommandContext; + return { ctx, notifications }; +} + +test("registers the beta command", () => { + const { commands } = extensionHarness(); + assert.deepEqual([...commands.keys()], ["beta"]); +}); + +test("forwards one canonical beta turn while Pi is idle", async () => { + const { commands, sent } = extensionHarness(); + const { ctx } = commandContext(); + await commands.get("beta")?.handler("", ctx); + assert.deepEqual(sent, [{ content: "/beta", options: undefined }]); +}); + +test("queues beta as a follow-up while Pi is busy", async () => { + const { commands, sent } = extensionHarness(); + const { ctx } = commandContext(false); + await commands.get("beta")?.handler("", ctx); + assert.deepEqual(sent, [{ content: "/beta", options: { deliverAs: "followUp" } }]); +}); + +test("rejects beta arguments without starting a turn", async () => { + const { commands, sent } = extensionHarness(); + const { ctx, notifications } = commandContext(); + await commands.get("beta")?.handler(" off ", ctx); + assert.deepEqual(sent, []); + assert.deepEqual(notifications, [{ message: "Usage: /beta", level: "warning" }]); +}); diff --git a/install/pi-router/test/e2e.sh b/install/pi-router/test/e2e.sh index 8cadca611..f1e2413dc 100755 --- a/install/pi-router/test/e2e.sh +++ b/install/pi-router/test/e2e.sh @@ -133,9 +133,9 @@ phase "Phase 2 — generated pricing + savings contract" if with_timeout 30 env PI_CODING_AGENT_DIR="$PI_DIR" \ pi -e "$UNIT_SUITE" --no-session --offline --model weave/claude-sonnet-4-6 \ -p "Run the unit suite." >"$WORK/unit.out" 2>&1 /dev/null)" "/beta \$ARGUMENTS" + # ---------- command baseline seeding ---------- # # The statusline's background wrapper refresh only swaps a file whose bytes @@ -183,11 +190,13 @@ check "update on an uninstalled target errors on the key" "$?" 1 # mktemp hands back /var/... for /private/var/.... cmd_dir_real="$(cd "$upd_home/.claude/commands" && pwd -P)" baseline="$upd_home/.cache/weave-router/commands$(printf '%s' "$cmd_dir_real" | tr -c 'A-Za-z0-9._-' '_')" -if [ -f "$baseline/force-model.md" ] && [ -f "$baseline/router-off.md" ]; then +if [ -f "$baseline/force-model.md" ] && [ -f "$baseline/beta.md" ] && [ -f "$baseline/router-off.md" ]; then ok "install seeds the slash-command baseline" else no "install seeds the slash-command baseline" "canonical wrappers cached" "missing under $baseline" fi +check "install writes the beta slash command" \ + "$(sed -n '5p' "$upd_home/.claude/commands/beta.md" 2>/dev/null)" "/beta \$ARGUMENTS" # Baselines are the UNRENDERED canonical files: the refresh renders {{SCOPE}} # per install, and a pre-rendered baseline would never match upstream. if grep -q '{{SCOPE}}' "$baseline/router-off.md" 2>/dev/null; then @@ -196,6 +205,17 @@ else no "seeded baseline keeps the {{SCOPE}} placeholder" "unrendered copy" "placeholder substituted" fi +# Claude uninstall removes only installer-owned wrappers, including /beta. +uninstall_home="$work/uninstall"; mkdir -p "$uninstall_home" +run "$uninstall_home" rk_uninstall -- --claude --scope user --quiet --non-interactive +HOME="$uninstall_home" XDG_CACHE_HOME="$uninstall_home/.cache" PATH="$test_path" NO_COLOR=1 \ + bash "$uninstaller" --claude --scope user >/dev/null 2>&1 +if [ ! -e "$uninstall_home/.claude/commands/beta.md" ]; then + ok "uninstall removes the beta slash command" +else + no "uninstall removes the beta slash command" "file absent" "file remains" +fi + # ---------- update after `off`: parked sidecar key + base URL carry-over ---------- # # `off` moves the router URL and key header out of settings.json/settings.local.json diff --git a/install/tests/registry_test.sh b/install/tests/registry_test.sh index 51718b0ca..7a95805c9 100644 --- a/install/tests/registry_test.sh +++ b/install/tests/registry_test.sh @@ -65,9 +65,10 @@ check "names shared across clients are the expected shared directives" \ check "an alias resolves to its canonical directive" "force-model" "$(weave_registry_canonical_for fm)" check "a canonical name resolves to itself" "router-feedback" "$(weave_registry_canonical_for router-feedback)" -# Pi implements /fm and /ufm in its extension rather than through installed -# files. Its registered names must still match the shared registry. -pi_registered="$(grep -oE '"(fm|force-model|ufm|unforce-model)"' "$install_dir/pi-router/src/force-model.ts" \ +# Pi implements /fm, /ufm and /beta in its extension rather than through +# installed files. Its registered names must still match the shared registry. +pi_registered="$(grep -hoE '"(fm|force-model|ufm|unforce-model|beta)"' \ + "$install_dir/pi-router/src/force-model.ts" "$install_dir/pi-router/src/beta.ts" \ | tr -d '"' | sort -u | tr '\n' ' ' | sed 's/ $//')" check "the pi extension registers exactly the registry's pi names" \ "$(weave_registry_names pi | sort -u | tr '\n' ' ' | sed 's/ $//')" "$pi_registered" diff --git a/install/uninstall.sh b/install/uninstall.sh index 97f1b69ee..2cd3c2988 100755 --- a/install/uninstall.sh +++ b/install/uninstall.sh @@ -48,6 +48,7 @@ router-status||local-toggle|yes|yes|no|no|manual|command,skill router-session||prompt|yes|no|no|no|manual|command router-models|models|local-toggle|yes|yes|no|no|manual|command,skill disable-routing||local-toggle|no|yes|no|no|manual|skill +beta||prompt|yes|no|no|yes|manual|command WEAVE_REGISTRY_EOF ) diff --git a/internal/api/admin/policy_catalog.go b/internal/api/admin/policy_catalog.go index cbecda338..fadfa8c5c 100644 --- a/internal/api/admin/policy_catalog.go +++ b/internal/api/admin/policy_catalog.go @@ -37,6 +37,11 @@ func PolicyCatalogHandler(service *proxy.Service, defaultStrategy router.Strateg }} if service != nil { for _, strategy := range service.RegisteredStrategies() { + // Beta is a session control, not an installation strategy. Keep it + // out of the control-plane catalog so /beta remains its only surface. + if strategy == router.StrategyHMMBeta { + continue + } capabilities, _ := service.PolicyCapabilities(strategy) // Derived, not separately negotiated: ranked fallback is the // precondition for cluster overrides taking effect. diff --git a/internal/api/admin/policy_catalog_test.go b/internal/api/admin/policy_catalog_test.go index 41ad4514b..81eac559a 100644 --- a/internal/api/admin/policy_catalog_test.go +++ b/internal/api/admin/policy_catalog_test.go @@ -35,6 +35,9 @@ func TestPolicyCatalogHandlerReportsDefaultAndCapabilities(t *testing.T) { HonorsPreferredModels: true, HonorsQualityPriceBias: true, }, + }).WithPolicyStrategy(policy.StrategySpec{ + Strategy: router.StrategyHMMBeta, + Router: policyCatalogRouter{}, }) engine := gin.New() engine.GET( @@ -67,4 +70,5 @@ func TestPolicyCatalogHandlerReportsDefaultAndCapabilities(t *testing.T) { assert.True(t, payload.Strategies[0].Capabilities.SupportsShadow) assert.Equal(t, "hmm", payload.Strategies[1].Strategy) assert.True(t, payload.Strategies[1].Available) + assert.NotContains(t, recorder.Body.String(), string(router.StrategyHMMBeta)) } diff --git a/internal/postgres/session_pin_repo.go b/internal/postgres/session_pin_repo.go index d1e07f05e..0c14549f5 100644 --- a/internal/postgres/session_pin_repo.go +++ b/internal/postgres/session_pin_repo.go @@ -6,6 +6,7 @@ import ( "errors" "time" + "workweave/router/internal/router" "workweave/router/internal/router/sessionpin" "workweave/router/internal/sqlc" @@ -40,11 +41,12 @@ func (r *SessionPinRepo) Get(ctx context.Context, sessionKey [sessionpin.Session } // Consume atomically removes and returns an unexpired one-shot pin. -func (r *SessionPinRepo) Consume(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string) (sessionpin.Pin, bool, error) { +func (r *SessionPinRepo) Consume(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string, expectedStrategy router.Strategy) (sessionpin.Pin, bool, error) { q := sqlc.New(r.tx) row, err := q.DeleteSessionPin(ctx, sqlc.DeleteSessionPinParams{ - SessionKey: sessionKey[:], - Role: role, + SessionKey: sessionKey[:], + Role: role, + ExpectedRoutingStrategy: string(expectedStrategy), }) if err != nil { if errors.Is(err, sql.ErrNoRows) { @@ -58,17 +60,18 @@ func (r *SessionPinRepo) Consume(ctx context.Context, sessionKey [sessionpin.Ses func (r *SessionPinRepo) Upsert(ctx context.Context, p sessionpin.Pin) error { q := sqlc.New(r.tx) return q.UpsertSessionPin(ctx, sqlc.UpsertSessionPinParams{ - SessionKey: p.SessionKey[:], - Role: p.Role, - InstallationID: p.InstallationID, - PinnedProvider: p.Provider, - PinnedModel: p.Model, - PairedProvider: p.PairedProvider, - PairedModel: p.PairedModel, - DecisionReason: p.Reason, - PolicyGroup: p.PolicyGroup, - TurnCount: int32(p.TurnCount), - PinnedUntil: pgtype.Timestamp{Time: p.PinnedUntil.UTC(), Valid: true}, + SessionKey: p.SessionKey[:], + Role: p.Role, + InstallationID: p.InstallationID, + PinnedProvider: p.Provider, + PinnedModel: p.Model, + PairedProvider: p.PairedProvider, + PairedModel: p.PairedModel, + DecisionReason: p.Reason, + RoutingStrategy: string(p.Strategy), + PolicyGroup: p.PolicyGroup, + TurnCount: int32(p.TurnCount), + PinnedUntil: pgtype.Timestamp{Time: p.PinnedUntil.UTC(), Valid: true}, }) } @@ -82,28 +85,30 @@ func (r *SessionPinRepo) UpdateUsage(ctx context.Context, sessionKey [sessionpin } q := sqlc.New(r.tx) return q.UpdateSessionPinUsage(ctx, sqlc.UpdateSessionPinUsageParams{ - SessionKey: sessionKey[:], - Role: role, - LastInputTokens: int32(usage.InputTokens), - LastCachedReadTokens: int32(usage.CachedReadTokens), - LastCachedWriteTokens: int32(usage.CachedWriteTokens), - LastOutputTokens: int32(usage.OutputTokens), - LastTurnEndedAt: pgtype.Timestamptz{Time: endedAt.UTC(), Valid: true}, - LastServedModel: usage.ServedModel, - LastServedProvider: usage.ServedProvider, - PriorServedModel: usage.PriorServedModel, - SessionEverSwitched: usage.SessionEverSwitched, + SessionKey: sessionKey[:], + Role: role, + LastInputTokens: int32(usage.InputTokens), + LastCachedReadTokens: int32(usage.CachedReadTokens), + LastCachedWriteTokens: int32(usage.CachedWriteTokens), + LastOutputTokens: int32(usage.OutputTokens), + LastTurnEndedAt: pgtype.Timestamptz{Time: endedAt.UTC(), Valid: true}, + LastServedModel: usage.ServedModel, + LastServedProvider: usage.ServedProvider, + PriorServedModel: usage.PriorServedModel, + SessionEverSwitched: usage.SessionEverSwitched, + ExpectedRoutingStrategy: string(usage.Strategy), }) } // IncrementUpstreamErrors atomically bumps the consecutive-error counter. // A missing pin (already evicted or never created) returns (0, nil): the // two-strike check treats it as a no-op since there's no row left to evict. -func (r *SessionPinRepo) IncrementUpstreamErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string) (int, error) { +func (r *SessionPinRepo) IncrementUpstreamErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string, expectedStrategy router.Strategy) (int, error) { q := sqlc.New(r.tx) count, err := q.IncrementSessionPinUpstreamErrors(ctx, sqlc.IncrementSessionPinUpstreamErrorsParams{ - SessionKey: sessionKey[:], - Role: role, + SessionKey: sessionKey[:], + Role: role, + ExpectedRoutingStrategy: string(expectedStrategy), }) if err != nil { if errors.Is(err, sql.ErrNoRows) { @@ -116,21 +121,23 @@ func (r *SessionPinRepo) IncrementUpstreamErrors(ctx context.Context, sessionKey // ResetUpstreamErrors clears the consecutive-error counter after a // successful turn. Missing pin is a no-op, same as UpdateUsage. -func (r *SessionPinRepo) ResetUpstreamErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string) error { +func (r *SessionPinRepo) ResetUpstreamErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string, expectedStrategy router.Strategy) error { q := sqlc.New(r.tx) return q.ResetSessionPinUpstreamErrors(ctx, sqlc.ResetSessionPinUpstreamErrorsParams{ - SessionKey: sessionKey[:], - Role: role, + SessionKey: sessionKey[:], + Role: role, + ExpectedRoutingStrategy: string(expectedStrategy), }) } // IncrementOverloadErrors atomically bumps the consecutive-529-exhaustion // counter. A missing pin returns (0, nil), mirroring IncrementUpstreamErrors. -func (r *SessionPinRepo) IncrementOverloadErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string) (int, error) { +func (r *SessionPinRepo) IncrementOverloadErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string, expectedStrategy router.Strategy) (int, error) { q := sqlc.New(r.tx) count, err := q.IncrementSessionPinOverloadErrors(ctx, sqlc.IncrementSessionPinOverloadErrorsParams{ - SessionKey: sessionKey[:], - Role: role, + SessionKey: sessionKey[:], + Role: role, + ExpectedRoutingStrategy: string(expectedStrategy), }) if err != nil { if errors.Is(err, sql.ErrNoRows) { @@ -143,23 +150,25 @@ func (r *SessionPinRepo) IncrementOverloadErrors(ctx context.Context, sessionKey // ResetOverloadErrors clears the consecutive-529-exhaustion counter after a // successful turn. Missing pin is a no-op, same as ResetUpstreamErrors. -func (r *SessionPinRepo) ResetOverloadErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string) error { +func (r *SessionPinRepo) ResetOverloadErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string, expectedStrategy router.Strategy) error { q := sqlc.New(r.tx) return q.ResetSessionPinOverloadErrors(ctx, sqlc.ResetSessionPinOverloadErrorsParams{ - SessionKey: sessionKey[:], - Role: role, + SessionKey: sessionKey[:], + Role: role, + ExpectedRoutingStrategy: string(expectedStrategy), }) } // DisableProvider appends provider to disabled_providers (deduped) and // resets the overload strike counter in the same write. Missing pin is a // no-op: the eviction that accompanies this call has nothing left to guard. -func (r *SessionPinRepo) DisableProvider(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role, provider string) error { +func (r *SessionPinRepo) DisableProvider(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role, provider string, expectedStrategy router.Strategy) error { q := sqlc.New(r.tx) return q.DisableSessionPinProvider(ctx, sqlc.DisableSessionPinProviderParams{ - SessionKey: sessionKey[:], - Role: role, - Provider: provider, + SessionKey: sessionKey[:], + Role: role, + Provider: provider, + ExpectedRoutingStrategy: string(expectedStrategy), }) } @@ -177,6 +186,7 @@ func toSessionPin(row sqlc.RouterSessionPin) sessionpin.Pin { PairedProvider: row.PairedProvider, PairedModel: row.PairedModel, Reason: row.DecisionReason, + Strategy: router.Strategy(row.RoutingStrategy), PolicyGroup: row.PolicyGroup, TurnCount: int(row.TurnCount), PinnedUntil: timestampOrZero(row.PinnedUntil), diff --git a/internal/postgres/session_strategy_repo.go b/internal/postgres/session_strategy_repo.go new file mode 100644 index 000000000..5ef568eda --- /dev/null +++ b/internal/postgres/session_strategy_repo.go @@ -0,0 +1,67 @@ +package postgres + +import ( + "context" + "database/sql" + "errors" + + "workweave/router/internal/router" + "workweave/router/internal/router/sessionstrategy" + "workweave/router/internal/sqlc" + + "github.com/google/uuid" +) + +// SessionStrategyRepo adapts sessionstrategy.Store to SQLC-generated queries. +type SessionStrategyRepo struct { + tx sqlc.DBTX +} + +// NewSessionStrategyRepo wires the adapter over a pgx pool or transaction. +func NewSessionStrategyRepo(tx sqlc.DBTX) *SessionStrategyRepo { + return &SessionStrategyRepo{tx: tx} +} + +var _ sessionstrategy.Store = (*SessionStrategyRepo)(nil) + +// Get reads the explicit beta preference. A missing row means stable routing. +func (r *SessionStrategyRepo) Get(ctx context.Context, installationID uuid.UUID, sessionKey [sessionstrategy.SessionKeyLen]byte) (sessionstrategy.Preference, bool, error) { + strategy, err := sqlc.New(r.tx).GetSessionStrategyPreference(ctx, sqlc.GetSessionStrategyPreferenceParams{ + InstallationID: installationID, + SessionKey: sessionKey[:], + }) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return sessionstrategy.Preference{}, false, nil + } + return sessionstrategy.Preference{}, false, err + } + return sessionstrategy.Preference{ + InstallationID: installationID, + SessionKey: sessionKey, + Strategy: router.Strategy(strategy), + }, true, nil +} + +// Toggle flips the explicit beta preference in one statement and returns the +// state now persisted. Disabled sessions keep a row and use stable routing. +func (r *SessionStrategyRepo) Toggle(ctx context.Context, preference sessionstrategy.Preference) (bool, error) { + if err := preference.Validate(); err != nil { + return false, err + } + return sqlc.New(r.tx).UpsertToggledSessionStrategyPreference(ctx, sqlc.UpsertToggledSessionStrategyPreferenceParams{ + InstallationID: preference.InstallationID, + SessionKey: preference.SessionKey[:], + Strategy: string(preference.Strategy), + }) +} + +// Disable turns the explicit beta preference off in one statement and reports +// whether it had been enabled. +func (r *SessionStrategyRepo) Disable(ctx context.Context, installationID uuid.UUID, sessionKey [sessionstrategy.SessionKeyLen]byte) (bool, error) { + disabled, err := sqlc.New(r.tx).UpdateSessionStrategyPreferenceDisabled(ctx, sqlc.UpdateSessionStrategyPreferenceDisabledParams{ + InstallationID: installationID, + SessionKey: sessionKey[:], + }) + return disabled > 0, err +} diff --git a/internal/postgres/session_strategy_repo_test.go b/internal/postgres/session_strategy_repo_test.go new file mode 100644 index 000000000..b22ab9f9f --- /dev/null +++ b/internal/postgres/session_strategy_repo_test.go @@ -0,0 +1,228 @@ +package postgres + +import ( + "context" + "database/sql" + "errors" + "strings" + "testing" + "time" + + "workweave/router/internal/router" + "workweave/router/internal/router/sessionpin" + "workweave/router/internal/router/sessionstrategy" + "workweave/router/internal/sqlc" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type sessionStrategyExec struct { + query string + args []any +} + +type sessionStrategyDB struct { + row pgx.Row + execTag pgconn.CommandTag + execErr error + execCalls []sessionStrategyExec + rowQuery string + rowArgs []any +} + +func (db *sessionStrategyDB) Exec(_ context.Context, query string, args ...any) (pgconn.CommandTag, error) { + db.execCalls = append(db.execCalls, sessionStrategyExec{query: query, args: args}) + return db.execTag, db.execErr +} + +func (*sessionStrategyDB) Query(context.Context, string, ...any) (pgx.Rows, error) { + return nil, errors.New("unexpected Query call") +} + +func (db *sessionStrategyDB) QueryRow(_ context.Context, query string, args ...any) pgx.Row { + db.rowQuery = query + db.rowArgs = args + return db.row +} + +type sessionStrategyRow struct { + strategy string + err error +} + +type sessionStrategyEnabledRow struct { + enabled bool + err error +} + +func (row sessionStrategyEnabledRow) Scan(dest ...any) error { + if row.err != nil { + return row.err + } + *dest[0].(*bool) = row.enabled + return nil +} + +func (row sessionStrategyRow) Scan(dest ...any) error { + if row.err != nil { + return row.err + } + *dest[0].(*string) = row.strategy + return nil +} + +func TestSessionStrategyRepoGet(t *testing.T) { + t.Parallel() + + installationID := uuid.New() + key := [sessionstrategy.SessionKeyLen]byte{1, 2, 3} + db := &sessionStrategyDB{row: sessionStrategyRow{strategy: string(router.StrategyHMMBeta)}} + repo := NewSessionStrategyRepo(db) + + preference, ok, err := repo.Get(context.Background(), installationID, key) + + require.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, sessionstrategy.Preference{ + InstallationID: installationID, + SessionKey: key, + Strategy: router.StrategyHMMBeta, + }, preference) + assert.Contains(t, db.rowQuery, "installation_id = $1::uuid") + assert.Contains(t, db.rowQuery, "session_key = $2::bytea") + require.Len(t, db.rowArgs, 2) + assert.Equal(t, installationID, db.rowArgs[0]) + assert.Equal(t, key[:], db.rowArgs[1]) +} + +func TestSessionStrategyRepoGetMissingMeansStable(t *testing.T) { + t.Parallel() + + db := &sessionStrategyDB{row: sessionStrategyRow{err: sql.ErrNoRows}} + repo := NewSessionStrategyRepo(db) + + preference, ok, err := repo.Get(context.Background(), uuid.New(), [sessionstrategy.SessionKeyLen]byte{}) + + require.NoError(t, err) + assert.False(t, ok) + assert.Empty(t, preference) +} + +func TestSessionStrategyRepoToggleAcceptsOnlyHMMBeta(t *testing.T) { + t.Parallel() + + installationID := uuid.New() + key := [sessionstrategy.SessionKeyLen]byte{4, 5, 6} + db := &sessionStrategyDB{row: sessionStrategyEnabledRow{enabled: true}} + repo := NewSessionStrategyRepo(db) + + _, err := repo.Toggle(context.Background(), sessionstrategy.Preference{ + InstallationID: installationID, + SessionKey: key, + Strategy: "stable", + }) + require.ErrorIs(t, err, sessionstrategy.ErrInvalidStrategy) + assert.Empty(t, db.rowQuery) + + enabled, err := repo.Toggle(context.Background(), sessionstrategy.Preference{ + InstallationID: installationID, + SessionKey: key, + Strategy: router.StrategyHMMBeta, + }) + require.NoError(t, err) + assert.True(t, enabled) + assert.True(t, strings.Contains(db.rowQuery, "ON CONFLICT (installation_id, session_key)")) + assert.Contains(t, db.rowQuery, "enabled = NOT router.session_strategy_preferences.enabled") + assert.Contains(t, db.rowQuery, "RETURNING enabled") + assert.Equal(t, []any{installationID, key[:], string(router.StrategyHMMBeta)}, db.rowArgs) +} + +func TestSessionStrategyRepoDisableReportsWhetherBetaWasOn(t *testing.T) { + t.Parallel() + + installationID := uuid.New() + key := [sessionstrategy.SessionKeyLen]byte{4, 5, 6} + db := &sessionStrategyDB{execTag: pgconn.NewCommandTag("UPDATE 1")} + repo := NewSessionStrategyRepo(db) + + wasEnabled, err := repo.Disable(context.Background(), installationID, key) + require.NoError(t, err) + assert.True(t, wasEnabled) + require.Len(t, db.execCalls, 1) + assert.Contains(t, db.execCalls[0].query, "SET enabled = FALSE") + assert.Contains(t, db.execCalls[0].query, "AND enabled") + assert.Equal(t, []any{installationID, key[:]}, db.execCalls[0].args) + + db.execTag = pgconn.NewCommandTag("UPDATE 0") + wasEnabled, err = repo.Disable(context.Background(), installationID, key) + require.NoError(t, err) + assert.False(t, wasEnabled, "a session already on stable routing was not enabled") +} + +func TestSessionStrategyRepoGetIgnoresDisabledRows(t *testing.T) { + t.Parallel() + + db := &sessionStrategyDB{row: sessionStrategyRow{err: sql.ErrNoRows}} + repo := NewSessionStrategyRepo(db) + + _, ok, err := repo.Get(context.Background(), uuid.New(), [sessionstrategy.SessionKeyLen]byte{7, 8, 9}) + + require.NoError(t, err) + assert.False(t, ok) + assert.Contains(t, db.rowQuery, "AND enabled") +} + +func TestSessionPinConversionPreservesRoutingStrategy(t *testing.T) { + t.Parallel() + + pin := toSessionPin(sqlc.RouterSessionPin{RoutingStrategy: "hmm_beta"}) + assert.Equal(t, router.StrategyHMMBeta, pin.Strategy) +} + +func TestSessionPinMutationsCarryExpectedRoutingStrategy(t *testing.T) { + t.Parallel() + + key := [sessionpin.SessionKeyLen]byte{10, 11, 12} + db := &sessionStrategyDB{row: sessionStrategyRow{err: sql.ErrNoRows}} + repo := NewSessionPinRepo(db) + ctx := context.Background() + + _, ok, err := repo.Consume(ctx, key, "default", router.StrategyCluster) + require.NoError(t, err) + assert.False(t, ok) + assert.Contains(t, db.rowQuery, "routing_strategy = ''") + assert.Contains(t, db.rowQuery, "<> 'hmm_beta'") + assert.Equal(t, []any{key[:], "default", string(router.StrategyCluster)}, db.rowArgs) + + _, ok, err = repo.Consume(ctx, key, "default", router.StrategyHMMBeta) + require.NoError(t, err) + assert.False(t, ok) + assert.Equal(t, []any{key[:], "default", string(router.StrategyHMMBeta)}, db.rowArgs) + + _, err = repo.IncrementUpstreamErrors(ctx, key, "default", router.StrategyHMMBeta) + require.NoError(t, err) + assert.Equal(t, []any{key[:], "default", string(router.StrategyHMMBeta)}, db.rowArgs) + + _, err = repo.IncrementOverloadErrors(ctx, key, "default", router.StrategyHMMBeta) + require.NoError(t, err) + assert.Equal(t, []any{key[:], "default", string(router.StrategyHMMBeta)}, db.rowArgs) + + require.NoError(t, repo.UpdateUsage(ctx, key, "default", sessionpin.Usage{ + Strategy: router.StrategyHMMBeta, + EndedAt: time.Unix(1, 0), + })) + require.NoError(t, repo.ResetUpstreamErrors(ctx, key, "default", router.StrategyHMMBeta)) + require.NoError(t, repo.ResetOverloadErrors(ctx, key, "default", router.StrategyHMMBeta)) + require.NoError(t, repo.DisableProvider(ctx, key, "default", "anthropic", router.StrategyHMMBeta)) + + require.Len(t, db.execCalls, 4) + for _, call := range db.execCalls { + assert.Contains(t, call.query, "routing_strategy = ''") + assert.Contains(t, call.query, "<> 'hmm_beta'") + assert.Equal(t, string(router.StrategyHMMBeta), call.args[len(call.args)-1]) + } +} diff --git a/internal/proxy/beta.go b/internal/proxy/beta.go new file mode 100644 index 000000000..2e1aad661 --- /dev/null +++ b/internal/proxy/beta.go @@ -0,0 +1,176 @@ +package proxy + +import ( + "context" + "fmt" + "net/http" + + "workweave/router/internal/observability" + "workweave/router/internal/router" + "workweave/router/internal/router/catalog" + "workweave/router/internal/router/sessionpin" + "workweave/router/internal/router/sessionstrategy" + "workweave/router/internal/translate" + + "github.com/google/uuid" +) + +const ( + betaEnabledMessage = "Beta enabled. Type /beta again to turn it off." + betaDisabledMessage = "Beta disabled. Stable routing restored." + betaUsageMessage = "Usage: /beta" + betaUnavailable = "Beta is unavailable for this session." +) + +type betaArtifactHistoryContextKey struct{} + +// Historical /beta control turns prove this transcript crossed a policy +// boundary even though its strategy-specific pin history was invalidated. +func withBetaArtifactHistory(ctx context.Context) context.Context { + return context.WithValue(ctx, betaArtifactHistoryContextKey{}, true) +} + +func betaArtifactHistoryFromContext(ctx context.Context) bool { + found, _ := ctx.Value(betaArtifactHistoryContextKey{}).(bool) + return found +} + +// WithSessionStrategyStore wires durable per-session routing preferences. +// A nil store leaves normal routing unchanged and makes /beta unavailable. +func (s *Service) WithSessionStrategyStore(store sessionstrategy.Store) *Service { + s.sessionStrategyStore = store + return s +} + +func (s *Service) applySessionStrategy( + ctx context.Context, + installationID uuid.UUID, + sessionKey [sessionpin.SessionKeyLen]byte, +) (context.Context, error) { + if s.sessionStrategyStore == nil || installationID == uuid.Nil || sessionKey == ([sessionpin.SessionKeyLen]byte{}) { + return ctx, nil + } + preference, found, err := s.sessionStrategyStore.Get(ctx, installationID, sessionKey) + if err != nil { + return ctx, fmt.Errorf("load session routing strategy: %w", err) + } + if !found { + return ctx, nil + } + if preference.Strategy != router.StrategyHMMBeta { + return ctx, fmt.Errorf("unsupported persisted session routing strategy %q", preference.Strategy) + } + return router.WithStrategy(ctx, preference.Strategy), nil +} + +func (s *Service) handleBetaCommand( + ctx context.Context, + w http.ResponseWriter, + env *translate.RequestEnvelope, + cmd translate.BetaCommandResult, + installationID uuid.UUID, + sessionKey [sessionpin.SessionKeyLen]byte, + inputTokens int, +) error { + if cmd.Invalid { + return writeBetaCommandResponse(w, env, betaUsageMessage, inputTokens) + } + if s.sessionStrategyStore == nil || installationID == uuid.Nil || sessionKey == ([sessionpin.SessionKeyLen]byte{}) || env.ClientSessionID() == "" { + return writeBetaCommandResponse(w, env, betaUnavailable, inputTokens) + } + + // Both branches decide from what the store persisted rather than a prior + // read, so overlapping /beta commands for one session cannot act on the + // same stale state: an unavailable beta policy can only ever be left. + nowEnabled := false + if s.PolicyStrategyAvailable(router.StrategyHMMBeta) { + enabled, err := s.sessionStrategyStore.Toggle(context.Background(), sessionstrategy.Preference{ + InstallationID: installationID, + SessionKey: sessionKey, + Strategy: router.StrategyHMMBeta, + }) + if err != nil { + return fmt.Errorf("toggle beta routing: %w", err) + } + nowEnabled = enabled + } else { + wasEnabled, err := s.sessionStrategyStore.Disable(context.Background(), installationID, sessionKey) + if err != nil { + return fmt.Errorf("disable beta routing: %w", err) + } + if !wasEnabled { + return writeBetaCommandResponse(w, env, betaUnavailable, inputTokens) + } + } + + message := betaEnabledMessage + previousStrategy := router.StrategyFromContext(ctx) + if !nowEnabled { + previousStrategy = router.StrategyHMMBeta + message = betaDisabledMessage + } + + log := observability.FromContext(ctx) + // Persist first: strategy-bound reads make old-strategy pins ineligible + // immediately, so cleanup cannot overwrite a concurrent new-mode pin. + if err := s.invalidateSessionRoutingState( + router.WithStrategy(ctx, previousStrategy), + sessionKey, + ); err != nil { + log.Error("post-toggle routing-state cleanup failed", "err", err) + } + + log.Info( + "session beta routing toggled", + "enabled", nowEnabled, + "session_key_prefix", shortSessionKey(sessionKey), + ) + return writeBetaCommandResponse(w, env, message, inputTokens) +} + +func (s *Service) invalidateSessionRoutingState( + ctx context.Context, + sessionKey [sessionpin.SessionKeyLen]byte, +) error { + if s.pinStore == nil { + return nil + } + roles := []string{ + roleForTier(catalog.TierUnknown), + roleForTier(catalog.TierLow), + roleForTier(catalog.TierMid), + roleForTier(catalog.TierHigh), + } + strategy := router.StrategyFromContext(ctx) + seen := make(map[string]struct{}, len(roles)*3) + var firstErr error + for _, role := range roles { + for _, stateRole := range []string{ + role, + hmmHistoryRole(role), + commandContinuationRole(role), + } { + if _, duplicate := seen[stateRole]; duplicate { + continue + } + seen[stateRole] = struct{}{} + if _, _, err := s.pinStore.Consume(context.Background(), sessionKey, stateRole, strategy); err != nil && firstErr == nil { + firstErr = err + } + } + } + return firstErr +} + +func writeBetaCommandResponse( + w http.ResponseWriter, + env *translate.RequestEnvelope, + message string, + inputTokens int, +) error { + text := "✦ **Weave Router** → " + message + "\n\n" + if env.SourceFormat() == translate.FormatOpenAI { + return writeSyntheticOpenAIResponse(w, env, text, inputTokens) + } + return writeSyntheticAnthropicResponse(w, env, text, inputTokens) +} diff --git a/internal/proxy/beta_internal_test.go b/internal/proxy/beta_internal_test.go new file mode 100644 index 000000000..744be51ab --- /dev/null +++ b/internal/proxy/beta_internal_test.go @@ -0,0 +1,686 @@ +package proxy + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "sync" + "testing" + + "workweave/router/internal/providers" + "workweave/router/internal/router" + "workweave/router/internal/router/policy" + "workweave/router/internal/router/sessionpin" + "workweave/router/internal/router/sessionstrategy" + "workweave/router/internal/translate" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +type betaTestRouter struct { + calls int + requests []router.Request + decision router.Decision +} + +type betaCaptureProvider struct { + body []byte +} + +func (p *betaCaptureProvider) Proxy( + _ context.Context, + _ router.Decision, + req providers.PreparedRequest, + _ http.ResponseWriter, + _ *http.Request, +) error { + p.body = append([]byte(nil), req.Body...) + return nil +} + +func (p *betaCaptureProvider) Passthrough( + _ context.Context, + req providers.PreparedRequest, + _ http.ResponseWriter, + _ *http.Request, +) error { + p.body = append([]byte(nil), req.Body...) + return nil +} + +func (r *betaTestRouter) Route(_ context.Context, req router.Request) (router.Decision, error) { + r.calls++ + r.requests = append(r.requests, req) + return r.decision, nil +} + +type betaTestPreferenceStore struct { + mu sync.Mutex + preference sessionstrategy.Preference + found bool + toggles int + disables int + getErr error + toggleErr error + beforeWrite func() +} + +func (s *betaTestPreferenceStore) Get( + _ context.Context, + installationID uuid.UUID, + sessionKey [sessionstrategy.SessionKeyLen]byte, +) (sessionstrategy.Preference, bool, error) { + s.mu.Lock() + preference, found, err := s.preference, s.found, s.getErr + s.mu.Unlock() + if err != nil { + return sessionstrategy.Preference{}, false, err + } + if !found || preference.InstallationID != installationID || preference.SessionKey != sessionKey { + return sessionstrategy.Preference{}, false, nil + } + return preference, true, nil +} + +func (s *betaTestPreferenceStore) Toggle(_ context.Context, preference sessionstrategy.Preference) (bool, error) { + if err := preference.Validate(); err != nil { + return false, err + } + if s.beforeWrite != nil { + s.beforeWrite() + } + s.mu.Lock() + defer s.mu.Unlock() + if s.toggleErr != nil { + return false, s.toggleErr + } + s.toggles++ + if s.found { + s.preference = sessionstrategy.Preference{} + s.found = false + return false, nil + } + s.preference = preference + s.found = true + return true, nil +} + +func (s *betaTestPreferenceStore) Disable( + _ context.Context, + installationID uuid.UUID, + sessionKey [sessionstrategy.SessionKeyLen]byte, +) (bool, error) { + if s.beforeWrite != nil { + s.beforeWrite() + } + s.mu.Lock() + defer s.mu.Unlock() + if s.toggleErr != nil { + return false, s.toggleErr + } + s.disables++ + if !s.found || s.preference.InstallationID != installationID || s.preference.SessionKey != sessionKey { + return false, nil + } + s.preference = sessionstrategy.Preference{} + s.found = false + return true, nil +} + +type betaCleanupPinStore struct { + consumeCalls int + consumedStrategy []router.Strategy + upsertCalls int +} + +func (s *betaCleanupPinStore) Get(context.Context, [sessionpin.SessionKeyLen]byte, string) (sessionpin.Pin, bool, error) { + return sessionpin.Pin{}, false, nil +} + +func (s *betaCleanupPinStore) Consume(_ context.Context, _ [sessionpin.SessionKeyLen]byte, _ string, strategy router.Strategy) (sessionpin.Pin, bool, error) { + s.consumeCalls++ + s.consumedStrategy = append(s.consumedStrategy, strategy) + return sessionpin.Pin{}, false, nil +} + +func (s *betaCleanupPinStore) Upsert(context.Context, sessionpin.Pin) error { + s.upsertCalls++ + return nil +} +func (*betaCleanupPinStore) UpdateUsage(context.Context, [sessionpin.SessionKeyLen]byte, string, sessionpin.Usage) error { + return nil +} +func (*betaCleanupPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { + return 0, nil +} +func (*betaCleanupPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { + return nil +} +func (*betaCleanupPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { + return 0, nil +} +func (*betaCleanupPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { + return nil +} +func (*betaCleanupPinStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string, router.Strategy) error { + return nil +} +func (*betaCleanupPinStore) SweepExpired(context.Context) error { return nil } + +func betaTestEnvelope(t *testing.T, text string, withSession bool) *translate.RequestEnvelope { + t.Helper() + metadata := "" + if withSession { + metadata = `,"metadata":{"user_id":"user_account__session_4dbee464-ebf7-437f-9f20-db5a6f7fe3b4"}` + } + env, err := translate.ParseAnthropic([]byte( + `{"model":"claude-sonnet-5","messages":[{"role":"user","content":` + + mustJSONQuote(t, text) + `}],"max_tokens":128` + metadata + `}`, + )) + require.NoError(t, err) + return env +} + +func mustJSONQuote(t *testing.T, value string) string { + t.Helper() + quoted, err := json.Marshal(value) + require.NoError(t, err) + return string(quoted) +} + +func TestHandleBetaCommandTogglesAndAcknowledges(t *testing.T) { + store := &betaTestPreferenceStore{} + pins := &betaCleanupPinStore{} + svc := (&Service{pinStore: pins}). + WithPolicyStrategy(policy.StrategySpec{Strategy: router.StrategyHMMBeta, Router: &betaTestRouter{}}). + WithSessionStrategyStore(store) + installationID := uuid.New() + var sessionKey [sessionstrategy.SessionKeyLen]byte + sessionKey[0] = 1 + + for _, want := range []string{betaEnabledMessage, betaDisabledMessage} { + env := betaTestEnvelope(t, "/beta", true) + cmd, found := env.ExtractBetaCommand() + require.True(t, found) + response := httptest.NewRecorder() + require.NoError(t, svc.handleBetaCommand( + context.Background(), response, env, cmd, installationID, sessionKey, 1, + )) + assert.Equal(t, "✦ **Weave Router** → "+want+"\n\n", gjson.Get(response.Body.String(), "content.0.text").String()) + } + + assert.Equal(t, 2, store.toggles) + assert.False(t, store.found) + require.NotEmpty(t, pins.consumedStrategy) + assert.Equal(t, router.StrategyCluster, pins.consumedStrategy[0]) + assert.Equal(t, router.StrategyHMMBeta, pins.consumedStrategy[len(pins.consumedStrategy)-1]) +} + +func TestHandleBetaCommandOverlappingTogglesAcknowledgeDistinctStates(t *testing.T) { + store := &betaTestPreferenceStore{} + var readers sync.WaitGroup + readers.Add(2) + store.beforeWrite = func() { + readers.Done() + readers.Wait() + } + svc := (&Service{}). + WithPolicyStrategy(policy.StrategySpec{Strategy: router.StrategyHMMBeta, Router: &betaTestRouter{}}). + WithSessionStrategyStore(store) + installationID := uuid.New() + var sessionKey [sessionstrategy.SessionKeyLen]byte + sessionKey[0] = 1 + + acknowledgements := make(chan string, 2) + var toggles sync.WaitGroup + for range 2 { + toggles.Add(1) + go func() { + defer toggles.Done() + env := betaTestEnvelope(t, "/beta", true) + cmd, found := env.ExtractBetaCommand() + assert.True(t, found) + response := httptest.NewRecorder() + assert.NoError(t, svc.handleBetaCommand( + context.Background(), response, env, cmd, installationID, sessionKey, 1, + )) + acknowledgements <- gjson.Get(response.Body.String(), "content.0.text").String() + }() + } + toggles.Wait() + close(acknowledgements) + + var acked []string + for text := range acknowledgements { + acked = append(acked, text) + } + assert.ElementsMatch(t, []string{ + "✦ **Weave Router** → " + betaEnabledMessage + "\n\n", + "✦ **Weave Router** → " + betaDisabledMessage + "\n\n", + }, acked) + assert.Equal(t, 2, store.toggles) + assert.False(t, store.found, "two overlapping toggles from stable must land back on stable") +} + +func TestHandleBetaCommandRejectsArgumentsWithoutStateChange(t *testing.T) { + store := &betaTestPreferenceStore{} + svc := (&Service{}). + WithPolicyStrategy(policy.StrategySpec{Strategy: router.StrategyHMMBeta, Router: &betaTestRouter{}}). + WithSessionStrategyStore(store) + env := betaTestEnvelope(t, "/beta status", true) + cmd, found := env.ExtractBetaCommand() + require.True(t, found) + response := httptest.NewRecorder() + var sessionKey [sessionstrategy.SessionKeyLen]byte + sessionKey[0] = 1 + + require.NoError(t, svc.handleBetaCommand( + context.Background(), response, env, cmd, uuid.New(), sessionKey, 1, + )) + assert.Contains(t, gjson.Get(response.Body.String(), "content.0.text").String(), betaUsageMessage) + assert.Zero(t, store.toggles) +} + +func TestHandleBetaCommandRequiresClientSessionAndAvailablePolicy(t *testing.T) { + for _, tt := range []struct { + name string + withSession bool + withPolicy bool + }{ + {name: "missing client session", withPolicy: true}, + {name: "missing beta policy", withSession: true}, + } { + t.Run(tt.name, func(t *testing.T) { + store := &betaTestPreferenceStore{} + svc := (&Service{}).WithSessionStrategyStore(store) + if tt.withPolicy { + svc.WithPolicyStrategy(policy.StrategySpec{Strategy: router.StrategyHMMBeta, Router: &betaTestRouter{}}) + } + env := betaTestEnvelope(t, "/beta", tt.withSession) + cmd, found := env.ExtractBetaCommand() + require.True(t, found) + response := httptest.NewRecorder() + var sessionKey [sessionstrategy.SessionKeyLen]byte + sessionKey[0] = 1 + + require.NoError(t, svc.handleBetaCommand( + context.Background(), response, env, cmd, uuid.New(), sessionKey, 1, + )) + assert.Contains(t, gjson.Get(response.Body.String(), "content.0.text").String(), betaUnavailable) + assert.Zero(t, store.toggles) + }) + } +} + +func TestHandleBetaCommandWriteFailureLeavesStablePinsUntouched(t *testing.T) { + store := &betaTestPreferenceStore{toggleErr: errors.New("write failed")} + pins := &betaCleanupPinStore{} + svc := (&Service{pinStore: pins}). + WithPolicyStrategy(policy.StrategySpec{Strategy: router.StrategyHMMBeta, Router: &betaTestRouter{}}). + WithSessionStrategyStore(store) + env := betaTestEnvelope(t, "/beta", true) + cmd, found := env.ExtractBetaCommand() + require.True(t, found) + var sessionKey [sessionstrategy.SessionKeyLen]byte + sessionKey[0] = 1 + + err := svc.handleBetaCommand( + context.Background(), httptest.NewRecorder(), env, cmd, uuid.New(), sessionKey, 1, + ) + + require.Error(t, err) + assert.False(t, store.found) + assert.Zero(t, pins.consumeCalls, "failed enable must not disturb stable routing state") +} + +func TestHandleBetaCommandCanDisablePersistedBetaWhilePolicyUnavailable(t *testing.T) { + installationID := uuid.New() + var sessionKey [sessionstrategy.SessionKeyLen]byte + sessionKey[0] = 1 + store := &betaTestPreferenceStore{ + preference: sessionstrategy.Preference{ + InstallationID: installationID, + SessionKey: sessionKey, + Strategy: router.StrategyHMMBeta, + }, + found: true, + } + svc := (&Service{}).WithSessionStrategyStore(store) + env := betaTestEnvelope(t, "/beta", true) + cmd, found := env.ExtractBetaCommand() + require.True(t, found) + response := httptest.NewRecorder() + + require.NoError(t, svc.handleBetaCommand( + context.Background(), response, env, cmd, installationID, sessionKey, 1, + )) + assert.False(t, store.found) + assert.Equal(t, 1, store.disables) + assert.Zero(t, store.toggles, "an unavailable policy must never flip the preference") + assert.Contains(t, gjson.Get(response.Body.String(), "content.0.text").String(), betaDisabledMessage) +} + +func TestHandleBetaCommandOverlappingTogglesCannotReenableUnavailableBeta(t *testing.T) { + installationID := uuid.New() + var sessionKey [sessionstrategy.SessionKeyLen]byte + sessionKey[0] = 1 + store := &betaTestPreferenceStore{ + preference: sessionstrategy.Preference{ + InstallationID: installationID, + SessionKey: sessionKey, + Strategy: router.StrategyHMMBeta, + }, + found: true, + } + var arrivals sync.WaitGroup + arrivals.Add(2) + store.beforeWrite = func() { + arrivals.Done() + arrivals.Wait() + } + svc := (&Service{}).WithSessionStrategyStore(store) + + acknowledgements := make(chan string, 2) + var commands sync.WaitGroup + for range 2 { + commands.Add(1) + go func() { + defer commands.Done() + env := betaTestEnvelope(t, "/beta", true) + cmd, found := env.ExtractBetaCommand() + assert.True(t, found) + response := httptest.NewRecorder() + assert.NoError(t, svc.handleBetaCommand( + context.Background(), response, env, cmd, installationID, sessionKey, 1, + )) + acknowledgements <- gjson.Get(response.Body.String(), "content.0.text").String() + }() + } + commands.Wait() + close(acknowledgements) + + var acked []string + for text := range acknowledgements { + acked = append(acked, text) + } + assert.ElementsMatch(t, []string{ + "✦ **Weave Router** → " + betaDisabledMessage + "\n\n", + "✦ **Weave Router** → " + betaUnavailable + "\n\n", + }, acked) + assert.False(t, store.found, "beta must stay off while its policy is unavailable") +} + +func TestApplySessionStrategyOnlyUsesPersistedPreference(t *testing.T) { + installationID := uuid.New() + var sessionKey [sessionstrategy.SessionKeyLen]byte + sessionKey[0] = 1 + stableCtx := router.WithStrategy(context.Background(), router.StrategyHMM) + store := &betaTestPreferenceStore{} + svc := (&Service{}).WithSessionStrategyStore(store) + + ctx, err := svc.applySessionStrategy(stableCtx, installationID, sessionKey) + require.NoError(t, err) + assert.Equal(t, router.StrategyHMM, router.StrategyFromContext(ctx)) + + store.preference = sessionstrategy.Preference{ + InstallationID: installationID, + SessionKey: sessionKey, + Strategy: router.StrategyHMMBeta, + } + store.found = true + ctx, err = svc.applySessionStrategy(stableCtx, installationID, sessionKey) + require.NoError(t, err) + assert.Equal(t, router.StrategyHMMBeta, router.StrategyFromContext(ctx)) +} + +func TestApplySessionStrategyReturnsOriginalContextOnStoreError(t *testing.T) { + marker := struct{}{} + ctx := context.WithValue(context.Background(), marker, "kept") + svc := (&Service{}).WithSessionStrategyStore(&betaTestPreferenceStore{getErr: errors.New("db down")}) + var sessionKey [sessionstrategy.SessionKeyLen]byte + sessionKey[0] = 1 + + got, err := svc.applySessionStrategy(ctx, uuid.New(), sessionKey) + + require.Error(t, err) + assert.Same(t, ctx, got) + assert.Equal(t, "kept", got.Value(marker)) +} + +func TestPersistedBetaPreferenceFailsClosedWhenBetaRouterUnavailable(t *testing.T) { + installationID := uuid.New() + var sessionKey [sessionstrategy.SessionKeyLen]byte + sessionKey[0] = 1 + store := &betaTestPreferenceStore{ + preference: sessionstrategy.Preference{ + InstallationID: installationID, + SessionKey: sessionKey, + Strategy: router.StrategyHMMBeta, + }, + found: true, + } + svc := (&Service{}). + WithPolicyStrategy(policy.StrategySpec{ + Strategy: router.StrategyHMMBeta, + Router: nil, + Unavailable: errors.New("beta unavailable"), + }). + WithSessionStrategyStore(store) + + ctx, err := svc.applySessionStrategy(context.Background(), installationID, sessionKey) + require.NoError(t, err) + assert.Equal(t, router.StrategyHMMBeta, router.StrategyFromContext(ctx)) + + _, err = svc.routeFor(ctx, router.Request{}) + require.Error(t, err) + assert.Contains(t, err.Error(), "beta unavailable") +} + +func TestProxyEntrypointsInterceptBetaWithoutRoutingUpstream(t *testing.T) { + for _, tt := range []struct { + name string + openAI bool + }{ + {name: "anthropic"}, + {name: "openai", openAI: true}, + } { + t.Run(tt.name, func(t *testing.T) { + stableRouter := &betaTestRouter{} + betaRouter := &betaTestRouter{} + store := &betaTestPreferenceStore{} + pins := &betaCleanupPinStore{} + svc := NewService(stableRouter, nil, nil, false, nil, pins, false, "", "", nil). + WithPolicyStrategy(policy.StrategySpec{Strategy: router.StrategyHMMBeta, Router: betaRouter}). + WithSessionStrategyStore(store) + ctx := context.WithValue(context.Background(), APIKeyIDContextKey{}, "beta-test-key") + ctx = context.WithValue(ctx, InstallationIDContextKey{}, uuid.NewString()) + body := []byte(`{"model":"claude-sonnet-5","messages":[{"role":"user","content":"/beta"}],"max_tokens":128,"metadata":{"user_id":"user_account__session_4dbee464-ebf7-437f-9f20-db5a6f7fe3b4"}}`) + request := httptest.NewRequest("POST", "/v1/messages", nil) + response := httptest.NewRecorder() + + var err error + if tt.openAI { + request = httptest.NewRequest("POST", "/v1/chat/completions", nil) + err = svc.ProxyOpenAIChatCompletion(ctx, body, response, request) + } else { + err = svc.ProxyMessages(ctx, body, response, request) + } + require.NoError(t, err) + if tt.openAI { + assert.Contains(t, gjson.Get(response.Body.String(), "choices.0.message.content").String(), betaEnabledMessage) + } else { + assert.Contains(t, gjson.Get(response.Body.String(), "content.0.text").String(), betaEnabledMessage) + } + assert.True(t, store.found) + assert.Zero(t, stableRouter.calls) + assert.Zero(t, betaRouter.calls) + assert.Zero(t, pins.upsertCalls, "/beta must not grant a continuation or write a routing pin") + }) + } +} + +func TestBetaToggleSelectsBetaThenRestoresStableRouter(t *testing.T) { + stableRouter := &betaTestRouter{} + betaRouter := &betaTestRouter{} + store := &betaTestPreferenceStore{} + svc := NewService(stableRouter, nil, nil, false, nil, nil, false, "", "", nil). + WithPolicyStrategy(policy.StrategySpec{Strategy: router.StrategyHMMBeta, Router: betaRouter}). + WithSessionStrategyStore(store) + installationID := uuid.New() + var sessionKey [sessionstrategy.SessionKeyLen]byte + sessionKey[0] = 1 + baseCtx := router.WithStrategy(context.Background(), router.StrategyCluster) + + enableEnv := betaTestEnvelope(t, "/beta", true) + enableCmd, found := enableEnv.ExtractBetaCommand() + require.True(t, found) + require.NoError(t, svc.handleBetaCommand( + baseCtx, httptest.NewRecorder(), enableEnv, enableCmd, installationID, sessionKey, 1, + )) + betaCtx, err := svc.applySessionStrategy(baseCtx, installationID, sessionKey) + require.NoError(t, err) + _, err = svc.routeFor(betaCtx, router.Request{}) + require.NoError(t, err) + assert.Equal(t, 1, betaRouter.calls) + assert.Zero(t, stableRouter.calls) + + disableEnv := betaTestEnvelope(t, "/beta", true) + disableCmd, found := disableEnv.ExtractBetaCommand() + require.True(t, found) + require.NoError(t, svc.handleBetaCommand( + baseCtx, httptest.NewRecorder(), disableEnv, disableCmd, installationID, sessionKey, 1, + )) + stableCtx, err := svc.applySessionStrategy(baseCtx, installationID, sessionKey) + require.NoError(t, err) + _, err = svc.routeFor(stableCtx, router.Request{}) + require.NoError(t, err) + assert.Equal(t, 1, betaRouter.calls) + assert.Equal(t, 1, stableRouter.calls) +} + +func TestProxyEntrypointsStripHistoricalBetaArtifactsBeforeRouting(t *testing.T) { + for _, tt := range []struct { + name string + openAI bool + tools []any + }{ + { + name: "anthropic", + tools: []any{ + map[string]any{"name": "Read", "input_schema": map[string]any{"type": "object"}}, + }, + }, + { + name: "openai", + openAI: true, + tools: []any{ + map[string]any{"type": "function", "function": map[string]any{"name": "Read", "parameters": map[string]any{"type": "object"}}}, + }, + }, + } { + t.Run(tt.name, func(t *testing.T) { + decisionProvider := providers.ProviderAnthropic + decisionModel := "claude-haiku-4-5" + if tt.openAI { + decisionProvider = providers.ProviderOpenAI + decisionModel = "gpt-5.5" + } + routing := &betaTestRouter{decision: router.Decision{ + Provider: decisionProvider, + Model: decisionModel, + Reason: "test", + }} + svc := NewService( + routing, + map[string]providers.Client{ + providers.ProviderAnthropic: embedTestProvider{}, + providers.ProviderOpenAI: embedTestProvider{}, + }, + nil, false, nil, nil, false, + providers.ProviderAnthropic, "claude-haiku-4-5", nil, + ) + body, err := json.Marshal(map[string]any{ + "model": "claude-opus-4-8", + "messages": []any{ + map[string]any{"role": "user", "content": "inspect this repository"}, + map[string]any{"role": "assistant", "content": "I will inspect it."}, + map[string]any{"role": "user", "content": "/beta"}, + map[string]any{"role": "assistant", "content": "✦ **Weave Router** → Beta enabled. Type /beta again to turn it off.\n\n"}, + map[string]any{"role": "user", "content": "continue with the implementation"}, + }, + "tools": tt.tools, + "max_tokens": 4096, + }) + require.NoError(t, err) + request := httptest.NewRequest("POST", "/v1/messages", nil) + response := httptest.NewRecorder() + if tt.openAI { + request = httptest.NewRequest("POST", "/v1/chat/completions", nil) + err = svc.ProxyOpenAIChatCompletion(context.Background(), body, response, request) + } else { + err = svc.ProxyMessages(context.Background(), body, response, request) + } + require.NoError(t, err) + require.Len(t, routing.requests, 1) + assert.NotContains(t, routing.requests[0].PromptText, "/beta") + assert.NotContains(t, routing.requests[0].PromptText, "Beta enabled") + assert.Contains(t, routing.requests[0].PromptText, "continue with the implementation") + assert.Len(t, routing.requests[0].ConversationMessages, 3, + "the command and acknowledgement must both be absent from routing features") + }) + } +} + +func TestHistoricalBetaArtifactsStripStaleThinkingSignatures(t *testing.T) { + for _, tt := range []struct { + name string + ack string + }{ + {name: "after enabling beta", ack: "✦ **Weave Router** → Beta enabled. Type /beta again to turn it off.\n\n"}, + {name: "after restoring stable", ack: "✦ **Weave Router** → Beta disabled. Stable routing restored.\n\n"}, + } { + t.Run(tt.name, func(t *testing.T) { + routing := &betaTestRouter{decision: router.Decision{ + Provider: providers.ProviderAnthropic, + Model: "claude-opus-4-7", + Reason: "test", + }} + provider := &betaCaptureProvider{} + svc := NewService( + routing, + map[string]providers.Client{providers.ProviderAnthropic: provider}, + nil, false, nil, nil, false, + providers.ProviderAnthropic, "claude-opus-4-7", nil, + ) + body := []byte(`{ + "model":"claude-opus-4-7", + "messages":[ + {"role":"user","content":"inspect this repository"}, + {"role":"assistant","content":[ + {"type":"thinking","thinking":"old thought","signature":"stale-signature"}, + {"type":"text","text":"I will inspect it."} + ]}, + {"role":"user","content":"/beta"}, + {"role":"assistant","content":` + mustJSONQuote(t, tt.ack) + `}, + {"role":"user","content":"continue with the implementation"} + ], + "max_tokens":4096, + "thinking":{"type":"adaptive"} + }`) + request := httptest.NewRequest("POST", "/v1/messages", nil) + response := httptest.NewRecorder() + + require.NoError(t, svc.ProxyMessages(context.Background(), body, response, request)) + require.NotEmpty(t, provider.body) + assert.NotContains(t, string(provider.body), "stale-signature") + assert.NotContains(t, string(provider.body), `"type":"thinking"`) + assert.NotContains(t, string(provider.body), "/beta") + assert.Contains(t, string(provider.body), "continue with the implementation") + }) + } +} diff --git a/internal/proxy/command_continuation.go b/internal/proxy/command_continuation.go index f2b389a0d..51605df38 100644 --- a/internal/proxy/command_continuation.go +++ b/internal/proxy/command_continuation.go @@ -4,6 +4,7 @@ import ( "context" "workweave/router/internal/observability" + "workweave/router/internal/router" "workweave/router/internal/router/sessionpin" "github.com/google/uuid" @@ -48,6 +49,7 @@ func (s *Service) grantPostCommandContinuation( } pin.Role = commandContinuationRole(role) pin.InstallationID = installationID + pin.Strategy = router.StrategyFromContext(ctx) pin.TurnCount = 1 // A slash command can arrive just before the active pin expires. Renew the // one-shot independently so its immediate follow-up remains eligible. @@ -66,11 +68,14 @@ func (s *Service) consumePostCommandContinuation( if s.pinStore == nil { return sessionpin.Pin{}, false } - pin, found, err := s.pinStore.Consume(ctx, sessionKey, commandContinuationRole(role)) + pin, found, err := s.pinStore.Consume(ctx, sessionKey, commandContinuationRole(role), router.StrategyFromContext(ctx)) if err != nil { observability.FromContext(ctx).Error("post-command continuation consume failed", "err", err) return sessionpin.Pin{}, false } + if found && !pinMatchesEffectiveStrategy(ctx, pin) { + return sessionpin.Pin{}, false + } return pin, found } @@ -85,7 +90,7 @@ func (s *Service) invalidatePostCommandContinuation( if s.pinStore == nil { return nil } - _, _, err := s.pinStore.Consume(context.Background(), sessionKey, commandContinuationRole(role)) + _, _, err := s.pinStore.Consume(context.Background(), sessionKey, commandContinuationRole(role), router.StrategyFromContext(ctx)) if err != nil { observability.FromContext(ctx).Error( "post-command continuation invalidation failed", diff --git a/internal/proxy/cyber_refusal_internal_test.go b/internal/proxy/cyber_refusal_internal_test.go index 136536eda..45a5f1f61 100644 --- a/internal/proxy/cyber_refusal_internal_test.go +++ b/internal/proxy/cyber_refusal_internal_test.go @@ -34,22 +34,22 @@ func (f *repinFakeStore) Upsert(_ context.Context, p sessionpin.Pin) error { func (f *repinFakeStore) UpdateUsage(context.Context, [sessionpin.SessionKeyLen]byte, string, sessionpin.Usage) error { return nil } -func (f *repinFakeStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (f *repinFakeStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (f *repinFakeStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (f *repinFakeStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (f *repinFakeStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (f *repinFakeStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (f *repinFakeStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (f *repinFakeStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (f *repinFakeStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string) error { +func (f *repinFakeStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string, router.Strategy) error { return nil } -func (f *repinFakeStore) Consume(context.Context, [sessionpin.SessionKeyLen]byte, string) (sessionpin.Pin, bool, error) { +func (f *repinFakeStore) Consume(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (sessionpin.Pin, bool, error) { return sessionpin.Pin{}, false, nil } func (f *repinFakeStore) SweepExpired(context.Context) error { return nil } diff --git a/internal/proxy/force_cluster_internal_test.go b/internal/proxy/force_cluster_internal_test.go index e38185009..aadabc863 100644 --- a/internal/proxy/force_cluster_internal_test.go +++ b/internal/proxy/force_cluster_internal_test.go @@ -33,7 +33,7 @@ func TestApplyForceClusterHeader_AbsentIsNoOp(t *testing.T) { } func TestApplyForceClusterHeader_ThreadsLabelOnHMMStrategy(t *testing.T) { - for _, strategy := range []router.Strategy{router.StrategyHMM, router.StrategyHMMEmbedding} { + for _, strategy := range []router.Strategy{router.StrategyHMM, router.StrategyHMMEmbedding, router.StrategyHMMBeta} { t.Run(string(strategy), func(t *testing.T) { ctx := router.WithStrategy(context.Background(), strategy) diff --git a/internal/proxy/force_model.go b/internal/proxy/force_model.go index 1e083c993..ade64a026 100644 --- a/internal/proxy/force_model.go +++ b/internal/proxy/force_model.go @@ -279,7 +279,7 @@ func (s *Service) setForceModelPin( existing, found, err := s.pinStore.Get(ctx, sessionKey, role) if err != nil { log.Error("force-model: prior pin lookup failed", "err", err) - } else if found { + } else if found && pinMatchesEffectiveStrategy(ctx, existing) { lastServedModel = existing.LastServedModel } forced := sessionpin.Pin{ @@ -289,6 +289,7 @@ func (s *Service) setForceModelPin( Provider: provider, Model: canonicalModel, Reason: translate.ReasonUserForceModel, + Strategy: router.StrategyFromContext(ctx), TurnCount: 1, PinnedUntil: pinNeverExpires, LastServedModel: lastServedModel, diff --git a/internal/proxy/force_model_pin_ttl_internal_test.go b/internal/proxy/force_model_pin_ttl_internal_test.go index 1f049e1a9..369bb3106 100644 --- a/internal/proxy/force_model_pin_ttl_internal_test.go +++ b/internal/proxy/force_model_pin_ttl_internal_test.go @@ -6,6 +6,7 @@ import ( "time" "workweave/router/internal/providers" + "workweave/router/internal/router" "workweave/router/internal/router/sessionpin" "workweave/router/internal/translate" @@ -31,22 +32,22 @@ func (s *recordingPinStore) Upsert(_ context.Context, p sessionpin.Pin) error { func (s *recordingPinStore) UpdateUsage(context.Context, [sessionpin.SessionKeyLen]byte, string, sessionpin.Usage) error { return nil } -func (s *recordingPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *recordingPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (s *recordingPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *recordingPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (s *recordingPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *recordingPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (s *recordingPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *recordingPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (s *recordingPinStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string) error { +func (s *recordingPinStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string, router.Strategy) error { return nil } -func (s *recordingPinStore) Consume(context.Context, [sessionpin.SessionKeyLen]byte, string) (sessionpin.Pin, bool, error) { +func (s *recordingPinStore) Consume(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (sessionpin.Pin, bool, error) { return sessionpin.Pin{}, false, nil } func (s *recordingPinStore) SweepExpired(context.Context) error { return nil } @@ -88,3 +89,17 @@ func TestSetForceModelPin_WritesNeverExpiresSentinel(t *testing.T) { assert.Equal(t, pinNeverExpires, store.upserts[0].PinnedUntil, "a /force-model pin must be written with the never-expires sentinel") } + +func TestSetForceModelPin_WritesEffectiveStrategy(t *testing.T) { + store := &recordingPinStore{} + svc := NewService(nil, nil, nil, false, nil, store, false, + providers.ProviderAnthropic, "claude-haiku-4-5", nil) + + ctx := router.WithStrategy(context.Background(), router.StrategyHMMBeta) + require.NoError(t, svc.setForceModelPin( + ctx, [sessionpin.SessionKeyLen]byte{}, sessionpin.DefaultRole, uuid.New(), + "claude-opus-4-8", providers.ProviderAnthropic)) + + require.Len(t, store.upserts, 1) + assert.Equal(t, router.StrategyHMMBeta, store.upserts[0].Strategy) +} diff --git a/internal/proxy/force_model_tier_fallback_internal_test.go b/internal/proxy/force_model_tier_fallback_internal_test.go index 9116bb48c..f75bf14d6 100644 --- a/internal/proxy/force_model_tier_fallback_internal_test.go +++ b/internal/proxy/force_model_tier_fallback_internal_test.go @@ -59,22 +59,22 @@ func (s *forcedPinStore) Upsert(context.Context, sessionpin.Pin) error { return func (s *forcedPinStore) UpdateUsage(context.Context, [sessionpin.SessionKeyLen]byte, string, sessionpin.Usage) error { return nil } -func (s *forcedPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *forcedPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (s *forcedPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *forcedPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (s *forcedPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *forcedPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (s *forcedPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *forcedPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (s *forcedPinStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string) error { +func (s *forcedPinStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string, router.Strategy) error { return nil } -func (s *forcedPinStore) Consume(context.Context, [sessionpin.SessionKeyLen]byte, string) (sessionpin.Pin, bool, error) { +func (s *forcedPinStore) Consume(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (sessionpin.Pin, bool, error) { return sessionpin.Pin{}, false, nil } func (s *forcedPinStore) SweepExpired(context.Context) error { return nil } @@ -102,22 +102,22 @@ func (s *overwritingPinStore) Upsert(_ context.Context, p sessionpin.Pin) error func (s *overwritingPinStore) UpdateUsage(context.Context, [sessionpin.SessionKeyLen]byte, string, sessionpin.Usage) error { return nil } -func (s *overwritingPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *overwritingPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (s *overwritingPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *overwritingPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (s *overwritingPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *overwritingPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (s *overwritingPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *overwritingPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (s *overwritingPinStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string) error { +func (s *overwritingPinStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string, router.Strategy) error { return nil } -func (s *overwritingPinStore) Consume(context.Context, [sessionpin.SessionKeyLen]byte, string) (sessionpin.Pin, bool, error) { +func (s *overwritingPinStore) Consume(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (sessionpin.Pin, bool, error) { return sessionpin.Pin{}, false, nil } func (s *overwritingPinStore) SweepExpired(context.Context) error { return nil } diff --git a/internal/proxy/loop_detection.go b/internal/proxy/loop_detection.go index 557679cbd..25daf49a1 100644 --- a/internal/proxy/loop_detection.go +++ b/internal/proxy/loop_detection.go @@ -7,6 +7,7 @@ import ( "workweave/router/internal/observability" "workweave/router/internal/providers" + "workweave/router/internal/router" "workweave/router/internal/router/sessionpin" "workweave/router/internal/translate" @@ -146,7 +147,7 @@ func (s *Service) handleLoopEscalation( existing, found, err := s.pinStore.Get(ctx, sessionKey, role) if err != nil { log.Error("loop-escalation: prior pin lookup failed", "err", err) - } else if found { + } else if found && pinMatchesEffectiveStrategy(ctx, existing) { if existing.Reason == translate.ReasonLoopEscalation { return // already rescued this session; don't re-pin or double-log } @@ -218,7 +219,7 @@ func (s *Service) handleLoopEscalation( return } var lastServed string - if existing, found, err := s.pinStore.Get(ctx, sessionKey, role); err == nil && found { + if existing, found, err := s.pinStore.Get(ctx, sessionKey, role); err == nil && found && pinMatchesEffectiveStrategy(ctx, existing) { lastServed = existing.LastServedModel } pin := sessionpin.Pin{ @@ -228,6 +229,7 @@ func (s *Service) handleLoopEscalation( Provider: providers.ProviderAnthropic, Model: escalateModel, Reason: translate.ReasonLoopEscalation, + Strategy: router.StrategyFromContext(ctx), TurnCount: 1, PinnedUntil: time.Now().Add(pinSessionTTL), LastServedModel: lastServed, diff --git a/internal/proxy/pin_eviction.go b/internal/proxy/pin_eviction.go index e3d80fe5e..8391c05f5 100644 --- a/internal/proxy/pin_eviction.go +++ b/internal/proxy/pin_eviction.go @@ -7,6 +7,7 @@ import ( "workweave/router/internal/observability" "workweave/router/internal/providers" + "workweave/router/internal/router" "workweave/router/internal/router/sessionpin" "workweave/router/internal/translate" @@ -65,6 +66,7 @@ func (s *Service) expireSessionPinRow( Provider: "", Model: "", Reason: reason, + Strategy: router.StrategyFromContext(ctx), TurnCount: 1, PinnedUntil: time.Now().Add(-time.Second), } @@ -192,7 +194,7 @@ func (s *Service) maybeEvictPinAfterUpstreamErr( if proxyErr == nil { // context.Background(): the request ctx is already canceled by the // time streaming finishes, but this reset must still go through. - if err := s.pinStore.ResetUpstreamErrors(context.Background(), sessionKey, role); err != nil { + if err := s.pinStore.ResetUpstreamErrors(context.Background(), sessionKey, role, router.StrategyFromContext(ctx)); err != nil { log.Error("pin error-counter reset failed", "err", err, "role", role) } return @@ -208,7 +210,7 @@ func (s *Service) maybeEvictPinAfterUpstreamErr( return } - count, err := s.pinStore.IncrementUpstreamErrors(context.Background(), sessionKey, role) + count, err := s.pinStore.IncrementUpstreamErrors(context.Background(), sessionKey, role, router.StrategyFromContext(ctx)) if err != nil { log.Error("pin error-counter increment failed", "err", err, "role", role, "upstream_status", status) return diff --git a/internal/proxy/pin_eviction_internal_test.go b/internal/proxy/pin_eviction_internal_test.go index ced926de9..351f01688 100644 --- a/internal/proxy/pin_eviction_internal_test.go +++ b/internal/proxy/pin_eviction_internal_test.go @@ -9,6 +9,7 @@ import ( "time" "workweave/router/internal/providers" + "workweave/router/internal/router" "workweave/router/internal/router/sessionpin" "workweave/router/internal/translate" @@ -44,7 +45,7 @@ func (s *evictionStubPinStore) UpdateUsage(context.Context, [sessionpin.SessionK return nil } -func (s *evictionStubPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *evictionStubPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { s.mu.Lock() defer s.mu.Unlock() s.incrementCalls++ @@ -56,26 +57,26 @@ func (s *evictionStubPinStore) IncrementUpstreamErrors(context.Context, [session return v, nil } -func (s *evictionStubPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *evictionStubPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { s.mu.Lock() defer s.mu.Unlock() s.resetCalls++ return nil } -func (s *evictionStubPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *evictionStubPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (s *evictionStubPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *evictionStubPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (s *evictionStubPinStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string) error { +func (s *evictionStubPinStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string, router.Strategy) error { return nil } -func (s *evictionStubPinStore) Consume(_ context.Context, key [sessionpin.SessionKeyLen]byte, role string) (sessionpin.Pin, bool, error) { +func (s *evictionStubPinStore) Consume(_ context.Context, key [sessionpin.SessionKeyLen]byte, role string, _ router.Strategy) (sessionpin.Pin, bool, error) { s.mu.Lock() defer s.mu.Unlock() s.consumeRoles = append(s.consumeRoles, role) diff --git a/internal/proxy/pin_strategy.go b/internal/proxy/pin_strategy.go new file mode 100644 index 000000000..735fa80fb --- /dev/null +++ b/internal/proxy/pin_strategy.go @@ -0,0 +1,35 @@ +package proxy + +import ( + "context" + + "workweave/router/internal/router" + "workweave/router/internal/router/sessionpin" +) + +// pinMatchesEffectiveStrategy reports whether a stored pin belongs to the +// strategy serving this request. Legacy (empty Strategy) rows remain eligible +// for non-beta strategies during rollout; beta never inherits a legacy pin. +func pinMatchesEffectiveStrategy(ctx context.Context, pin sessionpin.Pin) bool { + expected := router.StrategyFromContext(ctx) + if pin.Strategy == expected { + return true + } + return pin.Strategy == "" && expected != router.StrategyHMMBeta +} + +func strategyContext(strategy router.Strategy) context.Context { + return router.WithStrategy(context.Background(), strategy) +} + +func strategyForTurnLoopResult(res turnLoopResult) router.Strategy { + if res.Strategy != "" { + return res.Strategy + } + for _, decision := range []router.Decision{res.Decision, res.Fresh} { + if decision.Metadata != nil && decision.Metadata.Strategy != "" { + return router.Strategy(decision.Metadata.Strategy) + } + } + return router.StrategyCluster +} diff --git a/internal/proxy/pin_strategy_internal_test.go b/internal/proxy/pin_strategy_internal_test.go new file mode 100644 index 000000000..85fae7476 --- /dev/null +++ b/internal/proxy/pin_strategy_internal_test.go @@ -0,0 +1,38 @@ +package proxy + +import ( + "context" + "testing" + + "workweave/router/internal/router" + "workweave/router/internal/router/sessionpin" + + "github.com/stretchr/testify/assert" +) + +func TestPinMatchesEffectiveStrategy(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + request router.Strategy + stored router.Strategy + expected bool + }{ + {name: "stable exact", request: router.StrategyCluster, stored: router.StrategyCluster, expected: true}, + {name: "beta exact", request: router.StrategyHMMBeta, stored: router.StrategyHMMBeta, expected: true}, + {name: "beta rejects stable", request: router.StrategyHMMBeta, stored: router.StrategyCluster, expected: false}, + {name: "stable rejects beta", request: router.StrategyCluster, stored: router.StrategyHMMBeta, expected: false}, + {name: "other HMM rejects beta", request: router.StrategyHMM, stored: router.StrategyHMMBeta, expected: false}, + {name: "stable accepts legacy during rollout", request: router.StrategyCluster, stored: "", expected: true}, + {name: "non-beta opt-in accepts legacy during rollout", request: router.StrategyHMM, stored: "", expected: true}, + {name: "beta rejects legacy", request: router.StrategyHMMBeta, stored: "", expected: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := router.WithStrategy(context.Background(), tt.request) + assert.Equal(t, tt.expected, pinMatchesEffectiveStrategy(ctx, sessionpin.Pin{Strategy: tt.stored})) + }) + } +} diff --git a/internal/proxy/provider_overload.go b/internal/proxy/provider_overload.go index 607cffb65..ce0823f7a 100644 --- a/internal/proxy/provider_overload.go +++ b/internal/proxy/provider_overload.go @@ -5,6 +5,7 @@ import ( "strings" "workweave/router/internal/observability" + "workweave/router/internal/router" "workweave/router/internal/router/sessionpin" "workweave/router/internal/translate" @@ -52,7 +53,7 @@ func (s *Service) maybeDisableProviderAfterOverload( if proxyErr == nil { // context.Background(): the request ctx is already canceled by the // time streaming finishes, but this reset must still go through. - if err := s.pinStore.ResetOverloadErrors(context.Background(), sessionKey, role); err != nil { + if err := s.pinStore.ResetOverloadErrors(context.Background(), sessionKey, role, router.StrategyFromContext(ctx)); err != nil { log.Error("pin overload-counter reset failed", "err", err, "role", role) } return @@ -62,7 +63,7 @@ func (s *Service) maybeDisableProviderAfterOverload( return } - count, err := s.pinStore.IncrementOverloadErrors(context.Background(), sessionKey, role) + count, err := s.pinStore.IncrementOverloadErrors(context.Background(), sessionKey, role, router.StrategyFromContext(ctx)) if err != nil { log.Error("pin overload-counter increment failed", "err", err, "role", role, "provider", finalProvider) return @@ -77,7 +78,7 @@ func (s *Service) maybeDisableProviderAfterOverload( return } - if err := s.pinStore.DisableProvider(context.Background(), sessionKey, role, finalProvider); err != nil { + if err := s.pinStore.DisableProvider(context.Background(), sessionKey, role, finalProvider, router.StrategyFromContext(ctx)); err != nil { log.Error("pin provider-disable upsert failed", "err", err, "role", role, "provider", finalProvider) return } diff --git a/internal/proxy/provider_overload_internal_test.go b/internal/proxy/provider_overload_internal_test.go index e4453fd02..7302467d4 100644 --- a/internal/proxy/provider_overload_internal_test.go +++ b/internal/proxy/provider_overload_internal_test.go @@ -9,6 +9,7 @@ import ( "time" "workweave/router/internal/providers" + "workweave/router/internal/router" "workweave/router/internal/router/sessionpin" "workweave/router/internal/translate" @@ -44,15 +45,15 @@ func (s *overloadStubPinStore) UpdateUsage(context.Context, [sessionpin.SessionK return nil } -func (s *overloadStubPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *overloadStubPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (s *overloadStubPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *overloadStubPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (s *overloadStubPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *overloadStubPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { s.mu.Lock() defer s.mu.Unlock() s.incrementCalls++ @@ -64,21 +65,21 @@ func (s *overloadStubPinStore) IncrementOverloadErrors(context.Context, [session return v, nil } -func (s *overloadStubPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *overloadStubPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { s.mu.Lock() defer s.mu.Unlock() s.resetCalls++ return nil } -func (s *overloadStubPinStore) DisableProvider(_ context.Context, _ [sessionpin.SessionKeyLen]byte, _, provider string) error { +func (s *overloadStubPinStore) DisableProvider(_ context.Context, _ [sessionpin.SessionKeyLen]byte, _, provider string, _ router.Strategy) error { s.mu.Lock() defer s.mu.Unlock() s.disabledProviders = append(s.disabledProviders, provider) return nil } -func (s *overloadStubPinStore) Consume(context.Context, [sessionpin.SessionKeyLen]byte, string) (sessionpin.Pin, bool, error) { +func (s *overloadStubPinStore) Consume(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (sessionpin.Pin, bool, error) { return sessionpin.Pin{}, false, nil } diff --git a/internal/proxy/route_preview.go b/internal/proxy/route_preview.go index da9d0c7a0..5a1727b4f 100644 --- a/internal/proxy/route_preview.go +++ b/internal/proxy/route_preview.go @@ -12,14 +12,20 @@ import ( "workweave/router/internal/providers" "workweave/router/internal/router" "workweave/router/internal/router/policy" + "workweave/router/internal/router/sessionpin" "workweave/router/internal/translate" ) -func (s *Service) anthropicRoutingRequest(ctx context.Context, body []byte, headers http.Header) (router.Request, error) { +func (s *Service) anthropicRoutingRequest( + ctx context.Context, + body []byte, + headers http.Header, + ingress string, +) (context.Context, router.Request, error) { log := observability.FromContext(ctx) cleanBody, err := stripRoutingMarkerFromMessages(body) if err != nil { - return router.Request{}, fmt.Errorf("strip routing marker: %w", err) + return ctx, router.Request{}, fmt.Errorf("strip routing marker: %w", err) } if withoutFooter, footerErr := translate.StripFeedbackFooterFromMessages(cleanBody); footerErr != nil { log.Error("Failed to strip feedback footer from route preview", "err", footerErr) @@ -34,7 +40,21 @@ func (s *Service) anthropicRoutingRequest(ctx context.Context, body []byte, head env, err := translate.ParseAnthropic(cleanBody) if err != nil { - return router.Request{}, fmt.Errorf("parse request: %w", err) + return ctx, router.Request{}, fmt.Errorf("parse request: %w", err) + } + + apiKeyID, _ := ctx.Value(APIKeyIDContextKey{}).(string) + var sessionKey [sessionpin.SessionKeyLen]byte + ctx, log, sessionKey = bindRequestLogger(ctx, env, apiKeyID, "", ingress) + if removed := env.StripRouterFeedbackArtifacts(); removed > 0 { + log.Info("Stripped router-feedback artifacts from route preview", "removed_messages", removed) + } + if removed := env.StripBetaArtifacts(); removed > 0 { + log.Info("Stripped beta artifacts from route preview", "removed_messages", removed) + } + ctx, err = s.applySessionStrategy(ctx, installationIDFromContext(ctx), sessionKey) + if err != nil { + return ctx, router.Request{}, err } embedOnlyUser := s.ResolveEmbedOnlyUserMessage(ctx) features := env.RoutingFeatures(embedOnlyUser) @@ -67,7 +87,7 @@ func (s *Service) anthropicRoutingRequest(ctx context.Context, body []byte, head if id := installationIDFromContext(ctx); id != uuid.Nil { installationID = id.String() } - return router.Request{ + return ctx, router.Request{ RequestedModel: features.Model, EstimatedInputTokens: features.Tokens, HasTools: features.HasTools, @@ -94,7 +114,7 @@ func (s *Service) anthropicRoutingRequest(ctx context.Context, body []byte, head // PreviewAnthropicRoute evaluates an Anthropic request with the registered // policy preview contract without dispatching or invoking serving lifecycle state. func (s *Service) PreviewAnthropicRoute(ctx context.Context, body []byte, headers http.Header) (policy.PreviewResult, error) { - req, err := s.anthropicRoutingRequest(ctx, body, headers) + ctx, req, err := s.anthropicRoutingRequest(ctx, body, headers, "anthropic_route_preview") if err != nil { return policy.PreviewResult{}, err } diff --git a/internal/proxy/router_feedback.go b/internal/proxy/router_feedback.go index 5427aa29e..28c711150 100644 --- a/internal/proxy/router_feedback.go +++ b/internal/proxy/router_feedback.go @@ -133,7 +133,7 @@ func (s *Service) handleRouterFeedbackCommand( if servedModel == "" && s.pinStore != nil { if pin, found, err := s.pinStore.Get(ctx, sessionKey, role); err != nil { log.Error("/router-feedback: pin lookup failed", "err", err) - } else if found { + } else if found && pinMatchesEffectiveStrategy(ctx, pin) { servedModel = pin.LastServedModel if servedModel == "" { servedModel = pin.Model diff --git a/internal/proxy/service.go b/internal/proxy/service.go index d550331fd..cf53abe42 100644 --- a/internal/proxy/service.go +++ b/internal/proxy/service.go @@ -34,6 +34,7 @@ import ( "workweave/router/internal/router/policy" "workweave/router/internal/router/rl" "workweave/router/internal/router/sessionpin" + "workweave/router/internal/router/sessionstrategy" "workweave/router/internal/router/turntype" "workweave/router/internal/sse" "workweave/router/internal/timing" @@ -66,6 +67,9 @@ type Service struct { // pinStore persists session-sticky routing decisions. Nil when the feature // flag is off; the orchestrator then runs the scorer every turn. pinStore sessionpin.Store + // sessionStrategyStore persists the explicit per-session /beta selection. + // Stable routing is represented by no row. + sessionStrategyStore sessionstrategy.Store // noProgress tracks per-session dispatch fingerprints to catch the // cross-envelope subagent loop (parent agent re-spawning identical // sub-conversations). Nil disables the detector. @@ -2302,7 +2306,7 @@ func defaultStrategyUnavailable(strategy router.Strategy) error { switch strategy { case router.StrategyRL: return rl.ErrPolicyUnavailable - case router.StrategyHMM, router.StrategyHMMEmbedding: + case router.StrategyHMM, router.StrategyHMMEmbedding, router.StrategyHMMBeta: return hmm.ErrHMMUnavailable case router.StrategyBandit: return bandit.ErrBanditUnavailable @@ -2323,7 +2327,7 @@ func (s *Service) Route(ctx context.Context, req router.Request) (router.Decisio // callers in internal/api/* never import internal/translate directly, // matching ProxyMessages. func (s *Service) RouteAnthropicRequest(ctx context.Context, body []byte, headers http.Header) (decision router.Decision, err error) { - req, err := s.anthropicRoutingRequest(ctx, body, headers) + ctx, req, err := s.anthropicRoutingRequest(ctx, body, headers, "anthropic_route") if err != nil { return decision, err } @@ -2642,7 +2646,7 @@ func (s *Service) maybeRepinOnRefusal(ctx context.Context, obs *refusalObserver, // Prefer the scorer's runner-up (PairedModel); use context.Background() because // the request ctx may already be canceled when the response has been written. fbModel, fbProvider := s.ResolveCyberRefusalFallbackModel(ctx), "" - if existing, found, err := s.pinStore.Get(context.Background(), sessionKey, role); err == nil && found && existing.PairedModel != "" { + if existing, found, err := s.pinStore.Get(context.Background(), sessionKey, role); err == nil && found && pinMatchesEffectiveStrategy(ctx, existing) && existing.PairedModel != "" { fbModel, fbProvider = existing.PairedModel, existing.PairedProvider } if fbProvider == "" { @@ -2662,6 +2666,7 @@ func (s *Service) maybeRepinOnRefusal(ctx context.Context, obs *refusalObserver, Provider: fbProvider, Model: fbModel, Reason: "cyber-refusal-repin", + Strategy: router.StrategyFromContext(ctx), TurnCount: 1, PinnedUntil: pinExpiry("cyber-refusal-repin"), } @@ -2751,6 +2756,10 @@ func (s *Service) ProxyMessages(ctx context.Context, body []byte, w http.Respons if removed := env.StripRouterFeedbackArtifacts(); removed > 0 { log.Info("Stripped router-feedback artifacts from Anthropic history", "removed_messages", removed) } + if removed := env.StripBetaArtifacts(); removed > 0 { + ctx = withBetaArtifactHistory(ctx) + log.Info("Stripped beta artifacts from Anthropic history", "removed_messages", removed) + } embedFlag := s.ResolveEmbedOnlyUserMessage(ctx) feats := env.RoutingFeatures(embedFlag) @@ -2770,6 +2779,19 @@ func (s *Service) ProxyMessages(ctx context.Context, body []byte, w http.Respons "prompt_preview", observability.Preview(promptText, 200), ) + // /beta toggle: handled server-side, never forwarded upstream, no post-command continuation. + if !agentShadowMode { + if cmd, hasCmd := env.ExtractBetaCommand(); hasCmd { + log.Info("ProxyMessages beta command") + return s.handleBetaCommand(ctx, w, env, cmd, installationID, sessionKey, feats.Tokens) + } + ctx, err = s.applySessionStrategy(ctx, installationID, sessionKey) + if err != nil { + return err + } + *r = *r.WithContext(ctx) + } + // Handle /force-model and /unforce-model before routing (stripped from // env.body so the upstream never sees it). Session key is derived before // extraction: DeriveSessionKey can fall back to prompt text, and deriving @@ -4389,6 +4411,7 @@ func (s *Service) recordTurnUsage(res turnLoopResult, servedProvider, servedMode return } usage := sessionpin.Usage{ + Strategy: strategyForTurnLoopResult(res), InputTokens: in, CachedReadTokens: cacheRead, CachedWriteTokens: cacheCreation, @@ -4421,23 +4444,25 @@ func (s *Service) recordHMMTurnHistory(res turnLoopResult, servedProvider, serve return } hasUsage := in != 0 || out != 0 || cacheCreation != 0 || cacheRead != 0 + strategyCtx := strategyContext(strategyForTurnLoopResult(res)) historyProvider := servedProvider if !hasUsage { // A failed turn has no usage writeback; preserve the prior provider to // avoid an invalid model/provider pair on the next HMM stay. - if prior := s.loadHMMHistory(context.Background(), res.SessionKey, res.PinRole); prior.Provider != "" { + if prior := s.loadHMMHistory(strategyCtx, res.SessionKey, res.PinRole); prior.Provider != "" { historyProvider = prior.Provider } } role := hmmHistoryRole(res.PinRole) // The upsert only refreshes the row's TTL/turn_count/provider (ON CONFLICT // leaves the usage columns untouched), so it is always safe to run. - s.upsertPin(context.Background(), sessionpin.Pin{ + s.upsertPin(strategyCtx, sessionpin.Pin{ SessionKey: res.SessionKey, Role: role, InstallationID: res.InstallationID, Provider: historyProvider, Reason: hmmHistoryStoredReason(res), + Strategy: router.StrategyFromContext(strategyCtx), TurnCount: 1, PinnedUntil: pinExpiry(hmmHistoryReason), }) @@ -4448,6 +4473,7 @@ func (s *Service) recordHMMTurnHistory(res turnLoopResult, servedProvider, serve } now := time.Now() if err := s.pinStore.UpdateUsage(context.Background(), res.SessionKey, role, sessionpin.Usage{ + Strategy: router.StrategyFromContext(strategyCtx), InputTokens: in, CachedReadTokens: cacheRead, CachedWriteTokens: cacheCreation, @@ -5357,6 +5383,10 @@ func (s *Service) ProxyOpenAIChatCompletion(ctx context.Context, body []byte, w if removed := env.StripRouterFeedbackArtifacts(); removed > 0 { log.Info("Stripped router-feedback artifacts from OpenAI history", "removed_messages", removed) } + if removed := env.StripBetaArtifacts(); removed > 0 { + ctx = withBetaArtifactHistory(ctx) + log.Info("Stripped beta artifacts from OpenAI history", "removed_messages", removed) + } embedFlag := s.ResolveEmbedOnlyUserMessage(ctx) feats := env.RoutingFeatures(embedFlag) promptText := feats.PromptText @@ -5377,6 +5407,17 @@ func (s *Service) ProxyOpenAIChatCompletion(ctx context.Context, body []byte, w "prompt_preview", observability.Preview(promptText, 200), ) + // /beta toggle: handled server-side before other routing commands; no post-command continuation. + if cmd, hasCmd := env.ExtractBetaCommand(); hasCmd { + log.Info("ProxyOpenAIChatCompletion beta command") + return s.handleBetaCommand(ctx, w, env, cmd, installationID, sessionKey, feats.Tokens) + } + ctx, err = s.applySessionStrategy(ctx, installationID, sessionKey) + if err != nil { + return err + } + *r = *r.WithContext(ctx) + // Handle /force-model and /unforce-model before routing (stripped from // env.body so the upstream never sees it). Session key is derived before // extraction: DeriveSessionKey can fall back to prompt text, and deriving diff --git a/internal/proxy/service_session_pin_test.go b/internal/proxy/service_session_pin_test.go index 6b0f04042..517c78302 100644 --- a/internal/proxy/service_session_pin_test.go +++ b/internal/proxy/service_session_pin_test.go @@ -109,7 +109,7 @@ func (f *fakePinStore) Upsert(ctx context.Context, p sessionpin.Pin) error { return nil } -func (f *fakePinStore) Consume(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role string) (sessionpin.Pin, bool, error) { +func (f *fakePinStore) Consume(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role string, expected router.Strategy) (sessionpin.Pin, bool, error) { f.mu.Lock() defer f.mu.Unlock() pin, found := f.commandContinuations[role] @@ -133,7 +133,7 @@ func (f *fakePinStore) UpdateUsage(ctx context.Context, key [sessionpin.SessionK return nil } -func (f *fakePinStore) IncrementUpstreamErrors(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role string) (int, error) { +func (f *fakePinStore) IncrementUpstreamErrors(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role string, expected router.Strategy) (int, error) { f.mu.Lock() defer f.mu.Unlock() f.incrementCalls++ @@ -143,14 +143,14 @@ func (f *fakePinStore) IncrementUpstreamErrors(ctx context.Context, key [session return f.incrementReturns, nil } -func (f *fakePinStore) ResetUpstreamErrors(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role string) error { +func (f *fakePinStore) ResetUpstreamErrors(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role string, expected router.Strategy) error { f.mu.Lock() defer f.mu.Unlock() f.resetCalls++ return nil } -func (f *fakePinStore) IncrementOverloadErrors(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role string) (int, error) { +func (f *fakePinStore) IncrementOverloadErrors(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role string, expected router.Strategy) (int, error) { f.mu.Lock() defer f.mu.Unlock() f.overloadIncrementCalls++ @@ -163,7 +163,7 @@ func (f *fakePinStore) IncrementOverloadErrors(ctx context.Context, key [session return f.pin.ConsecutiveOverloadErrors, nil } -func (f *fakePinStore) ResetOverloadErrors(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role string) error { +func (f *fakePinStore) ResetOverloadErrors(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role string, expected router.Strategy) error { f.mu.Lock() defer f.mu.Unlock() f.overloadResetCalls++ @@ -173,7 +173,7 @@ func (f *fakePinStore) ResetOverloadErrors(ctx context.Context, key [sessionpin. return nil } -func (f *fakePinStore) DisableProvider(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role, provider string) error { +func (f *fakePinStore) DisableProvider(ctx context.Context, key [sessionpin.SessionKeyLen]byte, role, provider string, expected router.Strategy) error { f.mu.Lock() defer f.mu.Unlock() f.disabledProviders = append(f.disabledProviders, provider) diff --git a/internal/proxy/struggle_escalation.go b/internal/proxy/struggle_escalation.go index 0c9bb3293..5afec310a 100644 --- a/internal/proxy/struggle_escalation.go +++ b/internal/proxy/struggle_escalation.go @@ -5,6 +5,7 @@ import ( "time" "workweave/router/internal/observability" + "workweave/router/internal/router" "workweave/router/internal/router/catalog" "workweave/router/internal/router/sessionpin" "workweave/router/internal/translate" @@ -85,7 +86,7 @@ func (s *Service) handleStruggleEscalation( log.Error("struggle-escalation: pin lookup failed", "err", err) return } - if !found { + if !found || !pinMatchesEffectiveStrategy(ctx, pin) { return } @@ -189,6 +190,7 @@ func (s *Service) handleStruggleEscalation( Provider: m.Providers[0].Provider, Model: target, Reason: translate.ReasonStruggleEscalation, + Strategy: router.StrategyFromContext(ctx), TurnCount: 1, PinnedUntil: time.Now().Add(pinSessionTTL), PolicyGroup: targetCluster, diff --git a/internal/proxy/turnloop.go b/internal/proxy/turnloop.go index 341dde6ff..a987b495f 100644 --- a/internal/proxy/turnloop.go +++ b/internal/proxy/turnloop.go @@ -162,9 +162,12 @@ type turnLoopResult struct { Decision router.Decision SessionKey [sessionpin.SessionKeyLen]byte InstallationID uuid.UUID - TurnType turntype.TurnType - StickyHit bool - HardPinned bool + // Strategy is the effective request strategy, carried through the response + // path so async pin writes stay strategy-bound after ctx is cancelled. + Strategy router.Strategy + TurnType turntype.TurnType + StickyHit bool + HardPinned bool // AuthoritativePerTurn is true only for eligible main/tool-result turns // whose active policy declared model-authoritative dispatch. AuthoritativePerTurn bool @@ -553,10 +556,12 @@ func (s *Service) runTurnLoop( req.InstallationID = installationID.String() } res := turnLoopResult{ - InstallationID: installationID, - TurnType: turntype.DetectFromEnvelope(env, feats, subAgentHint), - PinTier: "miss", - RequestedTier: catalog.TierFor(feats.Model), + InstallationID: installationID, + Strategy: router.StrategyFromContext(ctx), + TurnType: turntype.DetectFromEnvelope(env, feats, subAgentHint), + PinTier: "miss", + RequestedTier: catalog.TierFor(feats.Model), + StripThinkingBlocks: betaArtifactHistoryFromContext(ctx), } res.AuthoritativePerTurn = authoritativePolicyTurn(res.TurnType) && s.authoritativePerTurnSelection(ctx) @@ -1817,6 +1822,9 @@ func (s *Service) loadPin(ctx context.Context, sessionKey [sessionpin.SessionKey if !found { return sessionpin.Pin{}, false } + if !pinMatchesEffectiveStrategy(ctx, pin) { + return sessionpin.Pin{}, false + } if !pin.PinnedUntil.After(time.Now()) { return pin, false } @@ -1833,6 +1841,9 @@ func (s *Service) loadHMMHistory(ctx context.Context, sessionKey [sessionpin.Ses if !found { return sessionpin.Pin{} } + if !pinMatchesEffectiveStrategy(ctx, pin) { + return sessionpin.Pin{} + } return pin } @@ -1916,6 +1927,7 @@ func (s *Service) refreshPin(ctx context.Context, installationID uuid.UUID, sess PairedProvider: existing.PairedProvider, PairedModel: existing.PairedModel, Reason: chosen.Reason, + Strategy: router.StrategyFromContext(ctx), // Same rationale as the pair above: a refresh runs no policy, so the // reconstructed decision carries no group. Carry the stored one forward. PolicyGroup: existing.PolicyGroup, @@ -1957,6 +1969,7 @@ func (s *Service) writeNewPin(ctx context.Context, installationID uuid.UUID, ses PairedProvider: pairedProvider, PairedModel: pairedModel, Reason: chosen.Reason, + Strategy: router.StrategyFromContext(ctx), PolicyGroup: decisionPolicyGroup(chosen), TurnCount: 1, PinnedUntil: pinExpiry(chosen.Reason), @@ -1969,6 +1982,9 @@ func (s *Service) writeNewPin(ctx context.Context, installationID uuid.UUID, ses // finished streaming. func (s *Service) upsertPin(ctx context.Context, p sessionpin.Pin) { log := observability.FromContext(ctx) + if p.Strategy == "" { + p.Strategy = router.StrategyFromContext(ctx) + } if err := s.pinStore.Upsert(context.Background(), p); err != nil { log.Error("session pin upsert failed", "err", err) return diff --git a/internal/proxy/turnloop_internal_test.go b/internal/proxy/turnloop_internal_test.go index ee0ac78d3..0dd7dd055 100644 --- a/internal/proxy/turnloop_internal_test.go +++ b/internal/proxy/turnloop_internal_test.go @@ -26,6 +26,9 @@ type stubPinStore struct { usageRoles []string getPin sessionpin.Pin getFound bool + consumePin sessionpin.Pin + consumeHit bool + consumeFor router.Strategy upserts []sessionpin.Pin upsertErr error } @@ -59,28 +62,31 @@ func (s *stubPinStore) UpdateUsage(_ context.Context, _ [sessionpin.SessionKeyLe return nil } -func (s *stubPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *stubPinStore) IncrementUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (s *stubPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *stubPinStore) ResetUpstreamErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (s *stubPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) (int, error) { +func (s *stubPinStore) IncrementOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) (int, error) { return 0, nil } -func (s *stubPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string) error { +func (s *stubPinStore) ResetOverloadErrors(context.Context, [sessionpin.SessionKeyLen]byte, string, router.Strategy) error { return nil } -func (s *stubPinStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string) error { +func (s *stubPinStore) DisableProvider(context.Context, [sessionpin.SessionKeyLen]byte, string, string, router.Strategy) error { return nil } -func (s *stubPinStore) Consume(context.Context, [sessionpin.SessionKeyLen]byte, string) (sessionpin.Pin, bool, error) { - return sessionpin.Pin{}, false, nil +func (s *stubPinStore) Consume(_ context.Context, _ [sessionpin.SessionKeyLen]byte, _ string, expected router.Strategy) (sessionpin.Pin, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + s.consumeFor = expected + return s.consumePin, s.consumeHit, nil } func (s *stubPinStore) SweepExpired(context.Context) error { return nil } @@ -100,6 +106,23 @@ func TestPinCacheCold_OrdinaryAndHMMShareTheSameRule(t *testing.T) { assert.True(t, pinCacheCold(cold, false), "an expired pin is cold without a prefix break") } +func TestConsumePostCommandContinuation_RequiresEffectiveStrategy(t *testing.T) { + store := newStubPinStore() + store.consumeHit = true + store.consumePin = sessionpin.Pin{Strategy: router.StrategyCluster} + svc := NewService(nil, nil, nil, false, nil, store, false, + providers.ProviderAnthropic, "claude-haiku-4-5", nil) + ctx := router.WithStrategy(context.Background(), router.StrategyHMMBeta) + + _, found := svc.consumePostCommandContinuation(ctx, [sessionpin.SessionKeyLen]byte{1}, sessionpin.DefaultRole) + assert.False(t, found, "a non-beta continuation must not cross into a beta session") + + store.mu.Lock() + defer store.mu.Unlock() + assert.Equal(t, router.StrategyHMMBeta, store.consumeFor, + "the atomic consume must be filtered by the effective strategy in storage too") +} + func TestApplyPinEvidence_UsesAvailablePriorTurnData(t *testing.T) { t.Parallel() @@ -250,12 +273,13 @@ func TestRecordTurnUsage_HMMDecisionWritesHistoryOnly(t *testing.T) { res := turnLoopResult{ InstallationID: uuid.New(), + Strategy: router.StrategyHMMBeta, Decision: router.Decision{ Provider: "anthropic", Model: "claude-sonnet-5", Reason: "hmm_policy(label=high)", Metadata: &router.RoutingMetadata{ - Strategy: string(router.StrategyHMM), + Strategy: string(router.StrategyHMMBeta), RouteID: "route-1", }, }, @@ -275,6 +299,7 @@ func TestRecordTurnUsage_HMMDecisionWritesHistoryOnly(t *testing.T) { assert.Equal(t, "hmm_policy(label=high)", store.upserts[0].Reason) assert.Equal(t, providers.ProviderAnthropic, store.upserts[0].Provider) assert.Empty(t, store.upserts[0].Model, "HMM history rows must not be routable pins") + assert.Equal(t, router.StrategyHMMBeta, store.upserts[0].Strategy) assert.Equal(t, []string{hmmHistoryRole(sessionpin.DefaultRole)}, store.usageRoles) assert.NotContains(t, store.usageRoles, sessionpin.DefaultRole, "HMM turns must not mutate the active routing pin role") assert.Equal(t, 1200, store.lastUsage.InputTokens) @@ -283,6 +308,7 @@ func TestRecordTurnUsage_HMMDecisionWritesHistoryOnly(t *testing.T) { assert.Equal(t, 80, store.lastUsage.OutputTokens) assert.Equal(t, "claude-sonnet-5", store.lastUsage.ServedModel) assert.Equal(t, "claude-haiku-4-5", store.lastUsage.PriorServedModel) + assert.Equal(t, router.StrategyHMMBeta, store.lastUsage.Strategy) } func TestRecordTurnUsage_HMMModelChangeWritesCurrentUsageOnly(t *testing.T) { @@ -1334,6 +1360,44 @@ func TestLoadPin_ServesFreshPostgresPin(t *testing.T) { assert.Equal(t, "anthropic", pin.Provider) } +func TestLoadPin_RequiresBetaStrategyMatch(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + request router.Strategy + stored router.Strategy + found bool + }{ + {name: "beta exact", request: router.StrategyHMMBeta, stored: router.StrategyHMMBeta, found: true}, + {name: "beta rejects stable", request: router.StrategyHMMBeta, stored: router.StrategyCluster, found: false}, + {name: "beta rejects legacy", request: router.StrategyHMMBeta, stored: "", found: false}, + {name: "stable rejects beta", request: router.StrategyCluster, stored: router.StrategyHMMBeta, found: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + store := newStubPinStore() + store.getFound = true + store.getPin = sessionpin.Pin{ + Provider: providers.ProviderAnthropic, + Model: "claude-opus-4-7", + Strategy: tt.stored, + PinnedUntil: time.Now().Add(time.Hour), + } + svc := NewService(nil, nil, nil, false, nil, store, false, + providers.ProviderAnthropic, "claude-haiku-4-5", nil) + ctx := router.WithStrategy(context.Background(), tt.request) + + pin, found := svc.loadPin(ctx, [sessionpin.SessionKeyLen]byte{1}, sessionpin.DefaultRole) + assert.Equal(t, tt.found, found) + if !tt.found { + assert.Equal(t, sessionpin.Pin{}, pin, "a mismatched pin must not leak reuse or history evidence") + } + }) + } +} + func TestSwitchHistoryFromPins_UsesHMMHistory(t *testing.T) { now := time.Now() active := sessionpin.Pin{ diff --git a/internal/router/hmm/hmm_test.go b/internal/router/hmm/hmm_test.go index 9e336b429..8dcf315ab 100644 --- a/internal/router/hmm/hmm_test.go +++ b/internal/router/hmm/hmm_test.go @@ -103,22 +103,26 @@ func TestRouterMapsSidecarRosterModelBackToCatalogDecision(t *testing.T) { assert.False(t, candidate.Capabilities.SupportsImages) } -func TestRouterUsesSeparatelySelectableEmbeddingStrategy(t *testing.T) { - decider := &fakeDecider{res: Result{Model: "moonshotai/kimi-k2.7-code"}} - r := newWithRoutingTargets( - router.StrategyHMMEmbedding, - decider, - map[string]struct{}{"moonshotai/kimi-k2.7": {}}, - map[string]struct{}{providers.ProviderFireworks: {}}, - ) +func TestRouterUsesSeparatelySelectableHMMStrategies(t *testing.T) { + for _, strategy := range []router.Strategy{router.StrategyHMMEmbedding, router.StrategyHMMBeta} { + t.Run(string(strategy), func(t *testing.T) { + decider := &fakeDecider{res: Result{Model: "moonshotai/kimi-k2.7-code"}} + r := newWithRoutingTargets( + strategy, + decider, + map[string]struct{}{"moonshotai/kimi-k2.7": {}}, + map[string]struct{}{providers.ProviderFireworks: {}}, + ) + + decision, err := r.Route(context.Background(), router.Request{PromptText: "hello"}) - decision, err := r.Route(context.Background(), router.Request{PromptText: "hello"}) - - require.NoError(t, err) - assert.Equal(t, router.StrategyHMMEmbedding, decider.query.Strategy) - require.NotNil(t, decision.Metadata) - assert.Equal(t, string(router.StrategyHMMEmbedding), decision.Metadata.Strategy) - assert.Contains(t, decision.Reason, "hmm_policy") + require.NoError(t, err) + assert.Equal(t, strategy, decider.query.Strategy) + require.NotNil(t, decision.Metadata) + assert.Equal(t, string(strategy), decision.Metadata.Strategy) + assert.Contains(t, decision.Reason, "hmm_policy") + }) + } } func TestRouterKeepsGeneratedRouteIDWhenSidecarOmitsIt(t *testing.T) { diff --git a/internal/router/sessionpin/store.go b/internal/router/sessionpin/store.go index 29b48389a..0806ac0d6 100644 --- a/internal/router/sessionpin/store.go +++ b/internal/router/sessionpin/store.go @@ -6,6 +6,8 @@ import ( "context" "time" + "workweave/router/internal/router" + "github.com/google/uuid" ) @@ -43,6 +45,9 @@ type Pin struct { PairedProvider string PairedModel string Reason string + // Strategy is the routing strategy that produced the pin. During rollout, + // non-beta stable strategies may accept an empty legacy value; beta never does. + Strategy router.Strategy // PolicyGroup is the HMM complexity cluster the pinned decision came from // (RoutingMetadata.PolicyGroup). Compared to the fresh decision's group to // distinguish a within-cluster reroute from a genuine cluster change. @@ -80,6 +85,9 @@ type Pin struct { // Usage captures the previous turn's upstream token accounting. type Usage struct { + // Strategy is the routing strategy whose pin may receive this usage. The + // update is ignored when the row has since been replaced by another strategy. + Strategy router.Strategy InputTokens int CachedReadTokens int CachedWriteTokens int @@ -100,8 +108,8 @@ type Usage struct { } // Store is the I/O surface for session pins. Get returns (zero, false, nil) -// when no row exists; UpdateUsage/IncrementUpstreamErrors/ResetUpstreamErrors -// are no-ops on an evicted or missing pin. +// when no row exists. In-place mutations are no-ops on an evicted, missing, or +// different-strategy pin so a late request cannot mutate a replacement row. // // IncrementUpstreamErrors atomically bumps the error counter and returns the // new count so the turn loop can two-strike-evict without a cross-pod @@ -111,20 +119,20 @@ type Store interface { // Consume atomically removes and returns one unexpired pin. It is used for // one-shot continuations, where a Get followed by an expiry write could let // two concurrent requests reuse the same pin. - Consume(ctx context.Context, sessionKey [SessionKeyLen]byte, role string) (Pin, bool, error) + Consume(ctx context.Context, sessionKey [SessionKeyLen]byte, role string, expectedStrategy router.Strategy) (Pin, bool, error) Upsert(ctx context.Context, p Pin) error UpdateUsage(ctx context.Context, sessionKey [SessionKeyLen]byte, role string, usage Usage) error - IncrementUpstreamErrors(ctx context.Context, sessionKey [SessionKeyLen]byte, role string) (int, error) - ResetUpstreamErrors(ctx context.Context, sessionKey [SessionKeyLen]byte, role string) error + IncrementUpstreamErrors(ctx context.Context, sessionKey [SessionKeyLen]byte, role string, expectedStrategy router.Strategy) (int, error) + ResetUpstreamErrors(ctx context.Context, sessionKey [SessionKeyLen]byte, role string, expectedStrategy router.Strategy) error // IncrementOverloadErrors atomically bumps ConsecutiveOverloadErrors and // returns the new count, mirroring IncrementUpstreamErrors but for // client-visible 529 exhaustion instead of non-retryable 4xx. - IncrementOverloadErrors(ctx context.Context, sessionKey [SessionKeyLen]byte, role string) (int, error) + IncrementOverloadErrors(ctx context.Context, sessionKey [SessionKeyLen]byte, role string, expectedStrategy router.Strategy) (int, error) // ResetOverloadErrors clears ConsecutiveOverloadErrors after a successful // turn, mirroring ResetUpstreamErrors. - ResetOverloadErrors(ctx context.Context, sessionKey [SessionKeyLen]byte, role string) error + ResetOverloadErrors(ctx context.Context, sessionKey [SessionKeyLen]byte, role string, expectedStrategy router.Strategy) error // DisableProvider appends provider to DisabledProviders (deduped) and // resets ConsecutiveOverloadErrors in the same write. - DisableProvider(ctx context.Context, sessionKey [SessionKeyLen]byte, role, provider string) error + DisableProvider(ctx context.Context, sessionKey [SessionKeyLen]byte, role, provider string, expectedStrategy router.Strategy) error SweepExpired(ctx context.Context) error } diff --git a/internal/router/sessionstrategy/store.go b/internal/router/sessionstrategy/store.go new file mode 100644 index 000000000..c2a38499b --- /dev/null +++ b/internal/router/sessionstrategy/store.go @@ -0,0 +1,45 @@ +// Package sessionstrategy defines the inner-ring contract for explicit +// per-session routing strategy preferences. +package sessionstrategy + +import ( + "context" + "errors" + + "workweave/router/internal/router" + "workweave/router/internal/router/sessionpin" + + "github.com/google/uuid" +) + +// SessionKeyLen is the shared sha256-truncated session key length. +const SessionKeyLen = sessionpin.SessionKeyLen + +// ErrInvalidStrategy is returned when a caller tries to persist any strategy +// other than the explicit beta override. Stable routing is represented by no row. +var ErrInvalidStrategy = errors.New("invalid session strategy preference") + +// Preference is the explicit strategy selected for one installation-scoped session. +type Preference struct { + InstallationID uuid.UUID + SessionKey [SessionKeyLen]byte + Strategy router.Strategy +} + +// Validate rejects values that are not valid explicit preferences. +func (p Preference) Validate() error { + if p.Strategy != router.StrategyHMMBeta { + return ErrInvalidStrategy + } + return nil +} + +// Store persists explicit session strategy preferences. Get returns +// (zero, false, nil) when the session uses stable routing. Toggle and Disable +// are atomic writes so overlapping /beta commands cannot both act on the same +// prior state; each caller sees its own persisted result. +type Store interface { + Get(ctx context.Context, installationID uuid.UUID, sessionKey [SessionKeyLen]byte) (Preference, bool, error) + Toggle(ctx context.Context, preference Preference) (bool, error) + Disable(ctx context.Context, installationID uuid.UUID, sessionKey [SessionKeyLen]byte) (bool, error) +} diff --git a/internal/router/sessionstrategy/store_test.go b/internal/router/sessionstrategy/store_test.go new file mode 100644 index 000000000..836fc09db --- /dev/null +++ b/internal/router/sessionstrategy/store_test.go @@ -0,0 +1,27 @@ +package sessionstrategy_test + +import ( + "errors" + "testing" + + "workweave/router/internal/router" + "workweave/router/internal/router/sessionstrategy" + + "github.com/stretchr/testify/assert" +) + +func TestPreferenceValidateAcceptsOnlyHMMBeta(t *testing.T) { + t.Parallel() + + assert.NoError(t, (sessionstrategy.Preference{Strategy: router.StrategyHMMBeta}).Validate()) + for _, strategy := range []router.Strategy{"", "stable", router.StrategyHMM, router.StrategyCluster} { + assert.ErrorIs(t, (sessionstrategy.Preference{Strategy: strategy}).Validate(), sessionstrategy.ErrInvalidStrategy) + } +} + +func TestInvalidStrategyErrorIsStable(t *testing.T) { + t.Parallel() + + err := (sessionstrategy.Preference{Strategy: "stable"}).Validate() + assert.True(t, errors.Is(err, sessionstrategy.ErrInvalidStrategy)) +} diff --git a/internal/router/strategy.go b/internal/router/strategy.go index ceb76d1c4..409a14e93 100644 --- a/internal/router/strategy.go +++ b/internal/router/strategy.go @@ -9,9 +9,8 @@ import ( // implementation. Strategy-specific sentinels may wrap this error. var ErrStrategyUnavailable = errors.New("router: strategy unavailable") -// Strategy names a routing strategy a request can opt into via the -// x-weave-router-strategy header. The zero value ("") means the deployment -// default (cluster). +// Strategy identifies the routing policy selected for a request. The zero +// value ("") means the deployment default (cluster). type Strategy string const ( @@ -25,6 +24,9 @@ const ( StrategyHMM Strategy = "hmm" // StrategyHMMEmbedding uses the HMM sidecar with the embedding-quality-seeded bandit prior. StrategyHMMEmbedding Strategy = "hmm_embedding" + // StrategyHMMBeta routes through an independently deployed beta HMM policy. + // It is selected only by the session-scoped /beta toggle. + StrategyHMMBeta Strategy = "hmm_beta" // StrategyBandit routes via Thompson sampling over a frozen // ts_posterior.json (cluster×model reward posterior). Opt-in only; wired // when ROUTER_BANDIT_POSTERIOR_FILE is set at boot. @@ -34,7 +36,7 @@ const ( // IsHMMStrategy reports whether strategy uses the HMM policy contract and // lifecycle semantics. func IsHMMStrategy(strategy Strategy) bool { - return strategy == StrategyHMM || strategy == StrategyHMMEmbedding + return strategy == StrategyHMM || strategy == StrategyHMMEmbedding || strategy == StrategyHMMBeta } type strategyContextKey struct{} diff --git a/internal/router/strategy_test.go b/internal/router/strategy_test.go index 57e52ea46..ac90fea6e 100644 --- a/internal/router/strategy_test.go +++ b/internal/router/strategy_test.go @@ -9,6 +9,7 @@ import ( func TestIsHMMStrategy(t *testing.T) { assert.True(t, IsHMMStrategy(StrategyHMM)) assert.True(t, IsHMMStrategy(StrategyHMMEmbedding)) + assert.True(t, IsHMMStrategy(StrategyHMMBeta)) assert.False(t, IsHMMStrategy(StrategyCluster)) assert.False(t, IsHMMStrategy(StrategyRL)) } diff --git a/internal/server/middleware/router_strategy_override.go b/internal/server/middleware/router_strategy_override.go index 54e0d1011..e898c6b76 100644 --- a/internal/server/middleware/router_strategy_override.go +++ b/internal/server/middleware/router_strategy_override.go @@ -36,6 +36,10 @@ func WithRouterStrategyDefault(defaultStrategy router.Strategy, available ...rou } } for _, strategy := range available { + // hmm_beta is session-only: never activated by header, installation, or deployment default. + if strategy == router.StrategyHMMBeta { + continue + } allowed[strategy] = struct{}{} } defaultStrategy = normalizeRouterStrategyDefault(defaultStrategy, allowed) @@ -84,6 +88,9 @@ func NormalizeRouterStrategyDefault(defaultStrategy router.Strategy, available . allowed := make(map[router.Strategy]struct{}, len(available)+1) allowed[router.StrategyCluster] = struct{}{} for _, strategy := range available { + if strategy == router.StrategyHMMBeta { + continue + } allowed[strategy] = struct{}{} } return normalizeRouterStrategyDefault(defaultStrategy, allowed) diff --git a/internal/server/middleware/router_strategy_override_test.go b/internal/server/middleware/router_strategy_override_test.go index 2ea107be2..202b0a644 100644 --- a/internal/server/middleware/router_strategy_override_test.go +++ b/internal/server/middleware/router_strategy_override_test.go @@ -122,6 +122,29 @@ func TestRouterStrategyOverride_ExplicitClusterWinsOverDeploymentDefault(t *test func TestNormalizeRouterStrategyDefault(t *testing.T) { assert.Equal(t, router.StrategyHMM, middleware.NormalizeRouterStrategyDefault(router.StrategyHMM, router.StrategyHMM)) assert.Equal(t, router.StrategyCluster, middleware.NormalizeRouterStrategyDefault(router.Strategy("typo"), router.StrategyHMM)) + assert.Equal(t, router.StrategyCluster, + middleware.NormalizeRouterStrategyDefault(router.StrategyHMMBeta, router.StrategyHMMBeta), + "beta cannot be activated by a deployment default") +} + +func TestRouterStrategyOverride_BetaHeaderCannotActivateBeta(t *testing.T) { + got := runStrategyOverride( + t, + &auth.Installation{ID: "inst-beta", PolicyHeaderOverridesEnabled: true}, + string(router.StrategyHMMBeta), + router.StrategyHMMBeta, + ) + assert.Equal(t, router.StrategyCluster, got) +} + +func TestRouterStrategyOverride_InstallationDefaultBetaCannotActivateBeta(t *testing.T) { + got := runStrategyOverride( + t, + &auth.Installation{ID: "inst-beta", RoutingStrategy: router.StrategyHMMBeta}, + "", + router.StrategyHMMBeta, + ) + assert.Equal(t, router.StrategyCluster, got) } func TestRouterStrategyOverride_UnknownValueIgnored(t *testing.T) { diff --git a/internal/sqlc/models.go b/internal/sqlc/models.go index a46a8f54b..4c09fee0f 100644 --- a/internal/sqlc/models.go +++ b/internal/sqlc/models.go @@ -537,6 +537,15 @@ type RouterSessionPin struct { ConsecutiveOverloadErrors int32 DisabledProviders []string PolicyGroup string + RoutingStrategy string +} + +// Explicit per-session router strategy preferences +type RouterSessionStrategyPreference struct { + InstallationID uuid.UUID + SessionKey []byte + Strategy string + Enabled bool } // Shadow-mode spiral (death-march) detections: log-only fire-rate corpus measured on live traffic before escalation is armed diff --git a/internal/sqlc/session_pins.sql.go b/internal/sqlc/session_pins.sql.go index 2199562c5..ac6ce2dad 100644 --- a/internal/sqlc/session_pins.sql.go +++ b/internal/sqlc/session_pins.sql.go @@ -16,25 +16,35 @@ const deleteSessionPin = `-- name: DeleteSessionPin :one DELETE FROM router.session_pins WHERE session_key = $1::bytea AND role = $2::varchar + AND ( + routing_strategy = $3::varchar + OR (routing_strategy = '' AND $3::varchar <> 'hmm_beta') + ) AND pinned_until > CURRENT_TIMESTAMP -RETURNING session_key, role, installation_id, pinned_provider, pinned_model, decision_reason, turn_count, pinned_until, first_pinned_at, last_seen_at, last_input_tokens, last_cached_read_tokens, last_cached_write_tokens, last_output_tokens, last_turn_ended_at, consecutive_upstream_errors, last_served_model, has_ever_switched, paired_provider, paired_model, consecutive_overload_errors, disabled_providers, policy_group +RETURNING session_key, role, installation_id, pinned_provider, pinned_model, decision_reason, turn_count, pinned_until, first_pinned_at, last_seen_at, last_input_tokens, last_cached_read_tokens, last_cached_write_tokens, last_output_tokens, last_turn_ended_at, consecutive_upstream_errors, last_served_model, has_ever_switched, paired_provider, paired_model, consecutive_overload_errors, disabled_providers, policy_group, routing_strategy ` type DeleteSessionPinParams struct { - SessionKey []byte - Role string + SessionKey []byte + Role string + ExpectedRoutingStrategy string } -// Atomically consumes one active pin so a one-shot continuation cannot be -// reused by concurrent requests. Expired rows remain for the normal sweep. +// Atomically consumes one active pin for the expected strategy so a stale +// continuation cannot delete a replacement strategy's pin. Expired rows +// remain for the normal sweep. // // DELETE FROM router.session_pins // WHERE session_key = $1::bytea // AND role = $2::varchar +// AND ( +// routing_strategy = $3::varchar +// OR (routing_strategy = '' AND $3::varchar <> 'hmm_beta') +// ) // AND pinned_until > CURRENT_TIMESTAMP -// RETURNING session_key, role, installation_id, pinned_provider, pinned_model, decision_reason, turn_count, pinned_until, first_pinned_at, last_seen_at, last_input_tokens, last_cached_read_tokens, last_cached_write_tokens, last_output_tokens, last_turn_ended_at, consecutive_upstream_errors, last_served_model, has_ever_switched, paired_provider, paired_model, consecutive_overload_errors, disabled_providers, policy_group +// RETURNING session_key, role, installation_id, pinned_provider, pinned_model, decision_reason, turn_count, pinned_until, first_pinned_at, last_seen_at, last_input_tokens, last_cached_read_tokens, last_cached_write_tokens, last_output_tokens, last_turn_ended_at, consecutive_upstream_errors, last_served_model, has_ever_switched, paired_provider, paired_model, consecutive_overload_errors, disabled_providers, policy_group, routing_strategy func (q *Queries) DeleteSessionPin(ctx context.Context, arg DeleteSessionPinParams) (RouterSessionPin, error) { - row := q.db.QueryRow(ctx, deleteSessionPin, arg.SessionKey, arg.Role) + row := q.db.QueryRow(ctx, deleteSessionPin, arg.SessionKey, arg.Role, arg.ExpectedRoutingStrategy) var i RouterSessionPin err := row.Scan( &i.SessionKey, @@ -60,6 +70,7 @@ func (q *Queries) DeleteSessionPin(ctx context.Context, arg DeleteSessionPinPara &i.ConsecutiveOverloadErrors, &i.DisabledProviders, &i.PolicyGroup, + &i.RoutingStrategy, ) return i, err } @@ -73,20 +84,24 @@ SET disabled_providers = CASE consecutive_overload_errors = 0 WHERE session_key = $2::bytea AND role = $3::varchar + AND ( + routing_strategy = $4::varchar + OR (routing_strategy = '' AND $4::varchar <> 'hmm_beta') + ) ` type DisableSessionPinProviderParams struct { - Provider string - SessionKey []byte - Role string + Provider string + SessionKey []byte + Role string + ExpectedRoutingStrategy string } // Appends a provider to disabled_providers (deduped) and resets the // overload strike counter in the same statement, fired once the -// two-strike threshold is reached. disabled_providers only grows for the -// life of this pin row -- UpsertSessionPin's ON CONFLICT update never -// touches it, so a struck-out provider stays disabled until the pin -// itself is evicted/expires, with no separate time-based cooldown. +// two-strike threshold is reached. disabled_providers only grows within one +// strategy's pin lifecycle; a strategy replacement resets it with the other +// strategy-bound evidence. There is no separate time-based cooldown. // // UPDATE router.session_pins // SET disabled_providers = CASE @@ -96,13 +111,22 @@ type DisableSessionPinProviderParams struct { // consecutive_overload_errors = 0 // WHERE session_key = $2::bytea // AND role = $3::varchar +// AND ( +// routing_strategy = $4::varchar +// OR (routing_strategy = '' AND $4::varchar <> 'hmm_beta') +// ) func (q *Queries) DisableSessionPinProvider(ctx context.Context, arg DisableSessionPinProviderParams) error { - _, err := q.db.Exec(ctx, disableSessionPinProvider, arg.Provider, arg.SessionKey, arg.Role) + _, err := q.db.Exec(ctx, disableSessionPinProvider, + arg.Provider, + arg.SessionKey, + arg.Role, + arg.ExpectedRoutingStrategy, + ) return err } const getSessionPin = `-- name: GetSessionPin :one -SELECT session_key, role, installation_id, pinned_provider, pinned_model, decision_reason, turn_count, pinned_until, first_pinned_at, last_seen_at, last_input_tokens, last_cached_read_tokens, last_cached_write_tokens, last_output_tokens, last_turn_ended_at, consecutive_upstream_errors, last_served_model, has_ever_switched, paired_provider, paired_model, consecutive_overload_errors, disabled_providers, policy_group +SELECT session_key, role, installation_id, pinned_provider, pinned_model, decision_reason, turn_count, pinned_until, first_pinned_at, last_seen_at, last_input_tokens, last_cached_read_tokens, last_cached_write_tokens, last_output_tokens, last_turn_ended_at, consecutive_upstream_errors, last_served_model, has_ever_switched, paired_provider, paired_model, consecutive_overload_errors, disabled_providers, policy_group, routing_strategy FROM router.session_pins WHERE session_key = $1::bytea AND role = $2::varchar @@ -120,7 +144,7 @@ type GetSessionPinParams struct { // last_turn_ended_at carry the previous turn's upstream usage; the // planner reads them to weigh switch EV against eviction cost. // -// SELECT session_key, role, installation_id, pinned_provider, pinned_model, decision_reason, turn_count, pinned_until, first_pinned_at, last_seen_at, last_input_tokens, last_cached_read_tokens, last_cached_write_tokens, last_output_tokens, last_turn_ended_at, consecutive_upstream_errors, last_served_model, has_ever_switched, paired_provider, paired_model, consecutive_overload_errors, disabled_providers, policy_group +// SELECT session_key, role, installation_id, pinned_provider, pinned_model, decision_reason, turn_count, pinned_until, first_pinned_at, last_seen_at, last_input_tokens, last_cached_read_tokens, last_cached_write_tokens, last_output_tokens, last_turn_ended_at, consecutive_upstream_errors, last_served_model, has_ever_switched, paired_provider, paired_model, consecutive_overload_errors, disabled_providers, policy_group, routing_strategy // FROM router.session_pins // WHERE session_key = $1::bytea // AND role = $2::varchar @@ -151,6 +175,7 @@ func (q *Queries) GetSessionPin(ctx context.Context, arg GetSessionPinParams) (R &i.ConsecutiveOverloadErrors, &i.DisabledProviders, &i.PolicyGroup, + &i.RoutingStrategy, ) return i, err } @@ -160,12 +185,17 @@ UPDATE router.session_pins SET consecutive_overload_errors = consecutive_overload_errors + 1 WHERE session_key = $1::bytea AND role = $2::varchar + AND ( + routing_strategy = $3::varchar + OR (routing_strategy = '' AND $3::varchar <> 'hmm_beta') + ) RETURNING consecutive_overload_errors ` type IncrementSessionPinOverloadErrorsParams struct { - SessionKey []byte - Role string + SessionKey []byte + Role string + ExpectedRoutingStrategy string } // Atomically increments consecutive_overload_errors and returns the new @@ -181,9 +211,13 @@ type IncrementSessionPinOverloadErrorsParams struct { // SET consecutive_overload_errors = consecutive_overload_errors + 1 // WHERE session_key = $1::bytea // AND role = $2::varchar +// AND ( +// routing_strategy = $3::varchar +// OR (routing_strategy = '' AND $3::varchar <> 'hmm_beta') +// ) // RETURNING consecutive_overload_errors func (q *Queries) IncrementSessionPinOverloadErrors(ctx context.Context, arg IncrementSessionPinOverloadErrorsParams) (int32, error) { - row := q.db.QueryRow(ctx, incrementSessionPinOverloadErrors, arg.SessionKey, arg.Role) + row := q.db.QueryRow(ctx, incrementSessionPinOverloadErrors, arg.SessionKey, arg.Role, arg.ExpectedRoutingStrategy) var consecutive_overload_errors int32 err := row.Scan(&consecutive_overload_errors) return consecutive_overload_errors, err @@ -194,12 +228,17 @@ UPDATE router.session_pins SET consecutive_upstream_errors = consecutive_upstream_errors + 1 WHERE session_key = $1::bytea AND role = $2::varchar + AND ( + routing_strategy = $3::varchar + OR (routing_strategy = '' AND $3::varchar <> 'hmm_beta') + ) RETURNING consecutive_upstream_errors ` type IncrementSessionPinUpstreamErrorsParams struct { - SessionKey []byte - Role string + SessionKey []byte + Role string + ExpectedRoutingStrategy string } // Atomically increments consecutive_upstream_errors and returns the @@ -213,9 +252,13 @@ type IncrementSessionPinUpstreamErrorsParams struct { // SET consecutive_upstream_errors = consecutive_upstream_errors + 1 // WHERE session_key = $1::bytea // AND role = $2::varchar +// AND ( +// routing_strategy = $3::varchar +// OR (routing_strategy = '' AND $3::varchar <> 'hmm_beta') +// ) // RETURNING consecutive_upstream_errors func (q *Queries) IncrementSessionPinUpstreamErrors(ctx context.Context, arg IncrementSessionPinUpstreamErrorsParams) (int32, error) { - row := q.db.QueryRow(ctx, incrementSessionPinUpstreamErrors, arg.SessionKey, arg.Role) + row := q.db.QueryRow(ctx, incrementSessionPinUpstreamErrors, arg.SessionKey, arg.Role, arg.ExpectedRoutingStrategy) var consecutive_upstream_errors int32 err := row.Scan(&consecutive_upstream_errors) return consecutive_upstream_errors, err @@ -226,12 +269,17 @@ UPDATE router.session_pins SET consecutive_overload_errors = 0 WHERE session_key = $1::bytea AND role = $2::varchar + AND ( + routing_strategy = $3::varchar + OR (routing_strategy = '' AND $3::varchar <> 'hmm_beta') + ) AND consecutive_overload_errors > 0 ` type ResetSessionPinOverloadErrorsParams struct { - SessionKey []byte - Role string + SessionKey []byte + Role string + ExpectedRoutingStrategy string } // Clears the overload strike counter after a successful turn. UPDATE @@ -242,9 +290,13 @@ type ResetSessionPinOverloadErrorsParams struct { // SET consecutive_overload_errors = 0 // WHERE session_key = $1::bytea // AND role = $2::varchar +// AND ( +// routing_strategy = $3::varchar +// OR (routing_strategy = '' AND $3::varchar <> 'hmm_beta') +// ) // AND consecutive_overload_errors > 0 func (q *Queries) ResetSessionPinOverloadErrors(ctx context.Context, arg ResetSessionPinOverloadErrorsParams) error { - _, err := q.db.Exec(ctx, resetSessionPinOverloadErrors, arg.SessionKey, arg.Role) + _, err := q.db.Exec(ctx, resetSessionPinOverloadErrors, arg.SessionKey, arg.Role, arg.ExpectedRoutingStrategy) return err } @@ -253,12 +305,17 @@ UPDATE router.session_pins SET consecutive_upstream_errors = 0 WHERE session_key = $1::bytea AND role = $2::varchar + AND ( + routing_strategy = $3::varchar + OR (routing_strategy = '' AND $3::varchar <> 'hmm_beta') + ) AND consecutive_upstream_errors > 0 ` type ResetSessionPinUpstreamErrorsParams struct { - SessionKey []byte - Role string + SessionKey []byte + Role string + ExpectedRoutingStrategy string } // Clears the two-strike counter after a successful turn. UPDATE @@ -269,9 +326,13 @@ type ResetSessionPinUpstreamErrorsParams struct { // SET consecutive_upstream_errors = 0 // WHERE session_key = $1::bytea // AND role = $2::varchar +// AND ( +// routing_strategy = $3::varchar +// OR (routing_strategy = '' AND $3::varchar <> 'hmm_beta') +// ) // AND consecutive_upstream_errors > 0 func (q *Queries) ResetSessionPinUpstreamErrors(ctx context.Context, arg ResetSessionPinUpstreamErrorsParams) error { - _, err := q.db.Exec(ctx, resetSessionPinUpstreamErrors, arg.SessionKey, arg.Role) + _, err := q.db.Exec(ctx, resetSessionPinUpstreamErrors, arg.SessionKey, arg.Role, arg.ExpectedRoutingStrategy) return err } @@ -307,20 +368,25 @@ SET last_input_tokens = $1::int, last_served_model = $8::varchar WHERE session_key = $10::bytea AND role = $11::varchar + AND ( + routing_strategy = $12::varchar + OR (routing_strategy = '' AND $12::varchar <> 'hmm_beta') + ) ` type UpdateSessionPinUsageParams struct { - LastInputTokens int32 - LastCachedReadTokens int32 - LastCachedWriteTokens int32 - LastOutputTokens int32 - LastTurnEndedAt pgtype.Timestamptz - LastServedProvider string - SessionEverSwitched bool - LastServedModel string - PriorServedModel string - SessionKey []byte - Role string + LastInputTokens int32 + LastCachedReadTokens int32 + LastCachedWriteTokens int32 + LastOutputTokens int32 + LastTurnEndedAt pgtype.Timestamptz + LastServedProvider string + SessionEverSwitched bool + LastServedModel string + PriorServedModel string + SessionKey []byte + Role string + ExpectedRoutingStrategy string } // Records the previous turn's upstream token usage on an existing pin @@ -329,7 +395,8 @@ type UpdateSessionPinUsageParams struct { // turn to compute switch EV against eviction cost. The UPDATE matches // by (session_key, role); if the pin has been evicted or never // existed, zero rows are affected and the adapter wraps that as a -// successful no-op. last_served_model records the model that actually +// successful no-op. A strategy mismatch is also a no-op, preventing a late +// response from mutating a replacement strategy's pin. last_served_model records the model that actually // served this turn; it lives here (not in UpsertSessionPin) so a // /force-model upsert cannot overwrite the genuinely-last-served model // before the next turn reads it to detect a mid-session model switch. @@ -355,6 +422,10 @@ type UpdateSessionPinUsageParams struct { // last_served_model = $8::varchar // WHERE session_key = $10::bytea // AND role = $11::varchar +// AND ( +// routing_strategy = $12::varchar +// OR (routing_strategy = '' AND $12::varchar <> 'hmm_beta') +// ) func (q *Queries) UpdateSessionPinUsage(ctx context.Context, arg UpdateSessionPinUsageParams) error { _, err := q.db.Exec(ctx, updateSessionPinUsage, arg.LastInputTokens, @@ -368,6 +439,7 @@ func (q *Queries) UpdateSessionPinUsage(ctx context.Context, arg UpdateSessionPi arg.PriorServedModel, arg.SessionKey, arg.Role, + arg.ExpectedRoutingStrategy, ) return err } @@ -376,33 +448,40 @@ const upsertSessionPin = `-- name: UpsertSessionPin :exec INSERT INTO router.session_pins ( session_key, role, installation_id, pinned_provider, pinned_model, paired_provider, paired_model, - decision_reason, policy_group, turn_count, pinned_until + decision_reason, routing_strategy, policy_group, turn_count, pinned_until ) VALUES ( $1::bytea, $2::varchar, $3::uuid, $4::varchar, $5::varchar, $6::varchar, $7::varchar, - $8::text, $9::varchar, - $10::int, $11::timestamp + $8::text, $9::varchar, $10::varchar, + $11::int, $12::timestamp ) ON CONFLICT (session_key, role) DO UPDATE SET pinned_provider = EXCLUDED.pinned_provider, pinned_model = EXCLUDED.pinned_model, decision_reason = EXCLUDED.decision_reason, - turn_count = router.session_pins.turn_count + 1, + routing_strategy = EXCLUDED.routing_strategy, + turn_count = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.turn_count + 1 + ELSE EXCLUDED.turn_count + END, pinned_until = EXCLUDED.pinned_until, + first_pinned_at = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.first_pinned_at + ELSE CURRENT_TIMESTAMP + END, last_seen_at = CURRENT_TIMESTAMP, -- Band pair maintenance, in priority order: -- 1. A fresh scorer decision supplies a non-empty pair -> take it. - -- 2. Empty incoming pair but the pinned model is unchanged (sticky refresh, - -- reconstructed re-anchor of the same model) -> preserve the stored pair. - -- 3. Empty incoming pair and the pinned model changed (force-model, - -- loop-escalation, eviction -- non-scorer writes) -> clear the pair, so a - -- model change never inherits a stale runner-up or collapses pinned_model - -- and paired_model onto the same slug. + -- 2. Empty incoming pair but model and strategy are unchanged -> preserve it. + -- 3. Empty incoming pair and either changed -> clear it. paired_provider = CASE WHEN EXCLUDED.paired_model <> '' THEN EXCLUDED.paired_provider WHEN EXCLUDED.pinned_model = router.session_pins.pinned_model + AND EXCLUDED.routing_strategy = router.session_pins.routing_strategy THEN router.session_pins.paired_provider ELSE '' END, @@ -410,6 +489,7 @@ ON CONFLICT (session_key, role) DO UPDATE SET WHEN EXCLUDED.paired_model <> '' THEN EXCLUDED.paired_model WHEN EXCLUDED.pinned_model = router.session_pins.pinned_model + AND EXCLUDED.routing_strategy = router.session_pins.routing_strategy THEN router.session_pins.paired_model ELSE '' END, @@ -417,11 +497,13 @@ ON CONFLICT (session_key, role) DO UPDATE SET WHEN EXCLUDED.policy_group <> '' THEN EXCLUDED.policy_group WHEN EXCLUDED.pinned_model = router.session_pins.pinned_model + AND EXCLUDED.routing_strategy = router.session_pins.routing_strategy THEN router.session_pins.policy_group ELSE '' END, consecutive_upstream_errors = CASE WHEN router.session_pins.pinned_model = EXCLUDED.pinned_model + AND router.session_pins.routing_strategy = EXCLUDED.routing_strategy THEN router.session_pins.consecutive_upstream_errors ELSE 0 END, @@ -432,23 +514,67 @@ ON CONFLICT (session_key, role) DO UPDATE SET -- requiring two genuine consecutive strikes on the SAME served provider. consecutive_overload_errors = CASE WHEN router.session_pins.pinned_model = EXCLUDED.pinned_model + AND router.session_pins.routing_strategy = EXCLUDED.routing_strategy THEN router.session_pins.consecutive_overload_errors ELSE 0 + END, + -- A strategy switch selects a different policy. Do not carry cache, + -- switch, or error evidence from the previous policy into it. + last_input_tokens = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_input_tokens + ELSE 0 + END, + last_cached_read_tokens = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_cached_read_tokens + ELSE 0 + END, + last_cached_write_tokens = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_cached_write_tokens + ELSE 0 + END, + last_output_tokens = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_output_tokens + ELSE 0 + END, + last_turn_ended_at = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_turn_ended_at + ELSE NULL + END, + last_served_model = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.last_served_model + ELSE '' + END, + has_ever_switched = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.has_ever_switched + ELSE FALSE + END, + disabled_providers = CASE + WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy + THEN router.session_pins.disabled_providers + ELSE '{}' END ` type UpsertSessionPinParams struct { - SessionKey []byte - Role string - InstallationID uuid.UUID - PinnedProvider string - PinnedModel string - PairedProvider string - PairedModel string - DecisionReason string - PolicyGroup string - TurnCount int32 - PinnedUntil pgtype.Timestamp + SessionKey []byte + Role string + InstallationID uuid.UUID + PinnedProvider string + PinnedModel string + PairedProvider string + PairedModel string + DecisionReason string + RoutingStrategy string + PolicyGroup string + TurnCount int32 + PinnedUntil pgtype.Timestamp } // Upserts a pin, refreshing pinned_until on every hit (sliding TTL). @@ -461,59 +587,63 @@ type UpsertSessionPinParams struct { // them, so the at-start-of-turn refresh here cannot clobber the // previous turn's usage with zeros before the planner reads it. // -// consecutive_upstream_errors is preserved on a same-model refresh (so -// the two-strike eviction counter accumulates across turns of the same -// sticky pin) but reset to 0 on a switch (different model = clean -// slate). The reset on switch also covers the loop-break / force-model -// pin-expiry writes, which set pinned_model to the empty string. +// consecutive_upstream_errors is preserved on a same-model, same-strategy +// refresh (so the two-strike eviction counter accumulates across turns of the +// same sticky pin) but reset to 0 on a model or strategy switch. The reset also +// covers loop-break / force-model pin-expiry writes, which set pinned_model to +// the empty string. // // paired_provider / paired_model hold the runner-up half of the band pair the // scorer picks. On the conflict update they refresh to a fresh scorer runner-up -// (non-empty incoming pair), are preserved when the pinned model is unchanged -// (sticky refresh / same-model re-anchor carry an empty pair), and are cleared -// when the pinned model changes without a fresh pair (force-model, -// loop-escalation, eviction -- non-scorer writes). This keeps the stored pair -// consistent with the live decision: it tracks genuine re-routes, never -// inherits a stale runner-up across a non-scorer model change, and never -// collapses pinned_model and paired_model onto the same slug. A later per-turn -// swap policy reads the pair that matches the active decision. +// (non-empty incoming pair), are preserved when both the pinned model and +// strategy are unchanged (sticky refresh / same-model re-anchor carry an empty +// pair), and are cleared when either changes without a fresh pair. This keeps +// the stored pair consistent with the live decision: it tracks genuine +// re-routes and never inherits a stale runner-up across a strategy change. // // policy_group follows the same three-way maintenance: a fresh policy decision -// supplies a non-empty group, a same-model refresh preserves the stored one, and -// a model change without a group (force-model, loop-break, eviction) clears it. +// supplies a non-empty group, a same-model same-strategy refresh preserves the +// stored one, and a model or strategy change without a group clears it. // The pin-sticky arm-selector guard compares it against the fresh decision's // group, so a stale group must never survive onto a different pinned model. // // INSERT INTO router.session_pins ( // session_key, role, installation_id, pinned_provider, // pinned_model, paired_provider, paired_model, -// decision_reason, policy_group, turn_count, pinned_until +// decision_reason, routing_strategy, policy_group, turn_count, pinned_until // ) VALUES ( // $1::bytea, $2::varchar, $3::uuid, // $4::varchar, $5::varchar, // $6::varchar, $7::varchar, -// $8::text, $9::varchar, -// $10::int, $11::timestamp +// $8::text, $9::varchar, $10::varchar, +// $11::int, $12::timestamp // ) // ON CONFLICT (session_key, role) DO UPDATE SET // pinned_provider = EXCLUDED.pinned_provider, // pinned_model = EXCLUDED.pinned_model, // decision_reason = EXCLUDED.decision_reason, -// turn_count = router.session_pins.turn_count + 1, +// routing_strategy = EXCLUDED.routing_strategy, +// turn_count = CASE +// WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy +// THEN router.session_pins.turn_count + 1 +// ELSE EXCLUDED.turn_count +// END, // pinned_until = EXCLUDED.pinned_until, +// first_pinned_at = CASE +// WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy +// THEN router.session_pins.first_pinned_at +// ELSE CURRENT_TIMESTAMP +// END, // last_seen_at = CURRENT_TIMESTAMP, // -- Band pair maintenance, in priority order: // -- 1. A fresh scorer decision supplies a non-empty pair -> take it. -// -- 2. Empty incoming pair but the pinned model is unchanged (sticky refresh, -// -- reconstructed re-anchor of the same model) -> preserve the stored pair. -// -- 3. Empty incoming pair and the pinned model changed (force-model, -// -- loop-escalation, eviction -- non-scorer writes) -> clear the pair, so a -// -- model change never inherits a stale runner-up or collapses pinned_model -// -- and paired_model onto the same slug. +// -- 2. Empty incoming pair but model and strategy are unchanged -> preserve it. +// -- 3. Empty incoming pair and either changed -> clear it. // paired_provider = CASE // WHEN EXCLUDED.paired_model <> '' // THEN EXCLUDED.paired_provider // WHEN EXCLUDED.pinned_model = router.session_pins.pinned_model +// AND EXCLUDED.routing_strategy = router.session_pins.routing_strategy // THEN router.session_pins.paired_provider // ELSE '' // END, @@ -521,6 +651,7 @@ type UpsertSessionPinParams struct { // WHEN EXCLUDED.paired_model <> '' // THEN EXCLUDED.paired_model // WHEN EXCLUDED.pinned_model = router.session_pins.pinned_model +// AND EXCLUDED.routing_strategy = router.session_pins.routing_strategy // THEN router.session_pins.paired_model // ELSE '' // END, @@ -528,11 +659,13 @@ type UpsertSessionPinParams struct { // WHEN EXCLUDED.policy_group <> '' // THEN EXCLUDED.policy_group // WHEN EXCLUDED.pinned_model = router.session_pins.pinned_model +// AND EXCLUDED.routing_strategy = router.session_pins.routing_strategy // THEN router.session_pins.policy_group // ELSE '' // END, // consecutive_upstream_errors = CASE // WHEN router.session_pins.pinned_model = EXCLUDED.pinned_model +// AND router.session_pins.routing_strategy = EXCLUDED.routing_strategy // THEN router.session_pins.consecutive_upstream_errors // ELSE 0 // END, @@ -543,8 +676,51 @@ type UpsertSessionPinParams struct { // -- requiring two genuine consecutive strikes on the SAME served provider. // consecutive_overload_errors = CASE // WHEN router.session_pins.pinned_model = EXCLUDED.pinned_model +// AND router.session_pins.routing_strategy = EXCLUDED.routing_strategy // THEN router.session_pins.consecutive_overload_errors // ELSE 0 +// END, +// -- A strategy switch selects a different policy. Do not carry cache, +// -- switch, or error evidence from the previous policy into it. +// last_input_tokens = CASE +// WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy +// THEN router.session_pins.last_input_tokens +// ELSE 0 +// END, +// last_cached_read_tokens = CASE +// WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy +// THEN router.session_pins.last_cached_read_tokens +// ELSE 0 +// END, +// last_cached_write_tokens = CASE +// WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy +// THEN router.session_pins.last_cached_write_tokens +// ELSE 0 +// END, +// last_output_tokens = CASE +// WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy +// THEN router.session_pins.last_output_tokens +// ELSE 0 +// END, +// last_turn_ended_at = CASE +// WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy +// THEN router.session_pins.last_turn_ended_at +// ELSE NULL +// END, +// last_served_model = CASE +// WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy +// THEN router.session_pins.last_served_model +// ELSE '' +// END, +// has_ever_switched = CASE +// WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy +// THEN router.session_pins.has_ever_switched +// ELSE FALSE +// END, +// disabled_providers = CASE +// WHEN router.session_pins.routing_strategy = EXCLUDED.routing_strategy +// THEN router.session_pins.disabled_providers +// ELSE '{}' // END func (q *Queries) UpsertSessionPin(ctx context.Context, arg UpsertSessionPinParams) error { _, err := q.db.Exec(ctx, upsertSessionPin, @@ -556,6 +732,7 @@ func (q *Queries) UpsertSessionPin(ctx context.Context, arg UpsertSessionPinPara arg.PairedProvider, arg.PairedModel, arg.DecisionReason, + arg.RoutingStrategy, arg.PolicyGroup, arg.TurnCount, arg.PinnedUntil, diff --git a/internal/sqlc/session_strategy_preferences.sql.go b/internal/sqlc/session_strategy_preferences.sql.go new file mode 100644 index 000000000..34c623b00 --- /dev/null +++ b/internal/sqlc/session_strategy_preferences.sql.go @@ -0,0 +1,111 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: session_strategy_preferences.sql + +package sqlc + +import ( + "context" + + "github.com/google/uuid" +) + +const getSessionStrategyPreference = `-- name: GetSessionStrategyPreference :one +SELECT strategy +FROM router.session_strategy_preferences +WHERE installation_id = $1::uuid + AND session_key = $2::bytea + AND enabled +` + +type GetSessionStrategyPreferenceParams struct { + InstallationID uuid.UUID + SessionKey []byte +} + +// Reads an explicit beta strategy for one installation-scoped session. A +// missing or disabled row means the session uses stable routing. +// +// SELECT strategy +// FROM router.session_strategy_preferences +// WHERE installation_id = $1::uuid +// AND session_key = $2::bytea +// AND enabled +func (q *Queries) GetSessionStrategyPreference(ctx context.Context, arg GetSessionStrategyPreferenceParams) (string, error) { + row := q.db.QueryRow(ctx, getSessionStrategyPreference, arg.InstallationID, arg.SessionKey) + var strategy string + err := row.Scan(&strategy) + return strategy, err +} + +const updateSessionStrategyPreferenceDisabled = `-- name: UpdateSessionStrategyPreferenceDisabled :execrows +UPDATE router.session_strategy_preferences +SET enabled = FALSE +WHERE installation_id = $1::uuid + AND session_key = $2::bytea + AND enabled +` + +type UpdateSessionStrategyPreferenceDisabledParams struct { + InstallationID uuid.UUID + SessionKey []byte +} + +// Turns the session's explicit override off and reports one affected row when +// beta had been enabled. Callers use this instead of the toggle when the beta +// policy is unavailable, so a concurrent command can never re-enable it. +// +// UPDATE router.session_strategy_preferences +// SET enabled = FALSE +// WHERE installation_id = $1::uuid +// AND session_key = $2::bytea +// AND enabled +func (q *Queries) UpdateSessionStrategyPreferenceDisabled(ctx context.Context, arg UpdateSessionStrategyPreferenceDisabledParams) (int64, error) { + result, err := q.db.Exec(ctx, updateSessionStrategyPreferenceDisabled, arg.InstallationID, arg.SessionKey) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const upsertToggledSessionStrategyPreference = `-- name: UpsertToggledSessionStrategyPreference :one +INSERT INTO router.session_strategy_preferences ( + installation_id, session_key, strategy, enabled +) VALUES ( + $1::uuid, $2::bytea, $3::varchar, TRUE +) +ON CONFLICT (installation_id, session_key) +DO UPDATE SET + strategy = EXCLUDED.strategy, + enabled = NOT router.session_strategy_preferences.enabled +RETURNING enabled +` + +type UpsertToggledSessionStrategyPreferenceParams struct { + InstallationID uuid.UUID + SessionKey []byte + Strategy string +} + +// Flips the session's explicit override and returns the state now persisted. +// The row lock taken on conflict serializes overlapping toggles across router +// instances, so each caller observes its own flip instead of a stale read. +// The database constraint rejects any strategy other than hmm_beta. +// +// INSERT INTO router.session_strategy_preferences ( +// installation_id, session_key, strategy, enabled +// ) VALUES ( +// $1::uuid, $2::bytea, $3::varchar, TRUE +// ) +// ON CONFLICT (installation_id, session_key) +// DO UPDATE SET +// strategy = EXCLUDED.strategy, +// enabled = NOT router.session_strategy_preferences.enabled +// RETURNING enabled +func (q *Queries) UpsertToggledSessionStrategyPreference(ctx context.Context, arg UpsertToggledSessionStrategyPreferenceParams) (bool, error) { + row := q.db.QueryRow(ctx, upsertToggledSessionStrategyPreference, arg.InstallationID, arg.SessionKey, arg.Strategy) + var enabled bool + err := row.Scan(&enabled) + return enabled, err +} diff --git a/internal/translate/beta.go b/internal/translate/beta.go new file mode 100644 index 000000000..84cf19f78 --- /dev/null +++ b/internal/translate/beta.go @@ -0,0 +1,205 @@ +package translate + +import ( + "strings" + + "github.com/tidwall/gjson" +) + +// BetaCommandResult describes a leading /beta directive. Invalid is true +// when the command has arguments or trailing prompt text; /beta intentionally +// has one toggle-only form. +type BetaCommandResult struct { + Invalid bool +} + +// ExtractBetaCommand scans the final user message for a leading /beta +// directive and strips it so the command is never forwarded upstream. +func (env *RequestEnvelope) ExtractBetaCommand() (BetaCommandResult, bool) { + var result BetaCommandResult + found := env.extractLeadingCommand(func(text string) (bool, string) { + parsed, ok, stripped := parseBetaCommand(text) + if ok { + result = parsed + } + return ok, stripped + }) + return result, found +} + +// StripBetaArtifacts removes prior command-only /beta turns and the router's +// synthetic acknowledgements from model-visible history. The trailing user +// command is preserved so ExtractBetaCommand can still toggle the session. +func (env *RequestEnvelope) StripBetaArtifacts() int { + switch env.format { + case FormatAnthropic, FormatOpenAI: + default: + return 0 + } + msgs := gjson.GetBytes(env.body, "messages") + if !msgs.IsArray() { + return 0 + } + + lastUserIdx := -1 + msgs.ForEach(func(key, msg gjson.Result) bool { + if msg.Get("role").String() == "user" { + lastUserIdx = int(key.Int()) + } + return true + }) + + removed := 0 + removedCommands := make(map[int]struct{}) + rebuilt := make([]string, 0, len(msgs.Array())) + msgs.ForEach(func(key, msg gjson.Result) bool { + idx := int(key.Int()) + role := msg.Get("role").String() + content := msg.Get("content") + if role == "user" && idx != lastUserIdx && isBetaCommandOnlyContent(content) { + removedCommands[idx] = struct{}{} + removed++ + return true + } + if role == "assistant" { + _, followsRemovedCommand := removedCommands[idx-1] + if isBetaAckOnlyContent(content) || (followsRemovedCommand && isEmptyTextContent(content)) { + removed++ + return true + } + } + rebuilt = append(rebuilt, msg.Raw) + return true + }) + if removed == 0 { + return 0 + } + return env.setMessages(rebuilt, removed) +} + +func isEmptyTextContent(content gjson.Result) bool { + switch { + case content.Type == gjson.String: + return strings.TrimSpace(content.String()) == "" + case content.Type == gjson.JSON && content.IsArray(): + empty := true + content.ForEach(func(_, block gjson.Result) bool { + if block.Get("type").String() != "text" || strings.TrimSpace(block.Get("text").String()) != "" { + empty = false + return false + } + return true + }) + return empty + default: + return false + } +} + +func parseBetaCommand(text string) (result BetaCommandResult, found bool, stripped string) { + prefixEnd := leadingInjectedPrefixEnd(text) + prefix := text[:prefixEnd] + body := text[prefixEnd:] + lines := strings.Split(body, "\n") + commandLine := -1 + + for i, line := range lines { + trimmed := strings.TrimSpace(line) + if trimmed == "" { + continue + } + fields := strings.Fields(trimmed) + if len(fields) == 0 || fields[0] != "/beta" { + return BetaCommandResult{}, false, text + } + commandLine = i + result.Invalid = len(fields) != 1 + break + } + if commandLine < 0 { + return BetaCommandResult{}, false, text + } + + remaining := make([]string, 0, len(lines)-1) + remaining = append(remaining, lines[:commandLine]...) + remaining = append(remaining, lines[commandLine+1:]...) + for _, line := range remaining { + if strings.TrimSpace(line) != "" { + result.Invalid = true + break + } + } + return result, true, strings.TrimSpace(prefix + strings.Join(remaining, "\n")) +} + +func isBetaCommandOnlyContent(content gjson.Result) bool { + switch { + case content.Type == gjson.String: + _, found, stripped := parseBetaCommand(content.String()) + return found && isOnlyInjectedCommandText(stripped) + case content.Type == gjson.JSON && content.IsArray(): + seenCommand := false + allSynthetic := true + content.ForEach(func(_, block gjson.Result) bool { + if block.Get("type").String() != "text" { + allSynthetic = false + return false + } + text := block.Get("text").String() + if strings.TrimSpace(text) == "" || isClaudeCodeInjectedBlock(text) { + return true + } + _, found, stripped := parseBetaCommand(text) + if found && isOnlyInjectedCommandText(stripped) { + seenCommand = true + return true + } + allSynthetic = false + return false + }) + return seenCommand && allSynthetic + default: + return false + } +} + +func isBetaAckOnlyContent(content gjson.Result) bool { + switch { + case content.Type == gjson.String: + return isBetaAckText(content.String()) + case content.Type == gjson.JSON && content.IsArray(): + seenAck := false + allSynthetic := true + content.ForEach(func(_, block gjson.Result) bool { + if block.Get("type").String() != "text" { + allSynthetic = false + return false + } + text := block.Get("text").String() + if strings.TrimSpace(text) == "" { + return true + } + if isBetaAckText(text) { + seenAck = true + return true + } + allSynthetic = false + return false + }) + return seenAck && allSynthetic + default: + return false + } +} + +func isBetaAckText(text string) bool { + switch strings.TrimSpace(text) { + case "✦ **Weave Router** → Beta enabled. Type /beta again to turn it off.", + "✦ **Weave Router** → Beta disabled. Stable routing restored.", + "✦ **Weave Router** → Beta is unavailable for this session.", + "✦ **Weave Router** → Usage: /beta": + return true + default: + return false + } +} diff --git a/internal/translate/beta_test.go b/internal/translate/beta_test.go new file mode 100644 index 000000000..c2b0982c9 --- /dev/null +++ b/internal/translate/beta_test.go @@ -0,0 +1,144 @@ +package translate_test + +import ( + "testing" + + "workweave/router/internal/translate" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestExtractBetaCommand(t *testing.T) { + tests := []struct { + name string + input string + found bool + invalid bool + wantContent string + }{ + {name: "toggle", input: "/beta", found: true}, + {name: "whitespace", input: " /beta ", found: true}, + {name: "arguments rejected", input: "/beta status", found: true, invalid: true}, + {name: "trailing prompt rejected", input: "/beta\nroute this", found: true, invalid: true, wantContent: "route this"}, + {name: "non-leading ignored", input: "route this\n/beta", found: false, wantContent: "route this\n/beta"}, + {name: "prefix boundary", input: "/betamax", found: false, wantContent: "/betamax"}, + { + name: "injected prefix", + input: "/beta\n/beta", + found: true, + wantContent: "/beta", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + env, err := translate.ParseAnthropic(buildAnthropicBody(t, tt.input)) + require.NoError(t, err) + + result, found := env.ExtractBetaCommand() + assert.Equal(t, tt.found, found) + assert.Equal(t, tt.invalid, result.Invalid) + assert.Equal(t, tt.wantContent, lastUserMessageText(t, env)) + }) + } +} + +func TestExtractBetaCommandOpenAI(t *testing.T) { + body := mustMarshalJSON(t, map[string]any{ + "model": "gpt-5.6-sol", + "messages": []any{ + map[string]any{"role": "user", "content": "/beta"}, + }, + }) + env, err := translate.ParseOpenAI(body) + require.NoError(t, err) + + result, found := env.ExtractBetaCommand() + require.True(t, found) + assert.False(t, result.Invalid) + assert.Empty(t, lastOpenAIUserMessageText(t, env)) +} + +func TestStripBetaArtifactsAnthropicPreservesCurrentToggle(t *testing.T) { + body := mustMarshalJSON(t, map[string]any{ + "model": "claude-sonnet-5", + "messages": []any{ + map[string]any{"role": "user", "content": "inspect this repository"}, + map[string]any{"role": "assistant", "content": "I will inspect it."}, + map[string]any{ + "role": "user", + "content": []any{ + map[string]any{"type": "text", "text": "beta\n/beta"}, + map[string]any{"type": "text", "text": "/beta"}, + }, + }, + map[string]any{ + "role": "assistant", + "content": []any{ + map[string]any{"type": "text", "text": "✦ **Weave Router** → Beta enabled. Type /beta again to turn it off.\n\n"}, + }, + }, + map[string]any{"role": "user", "content": "/beta"}, + }, + "max_tokens": 128, + }) + env, err := translate.ParseAnthropic(body) + require.NoError(t, err) + + assert.Equal(t, 2, env.StripBetaArtifacts()) + prepared, err := env.PrepareAnthropic(nil, translate.EmitOptions{TargetModel: "claude-sonnet-5"}) + require.NoError(t, err) + assert.Equal(t, int64(3), gjson.GetBytes(prepared.Body, "messages.#").Int()) + assert.Equal(t, "/beta", gjson.GetBytes(prepared.Body, "messages.2.content.0.text").String(), + "the current trailing toggle must survive until interception") + + result, found := env.ExtractBetaCommand() + require.True(t, found) + assert.False(t, result.Invalid) +} + +func TestStripBetaArtifactsOpenAI(t *testing.T) { + body := mustMarshalJSON(t, map[string]any{ + "model": "gpt-5.6-sol", + "messages": []any{ + map[string]any{"role": "user", "content": "inspect this repository"}, + map[string]any{"role": "assistant", "content": "I will inspect it."}, + map[string]any{"role": "user", "content": "/beta"}, + map[string]any{"role": "assistant", "content": "✦ **Weave Router** → Beta is unavailable for this session.\n\n"}, + map[string]any{"role": "user", "content": "continue with the implementation"}, + }, + }) + env, err := translate.ParseOpenAI(body) + require.NoError(t, err) + + assert.Equal(t, 2, env.StripBetaArtifacts()) + prepared, err := env.PrepareOpenAI(nil, translate.EmitOptions{TargetModel: "gpt-5.6-sol"}) + require.NoError(t, err) + assert.Equal(t, int64(3), gjson.GetBytes(prepared.Body, "messages.#").Int()) + assert.Equal(t, "continue with the implementation", gjson.GetBytes(prepared.Body, "messages.2.content").String()) +} + +func TestStripBetaArtifactsRemovesInvalidControlTurnButKeepsDiscussion(t *testing.T) { + body := mustMarshalJSON(t, map[string]any{ + "model": "claude-sonnet-5", + "messages": []any{ + map[string]any{"role": "user", "content": "/beta status"}, + map[string]any{"role": "assistant", "content": "✦ **Weave Router** → Usage: /beta\n\n"}, + map[string]any{"role": "user", "content": "Please explain /beta instead of toggling it."}, + map[string]any{"role": "assistant", "content": "✦ **Weave Router** → Beta enabled. Type /beta again to turn it off. This is quoted documentation."}, + map[string]any{"role": "user", "content": "continue"}, + }, + "max_tokens": 128, + }) + env, err := translate.ParseAnthropic(body) + require.NoError(t, err) + + assert.Equal(t, 2, env.StripBetaArtifacts()) + prepared, err := env.PrepareAnthropic(nil, translate.EmitOptions{TargetModel: "claude-sonnet-5"}) + require.NoError(t, err) + assert.Equal(t, int64(3), gjson.GetBytes(prepared.Body, "messages.#").Int()) + assert.Equal(t, "Please explain /beta instead of toggling it.", gjson.GetBytes(prepared.Body, "messages.0.content").String()) + assert.Contains(t, gjson.GetBytes(prepared.Body, "messages.1.content").String(), "quoted documentation") +}