From ea0ce9d15f6749d32b047556cf58a8c1e9d4dc06 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Fri, 17 Jul 2026 16:05:51 -0400 Subject: [PATCH] F3: Kimi inference pipeline + full grok cloud-surface excision MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sampler / inference (PRD F3): - kimi_compat.rs: single adaptation point for the Kimi chat/completions dialect (thinking-field mapping, model_id stripping, empty-content tool-call message fix, stream_options.include_usage), with kimi-cli source citations - Rate-limit handling reworked for Kimi/Moonshot semantics; UA kigi/{version} - /models replaces the xAI models-v2 endpoint everywhere; idle model refresh carries X-Msh-* device headers only (X-XAI-Token-Auth and x-grok-client-mode/CLIENT_MODE_HEADER machinery deleted) Cloud-surface excision (PRD §5, zero-egress): - remote/ conversations lane, cli-chat-proxy-types crate, prod/ dir, share command, credit bar: deleted (single local session lane; paginate() replaces merge_and_paginate) - Subscription/tier gate stack deleted end-to-end: AppView gate/tier/team/ZDR fields, app/subscription.rs watch loop, dispatch/billing.rs paywall + SuperGrok upsell, free-usage-exhausted chain, tier-restricted commands, GateInfo, RemoteSettings gate fields, SettingsUpdateNotification gate fields - /privacy + coding-data-sharing setting deleted (backed by a dead xAI RPC; Kigi is zero-egress — nothing to share or retain remotely) Auth UX correctness (user-reported): - Device-flow fixtures now mirror the live Kimi payload shape (https://www.kimi.com/code/authorize_device?user_code=..., verified against auth.kimi.com); the fabricated auth.kimi.com/device?code=... URLs are gone - open_browser_detached is a no-op under cfg(test): unit tests drove wiremock fixture URLs into the real browser (root cause of the "garbage mock link" ABCD-1234 tabs) - Welcome/pager-minimal rebrand: Grok Build -> Kigi, grok.com -> kimi.com, "Sign in to Grok" -> "Sign in to Kimi" --- Cargo.lock | 13 - Cargo.toml | 16 +- crates/codegen/kigi-bin/src/main.rs | 161 +- crates/codegen/kigi-chat-state/src/types.rs | 2 - crates/codegen/kigi-config-types/src/lib.rs | 33 - crates/codegen/kigi-config/Cargo.toml | 2 - .../codegen/kigi-config/src/signed_policy.rs | 59 +- .../kigi-config/src/signed_policy/tests.rs | 25 + crates/codegen/kigi-config/src/validation.rs | 20 +- crates/codegen/kigi-http/src/lib.rs | 38 +- crates/codegen/kigi-pager-minimal/src/auth.rs | 20 +- .../kigi-pager-pty-harness/src/content.rs | 4 +- crates/codegen/kigi-sampler/Cargo.toml | 2 +- .../kigi-sampler/src/actor/request_task.rs | 16 +- .../codegen/kigi-sampler/src/actor/state.rs | 4 - crates/codegen/kigi-sampler/src/client.rs | 321 +-- crates/codegen/kigi-sampler/src/config.rs | 14 +- crates/codegen/kigi-sampler/src/events.rs | 3 - .../codegen/kigi-sampler/src/kimi_compat.rs | 424 +++ crates/codegen/kigi-sampler/src/lib.rs | 3 +- crates/codegen/kigi-sampler/src/retry.rs | 87 +- .../src/stream/chat_completions.rs | 41 +- .../kigi-sampler/src/stream/collect.rs | 2 + .../kigi-sampler/src/stream/messages.rs | 1 - .../kigi-sampler/src/stream/responses.rs | 3 - .../codegen/kigi-sampler/tests/test_actor.rs | 4 - .../kigi-sampler/tests/test_kimi_wire.rs | 564 ++++ .../kigi-sampling-types/src/conversation.rs | 13 +- .../codegen/kigi-sampling-types/src/error.rs | 109 +- .../codegen/kigi-sampling-types/src/types.rs | 17 + crates/codegen/kigi-shared/Cargo.toml | 1 - crates/codegen/kigi-shared/src/session/mod.rs | 38 +- .../codegen/kigi-shell-base/src/util/mod.rs | 64 +- crates/codegen/kigi-shell/Cargo.toml | 1 - .../kigi-shell/benches/session_list.rs | 2 - crates/codegen/kigi-shell/src/agent/app.rs | 25 +- .../kigi-shell/src/agent/chat_modes.rs | 342 +-- crates/codegen/kigi-shell/src/agent/config.rs | 195 +- .../kigi-shell/src/agent/feedback_client.rs | 1458 ++--------- .../kigi-shell/src/agent/handlers/mod.rs | 1 - .../kigi-shell/src/agent/handlers/session.rs | 13 +- .../src/agent/handlers/workspaces.rs | 174 -- crates/codegen/kigi-shell/src/agent/init.rs | 21 - crates/codegen/kigi-shell/src/agent/mod.rs | 3 +- crates/codegen/kigi-shell/src/agent/models.rs | 101 +- .../kigi-shell/src/agent/models_fetch.rs | 1084 ++++++++ .../src/agent/mvp_agent/acp_agent.rs | 252 +- .../src/agent/mvp_agent/agent_ops.rs | 228 +- .../kigi-shell/src/agent/mvp_agent/mod.rs | 188 -- .../agent/mvp_agent/subagent_coordinator.rs | 2 +- .../kigi-shell/src/agent/mvp_agent/tests.rs | 12 - .../src/agent/subagent/handle_request.rs | 5 - .../kigi-shell/src/agent/subagent/mod.rs | 13 +- crates/codegen/kigi-shell/src/auth/device.rs | 8 +- .../kigi-shell/src/auth/device_code.rs | 16 +- .../codegen/kigi-shell/src/auth/kimi_oauth.rs | 18 +- crates/codegen/kigi-shell/src/auth/meta.rs | 12 - crates/codegen/kigi-shell/src/auth/mod.rs | 3 +- crates/codegen/kigi-shell/src/bundle.rs | 1512 +---------- crates/codegen/kigi-shell/src/config/mod.rs | 8 +- crates/codegen/kigi-shell/src/config/tests.rs | 10 +- .../kigi-shell/src/extensions/billing.rs | 882 +++---- .../kigi-shell/src/extensions/bundle.rs | 1204 --------- .../kigi-shell/src/extensions/feedback.rs | 77 +- .../codegen/kigi-shell/src/extensions/mod.rs | 2 - .../src/extensions/session_admin.rs | 126 +- .../kigi-shell/src/extensions/share.rs | 204 -- crates/codegen/kigi-shell/src/lib.rs | 1 - crates/codegen/kigi-shell/src/remote/agent.rs | 326 --- .../src/remote/chat_models_client.rs | 202 -- .../codegen/kigi-shell/src/remote/client.rs | 2268 ----------------- .../src/remote/conversations_client.rs | 305 --- crates/codegen/kigi-shell/src/remote/mod.rs | 37 - crates/codegen/kigi-shell/src/remote/pull.rs | 772 ------ .../kigi-shell/src/remote/pull_smoke_test.rs | 127 - crates/codegen/kigi-shell/src/remote/sync.rs | 167 -- .../src/remote/workspaces_client.rs | 176 -- .../codegen/kigi-shell/src/sampling/error.rs | 55 +- .../kigi-shell/src/session/acp_session.rs | 2 - .../session/acp_session_impl/model_switch.rs | 3 +- .../session/acp_session_impl/prompt_build.rs | 1 - .../src/session/acp_session_impl/recap.rs | 2 +- .../src/session/acp_session_impl/reminders.rs | 2 +- .../src/session/acp_session_impl/run_loop.rs | 14 +- .../session/acp_session_impl/sampler_turn.rs | 15 +- .../session/acp_session_impl/session_setup.rs | 20 +- .../session/acp_session_impl/slash_exec.rs | 2 +- .../src/session/acp_session_impl/spawn.rs | 72 +- .../auth_error_no_retry_tests.rs | 4 - .../cancel_running_task_tests.rs | 20 - .../acp_session_tests/idle_resume_tests.rs | 10 +- .../inline_auto_compact_flow_tests.rs | 12 +- .../acp_session_tests/memory_config_tests.rs | 1 - .../reactive_managed_reauth_e2e_tests.rs | 2 +- .../replay_buffer_send_update_tests.rs | 1 - .../src/session/acp_session_tests/support.rs | 1 - .../acp_session_tests/web_search_e2e_tests.rs | 3 - .../kigi-shell/src/session/acp_types.rs | 22 +- .../kigi-shell/src/session/compaction.rs | 3 +- .../kigi-shell/src/session/feedback.rs | 8 +- .../src/session/feedback_manager.rs | 655 +---- .../kigi-shell/src/session/feedback_types.rs | 763 ++++++ crates/codegen/kigi-shell/src/session/fork.rs | 75 +- .../src/session/helpers/session_compact.rs | 8 - crates/codegen/kigi-shell/src/session/mod.rs | 18 +- .../kigi-shell/src/session/persistence.rs | 351 +-- .../src/session/storage/jsonl/tests.rs | 13 +- .../src/session/unified_list/cursor.rs | 470 +--- .../src/session/unified_list/facets.rs | 248 +- .../src/session/unified_list/mod.rs | 443 +--- .../src/session/unified_list/row.rs | 82 - .../kigi-shell/src/session/worktree.rs | 2 +- .../src/test_support/lsp_runtime.rs | 4 - crates/codegen/kigi-shell/src/tools/config.rs | 6 +- .../src/util/config/resolve/auto_mode.rs | 3 +- .../src/util/config/resolve/features.rs | 22 - crates/codegen/kigi-shell/tests/common/mod.rs | 4 - .../kigi-shell/tests/git_contention_e2e.rs | 2 +- .../kigi-shell/tests/session_load_perf.rs | 2 +- .../tests/signed_managed_config/common.rs | 2 +- .../tests/signed_managed_config_extended.rs | 2 +- .../kigi-shell/tests/test_built_binary_e2e.rs | 2 +- .../kigi-shell/tests/test_leader_soak.rs | 2 +- .../kigi-shell/tests/test_settings_refresh.rs | 258 -- crates/codegen/kigi-test-support/src/env.rs | 2 +- .../codegen/kigi-test-support/src/leader.rs | 2 +- .../src/app/acp_handler/interactions.rs | 2 - .../kigi-tui/src/app/acp_handler/mod.rs | 2 - .../app/acp_handler/session_notification.rs | 33 +- .../kigi-tui/src/app/acp_handler/settings.rs | 58 - .../kigi-tui/src/app/acp_handler/tests/mod.rs | 13 - .../src/app/acp_handler/tests/plan_mode.rs | 2 +- .../acp_handler/tests/queue_and_adoption.rs | 19 +- .../app/acp_handler/tests/session_events.rs | 143 +- .../src/app/acp_handler/tests/settings.rs | 35 - crates/codegen/kigi-tui/src/app/actions.rs | 120 +- crates/codegen/kigi-tui/src/app/agent.rs | 14 - .../kigi-tui/src/app/agent_view/input.rs | 7 +- .../src/app/agent_view/interactions.rs | 2 - .../kigi-tui/src/app/agent_view/mod.rs | 34 - .../kigi-tui/src/app/agent_view/paste.rs | 2 - .../kigi-tui/src/app/agent_view/plan.rs | 2 - .../kigi-tui/src/app/agent_view/render.rs | 26 - .../kigi-tui/src/app/agent_view/rewind.rs | 2 - .../kigi-tui/src/app/agent_view/session.rs | 36 +- crates/codegen/kigi-tui/src/app/app_view.rs | 367 +-- crates/codegen/kigi-tui/src/app/cli.rs | 7 +- .../codegen/kigi-tui/src/app/dispatch/auth.rs | 23 +- .../kigi-tui/src/app/dispatch/billing.rs | 532 ---- .../kigi-tui/src/app/dispatch/dashboard.rs | 23 +- .../codegen/kigi-tui/src/app/dispatch/mod.rs | 5 - .../kigi-tui/src/app/dispatch/prompt.rs | 133 +- .../kigi-tui/src/app/dispatch/queue.rs | 24 +- .../kigi-tui/src/app/dispatch/router.rs | 31 +- .../src/app/dispatch/session/foreign.rs | 25 +- .../kigi-tui/src/app/dispatch/session/fork.rs | 13 +- .../src/app/dispatch/session/lifecycle.rs | 34 +- .../kigi-tui/src/app/dispatch/session/load.rs | 27 +- .../kigi-tui/src/app/dispatch/settings/ui.rs | 15 +- .../kigi-tui/src/app/dispatch/status.rs | 309 +-- .../kigi-tui/src/app/dispatch/task_result.rs | 96 +- .../src/app/dispatch/tests/billing.rs | 990 +------ .../src/app/dispatch/tests/dashboard.rs | 31 - .../kigi-tui/src/app/dispatch/tests/mod.rs | 46 - .../kigi-tui/src/app/dispatch/tests/prompt.rs | 91 +- .../src/app/dispatch/tests/session/foreign.rs | 4 - .../src/app/dispatch/tests/session/fork.rs | 23 +- .../app/dispatch/tests/session/lifecycle.rs | 13 +- .../src/app/dispatch/tests/session/load.rs | 13 - .../src/app/dispatch/tests/settings.rs | 5 - .../kigi-tui/src/app/dispatch/tests/status.rs | 701 +---- .../src/app/dispatch/tests/task_result.rs | 505 +--- .../kigi-tui/src/app/dispatch/transcript.rs | 3 +- .../kigi-tui/src/app/effects/helpers.rs | 234 +- .../codegen/kigi-tui/src/app/effects/mod.rs | 286 +-- .../codegen/kigi-tui/src/app/effects/tests.rs | 325 +-- crates/codegen/kigi-tui/src/app/event_loop.rs | 162 +- .../kigi-tui/src/app/foreign_sessions.rs | 1 - .../kigi-tui/src/app/leader_cluster/mod.rs | 2 +- crates/codegen/kigi-tui/src/app/mod.rs | 126 +- crates/codegen/kigi-tui/src/app/modals.rs | 13 +- crates/codegen/kigi-tui/src/app/mouse.rs | 26 - crates/codegen/kigi-tui/src/app/queue_edit.rs | 2 +- crates/codegen/kigi-tui/src/app/subagent.rs | 2 - .../codegen/kigi-tui/src/app/subscription.rs | 487 ---- .../codegen/kigi-tui/src/client_identity.rs | 2 +- crates/codegen/kigi-tui/src/headless.rs | 9 +- crates/codegen/kigi-tui/src/lib.rs | 1 - .../codegen/kigi-tui/src/scrollback/block.rs | 48 +- .../src/scrollback/blocks/credit_limit.rs | 253 -- .../kigi-tui/src/scrollback/blocks/mod.rs | 2 - .../codegen/kigi-tui/src/scrollback/export.rs | 2 +- .../src/scrollback/wrappers/entry_renderer.rs | 2 +- crates/codegen/kigi-tui/src/sessions_cmd.rs | 108 +- crates/codegen/kigi-tui/src/settings/defs.rs | 53 - .../codegen/kigi-tui/src/settings/registry.rs | 25 - crates/codegen/kigi-tui/src/share_cmd.rs | 49 - .../kigi-tui/src/slash/commands/mod.rs | 71 +- .../kigi-tui/src/slash/commands/privacy.rs | 212 -- .../kigi-tui/src/slash/commands/share.rs | 34 - .../kigi-tui/src/slash/commands/usage.rs | 58 +- crates/codegen/kigi-tui/src/slash/mod.rs | 1 - crates/codegen/kigi-tui/src/slash/registry.rs | 50 +- crates/codegen/kigi-tui/src/test_util.rs | 2 - crates/codegen/kigi-tui/src/views/agent.rs | 14 - .../codegen/kigi-tui/src/views/credit_bar.rs | 817 ------ .../kigi-tui/src/views/dashboard/peek.rs | 4 - .../kigi-tui/src/views/dashboard/render.rs | 2 - .../kigi-tui/src/views/dashboard/row.rs | 2 - crates/codegen/kigi-tui/src/views/mod.rs | 1 - crates/codegen/kigi-tui/src/views/modal.rs | 73 +- .../kigi-tui/src/views/prompt_widget/mod.rs | 15 - .../kigi-tui/src/views/question_view.rs | 8 - .../kigi-tui/src/views/settings_modal.rs | 66 +- .../codegen/kigi-tui/src/views/welcome/mod.rs | 311 +-- .../kigi-tui/src/views/welcome/prompt.rs | 2 - ...licy_gate_refusal_reaches_real_terminal.rs | 2 +- crates/codegen/kigi-tui/tests/settings_e2e.rs | 427 ---- crates/codegen/kigi-workspace/src/handle.rs | 3 +- .../kigi-workspace/src/permission/types.rs | 16 +- prod/mc/cli-chat-proxy-types/Cargo.toml | 20 - .../src/client_metrics_types.rs | 55 - .../src/deployment_config_types.rs | 105 - .../src/feedback_types.rs | 2024 --------------- prod/mc/cli-chat-proxy-types/src/lib.rs | 23 - .../src/metadata_types.rs | 219 -- .../cli-chat-proxy-types/src/sandbox_types.rs | 715 ------ .../cli-chat-proxy-types/src/serde_helpers.rs | 9 - .../cli-chat-proxy-types/src/session_types.rs | 116 - .../cli-chat-proxy-types/src/storage_types.rs | 226 -- .../src/subagent_bundle.rs | 101 - 231 files changed, 4730 insertions(+), 26358 deletions(-) create mode 100644 crates/codegen/kigi-sampler/src/kimi_compat.rs create mode 100644 crates/codegen/kigi-sampler/tests/test_kimi_wire.rs delete mode 100644 crates/codegen/kigi-shell/src/agent/handlers/workspaces.rs create mode 100644 crates/codegen/kigi-shell/src/agent/models_fetch.rs delete mode 100644 crates/codegen/kigi-shell/src/extensions/bundle.rs delete mode 100644 crates/codegen/kigi-shell/src/extensions/share.rs delete mode 100644 crates/codegen/kigi-shell/src/remote/agent.rs delete mode 100644 crates/codegen/kigi-shell/src/remote/chat_models_client.rs delete mode 100644 crates/codegen/kigi-shell/src/remote/client.rs delete mode 100644 crates/codegen/kigi-shell/src/remote/conversations_client.rs delete mode 100644 crates/codegen/kigi-shell/src/remote/mod.rs delete mode 100644 crates/codegen/kigi-shell/src/remote/pull.rs delete mode 100644 crates/codegen/kigi-shell/src/remote/pull_smoke_test.rs delete mode 100644 crates/codegen/kigi-shell/src/remote/sync.rs delete mode 100644 crates/codegen/kigi-shell/src/remote/workspaces_client.rs create mode 100644 crates/codegen/kigi-shell/src/session/feedback_types.rs delete mode 100644 crates/codegen/kigi-shell/tests/test_settings_refresh.rs delete mode 100644 crates/codegen/kigi-tui/src/app/dispatch/billing.rs delete mode 100644 crates/codegen/kigi-tui/src/app/subscription.rs delete mode 100644 crates/codegen/kigi-tui/src/scrollback/blocks/credit_limit.rs delete mode 100644 crates/codegen/kigi-tui/src/share_cmd.rs delete mode 100644 crates/codegen/kigi-tui/src/slash/commands/privacy.rs delete mode 100644 crates/codegen/kigi-tui/src/slash/commands/share.rs delete mode 100644 crates/codegen/kigi-tui/src/views/credit_bar.rs delete mode 100644 prod/mc/cli-chat-proxy-types/Cargo.toml delete mode 100644 prod/mc/cli-chat-proxy-types/src/client_metrics_types.rs delete mode 100644 prod/mc/cli-chat-proxy-types/src/deployment_config_types.rs delete mode 100644 prod/mc/cli-chat-proxy-types/src/feedback_types.rs delete mode 100644 prod/mc/cli-chat-proxy-types/src/lib.rs delete mode 100644 prod/mc/cli-chat-proxy-types/src/metadata_types.rs delete mode 100644 prod/mc/cli-chat-proxy-types/src/sandbox_types.rs delete mode 100644 prod/mc/cli-chat-proxy-types/src/serde_helpers.rs delete mode 100644 prod/mc/cli-chat-proxy-types/src/session_types.rs delete mode 100644 prod/mc/cli-chat-proxy-types/src/storage_types.rs delete mode 100644 prod/mc/cli-chat-proxy-types/src/subagent_bundle.rs diff --git a/Cargo.lock b/Cargo.lock index 6920eb6..13f450e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5785,7 +5785,6 @@ dependencies = [ "dunce", "kigi-tty-utils", "kigi-version", - "prod-mc-cli-chat-proxy-types", "ring", "semver", "serde", @@ -6385,7 +6384,6 @@ dependencies = [ "libc", "objc2 0.6.4", "parking_lot", - "prod-mc-cli-chat-proxy-types", "regex", "serde", "serde_json", @@ -6490,7 +6488,6 @@ dependencies = [ "parking_lot", "portable-pty", "process-wrap", - "prod-mc-cli-chat-proxy-types", "prost", "rand 0.9.5", "regex", @@ -9047,16 +9044,6 @@ dependencies = [ "hex", ] -[[package]] -name = "prod-mc-cli-chat-proxy-types" -version = "0.1.0" -dependencies = [ - "chrono", - "serde", - "serde_json", - "toml", -] - [[package]] name = "prodash" version = "31.0.0" diff --git a/Cargo.toml b/Cargo.toml index 4d7c4dc..b0f9762 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -75,7 +75,6 @@ members = [ "crates/common/kigi-tool-runtime", "crates/common/kigi-tool-types", "crates/common/kigi-tracing", - "prod/mc/cli-chat-proxy-types", "third_party/dagre_rust", "third_party/graphlib_rust", "third_party/mermaid-to-svg", @@ -337,17 +336,6 @@ strip = false debug = 1 split-debuginfo = "off" -# Production profile for latency-sensitive x-product services (VF, home-mixer). -# Thin LTO gives ~90% of full LTO benefit at significantly faster link time. -# Keeps symbols + line tables for prod debuggability (perf, flamegraph, stack traces). -[profile.x-prod] -inherits = "release" -lto = "thin" -strip = false -codegen-units = 1 -debug = "line-tables-only" -panic = "unwind" - # Desktop release profile. Functionally identical to release-dist — kept as a # named alias so the desktop workflow can reference it without coupling to the # CLI pipeline's profile name. Alpha and stable share a single release-dist @@ -365,7 +353,9 @@ codegen-units = 128 debug = "line-tables-only" opt-level = 0 lto = false -incremental = true +# Incremental caches balloon to tens of GB across this workspace's 60+ crates +# under iterative full-workspace builds; rebuild cost is the cheaper trade. +incremental = false [profile.bench] debug = true diff --git a/crates/codegen/kigi-bin/src/main.rs b/crates/codegen/kigi-bin/src/main.rs index 7d97726..082754b 100644 --- a/crates/codegen/kigi-bin/src/main.rs +++ b/crates/codegen/kigi-bin/src/main.rs @@ -35,7 +35,7 @@ use kigi_shell::leader::{ use kigi_shell::leader::{ControlPayload, LeaderClient, connect_or_spawn, default_socket_path}; use kigi_tui::app::{ AgentCmd, Command, LeaderMgmtArgs, LeaderMgmtCommand, LeaderTargetArgs, PagerArgs, - join_early_prefetch, resolve_use_leader, + resolve_use_leader, }; use kigi_tui::app::{WorkspaceMgmtArgs, WorkspaceMgmtCommand, WorkspaceStartArgs}; use kigi_tui::client_identity::PAGER_CLIENT_VERSION; @@ -45,8 +45,8 @@ use std::net::SocketAddr; use tokio_util::sync::CancellationToken; /// Apply global endpoint CLI args to an existing config. fn apply_agent_endpoint_args(agent_args: &kigi_tui::app::AgentArgs, config: &mut AgentConfig) { - if let Some(v) = &agent_args.cli_chat_proxy_base_url { - config.endpoints.cli_chat_proxy_base_url = Some(v.clone()); + if let Some(v) = &agent_args.coding_api_base_url { + config.endpoints.coding_api_base_url = Some(v.clone()); } if let Some(v) = &agent_args.xai_api_base_url { config.endpoints.xai_api_base_url = v.clone(); @@ -324,43 +324,20 @@ fn ensure_control_caps(reg: &LeaderRegistration) -> Result<&LeaderCapabilities> .as_ref() .ok_or_else(|| anyhow::anyhow!("Leader does not advertise capabilities (legacy version)")) } -/// Env override for the `grok workspace` gate: any truthy value enables the -/// command locally, a falsy one disables it — bypassing the remote settings flag. +/// Env override for the `kigi workspace` gate: any truthy value enables the +/// command locally, a falsy one disables it. This is the only gate now that +/// the server-side feature flag (xAI remote settings) is gone. const WORKSPACE_COMMAND_ENV: &str = "KIGI_WORKSPACE_COMMAND"; -/// Resolution of the `grok workspace` gate. `Unknown` is kept separate from -/// `Disabled` so we don't tell the user the flag is off when the settings were -/// simply never read (both fail closed, but `Unknown` earns an honest message). -#[derive(Debug, PartialEq, Eq)] -enum WorkspaceGate { - Enabled, - Disabled, - Unknown, -} /// The `KIGI_WORKSPACE_COMMAND` override, if set (`Some(true)`/`Some(false)`); -/// `None` defers to the remote settings flag. +/// `None` means unset (the command stays disabled by default). fn workspace_command_env_override() -> Option { std::env::var(WORKSPACE_COMMAND_ENV) .ok() .map(|v| env_flag_enabled(&v)) } -/// Resolve the gate. Precedence: env override > remote `Some(true)` > -/// loaded-but-off (`Disabled`) > settings-not-loaded (`Unknown`). -fn workspace_command_gate( - env_override: Option, - remote_settings: Option<&kigi_shell::util::config::RemoteSettings>, -) -> WorkspaceGate { - if let Some(enabled) = env_override { - return if enabled { - WorkspaceGate::Enabled - } else { - WorkspaceGate::Disabled - }; - } - match remote_settings { - Some(rs) if rs.workspace_command_enabled.unwrap_or(false) => WorkspaceGate::Enabled, - Some(_) => WorkspaceGate::Disabled, - None => WorkspaceGate::Unknown, - } +/// Resolve the gate: enabled exactly when the env override says so. +fn workspace_command_gate(env_override: Option) -> bool { + env_override.unwrap_or(false) } /// Truthy parse for grok on/off env vars: everything enables except the common /// falsy spellings (`0`, `false`, `off`, `no`, empty). @@ -370,40 +347,16 @@ fn env_flag_enabled(value: &str) -> bool { "" | "0" | "false" | "off" | "no" ) } -/// Blocking fetch of remote settings via the startup prefetch path. -fn fetch_remote_settings() -> Option { - join_early_prefetch(kigi_shell::agent::models::start_early_prefetch(None)) -} async fn run_workspace_mgmt(args: WorkspaceMgmtArgs) -> Result<()> { - let env_override = workspace_command_env_override(); - let remote_settings = if env_override.is_none() { - fetch_remote_settings() - } else { - None - }; - match workspace_command_gate(env_override, remote_settings.as_ref()) { - WorkspaceGate::Enabled => {} - WorkspaceGate::Disabled => { - anyhow::bail!( - "`grok workspace` is not enabled for this account \ - (gated by a server-side feature flag that is currently off)." - ) - } - WorkspaceGate::Unknown => { - anyhow::bail!( - "Could not load your settings for `grok workspace`. Check your \ - network connection (run `grok login` if you are signed out), then \ - try again." - ) - } + if !workspace_command_gate(workspace_command_env_override()) { + anyhow::bail!( + "`kigi workspace` is experimental and disabled by default. \ + Set {WORKSPACE_COMMAND_ENV}=1 to enable it." + ) } match args.command { - WorkspaceMgmtCommand::Start(a) => { - workspace_start(a, false, remote_settings.or_else(fetch_remote_settings)).await - } - WorkspaceMgmtCommand::Restart(a) => { - workspace_start(a, true, remote_settings.or_else(fetch_remote_settings)).await - } + WorkspaceMgmtCommand::Start(a) => workspace_start(a, false).await, + WorkspaceMgmtCommand::Restart(a) => workspace_start(a, true).await, WorkspaceMgmtCommand::Pause { target, json } => { workspace_control(&target, json, ControlCommand::WorkspacePause).await } @@ -467,24 +420,13 @@ async fn workspace_control( client.cancel(); Ok(()) } -async fn workspace_start( - args: WorkspaceStartArgs, - restart: bool, - remote_settings: Option, -) -> Result<()> { +async fn workspace_start(args: WorkspaceStartArgs, restart: bool) -> Result<()> { use kigi_shell::auth::ensure_authenticated; - kigi_shell::util::config::set_remote_campaigns_from_settings(remote_settings.as_ref()); let raw_config = kigi_shell::config::load_effective_config() .map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?; let agent_config = AgentConfig::new_from_toml_cfg(&raw_config) .map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?; - let (use_leader, _) = resolve_use_leader( - args.leader, - args.no_leader, - &raw_config, - remote_settings.as_ref(), - true, - ); + let (use_leader, _) = resolve_use_leader(args.leader, args.no_leader, &raw_config, true); if !use_leader { anyhow::bail!( "`grok workspace` requires leader mode (the workspace is shared via the leader).\n\ @@ -949,7 +891,9 @@ async fn run_agent_command( } } } - let early_prefetch = kigi_shell::agent::models::start_early_prefetch(None); + // Fire-and-forget model-catalog warmup (nothing joins the handle now that + // the xAI settings fetch it used to carry is gone). + drop(kigi_shell::agent::models::start_early_prefetch(None)); kigi_shell::agent::mvp_agent::warm_async_http_client(); tokio::task::spawn_blocking(|| {}); let is_stdio = matches!(agent_args.mode, AgentCmd::Stdio); @@ -972,8 +916,6 @@ async fn run_agent_command( .ok(); } } - let remote_settings = join_early_prefetch(early_prefetch); - kigi_shell::util::config::set_remote_campaigns_from_settings(remote_settings.as_ref()); let raw_config = kigi_shell::config::load_effective_config() .map_err(|e| anyhow::anyhow!("Failed to load config: {}", e))?; let mut agent_config = AgentConfig::new_from_toml_cfg(&raw_config) @@ -1008,10 +950,9 @@ async fn run_agent_command( agent_config.plugins.cli_plugin_dirs = agent_args.canonical_plugin_dirs(); } apply_agent_endpoint_args(&agent_args, &mut agent_config); - agent_config.remote_settings = remote_settings.clone(); agent_config.resolve_runtime_fields(&kigi_shell::agent::config::RuntimeResolutionContext { raw_config: &raw_config, - remote_settings: remote_settings.as_ref(), + remote_settings: None, cwd: None, is_headless: !is_leader, cli_subagents: None, @@ -1030,7 +971,6 @@ async fn run_agent_command( agent_args.leader, agent_args.no_leader, &raw_config, - remote_settings.as_ref(), leader_eligible, ); tracing::info!(use_leader, ?policy_disable_reason, "leader mode resolved"); @@ -1492,9 +1432,10 @@ async fn async_main() -> Result<()> { unsafe { std::env::set_var("KIGI_COMPACTION_DETAIL", detail) }; } if args.chat() { - unsafe { - std::env::set_var(kigi_shell::agent::chat_modes::KIGI_CHAT_MODE_ENV, "1"); - } + anyhow::bail!( + "--chat is no longer supported: the grok.com chat frontend it drove was \ + removed along with the xAI backend." + ); } if let Some(ref socket) = args.leader_socket { unsafe { std::env::set_var(kigi_shell::leader::LEADER_SOCKET_ENV, socket) }; @@ -1632,19 +1573,7 @@ async fn async_main() -> Result<()> { } Command::Sessions(sessions_args) => { init_tracing_simple("cli"); - let config = kigi_shell::config::load_effective_config_disk_only() - .map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?; - let agent_config = AgentConfig::new_from_toml_cfg(&config) - .map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?; - return kigi_tui::sessions_cmd::run(sessions_args, &agent_config).await; - } - Command::Share(ref share_args) => { - init_tracing_simple("cli"); - let config = kigi_shell::config::load_effective_config_disk_only() - .map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?; - let agent_config = AgentConfig::new_from_toml_cfg(&config) - .map_err(|e| anyhow::anyhow!("Failed to create agent config: {e}"))?; - return kigi_tui::share_cmd::run(share_args, &agent_config).await; + return kigi_tui::sessions_cmd::run(sessions_args).await; } Command::Export(export_args) => { init_tracing_simple("cli"); @@ -2130,37 +2059,9 @@ mod tests { } #[test] fn workspace_command_gate_resolution() { - use kigi_shell::util::config::RemoteSettings; - let on = RemoteSettings { - workspace_command_enabled: Some(true), - ..RemoteSettings::default() - }; - let off = RemoteSettings::default(); - assert_eq!( - workspace_command_gate(None, Some(&on)), - WorkspaceGate::Enabled - ); - assert_eq!( - workspace_command_gate(None, Some(&off)), - WorkspaceGate::Disabled - ); - assert_eq!(workspace_command_gate(None, None), WorkspaceGate::Unknown); - assert_eq!( - workspace_command_gate(Some(true), Some(&off)), - WorkspaceGate::Enabled - ); - assert_eq!( - workspace_command_gate(Some(true), None), - WorkspaceGate::Enabled - ); - assert_eq!( - workspace_command_gate(Some(false), Some(&on)), - WorkspaceGate::Disabled - ); - assert_eq!( - workspace_command_gate(Some(false), None), - WorkspaceGate::Disabled - ); + assert!(workspace_command_gate(Some(true))); + assert!(!workspace_command_gate(Some(false))); + assert!(!workspace_command_gate(None), "unset env defaults to off"); } #[serial_test::serial(KIGI_WORKSPACE_COMMAND)] #[test] diff --git a/crates/codegen/kigi-chat-state/src/types.rs b/crates/codegen/kigi-chat-state/src/types.rs index 17b233b..68e7dca 100644 --- a/crates/codegen/kigi-chat-state/src/types.rs +++ b/crates/codegen/kigi-chat-state/src/types.rs @@ -123,8 +123,6 @@ pub struct Credentials { pub auth_type: AuthType, /// Optional extra auth material forwarded with requests when present. pub alpha_test_key: Option, - /// Client version string. - pub client_version: Option, } /// The messages captured during a single conversation turn. diff --git a/crates/codegen/kigi-config-types/src/lib.rs b/crates/codegen/kigi-config-types/src/lib.rs index 447324c..c8d8f49 100644 --- a/crates/codegen/kigi-config-types/src/lib.rs +++ b/crates/codegen/kigi-config-types/src/lib.rs @@ -282,12 +282,6 @@ pub struct RemoteSettings { pub dream_min_sessions: Option, #[serde(default)] pub dream_check_interval_secs: Option, - /// Cadence (seconds) of the pager's free→paid subscription watch. - /// `0` disables it; the pager clamps and defaults (see its - /// `app::subscription` module). Forwarded from the `grok_build_settings` - /// remote settings flag via the CCP `/settings` flatten catch-all. - #[serde(default)] - pub subscription_watch_interval_secs: Option, #[serde(default)] pub writeback_enabled: Option, /// OAuth2 provider issuer URL (e.g., "https://auth.x.ai"). When present @@ -601,10 +595,6 @@ pub struct RemoteSettings { /// is a separate client tier gate. #[serde(default)] pub voice_mode_enabled: Option, - /// Whether ZDR (Zero Data Retention) users are allowed to use the product. - /// Controlled via remote settings. Default `false` (blocked) during beta. - #[serde(default)] - pub zdr_access_enabled: Option, /// remote settings tier of the `remember_tool_approvals` gate (whether per-tool /// "Always allow …" prompt options are shown). Lowest precedence; typically /// targeted per-org. Default `false`. @@ -652,33 +642,10 @@ pub struct RemoteSettings { /// `"default"`). Used only when no effective TOML permission key is set. #[serde(default)] pub permission_mode: Option, - /// User's subscription tier from remote settings `grok_build_access_gate`. - /// E.g. "free", "premium", "supergrok", "supergrok_heavy". - /// Stamped on analytics events + user profile for filtering. - #[serde(default)] - pub subscription_tier: Option, - #[serde(default)] - pub gate_message: Option, - #[serde(default)] - pub gate_url: Option, - #[serde(default)] - pub gate_label: Option, /// Whether the session picker groups entries by repo name. /// When `None` or `Some(false)`, sessions are shown in a flat list. #[serde(default)] pub session_picker_grouped: Option, - /// Whether the user is allowed to use Grok Build. Set by remote settings - /// `grok_build_access_gate` targeting rules. `None` = no server response - /// yet (client uses own fallback check). `Some(false)` = blocked. - #[serde(default)] - pub allow_access: Option, - /// User-friendly display name for the current subscription tier - /// (e.g. "SuperGrok", "X Premium+", "Free", "API Key"). Set by CCP - /// from the JWT tier claim (OAuth) or credential kind (API key). - /// Free/Invalid OAuth → `"Free"`; API keys → `"API Key"` (Mixpanel - /// `api_key`, never free). - #[serde(default)] - pub subscription_tier_display: Option, /// Whether on-demand credit usage is enabled. When `Some(false)`, the /// billing extension blocks on-demand cap changes. #[serde(default)] diff --git a/crates/codegen/kigi-config/Cargo.toml b/crates/codegen/kigi-config/Cargo.toml index d2c58d4..d065258 100644 --- a/crates/codegen/kigi-config/Cargo.toml +++ b/crates/codegen/kigi-config/Cargo.toml @@ -9,8 +9,6 @@ description = "Shared config loading for Grok — kigi_home, effective config (r base64 = { workspace = true } blake3 = { workspace = true } dunce = { workspace = true } -# Shared signed deployment-config envelope contract with the cli-chat-proxy signer. -prod-mc-cli-chat-proxy-types = { path = "../../../prod/mc/cli-chat-proxy-types" } ring = { workspace = true } semver = { workspace = true } serde = { workspace = true } diff --git a/crates/codegen/kigi-config/src/signed_policy.rs b/crates/codegen/kigi-config/src/signed_policy.rs index f8d6e48..f661416 100644 --- a/crates/codegen/kigi-config/src/signed_policy.rs +++ b/crates/codegen/kigi-config/src/signed_policy.rs @@ -8,7 +8,64 @@ //! Inert until a public key is provisioned: with no embedded keys the cache //! marker stays the (best-effort) authority. use base64::Engine; -pub use prod_mc_cli_chat_proxy_types::{SignatureEnvelope, SignedPayload, now_unix}; +use serde::{Deserialize, Serialize}; + +/// The payload format version the server currently signs. Bump when the payload +/// gains semantics (e.g. an anti-replay counter or a key-fingerprint binding) so +/// verifiers can distinguish generations; `0` means a pre-versioned payload. +pub const SIGNED_PAYLOAD_VERSION: u32 = 1; + +/// The exact bytes the server signs: the served policy, the principal it is +/// bound to, and an expiry. Serialized once on the server and shipped verbatim +/// as `signed_payload`, so the client verifies the received bytes directly +/// instead of re-canonicalizing (no cross-language serialization drift). +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct SignedPayload { + /// Payload format version ([`SIGNED_PAYLOAD_VERSION`]); `default` 0 so + /// pre-versioned sidecars parse and verify unchanged. + #[serde(default)] + pub version: u32, + #[serde(default)] + pub deployment_id: Option, + #[serde(default)] + pub team_id: Option, + #[serde(default)] + pub managed_config: Option, + #[serde(default)] + pub requirements: Option, + /// Strict (fail-closed) opt-in, carried in the SIGNED bytes so a local actor can't + /// flip enforcement. `default` false so an older/unsigned payload stays lenient. + #[serde(default)] + pub fail_closed: bool, + /// Unix seconds after which the signature is no longer trusted. + pub expires_at: u64, + /// Identifies the signing key, so a rotation can be distinguished. + pub key_id: String, +} + +/// One signed envelope carried alongside the legacy policy fields in the +/// deployment-config response (additive: old clients ignore it). Also the +/// shape the client persists as its on-disk signature sidecar. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SignatureEnvelope { + /// The exact JSON string that was signed (a serialized [`SignedPayload`]). + pub signed_payload: String, + /// Base64 (standard) Ed25519 signature over `signed_payload`'s UTF-8 bytes. + pub signature: String, + /// Untrusted (outside the signed bytes): a hint for picking among multiple + /// envelopes, never for selecting the verifying key — only the signed + /// payload's `key_id` is authoritative. + #[serde(default)] + pub key_id: String, +} + +/// Unix seconds now (saturating to 0 on a pre-epoch clock). +pub fn now_unix() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} /// Compiled-in trusted Ed25519 public keys, `(key_id, raw 32 bytes)`; more than one /// entry only during a rotation. Empty ships dark (see [`verification_active`]). /// Compile-time, not an env flag: the local attacker controls their env. diff --git a/crates/codegen/kigi-config/src/signed_policy/tests.rs b/crates/codegen/kigi-config/src/signed_policy/tests.rs index 5a2a997..0607400 100644 --- a/crates/codegen/kigi-config/src/signed_policy/tests.rs +++ b/crates/codegen/kigi-config/src/signed_policy/tests.rs @@ -956,3 +956,28 @@ fn rotation_selects_the_trusted_key_by_signed_key_id() { Err(SigError::SignatureMismatch) ); } + +/// The version field round-trips, and a pre-versioned payload (no `version` +/// key) defaults to 0 — old sidecars keep parsing. +#[test] +fn signed_payload_version_round_trips_and_defaults() { + let versioned = SignedPayload { + version: SIGNED_PAYLOAD_VERSION, + deployment_id: None, + team_id: Some("team-007".into()), + managed_config: None, + requirements: None, + fail_closed: false, + expires_at: 4_000_000_000, + key_id: "v1".into(), + }; + let json = serde_json::to_string(&versioned).unwrap(); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + versioned + ); + + let legacy: SignedPayload = + serde_json::from_str(r#"{"expires_at": 1, "key_id": "v1"}"#).unwrap(); + assert_eq!(legacy.version, 0, "pre-versioned payloads default to 0"); +} diff --git a/crates/codegen/kigi-config/src/validation.rs b/crates/codegen/kigi-config/src/validation.rs index bf87c09..cfdea5c 100644 --- a/crates/codegen/kigi-config/src/validation.rs +++ b/crates/codegen/kigi-config/src/validation.rs @@ -7,11 +7,21 @@ use crate::loader::{apply_version_overrides_with_registered, load_toml_file}; use crate::paths::{system_config_dir, user_kigi_home}; use crate::version_overrides::{VersionOverrideError, apply_version_overrides}; -use prod_mc_cli_chat_proxy_types::FAIL_CLOSED_KEY; -/// The canonical opt-in key + string parse live in the shared types crate, next to -/// the signed payload that carries the flag, so the server-side signer and this -/// client parse the same semantics. -pub use prod_mc_cli_chat_proxy_types::fail_closed_flag_from_str; +/// The `requirements.toml` opt-in key for strict (fail-closed) enforcement. +/// Lives next to the signed payload that carries the flag +/// ([`crate::signed_policy::SignedPayload::fail_closed`]) so the two sides +/// can't drift. +pub const FAIL_CLOSED_KEY: &str = "fail_closed"; + +/// Read the `fail_closed` opt-in from a requirements-TOML string — THE canonical +/// parse shared by every caller so the semantics can't drift. +/// Invalid TOML or a non-bool value → `false`. +pub fn fail_closed_flag_from_str(requirements: &str) -> bool { + toml::from_str::(requirements) + .ok() + .and_then(|v| v.get(FAIL_CLOSED_KEY).and_then(toml::Value::as_bool)) + .unwrap_or(false) +} /// Read the `fail_closed` opt-in from a parsed requirements layer — same semantics as /// [`fail_closed_flag_from_str`]. Env tightening (file vs `KIGI_MANAGED_CONFIG_FAIL_CLOSED`) diff --git a/crates/codegen/kigi-http/src/lib.rs b/crates/codegen/kigi-http/src/lib.rs index 783c961..68cd513 100644 --- a/crates/codegen/kigi-http/src/lib.rs +++ b/crates/codegen/kigi-http/src/lib.rs @@ -176,7 +176,7 @@ pub fn process_user_agent_string() -> String { UserAgent { origin, - agent_product: "grok-shell", + agent_product: "kigi", agent_version, platform: PlatformInfo::current(), } @@ -186,7 +186,7 @@ pub fn process_user_agent_string() -> String { pub fn session_user_agent_string(origin: &OriginClientInfo) -> String { UserAgent { origin: origin.clone(), - agent_product: "grok-shell", + agent_product: "kigi", agent_version: agent_version(), platform: PlatformInfo::current(), } @@ -234,29 +234,9 @@ pub fn client_type_from_origin(origin: Option<&OriginClientInfo>) -> ClientType ClientType::from_client_identifier(origin.map(|o| o.product.as_str())) } -/// Process-level client identifier (`KIGI_CLIENT_NAME` env var, default `"grok-shell"`). +/// Process-level client identifier (`KIGI_CLIENT_NAME` env var, default `"kigi"`). pub fn process_client_identifier() -> String { - std::env::var("KIGI_CLIENT_NAME").unwrap_or_else(|_| "grok-shell".to_string()) -} - -/// Header telling cli-chat-proxy whether this process is a single-prompt -/// (`grok -p`) run or an interactive session; feeds the `client_mode` -/// metric label. -pub const CLIENT_MODE_HEADER: &str = "x-grok-client-mode"; - -/// One-way latch: set to `"headless"` at startup by the non-TUI entry points -/// (`run_single_turn` for `grok -p`, `run_headless_inner` for -/// `grok agent [headless]`), `"interactive"` otherwise. -static CLIENT_MODE: OnceLock<&'static str> = OnceLock::new(); - -/// Mark this process as headless (single-prompt). No-op if already set. -pub fn set_process_client_mode_headless() { - let _ = CLIENT_MODE.set("headless"); -} - -/// The mode sent in [`CLIENT_MODE_HEADER`]; defaults to `"interactive"`. -pub fn process_client_mode() -> &'static str { - CLIENT_MODE.get().copied().unwrap_or("interactive") + std::env::var("KIGI_CLIENT_NAME").unwrap_or_else(|_| "kigi".to_string()) } pub fn user_agent_string_for(origin: &OriginClientInfo) -> String { @@ -568,14 +548,14 @@ mod tests { product: "grok-desktop".to_string(), version: Some("1.2.3".to_string()), }); - assert!(with_version.starts_with("grok-desktop/1.2.3 grok-shell/")); + assert!(with_version.starts_with("grok-desktop/1.2.3 kigi/")); assert!(with_version.contains(" (")); let without_version = session_user_agent_string(&OriginClientInfo { product: "grok-web".to_string(), version: None, }); - assert!(without_version.starts_with("grok-web grok-shell/")); + assert!(without_version.starts_with("grok-web kigi/")); assert!(!without_version.starts_with("grok-web/")); } @@ -583,10 +563,10 @@ mod tests { fn user_agent_render_collapses_duplicate_origin_and_agent_identity() { let ua = UserAgent { origin: OriginClientInfo { - product: "grok-shell".to_string(), + product: "kigi".to_string(), version: Some("0.1.171".to_string()), }, - agent_product: "grok-shell", + agent_product: "kigi", agent_version: "0.1.171".to_string(), platform: PlatformInfo { os: "macos".to_string(), @@ -594,6 +574,6 @@ mod tests { }, }; - assert_eq!(ua.render(), "grok-shell/0.1.171 (macos; aarch64)"); + assert_eq!(ua.render(), "kigi/0.1.171 (macos; aarch64)"); } } diff --git a/crates/codegen/kigi-pager-minimal/src/auth.rs b/crates/codegen/kigi-pager-minimal/src/auth.rs index 3ece0ff..7ee0f77 100644 --- a/crates/codegen/kigi-pager-minimal/src/auth.rs +++ b/crates/codegen/kigi-pager-minimal/src/auth.rs @@ -141,7 +141,7 @@ pub(super) fn render_auth(buf: &mut Buffer, area: Rect, theme: &Theme, hint: &Mi area, y, bottom, - Line::from(Span::styled("Sign in to Grok", bold)), + Line::from(Span::styled("Sign in to Kimi", bold)), ); y = put_line(buf, area, y, bottom, Line::default()); match url { @@ -242,12 +242,13 @@ mod tests { #[test] fn device_user_code_parses_verification_url() { + // The live Kimi device flow returns this exact URL shape. assert_eq!( - device_user_code("https://accounts.x.ai/oauth2/device?user_code=ABCD-EFGH"), + device_user_code("https://www.kimi.com/code/authorize_device?user_code=ABCD-EFGH"), Some("ABCD-EFGH") ); assert_eq!( - device_user_code("https://accounts.x.ai/oauth2/device"), + device_user_code("https://www.kimi.com/code/authorize_device"), None ); assert_eq!(device_user_code("https://x/device?other=1"), None); @@ -261,14 +262,14 @@ mod tests { let st = AuthState::Authenticating { request_seq: 1, handle: None, - auth_url: Some("https://accounts.x.ai/device?user_code=ABCD-EFGH".into()), + auth_url: Some("https://www.kimi.com/code/authorize_device?user_code=ABCD-EFGH".into()), mode: AuthMode::Device, }; match minimal_auth_hint(&st) { MinimalAuthHint::SigningIn { url, code } => { assert_eq!( url.as_deref(), - Some("https://accounts.x.ai/device?user_code=ABCD-EFGH") + Some("https://www.kimi.com/code/authorize_device?user_code=ABCD-EFGH") ); assert_eq!(code.as_deref(), Some("ABCD-EFGH")); } @@ -308,7 +309,7 @@ mod tests { let area = Rect::new(0, 0, 80, 12); let mut buf = Buffer::empty(area); let hint = MinimalAuthHint::SigningIn { - url: Some("https://accounts.x.ai/device?user_code=ABCD-EFGH".into()), + url: Some("https://www.kimi.com/code/authorize_device?user_code=ABCD-EFGH".into()), code: Some("ABCD-EFGH".into()), }; render_auth(&mut buf, area, &theme, &hint); @@ -320,8 +321,11 @@ mod tests { } } } - assert!(text.contains("Sign in to Grok"), "header: {text:?}"); - assert!(text.contains("accounts.x.ai/device"), "url: {text:?}"); + assert!(text.contains("Sign in to Kimi"), "header: {text:?}"); + assert!( + text.contains("www.kimi.com/code/authorize_device"), + "url: {text:?}" + ); assert!(text.contains("ABCD-EFGH"), "device code: {text:?}"); assert!( text.contains("Waiting for approval"), diff --git a/crates/codegen/kigi-pager-pty-harness/src/content.rs b/crates/codegen/kigi-pager-pty-harness/src/content.rs index 2103ce8..d1daf4f 100644 --- a/crates/codegen/kigi-pager-pty-harness/src/content.rs +++ b/crates/codegen/kigi-pager-pty-harness/src/content.rs @@ -90,7 +90,7 @@ impl ContentController { // config.toml when $HOME alone isn't sufficient (e.g. if // KIGI_SHARE_DIR is set in the test runner's env). ("KIGI_SHARE_DIR".into(), kigi_home), - ("KIGI_CLI_CHAT_PROXY_BASE_URL".into(), self.url()), + ("KIGI_CODE_BASE_URL".into(), self.url()), ("KIGI_XAI_API_BASE_URL".into(), self.url()), ("XAI_API_KEY".into(), "test-key-for-ci".into()), ("KIGI_TELEMETRY_ENABLED".into(), "false".into()), @@ -278,7 +278,7 @@ mod tests { get("KIGI_SHARE_DIR").as_deref(), content.home().join(".kigi").to_str() ); - assert_eq!(get("KIGI_CLI_CHAT_PROXY_BASE_URL"), Some(content.url())); + assert_eq!(get("KIGI_CODE_BASE_URL"), Some(content.url())); assert_eq!(get("KIGI_XAI_API_BASE_URL"), Some(content.url())); assert_eq!(get("XAI_API_KEY").as_deref(), Some("test-key-for-ci")); assert_eq!(get("KIGI_TELEMETRY_ENABLED").as_deref(), Some("false")); diff --git a/crates/codegen/kigi-sampler/Cargo.toml b/crates/codegen/kigi-sampler/Cargo.toml index 7e34439..cd19866 100644 --- a/crates/codegen/kigi-sampler/Cargo.toml +++ b/crates/codegen/kigi-sampler/Cargo.toml @@ -3,7 +3,7 @@ license = "Apache-2.0" name = "kigi-sampler" version.workspace = true edition.workspace = true -description = "Actor-based sampling/inference layer for xAI grok (HTTP streaming + retry, no shell coupling)" +description = "Actor-based sampling/inference layer for the Kimi APIs (HTTP streaming + retry, no shell coupling)" [dependencies] # Internal diff --git a/crates/codegen/kigi-sampler/src/actor/request_task.rs b/crates/codegen/kigi-sampler/src/actor/request_task.rs index 6ff81cb..a4b74ca 100644 --- a/crates/codegen/kigi-sampler/src/actor/request_task.rs +++ b/crates/codegen/kigi-sampler/src/actor/request_task.rs @@ -373,16 +373,13 @@ async fn apply_retry_decision( RetryDecision::Fatal(fatal_err) => { // Emit only on true budget exhaustion (hit the retry / rate-limit // cap), mirroring `classify_error`'s Fatal conditions — NOT on a - // server `x-should-retry: false` or a non-retryable error, which - // are also Fatal but are not "exhausted". + // non-retryable error, which is also Fatal but not "exhausted". let next_attempt = *retry_count + 1; - let server_said_stop = matches!(err.should_retry_header(), Some(false)); - let budget_exhausted = !server_said_stop - && if err.is_rate_limited() { - next_attempt >= max_retries.min(rate_limit_threshold) - } else { - err.is_retryable() && next_attempt >= max_retries - }; + let budget_exhausted = if err.is_rate_limited() { + next_attempt >= max_retries.min(rate_limit_threshold) + } else { + err.is_retryable() && next_attempt >= max_retries + }; if budget_exhausted { let exhausted_span = tracing::info_span!( "http.retries_exhausted", @@ -609,7 +606,6 @@ fn synthesize_from_info(info: &SamplingErrorInfo) -> SamplingError { message: info.message.clone(), model_metadata: info.model_metadata.clone(), retry_after_secs: info.retry_after_secs, - should_retry: None, } } SamplingErrorKind::EmptyResponse => { diff --git a/crates/codegen/kigi-sampler/src/actor/state.rs b/crates/codegen/kigi-sampler/src/actor/state.rs index 10e49f7..cfa6685 100644 --- a/crates/codegen/kigi-sampler/src/actor/state.rs +++ b/crates/codegen/kigi-sampler/src/actor/state.rs @@ -97,10 +97,6 @@ mod tests { idle_timeout_secs: None, reasoning_effort: None, origin_client: None, - client_identifier: None, - deployment_id: None, - user_id: None, - client_version: None, attribution_callback: None, bearer_resolver: None, supports_backend_search: false, diff --git a/crates/codegen/kigi-sampler/src/client.rs b/crates/codegen/kigi-sampler/src/client.rs index a2338af..8b00de1 100644 --- a/crates/codegen/kigi-sampler/src/client.rs +++ b/crates/codegen/kigi-sampler/src/client.rs @@ -1,16 +1,20 @@ -//! HTTP client for the xAI sampling APIs. +//! HTTP client for the sampling APIs. //! //! Owns the `reqwest::Client`, default request headers, and per-method //! defaults. Talks to three backend shapes: //! -//! * Chat Completions (`/chat/completions`) -//! * Responses API (`/responses`) -//! * Anthropic Messages API (`/messages`) +//! * Chat Completions (`/chat/completions`) — the Kimi dialect; both +//! product channels (subscription OAuth, Moonshot API key) ride it. +//! Kimi-specific request deviations are absorbed by [`crate::kimi_compat`]. +//! * Responses API (`/responses`) — kept for custom providers. +//! * Anthropic Messages API (`/messages`) — kept for custom providers. //! -//! All trace-upload and URL-based header injection is intentionally -//! *not* here. The session is responsible for putting any per-request -//! headers (proxy auth, OTel context, etc.) -//! into [`SamplerConfig::extra_headers`] before constructing the client. +//! Auth on the wire is a plain `Authorization: Bearer {token}` (or +//! `x-api-key` for Anthropic-scheme custom providers). All URL-based +//! header injection is intentionally *not* here. The session is +//! responsible for putting any per-request headers (device identity, +//! OTel context, etc.) into [`SamplerConfig::extra_headers`] before +//! constructing the client. use eventsource_stream::Eventsource; use futures_util::StreamExt; @@ -33,46 +37,10 @@ use crate::config::{AuthScheme, OriginClientInfo, SamplerConfig}; // Re-export ApiBackend from the shared types crate for downstream callers. pub use kigi_sampling_types::ApiBackend; -/// Process-level fallback for the `x-grok-client-identifier` header. -const DEFAULT_CLIENT_IDENTIFIER: &str = "grok-shell"; - -/// Product identifier baked into User-Agent strings. -const AGENT_PRODUCT: &str = "grok-shell"; +/// Product identifier baked into User-Agent strings (PRD F3: `kigi/{version}`). +const AGENT_PRODUCT: &str = "kigi"; const ANTHROPIC_DEFAULT_MAX_TOKENS: u32 = 128_000; -/// Per-request `x-grok-*` headers. Optional fields are skipped when empty/`None`. -struct GrokRequestHeaders<'a> { - conv_id: &'a str, - req_id: &'a str, - model_id: &'a str, - session_id: &'a str, - turn_idx: Option<&'a str>, - agent_id: &'a str, - deployment_id: Option<&'a str>, - user_id: Option<&'a str>, -} - -impl GrokRequestHeaders<'_> { - fn apply(&self, builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder { - let mut b = builder - .header("x-grok-conv-id", self.conv_id) - .header("x-grok-req-id", self.req_id) - .header("x-grok-model-override", self.model_id) - .header("x-grok-session-id", self.session_id) - .header("x-grok-agent-id", self.agent_id); - if let Some(idx) = self.turn_idx { - b = b.header("x-grok-turn-idx", idx); - } - if let Some(id) = self.deployment_id.filter(|s| !s.is_empty()) { - b = b.header("x-grok-deployment-id", id); - } - if let Some(id) = self.user_id.filter(|s| !s.is_empty()) { - b = b.header("x-grok-user-id", id); - } - b - } -} - /// Parse the `Retry-After` response header as delta-seconds. /// Our inference backends only emit integer seconds (never HTTP-date), /// so we only handle that form. HTTP-dates silently return `None` and @@ -211,21 +179,6 @@ fn extract_retry_after(headers: &reqwest::header::HeaderMap) -> Option { .map(|s| s.min(120)) } -fn extract_should_retry(headers: &reqwest::header::HeaderMap) -> Option { - headers - .get("x-should-retry") - .and_then(|v| v.to_str().ok()) - .and_then(|s| { - if s.eq_ignore_ascii_case("true") { - Some(true) - } else if s.eq_ignore_ascii_case("false") { - Some(false) - } else { - None // unknown value — treat as absent - } - }) -} - fn extract_model_metadata(headers: &reqwest::header::HeaderMap) -> Option { let context_window = headers .get("x-grok-context-window") @@ -444,45 +397,11 @@ impl SamplingClient { headers.insert(header_name, header_value); } - // Add x-grok-client-version header for version gating at the proxy. - if let Some(client_version) = config.client_version.as_ref() - && let Ok(header_value) = HeaderValue::from_str(client_version) - { - headers.insert( - HeaderName::from_static("x-grok-client-version"), - header_value, - ); - } - - if let Some(deployment_id) = config.deployment_id.as_ref() - && let Ok(header_value) = HeaderValue::from_str(deployment_id) - { - headers.insert( - HeaderName::from_static("x-grok-deployment-id"), - header_value, - ); - } - - if let Some(user_id) = config.user_id.as_ref() - && let Ok(header_value) = HeaderValue::from_str(user_id) - { - headers.insert(HeaderName::from_static("x-grok-user-id"), header_value); - } - - { - let client_id = config - .client_identifier - .clone() - .unwrap_or_else(|| DEFAULT_CLIENT_IDENTIFIER.to_string()); - if let Ok(header_value) = HeaderValue::from_str(&client_id) { - headers.insert( - HeaderName::from_static("x-grok-client-identifier"), - header_value, - ); - } - } - // Always set User-Agent: per-session origin if available, else fallback. + // This and `extra_headers` are the only client-identity signals on the + // wire — the old xAI proxy's `x-grok-*` marker headers are gone + // (PRD F3: auth is a plain bearer; kimi-cli sends only User-Agent + // plus the OAuth device headers, src/kimi_cli/llm.py:317-323). { let ua_string = match config.origin_client.as_ref() { Some(origin) => user_agent_string_for(origin), @@ -689,13 +608,7 @@ impl SamplingClient { } /// Build request headers string for error messages (redacting sensitive values). - fn format_request_headers( - &self, - x_grok_conv_id: &str, - x_grok_req_id: &str, - model_id: &str, - include_accept: bool, - ) -> Vec { + fn format_request_headers(&self, include_accept: bool) -> Vec { let mut req_headers: Vec = self .default_headers .iter() @@ -704,9 +617,6 @@ impl SamplingClient { }) .collect(); - req_headers.push(Self::format_header("x-grok-conv-id", x_grok_conv_id)); - req_headers.push(Self::format_header("x-grok-req-id", x_grok_req_id)); - req_headers.push(Self::format_header("x-grok-model-override", model_id)); if include_accept { req_headers.push(Self::format_header("accept", "text/event-stream")); } @@ -799,7 +709,6 @@ impl SamplingClient { let status = response.status(); let model_metadata = extract_model_metadata(response.headers()); let retry_after_secs = extract_retry_after(response.headers()); - let should_retry = extract_should_retry(response.headers()); let bytes = response.bytes().await?; if !status.is_success() { @@ -816,7 +725,6 @@ impl SamplingClient { message, model_metadata, retry_after_secs, - should_retry, }); } @@ -841,8 +749,6 @@ impl SamplingClient { request: ChatCompletionRequest, ) -> Result { let payload = self.apply_defaults(request)?; - let x_grok_conv_id = &payload.x_grok_conv_id.clone().unwrap_or_default(); - let x_grok_req_id = &payload.x_grok_req_id.clone().unwrap_or_default(); let model_id = payload.model.clone().unwrap_or_default(); tracing::debug!( @@ -851,19 +757,18 @@ impl SamplingClient { "Sending chat completion request" ); - let grok_headers = GrokRequestHeaders { - conv_id: x_grok_conv_id, - req_id: x_grok_req_id, - model_id: &model_id, - session_id: payload.x_grok_session_id.as_deref().unwrap_or_default(), - turn_idx: payload.x_grok_turn_idx.as_deref(), - agent_id: payload.x_grok_agent_id.as_deref().unwrap_or_default(), - deployment_id: payload.x_grok_deployment_id.as_deref(), - user_id: payload.x_grok_user_id.as_deref(), - }; - let http_request = grok_headers - .apply(self.post(self.endpoint("chat/completions"))) - .json(&payload); + // Serialize, then run the Kimi dialect adaptations (single + // adaptation point for every request-side deviation; see + // `crate::kimi_compat`). + let mut request_body = serde_json::to_value(&payload).map_err(|e| { + tracing::error!("Failed to serialize chat/completions request: {}", e); + SamplingError::Serialization(e) + })?; + crate::kimi_compat::adapt_chat_completions_body(&mut request_body); + + let http_request = self + .post(self.endpoint("chat/completions")) + .json(&request_body); let response = http_request.send().await.map_err(|e| { // Log at debug level; errors are surfaced to the caller. @@ -894,13 +799,14 @@ impl SamplingClient { Option, )> { let payload = self.apply_defaults(request)?; - let x_grok_conv_id = &payload.x_grok_conv_id.clone().unwrap_or_default(); - let x_grok_req_id = &payload.x_grok_req_id.clone().unwrap_or_default(); let model_id = payload.model.clone().unwrap_or_default(); - // Wrap the request with streaming fields and serialize once. - // Previously this path serialized twice: first to serde_json::Value - // (to inject `stream` and `stream_options`), then to HTTP body bytes. + // Wrap the request with the streaming fields the Kimi API expects + // (`stream: true` + `stream_options.include_usage: true`, exactly + // what kimi-cli sends — + // packages/kosong/src/kosong/chat_provider/kimi.py:174-181), then + // serialize and run the Kimi dialect adaptations (single adaptation + // point for every request-side deviation; see `crate::kimi_compat`). let streaming_request = StreamingChatRequest { inner: &payload, stream: true, @@ -908,21 +814,16 @@ impl SamplingClient { include_usage: true, }, }; + let mut request_body = serde_json::to_value(&streaming_request).map_err(|e| { + tracing::error!("Failed to serialize chat/completions request: {}", e); + SamplingError::Serialization(e) + })?; + crate::kimi_compat::adapt_chat_completions_body(&mut request_body); - let grok_headers = GrokRequestHeaders { - conv_id: x_grok_conv_id, - req_id: x_grok_req_id, - model_id: &model_id, - session_id: payload.x_grok_session_id.as_deref().unwrap_or_default(), - turn_idx: payload.x_grok_turn_idx.as_deref(), - agent_id: payload.x_grok_agent_id.as_deref().unwrap_or_default(), - deployment_id: payload.x_grok_deployment_id.as_deref(), - user_id: payload.x_grok_user_id.as_deref(), - }; - let http_request = grok_headers - .apply(self.post(self.endpoint("chat/completions"))) + let http_request = self + .post(self.endpoint("chat/completions")) .header(ACCEPT, HeaderValue::from_static("text/event-stream")) - .json(&streaming_request); + .json(&request_body); let built_request = http_request.build().map_err(|e| { tracing::error!("Failed to build HTTP request: {}", e); @@ -948,7 +849,6 @@ impl SamplingClient { span.record("success", status.is_success()); let model_metadata = extract_model_metadata(response.headers()); let retry_after_secs = extract_retry_after(response.headers()); - let should_retry = extract_should_retry(response.headers()); if !status.is_success() { if status == reqwest::StatusCode::UNAUTHORIZED { span.record("error", "unauthorized (401)"); @@ -962,8 +862,7 @@ impl SamplingClient { ))); } - let req_headers = - self.format_request_headers(x_grok_conv_id, x_grok_req_id, &model_id, true); + let req_headers = self.format_request_headers(true); let resp_headers = Self::format_response_headers(&response); let bytes = response.bytes().await?; let server_message = parse_error_bytes(bytes.as_ref()); @@ -988,7 +887,6 @@ impl SamplingClient { message, model_metadata, retry_after_secs, - should_retry, }); } @@ -1111,8 +1009,6 @@ impl SamplingClient { ) -> Result { self.apply_response_defaults(&mut request)?; - let x_grok_conv_id = request.x_grok_conv_id.as_deref().unwrap_or_default(); - let x_grok_req_id = request.x_grok_req_id.as_deref().unwrap_or_default(); let model_id = request.inner.model.clone().unwrap_or_default(); // The trace field is process-local: it is consumed by upstream @@ -1123,16 +1019,6 @@ impl SamplingClient { tracing::debug!("create_response: {:?}", &request); tracing::debug!("endpoint: {:?}", self.endpoint("responses")); - let grok_headers = GrokRequestHeaders { - conv_id: x_grok_conv_id, - req_id: x_grok_req_id, - model_id: &model_id, - session_id: request.x_grok_session_id.as_deref().unwrap_or_default(), - turn_idx: request.x_grok_turn_idx.as_deref(), - agent_id: request.x_grok_agent_id.as_deref().unwrap_or_default(), - deployment_id: request.x_grok_deployment_id.as_deref(), - user_id: request.x_grok_user_id.as_deref(), - }; let mut request_body = serde_json::to_value(&request.inner).map_err(|e| { tracing::error!("Failed to serialize responses request: {}", e); SamplingError::Serialization(e) @@ -1142,9 +1028,7 @@ impl SamplingClient { // it in post-serialize. This is the last surviving piece of the // old raw_output machinery. kigi_sampling_types::patch_reasoning_text_types(&mut request_body); - let http_request = grok_headers - .apply(self.post(self.endpoint("responses"))) - .json(&request_body); + let http_request = self.post(self.endpoint("responses")).json(&request_body); let response = http_request.send().await.map_err(|e| { tracing::debug!("HTTP request failed: {}", e); @@ -1154,7 +1038,6 @@ impl SamplingClient { let status = response.status(); let model_metadata = extract_model_metadata(response.headers()); let retry_after_secs = extract_retry_after(response.headers()); - let should_retry = extract_should_retry(response.headers()); let bytes = response.bytes().await?; if !status.is_success() { @@ -1167,8 +1050,7 @@ impl SamplingClient { ))); } - let req_headers = - self.format_request_headers(x_grok_conv_id, x_grok_req_id, &model_id, false); + let req_headers = self.format_request_headers(false); let server_message = parse_error_bytes(bytes.as_ref()); let message = self.build_api_error_message( @@ -1189,7 +1071,6 @@ impl SamplingClient { message, model_metadata, retry_after_secs, - should_retry, }); } @@ -1245,8 +1126,6 @@ impl SamplingClient { // Enable streaming request.inner.stream = Some(true); - let x_grok_conv_id = request.x_grok_conv_id.as_deref().unwrap_or_default(); - let x_grok_req_id = request.x_grok_req_id.as_deref().unwrap_or_default(); let model_id = request.inner.model.clone().unwrap_or_default(); // Drop process-local trace data (see note in `create_response`). @@ -1258,16 +1137,6 @@ impl SamplingClient { "Sending responses API stream request" ); - let grok_headers = GrokRequestHeaders { - conv_id: x_grok_conv_id, - req_id: x_grok_req_id, - model_id: &model_id, - session_id: request.x_grok_session_id.as_deref().unwrap_or_default(), - turn_idx: request.x_grok_turn_idx.as_deref(), - agent_id: request.x_grok_agent_id.as_deref().unwrap_or_default(), - deployment_id: request.x_grok_deployment_id.as_deref(), - user_id: request.x_grok_user_id.as_deref(), - }; let extra_raw_tools = std::mem::take(&mut request.extra_raw_tools); let mut request_body = serde_json::to_value(&request.inner).map_err(|e| { tracing::error!("Failed to serialize responses request: {}", e); @@ -1293,8 +1162,8 @@ impl SamplingClient { .defaults .doom_loop_recovery .map(crate::doom_loop::DoomLoopSignalCollector::new); - let mut http_request = grok_headers - .apply(self.post(self.endpoint("responses"))) + let mut http_request = self + .post(self.endpoint("responses")) .header(ACCEPT, HeaderValue::from_static("text/event-stream")); if doom_loop.is_some() { // Presence opts in; the server ignores the value. @@ -1336,9 +1205,7 @@ impl SamplingClient { } let model_metadata = extract_model_metadata(response.headers()); let retry_after_secs = extract_retry_after(response.headers()); - let should_retry = extract_should_retry(response.headers()); - let req_headers = - self.format_request_headers(x_grok_conv_id, x_grok_req_id, &model_id, true); + let req_headers = self.format_request_headers(true); let resp_headers = Self::format_response_headers(&response); let bytes = response.bytes().await?; let server_message = parse_error_bytes(bytes.as_ref()); @@ -1363,7 +1230,6 @@ impl SamplingClient { message, model_metadata, retry_after_secs, - should_retry, }); } @@ -1480,8 +1346,6 @@ impl SamplingClient { ) -> Result { self.apply_message_defaults(&mut request)?; - let x_grok_conv_id = request.x_grok_conv_id.as_deref().unwrap_or_default(); - let x_grok_req_id = request.x_grok_req_id.as_deref().unwrap_or_default(); let model_id = request.inner.model.clone(); // Drop process-local trace data. @@ -1490,19 +1354,7 @@ impl SamplingClient { tracing::debug!("create_message: {:?}", &request.inner); tracing::debug!("endpoint: {:?}", self.endpoint("messages")); - let grok_headers = GrokRequestHeaders { - conv_id: x_grok_conv_id, - req_id: x_grok_req_id, - model_id: &model_id, - session_id: request.x_grok_session_id.as_deref().unwrap_or_default(), - turn_idx: request.x_grok_turn_idx.as_deref(), - agent_id: request.x_grok_agent_id.as_deref().unwrap_or_default(), - deployment_id: request.x_grok_deployment_id.as_deref(), - user_id: request.x_grok_user_id.as_deref(), - }; - let http_request = grok_headers - .apply(self.post(self.endpoint("messages"))) - .json(&request.inner); + let http_request = self.post(self.endpoint("messages")).json(&request.inner); let response = http_request.send().await.map_err(|e| { tracing::debug!("HTTP request failed: {}", e); @@ -1512,7 +1364,6 @@ impl SamplingClient { let status = response.status(); let model_metadata = extract_model_metadata(response.headers()); let retry_after_secs = extract_retry_after(response.headers()); - let should_retry = extract_should_retry(response.headers()); let bytes = response.bytes().await?; if !status.is_success() { @@ -1525,8 +1376,7 @@ impl SamplingClient { ))); } - let req_headers = - self.format_request_headers(x_grok_conv_id, x_grok_req_id, &model_id, false); + let req_headers = self.format_request_headers(false); let server_message = parse_error_bytes(bytes.as_ref()); let message = self.build_api_error_message( @@ -1547,7 +1397,6 @@ impl SamplingClient { message, model_metadata, retry_after_secs, - should_retry, }); } @@ -1593,8 +1442,6 @@ impl SamplingClient { // Enable streaming request.inner.stream = Some(true); - let x_grok_conv_id = request.x_grok_conv_id.as_deref().unwrap_or_default(); - let x_grok_req_id = request.x_grok_req_id.as_deref().unwrap_or_default(); let model_id = request.inner.model.clone(); // Drop process-local trace data. @@ -1606,18 +1453,8 @@ impl SamplingClient { "Sending Messages API stream request" ); - let grok_headers = GrokRequestHeaders { - conv_id: x_grok_conv_id, - req_id: x_grok_req_id, - model_id: &model_id, - session_id: request.x_grok_session_id.as_deref().unwrap_or_default(), - turn_idx: request.x_grok_turn_idx.as_deref(), - agent_id: request.x_grok_agent_id.as_deref().unwrap_or_default(), - deployment_id: request.x_grok_deployment_id.as_deref(), - user_id: request.x_grok_user_id.as_deref(), - }; - let http_request = grok_headers - .apply(self.post(self.endpoint("messages"))) + let http_request = self + .post(self.endpoint("messages")) .header(ACCEPT, HeaderValue::from_static("text/event-stream")) .json(&request.inner); @@ -1655,9 +1492,7 @@ impl SamplingClient { } let model_metadata = extract_model_metadata(response.headers()); let retry_after_secs = extract_retry_after(response.headers()); - let should_retry = extract_should_retry(response.headers()); - let req_headers = - self.format_request_headers(x_grok_conv_id, x_grok_req_id, &model_id, true); + let req_headers = self.format_request_headers(true); let resp_headers = Self::format_response_headers(&response); let bytes = response.bytes().await?; let server_message = parse_error_bytes(bytes.as_ref()); @@ -1682,7 +1517,6 @@ impl SamplingClient { message, model_metadata, retry_after_secs, - should_retry, }); } @@ -2002,7 +1836,6 @@ impl SamplingClient { message: info.message, model_metadata: info.model_metadata, retry_after_secs: info.retry_after_secs, - should_retry: None, }) } } @@ -2031,10 +1864,6 @@ mod tests { idle_timeout_secs: None, reasoning_effort: None, origin_client: None, - client_identifier: None, - deployment_id: None, - user_id: None, - client_version: None, attribution_callback: None, bearer_resolver: None, supports_backend_search: false, @@ -2146,40 +1975,6 @@ mod tests { assert_eq!(extract_retry_after(&headers), None); } - #[test] - fn extract_should_retry_true() { - let mut headers = reqwest::header::HeaderMap::new(); - headers.insert("x-should-retry", "true".parse().unwrap()); - assert_eq!(extract_should_retry(&headers), Some(true)); - } - - #[test] - fn extract_should_retry_true_case_insensitive() { - let mut headers = reqwest::header::HeaderMap::new(); - headers.insert("x-should-retry", "TRUE".parse().unwrap()); - assert_eq!(extract_should_retry(&headers), Some(true)); - } - - #[test] - fn extract_should_retry_false() { - let mut headers = reqwest::header::HeaderMap::new(); - headers.insert("x-should-retry", "false".parse().unwrap()); - assert_eq!(extract_should_retry(&headers), Some(false)); - } - - #[test] - fn extract_should_retry_unknown_value_is_none() { - let mut headers = reqwest::header::HeaderMap::new(); - headers.insert("x-should-retry", "banana".parse().unwrap()); - assert_eq!(extract_should_retry(&headers), None); - } - - #[test] - fn extract_should_retry_absent_is_none() { - let headers = reqwest::header::HeaderMap::new(); - assert_eq!(extract_should_retry(&headers), None); - } - #[test] fn new_with_minimal_config_succeeds() { let client = SamplingClient::new(minimal_config()).expect("client should construct"); @@ -2286,8 +2081,8 @@ mod tests { version: None, }; let ua = user_agent_string_for(&origin); - // No slash between product and the grok-shell agent product. - assert!(ua.starts_with("my-client grok-shell/")); + // No slash between product and the kigi agent product. + assert!(ua.starts_with("my-client kigi/")); } #[test] diff --git a/crates/codegen/kigi-sampler/src/config.rs b/crates/codegen/kigi-sampler/src/config.rs index 29e0089..872ce24 100644 --- a/crates/codegen/kigi-sampler/src/config.rs +++ b/crates/codegen/kigi-sampler/src/config.rs @@ -71,12 +71,12 @@ pub struct SamplerConfig { // Reasoning effort pub reasoning_effort: Option, - // Client identity + /// Client identity for the User-Agent header (`kigi/{version}` plus an + /// optional origin product). The old xAI proxy's identity headers + /// (`x-grok-client-identifier` / `-client-version` / `-deployment-id` / + /// `-user-id`) are gone — User-Agent and `extra_headers` are the only + /// identity signals on the wire. pub origin_client: Option, - pub client_identifier: Option, - pub deployment_id: Option, - pub user_id: Option, - pub client_version: Option, /// Optional hook invoked at every UNAUTHORIZED (401) response /// site. The sampler passes the bearer that was actually sent on @@ -146,10 +146,6 @@ impl Default for SamplerConfig { idle_timeout_secs: None, reasoning_effort: None, origin_client: None, - client_identifier: None, - deployment_id: None, - user_id: None, - client_version: None, attribution_callback: None, bearer_resolver: None, supports_backend_search: false, diff --git a/crates/codegen/kigi-sampler/src/events.rs b/crates/codegen/kigi-sampler/src/events.rs index 8636363..2f90e05 100644 --- a/crates/codegen/kigi-sampler/src/events.rs +++ b/crates/codegen/kigi-sampler/src/events.rs @@ -292,7 +292,6 @@ mod tests { message: "boom".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; let info = SamplingErrorInfo::from(&err); assert_eq!(info.kind, SamplingErrorKind::Api); @@ -307,7 +306,6 @@ mod tests { message: "slow down".into(), model_metadata: None, retry_after_secs: Some(15), - should_retry: None, }; let info = SamplingErrorInfo::from(&err); assert_eq!(info.kind, SamplingErrorKind::RateLimited); @@ -326,7 +324,6 @@ mod tests { ..Default::default() }), retry_after_secs: None, - should_retry: None, }; let info = SamplingErrorInfo::from(&err); assert_eq!(info.kind, SamplingErrorKind::Api); diff --git a/crates/codegen/kigi-sampler/src/kimi_compat.rs b/crates/codegen/kigi-sampler/src/kimi_compat.rs new file mode 100644 index 0000000..6bc52eb --- /dev/null +++ b/crates/codegen/kigi-sampler/src/kimi_compat.rs @@ -0,0 +1,424 @@ +//! Kimi (Moonshot) chat/completions request adaptations. +//! +//! The Kimi endpoints are OpenAI-compatible but deviate in a handful of +//! places (PRD F3 Q1). Every request-side deviation is absorbed HERE, in a +//! single adaptation point applied to the serialized chat/completions body +//! just before it is sent — never as scattered special-cases at call sites. +//! Each adaptation cites the kimi-cli source it was derived from +//! (kimi-cli == the authoritative official client; paths are relative to +//! that repository). +//! +//! Response-side deviations live with the wire types themselves +//! (`kigi_sampling_types::Usage::cached_tokens`, +//! `ChatChunkChoice::usage`) and the L2 stream transform +//! (`stream::chat_completions` synthesizes missing tool-call ids). +//! +//! The `ApiBackend::ChatCompletions` backend is the Kimi dialect: both +//! product channels (subscription OAuth and Moonshot API keys) ride it. +//! Custom providers that need vanilla OpenAI semantics for reasoning use +//! the `Responses` backend, which stays available in model configuration. + +use serde_json::Value; + +/// Adapt a fully-serialized chat/completions request body to the Kimi +/// dialect, in place. Applied by [`crate::SamplingClient`] to both the +/// streaming and non-streaming chat/completions paths. +pub(crate) fn adapt_chat_completions_body(body: &mut Value) { + adapt_thinking(body); + adapt_messages(body); + adapt_tool_schemas(body); +} + +/// Map the OpenAI-style `reasoning_effort` knob onto Kimi's `thinking` +/// request field and drop `reasoning_effort` from the wire. +/// +/// kimi-cli controls thinking exclusively through the request body's +/// `thinking: {"type": "enabled" | "disabled"}` field +/// (packages/kosong/src/kosong/chat_provider/kimi.py:214-223 `with_thinking`: +/// `"enabled" if effort != "off" else "disabled"`; wired by +/// src/kimi_cli/llm.py:475-481). When no effort is configured, nothing is +/// sent and the server default applies (llm.py:482 "leave as-is"). +fn adapt_thinking(body: &mut Value) { + let Some(obj) = body.as_object_mut() else { + return; + }; + let Some(effort) = obj.remove("reasoning_effort") else { + return; + }; + let enabled = effort.as_str() != Some("none"); + obj.insert( + "thinking".to_owned(), + serde_json::json!({ "type": if enabled { "enabled" } else { "disabled" } }), + ); +} + +/// Message-level adaptations: +/// +/// * Drop `model_id` — a grok-build extension recorded on assistant turns; +/// kimi-cli's message serializer sends no such field +/// (packages/kosong/src/kosong/chat_provider/kimi.py:326-353). +/// * Drop `content` from assistant tool-call messages whose visible content +/// is effectively empty. The Kimi-for-Coding compat layer rejects an +/// empty text content part with 400 "text content is empty"; omitting +/// `content` entirely is always accepted +/// (packages/kosong/src/kosong/chat_provider/kimi.py:339-350, with the +/// "effectively empty" predicate at kimi.py:356-362). +fn adapt_messages(body: &mut Value) { + let Some(messages) = body.get_mut("messages").and_then(Value::as_array_mut) else { + return; + }; + for message in messages { + let Some(obj) = message.as_object_mut() else { + continue; + }; + obj.remove("model_id"); + let is_assistant = obj.get("role").and_then(Value::as_str) == Some("assistant"); + let has_tool_calls = obj + .get("tool_calls") + .and_then(Value::as_array) + .is_some_and(|calls| !calls.is_empty()); + if is_assistant + && has_tool_calls + && obj.get("content").is_some_and(is_effectively_empty_content) + { + obj.remove("content"); + } + } +} + +/// Port of kimi-cli `_is_effectively_empty_content_parts` +/// (packages/kosong/src/kosong/chat_provider/kimi.py:356-362): a bare +/// whitespace-only string, or a block list whose entries are all +/// whitespace-only text blocks. Any non-text block (e.g. an image) makes +/// the content non-empty. +fn is_effectively_empty_content(content: &Value) -> bool { + match content { + Value::String(s) => s.trim().is_empty(), + Value::Array(blocks) => blocks.iter().all(|block| { + block.get("type").and_then(Value::as_str) == Some("text") + && block + .get("text") + .and_then(Value::as_str) + .is_some_and(|t| t.trim().is_empty()) + }), + Value::Null => true, + _ => false, + } +} + +/// Moonshot's schema validator rejects tool parameter schemas whose +/// property schemas omit `type` (e.g. enum-only properties exposed by some +/// MCP servers): HTTP 400 "At path 'properties.X': type is not defined". +/// Fill in an inferred `type` locally so such tools keep working. Port of +/// kimi-cli `ensure_property_types` +/// (packages/kosong/src/kosong/utils/jsonschema.py:88-142, applied per tool +/// at packages/kosong/src/kosong/chat_provider/kimi.py:378-388). +fn adapt_tool_schemas(body: &mut Value) { + let Some(tools) = body.get_mut("tools").and_then(Value::as_array_mut) else { + return; + }; + for tool in tools { + if let Some(parameters) = tool.pointer_mut("/function/parameters") { + recurse_schema(parameters); + } + } +} + +/// JSON Schema keywords that describe a property's shape without a `type` +/// keyword; nodes carrying one are left alone +/// (kosong/utils/jsonschema.py:15-24 `_COMBINATOR_KEYS`). +const COMBINATOR_KEYS: [&str; 8] = [ + "anyOf", "oneOf", "allOf", "not", "if", "then", "else", "$ref", +]; + +/// Walk property-schema positions under `node` (`properties`, `items`, +/// `additionalProperties`, `anyOf`/`oneOf`/`allOf`); `node` itself is a +/// container and is not normalized (kosong/utils/jsonschema.py:114-142). +fn recurse_schema(node: &mut Value) { + let Some(obj) = node.as_object_mut() else { + return; + }; + if let Some(props) = obj.get_mut("properties").and_then(Value::as_object_mut) { + for value in props.values_mut() { + normalize_property(value); + } + } + match obj.get_mut("items") { + Some(items @ Value::Object(_)) => normalize_property(items), + Some(Value::Array(items)) => { + for value in items { + normalize_property(value); + } + } + _ => {} + } + if let Some(additional @ Value::Object(_)) = obj.get_mut("additionalProperties") { + normalize_property(additional); + } + for key in ["anyOf", "oneOf", "allOf"] { + if let Some(branches) = obj.get_mut(key).and_then(Value::as_array_mut) { + for value in branches { + normalize_property(value); + } + } + } +} + +/// Ensure a property schema declares `type`, then recurse into it +/// (kosong/utils/jsonschema.py:145-162 `_normalize_property`). +fn normalize_property(node: &mut Value) { + let Some(obj) = node.as_object_mut() else { + return; + }; + if !obj.contains_key("type") && !COMBINATOR_KEYS.iter().any(|k| obj.contains_key(*k)) { + let inferred = if let Some(Value::Array(values)) = obj.get("enum") { + if values.is_empty() { + infer_type_from_structure(obj) + } else { + infer_type_from_values(values) + } + } else if let Some(constant) = obj.get("const") { + infer_type_from_values(std::slice::from_ref(constant)) + } else { + infer_type_from_structure(obj) + }; + obj.insert("type".to_owned(), Value::String(inferred.to_owned())); + } + recurse_schema(node); +} + +/// Infer `type` from structural keywords when no enum/const is present; +/// defaults to `"string"` only with no structural hints at all +/// (kosong/utils/jsonschema.py:165-215 `_infer_type_from_structure`). +fn infer_type_from_structure(obj: &serde_json::Map) -> &'static str { + const OBJECT_KEYWORDS: [&str; 7] = [ + "properties", + "additionalProperties", + "patternProperties", + "propertyNames", + "required", + "minProperties", + "maxProperties", + ]; + const ARRAY_KEYWORDS: [&str; 6] = [ + "items", + "prefixItems", + "minItems", + "maxItems", + "uniqueItems", + "contains", + ]; + const STRING_KEYWORDS: [&str; 4] = ["minLength", "maxLength", "pattern", "format"]; + const NUMERIC_KEYWORDS: [&str; 5] = [ + "minimum", + "maximum", + "multipleOf", + "exclusiveMinimum", + "exclusiveMaximum", + ]; + if OBJECT_KEYWORDS.iter().any(|k| obj.contains_key(*k)) { + "object" + } else if ARRAY_KEYWORDS.iter().any(|k| obj.contains_key(*k)) { + "array" + } else if STRING_KEYWORDS.iter().any(|k| obj.contains_key(*k)) { + "string" + } else if NUMERIC_KEYWORDS.iter().any(|k| obj.contains_key(*k)) { + "number" + } else { + "string" + } +} + +/// Infer a `type` from concrete enum/const values: single JSON type wins, +/// `{integer, number}` collapses to `"number"`, any other mix falls back to +/// `"string"` (kosong/utils/jsonschema.py:218-247 `_infer_type_from_values`). +fn infer_type_from_values(values: &[Value]) -> &'static str { + let mut inferred = std::collections::BTreeSet::new(); + for value in values { + let ty = match value { + Value::Bool(_) => "boolean", + Value::Number(n) if n.is_i64() || n.is_u64() => "integer", + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Null => "null", + Value::Object(_) => "object", + Value::Array(_) => "array", + }; + inferred.insert(ty); + } + if inferred.len() == 1 { + return inferred.pop_first().expect("non-empty set"); + } + if inferred == std::collections::BTreeSet::from(["integer", "number"]) { + return "number"; + } + "string" +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn reasoning_effort_maps_to_kimi_thinking_field() { + let mut body = json!({ "model": "kimi-for-coding", "reasoning_effort": "high" }); + adapt_chat_completions_body(&mut body); + assert_eq!(body.get("reasoning_effort"), None); + assert_eq!(body["thinking"], json!({ "type": "enabled" })); + + // kimi.py:218: "off" (our ReasoningEffort::None) → disabled. + let mut body = json!({ "reasoning_effort": "none" }); + adapt_chat_completions_body(&mut body); + assert_eq!(body["thinking"], json!({ "type": "disabled" })); + + // llm.py:482: unset → leave as-is (no `thinking` at all). + let mut body = json!({ "model": "kimi-for-coding" }); + adapt_chat_completions_body(&mut body); + assert_eq!(body.get("thinking"), None); + } + + #[test] + fn assistant_tool_call_with_empty_content_drops_content() { + let mut body = json!({ + "messages": [ + { "role": "user", "content": "hi" }, + { + "role": "assistant", + "content": "", + "model_id": "kimi-for-coding", + "tool_calls": [{ "id": "c1", "type": "function", + "function": { "name": "f", "arguments": "{}" } }] + }, + ] + }); + adapt_chat_completions_body(&mut body); + let assistant = &body["messages"][1]; + assert_eq!(assistant.get("content"), None, "empty content dropped"); + assert_eq!(assistant.get("model_id"), None, "grok extension dropped"); + assert!(assistant.get("tool_calls").is_some()); + // The user message keeps its content. + assert_eq!(body["messages"][0]["content"], json!("hi")); + } + + #[test] + fn assistant_tool_call_with_real_content_keeps_content() { + let mut body = json!({ + "messages": [{ + "role": "assistant", + "content": "let me check", + "tool_calls": [{ "id": "c1", "type": "function", + "function": { "name": "f", "arguments": "{}" } }] + }] + }); + adapt_chat_completions_body(&mut body); + assert_eq!(body["messages"][0]["content"], json!("let me check")); + } + + #[test] + fn assistant_without_tool_calls_keeps_empty_content() { + // Only tool-call turns drop content (kimi.py:339-350 guards on + // `message.tool_calls`); a plain empty assistant turn is left alone. + let mut body = json!({ + "messages": [{ "role": "assistant", "content": "" }] + }); + adapt_chat_completions_body(&mut body); + assert_eq!(body["messages"][0]["content"], json!("")); + } + + #[test] + fn empty_text_block_list_counts_as_empty_content() { + let mut body = json!({ + "messages": [{ + "role": "assistant", + "content": [{ "type": "text", "text": " " }], + "tool_calls": [{ "id": "c1", "type": "function", + "function": { "name": "f", "arguments": "{}" } }] + }] + }); + adapt_chat_completions_body(&mut body); + assert_eq!(body["messages"][0].get("content"), None); + } + + #[test] + fn image_block_is_not_empty_content() { + let mut body = json!({ + "messages": [{ + "role": "assistant", + "content": [{ "type": "image_url", "image_url": { "url": "data:x" } }], + "tool_calls": [{ "id": "c1", "type": "function", + "function": { "name": "f", "arguments": "{}" } }] + }] + }); + adapt_chat_completions_body(&mut body); + assert!(body["messages"][0].get("content").is_some()); + } + + #[test] + fn enum_only_property_gains_inferred_type() { + // The Moonshot validator 400s on `{"enum": [...]}` without `type` + // (kosong/utils/jsonschema.py:91-96). + let mut body = json!({ + "tools": [{ + "type": "function", + "function": { + "name": "search", + "parameters": { + "type": "object", + "properties": { + "mode": { "enum": ["smart", "full"] }, + "count": { "enum": [1, 2, 3] }, + "ratio": { "enum": [1, 2.5] }, + "nested": { + "type": "object", + "properties": { "inner": { "enum": ["a"] } } + }, + "combined": { "anyOf": [{ "type": "string" }] } + } + } + } + }] + }); + adapt_chat_completions_body(&mut body); + let props = &body["tools"][0]["function"]["parameters"]["properties"]; + assert_eq!(props["mode"]["type"], json!("string")); + assert_eq!(props["count"]["type"], json!("integer")); + assert_eq!(props["ratio"]["type"], json!("number")); + assert_eq!( + props["nested"]["properties"]["inner"]["type"], + json!("string") + ); + assert_eq!( + props["combined"].get("type"), + None, + "combinator nodes are left alone" + ); + } + + #[test] + fn structural_keywords_infer_shape_not_string() { + let mut body = json!({ + "tools": [{ + "type": "function", + "function": { + "name": "t", + "parameters": { + "type": "object", + "properties": { + "obj": { "properties": { "x": { "type": "string" } } }, + "arr": { "items": { "type": "string" } }, + "num": { "minimum": 0 }, + "free": {} + } + } + } + }] + }); + adapt_chat_completions_body(&mut body); + let props = &body["tools"][0]["function"]["parameters"]["properties"]; + assert_eq!(props["obj"]["type"], json!("object")); + assert_eq!(props["arr"]["type"], json!("array")); + assert_eq!(props["num"]["type"], json!("number")); + assert_eq!(props["free"]["type"], json!("string")); + } +} diff --git a/crates/codegen/kigi-sampler/src/lib.rs b/crates/codegen/kigi-sampler/src/lib.rs index 3933936..fe09553 100644 --- a/crates/codegen/kigi-sampler/src/lib.rs +++ b/crates/codegen/kigi-sampler/src/lib.rs @@ -1,4 +1,4 @@ -//! kigi-sampler - Actor-based sampling layer for xAI grok. +//! kigi-sampler - Actor-based sampling layer for the Kimi inference APIs. //! //! This crate extracts the HTTP streaming + retry logic out of //! `kigi-shell`'s session actor into a standalone, reusable @@ -24,6 +24,7 @@ pub mod config; pub mod doom_loop; pub mod events; pub mod handle; +mod kimi_compat; pub mod metrics; pub mod retry; pub mod sampling_log; diff --git a/crates/codegen/kigi-sampler/src/retry.rs b/crates/codegen/kigi-sampler/src/retry.rs index 7ba65b2..95a2b57 100644 --- a/crates/codegen/kigi-sampler/src/retry.rs +++ b/crates/codegen/kigi-sampler/src/retry.rs @@ -24,14 +24,10 @@ //! - `Serialization` (response parsing failure) //! - `MaxTokensTruncation` (by design) //! -//! **Server hint** (`x-should-retry` header from CCP): -//! - `false` → Fatal immediately, regardless of status code -//! - `true` / absent → falls through to status-code logic above -//! -//! Today CCP's header mirrors the client's `is_retryable()` logic -//! (4xx except 429 = false, 5xx + 429 = true), so no behavior changes -//! on merge. The header enables future CCP-side refinements (e.g. -//! marking content-caused 500s as non-retryable) without client updates. +//! 429 handling honors the standard `Retry-After` response header when +//! present (delta-seconds; see `client::extract_retry_after`), matching +//! the Kimi/Moonshot API. The old xAI proxy's `x-should-retry` hint +//! header was removed with the proxy. use std::time::Duration; @@ -170,22 +166,6 @@ pub fn classify_error( return RetryDecision::RetryWithImageStrip; } - // Server explicitly said don't retry (x-should-retry: false). - // Trust the server — it knows if the error is request-content-caused - // (e.g. malformed tool call in conversation history) vs transient. - // - // x-should-retry: true is intentionally NOT handled here — we only - // use the header to suppress retries (false), not to force them - // (true). Forcing retries on non-retryable status codes could - // amplify failures. true falls through to existing status-code logic. - // - // Checked AFTER image-strip guards: image stripping changes the - // request payload, so a server "don't retry" on the original - // request doesn't apply to the stripped request. - if let Some(false) = err.should_retry_header() { - return RetryDecision::Fatal(clone_error(err)); - } - // Context-window / size overflow is deterministic — re-sending the same (or // larger) payload always fails — so never retry it, whatever status the backend // used (in-stream `ResponseError`→500, HTTP 400/500, OpenAI/Anthropic variants). @@ -401,13 +381,11 @@ pub(crate) fn clone_error(err: &SamplingError) -> SamplingError { message, model_metadata, retry_after_secs, - should_retry, } => SamplingError::Api { status: *status, message: message.clone(), model_metadata: model_metadata.clone(), retry_after_secs: *retry_after_secs, - should_retry: *should_retry, }, SamplingError::EventStreamError(msg) => SamplingError::EventStreamError(msg.clone()), SamplingError::StreamError { @@ -445,7 +423,6 @@ mod tests { message: message.to_string(), model_metadata: None, retry_after_secs: None, - should_retry: None, } } @@ -455,7 +432,6 @@ mod tests { message: "x".to_string(), model_metadata: None, retry_after_secs: Some(retry_after), - should_retry: None, } } @@ -757,31 +733,15 @@ mod tests { assert!(s.contains("240s")); } - #[test] - fn should_retry_false_overrides_retryable_status() { - let err = SamplingError::Api { - status: StatusCode::INTERNAL_SERVER_ERROR, - message: "boom".into(), - model_metadata: None, - retry_after_secs: None, - should_retry: Some(false), - }; - assert!(matches!( - classify_error(&err, 0, 15, RATE_LIMIT_RETRY_THRESHOLD), - RetryDecision::Fatal(_) - )); - } - #[test] fn context_length_overflow_is_fatal_even_as_500() { - // The backend streams a size overflow as a ResponseError that becomes a 500 with no - // should_retry hint; without the context-length check it would retry the full budget. + // The backend streams a size overflow as a ResponseError that becomes a 500; + // without the context-length check it would retry the full budget. let err = SamplingError::Api { status: StatusCode::INTERNAL_SERVER_ERROR, message: "none: The prompt is too long for this model's context window.".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!(matches!( classify_error(&err, 0, 15, RATE_LIMIT_RETRY_THRESHOLD), @@ -790,28 +750,12 @@ mod tests { } #[test] - fn should_retry_true_falls_through_to_existing_logic() { + fn api_500_first_failure_retries_with_client_rebuild() { let err = SamplingError::Api { status: StatusCode::INTERNAL_SERVER_ERROR, message: "boom".into(), model_metadata: None, retry_after_secs: None, - should_retry: Some(true), - }; - assert!(matches!( - classify_error(&err, 0, 15, RATE_LIMIT_RETRY_THRESHOLD), - RetryDecision::RetryWithClientRebuild { .. } - )); - } - - #[test] - fn should_retry_absent_falls_through() { - let err = SamplingError::Api { - status: StatusCode::INTERNAL_SERVER_ERROR, - message: "boom".into(), - model_metadata: None, - retry_after_secs: None, - should_retry: None, }; assert!(matches!( classify_error(&err, 0, 15, RATE_LIMIT_RETRY_THRESHOLD), @@ -836,21 +780,4 @@ mod tests { } } } - - #[test] - fn should_retry_false_on_429_is_fatal() { - // Server says don't retry, even though 429 is normally retryable. - // should_retry check runs before rate-limit check. - let err = SamplingError::Api { - status: StatusCode::TOO_MANY_REQUESTS, - message: "rate limited".into(), - model_metadata: None, - retry_after_secs: Some(10), - should_retry: Some(false), - }; - assert!(matches!( - classify_error(&err, 0, 15, RATE_LIMIT_RETRY_THRESHOLD), - RetryDecision::Fatal(_) - )); - } } diff --git a/crates/codegen/kigi-sampler/src/stream/chat_completions.rs b/crates/codegen/kigi-sampler/src/stream/chat_completions.rs index 96b6d3f..bc7ff67 100644 --- a/crates/codegen/kigi-sampler/src/stream/chat_completions.rs +++ b/crates/codegen/kigi-sampler/src/stream/chat_completions.rs @@ -128,7 +128,16 @@ pub fn stream_chat_completions<'a>( first_chunk_seen = true; } - if let Some(u) = chunk.usage.clone() { + // Kimi/Moonshot deviation: usage may ride inside a choice instead + // of (or in addition to) the chunk's top-level `usage`. Same + // fallback as kimi-cli's `extract_usage_from_chunk` + // (packages/kosong/src/kosong/chat_provider/kimi.py:522-533): + // top-level wins, else the first choice carrying one. + let chunk_usage = chunk + .usage + .clone() + .or_else(|| chunk.choices.iter().find_map(|c| c.usage.clone())); + if let Some(u) = chunk_usage { // Wire cost is cumulative for the response, so last-write-wins. // Never clobber a known cost with missing/unreported. let chunk_cost = kigi_sampling_types::reported_cost_ticks(u.cost_in_usd_ticks); @@ -247,10 +256,27 @@ pub fn stream_chat_completions<'a>( // ── Build the final response ───────────────────────────────── let tool_calls: Vec = tool_call_acc .into_values() - .map(|(id, name, arguments)| ToolCall { - id: std::sync::Arc::::from(id), - name, - arguments: std::sync::Arc::::from(arguments), + .map(|(id, name, arguments)| { + // Kimi/Moonshot deviation: tool-call deltas may omit `id`. + // Synthesize one so the tool-result round-trip stays keyed, + // exactly like kimi-cli (`id=tool_call.id or str(uuid.uuid4())`, + // packages/kosong/src/kosong/chat_provider/kimi.py:505). + let id = if id.is_empty() { + let synthesized = uuid::Uuid::new_v4().to_string(); + tracing::debug!( + tool_name = %name, + synthesized_id = %synthesized, + "tool-call delta carried no id; synthesized one" + ); + synthesized + } else { + id + }; + ToolCall { + id: std::sync::Arc::::from(id), + name, + arguments: std::sync::Arc::::from(arguments), + } }) .collect(); @@ -329,6 +355,7 @@ mod tests { index: i as u32, delta, finish_reason: None, + usage: None, }) .collect(), usage: None, @@ -664,6 +691,7 @@ mod tests { prompt_tokens: 100, completion_tokens: 50, total_tokens: 150, + cached_tokens: None, prompt_tokens_details: None, completion_tokens_details: None, cost_in_usd_ticks: None, @@ -704,6 +732,7 @@ mod tests { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15, + cached_tokens: None, prompt_tokens_details: None, completion_tokens_details: None, cost_in_usd_ticks: wire, @@ -737,6 +766,7 @@ mod tests { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15, + cached_tokens: None, prompt_tokens_details: None, completion_tokens_details: None, cost_in_usd_ticks: Some(99), @@ -746,6 +776,7 @@ mod tests { prompt_tokens: 12, completion_tokens: 6, total_tokens: 18, + cached_tokens: None, prompt_tokens_details: None, completion_tokens_details: None, cost_in_usd_ticks: Some(0), diff --git a/crates/codegen/kigi-sampler/src/stream/collect.rs b/crates/codegen/kigi-sampler/src/stream/collect.rs index 78e33ef..1c4caf1 100644 --- a/crates/codegen/kigi-sampler/src/stream/collect.rs +++ b/crates/codegen/kigi-sampler/src/stream/collect.rs @@ -90,6 +90,7 @@ mod tests { tool_call_id: None, }, finish_reason: None, + usage: None, }], usage: None, system_fingerprint: None, @@ -106,6 +107,7 @@ mod tests { index: 0, delta: ChatChunkDelta::default(), finish_reason: Some(FinishReason::Stop), + usage: None, }], usage: None, system_fingerprint: None, diff --git a/crates/codegen/kigi-sampler/src/stream/messages.rs b/crates/codegen/kigi-sampler/src/stream/messages.rs index 1ebdceb..33b3f82 100644 --- a/crates/codegen/kigi-sampler/src/stream/messages.rs +++ b/crates/codegen/kigi-sampler/src/stream/messages.rs @@ -429,7 +429,6 @@ pub fn stream_messages<'a>( message: error_message, model_metadata: None, retry_after_secs: None, - should_retry: None, }; yield SamplingEvent::Failed { request_id: request_id.clone(), diff --git a/crates/codegen/kigi-sampler/src/stream/responses.rs b/crates/codegen/kigi-sampler/src/stream/responses.rs index 881b671..3cf2e27 100644 --- a/crates/codegen/kigi-sampler/src/stream/responses.rs +++ b/crates/codegen/kigi-sampler/src/stream/responses.rs @@ -299,7 +299,6 @@ pub fn stream_responses<'a>( message: error_message, model_metadata: None, retry_after_secs: None, - should_retry: None, }; yield SamplingEvent::Failed { request_id: request_id.clone(), @@ -316,7 +315,6 @@ pub fn stream_responses<'a>( message: error_message, model_metadata: None, retry_after_secs: None, - should_retry: None, }; yield SamplingEvent::Failed { request_id: request_id.clone(), @@ -419,7 +417,6 @@ pub fn stream_responses<'a>( .to_string(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; yield SamplingEvent::Failed { request_id: request_id.clone(), diff --git a/crates/codegen/kigi-sampler/tests/test_actor.rs b/crates/codegen/kigi-sampler/tests/test_actor.rs index 3b26175..5dd690f 100644 --- a/crates/codegen/kigi-sampler/tests/test_actor.rs +++ b/crates/codegen/kigi-sampler/tests/test_actor.rs @@ -87,10 +87,6 @@ fn test_config(base_url: String, model: &str) -> SamplerConfig { idle_timeout_secs: Some(30), reasoning_effort: None, origin_client: None, - client_identifier: None, - deployment_id: None, - user_id: None, - client_version: None, attribution_callback: None, bearer_resolver: None, supports_backend_search: false, diff --git a/crates/codegen/kigi-sampler/tests/test_kimi_wire.rs b/crates/codegen/kigi-sampler/tests/test_kimi_wire.rs new file mode 100644 index 0000000..11d2013 --- /dev/null +++ b/crates/codegen/kigi-sampler/tests/test_kimi_wire.rs @@ -0,0 +1,564 @@ +//! Kimi chat/completions wire tests (PRD F3 acceptance). +//! +//! Exercises the sampler end-to-end against a mock HTTP server: +//! * streaming happy path with `reasoning_content` deltas, tool-call deltas, +//! and a Kimi-shaped usage chunk (usage riding inside the choice, cache +//! hits as top-level `cached_tokens`), +//! * the request the wire actually carries: plain `Authorization: Bearer`, +//! `User-Agent: kigi/{version}`, no xAI proxy marker headers, and the +//! `crate::kimi_compat` body adaptations, +//! * 429 honoring the standard `Retry-After` header, +//! * mid-stream network drop recovering through the retry loop. +//! +//! 401-no-retry and the rate-limit retry threshold are covered by +//! `test_actor.rs`; this file does not duplicate them. + +use std::net::SocketAddr; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use axum::Router; +use axum::body::Bytes; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::sse::{Event, Sse}; +use axum::routing::post; +use futures_util::stream::{self}; +use indexmap::IndexMap; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tokio::sync::{mpsc, oneshot}; + +use kigi_sampler::{ + ApiBackend, RequestId, RetryPolicy, SamplerActor, SamplerConfig, SamplingChannel, SamplingEvent, +}; +use kigi_sampling_types::{ + AssistantItem, ContentPart, ConversationItem, ConversationRequest, ReasoningEffort, ToolCall, + ToolResultItem, ToolSpec, UserItem, synthesized_reasoning_item, +}; + +// --------------------------------------------------------------------------- +// Mock server harness (same shape as test_actor.rs) +// --------------------------------------------------------------------------- + +struct MockServer { + addr: SocketAddr, + shutdown_tx: oneshot::Sender<()>, +} + +impl MockServer { + async fn spawn(app: Router) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + tokio::spawn(async move { + let _ = axum::serve(listener, app) + .with_graceful_shutdown(async move { + let _ = shutdown_rx.await; + }) + .await; + }); + tokio::time::sleep(Duration::from_millis(20)).await; + Self { addr, shutdown_tx } + } + + fn base_url(&self) -> String { + format!("http://{}/v1", self.addr) + } + + fn shutdown(self) { + let _ = self.shutdown_tx.send(()); + } +} + +fn test_config(base_url: String) -> SamplerConfig { + SamplerConfig { + api_key: Some("test-kimi-key".into()), + base_url, + model: "kimi-for-coding".into(), + max_completion_tokens: Some(1024), + api_backend: ApiBackend::ChatCompletions, + extra_headers: IndexMap::new(), + context_window: 128_000, + max_retries: Some(3), + idle_timeout_secs: Some(30), + ..Default::default() + } +} + +fn user_request(text: &str) -> ConversationRequest { + ConversationRequest { + items: vec![ConversationItem::User(UserItem { + content: vec![ContentPart::Text { + text: Arc::::from(text), + }], + synthetic_reason: None, + ..Default::default() + })], + ..Default::default() + } +} + +fn chunk(delta: Value, finish: Option<&str>, usage: Option) -> Event { + let mut choice = json!({ "index": 0, "delta": delta }); + choice["finish_reason"] = finish.map(Value::from).unwrap_or(Value::Null); + if let Some(u) = usage { + // Kimi deviation under test: usage rides INSIDE the choice + // (kimi-cli kimi.py:522-533 `extract_usage_from_chunk`). + choice["usage"] = u; + } + let body = json!({ + "id": "chatcmpl-kimi", + "object": "chat.completion.chunk", + "created": 0, + "model": "kimi-for-coding", + "choices": [choice] + }); + Event::default().data(body.to_string()) +} + +async fn drain_until_terminal( + rx: &mut mpsc::UnboundedReceiver, + timeout: Duration, +) -> Vec { + let mut out = Vec::new(); + let deadline = tokio::time::Instant::now() + timeout; + loop { + let ev = tokio::time::timeout_at(deadline, rx.recv()) + .await + .expect("timed out waiting for terminal event") + .expect("event channel closed before terminal event"); + let terminal = matches!( + ev, + SamplingEvent::Completed { .. } | SamplingEvent::Failed { .. } + ); + out.push(ev); + if terminal { + return out; + } + } +} + +// --------------------------------------------------------------------------- +// Streaming happy path: reasoning + tool calls + Kimi usage shapes +// --------------------------------------------------------------------------- + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn kimi_stream_reasoning_tool_calls_and_choice_usage() { + let app = Router::new().route( + "/v1/chat/completions", + post(|| async { + let events = vec![ + chunk( + json!({ "role": "assistant", "reasoning_content": "let me think" }), + None, + None, + ), + chunk(json!({ "reasoning_content": " harder" }), None, None), + chunk(json!({ "content": "Running the tool." }), None, None), + // Tool call split across chunks: id+name first, args continue. + chunk( + json!({ "tool_calls": [{ "index": 0, "id": "call_1", "type": "function", + "function": { "name": "read_file", "arguments": "{\"path\":" } }] }), + None, + None, + ), + chunk( + json!({ "tool_calls": [{ "index": 0, + "function": { "arguments": "\"a.rs\"}" } }] }), + None, + None, + ), + // Terminal chunk: finish_reason + usage inside the choice with + // Moonshot's top-level `cached_tokens` (kimi.py:427-431). + chunk( + json!({}), + Some("tool_calls"), + Some(json!({ + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + "cached_tokens": 60 + })), + ), + ]; + Sse::new(stream::iter( + events.into_iter().map(Ok::<_, std::convert::Infallible>), + )) + }), + ); + let server = MockServer::spawn(app).await; + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let handle = SamplerActor::spawn( + test_config(server.base_url()), + RetryPolicy::default(), + event_tx, + ); + + handle.submit(RequestId::from("req-kimi"), user_request("hi")); + let events = drain_until_terminal(&mut event_rx, Duration::from_secs(30)).await; + server.shutdown(); + + // Reasoning tokens stream on the Reasoning channel, text on Text. + let reasoning: String = events + .iter() + .filter_map(|e| match e { + SamplingEvent::ChannelToken { + channel: SamplingChannel::Reasoning, + text, + .. + } => Some(text.as_str()), + _ => None, + }) + .collect(); + assert_eq!(reasoning, "let me think harder"); + let text: String = events + .iter() + .filter_map(|e| match e { + SamplingEvent::ChannelToken { + channel: SamplingChannel::Text, + text, + .. + } => Some(text.as_str()), + _ => None, + }) + .collect(); + assert_eq!(text, "Running the tool."); + + // Tool-call deltas surfaced incrementally. + assert!(events.iter().any(|e| matches!( + e, + SamplingEvent::ToolCallDelta { id: Some(id), .. } if id == "call_1" + ))); + + match events.last().unwrap() { + SamplingEvent::Completed { response, .. } => { + let calls = response.tool_calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].id.as_ref(), "call_1"); + assert_eq!(calls[0].name, "read_file"); + assert_eq!(calls[0].arguments.as_ref(), "{\"path\":\"a.rs\"}"); + let reasoning_item = response + .reasoning_items() + .next() + .expect("reasoning sibling preserved"); + let kigi_sampling_types::rs::SummaryPart::SummaryText(t) = &reasoning_item.summary[0]; + assert_eq!(t.text, "let me think harder"); + // Choice-level usage + top-level cached_tokens both absorbed. + let usage = response.usage.as_ref().expect("usage from choice"); + assert_eq!(usage.prompt_tokens, 100); + assert_eq!(usage.completion_tokens, 20); + assert_eq!(usage.cached_prompt_tokens, 60); + } + other => panic!("expected Completed, got {other:?}"), + } +} + +// --------------------------------------------------------------------------- +// Request surface: bearer auth, kigi UA, kimi_compat body adaptations +// --------------------------------------------------------------------------- + +type Captured = Arc>>; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn request_carries_bearer_kigi_ua_and_kimi_dialect_body() { + let captured: Captured = Arc::new(Mutex::new(None)); + let captured_handler = Arc::clone(&captured); + let app = Router::new().route( + "/v1/chat/completions", + post(move |headers: HeaderMap, body: Bytes| { + let captured = Arc::clone(&captured_handler); + async move { + let body: Value = serde_json::from_slice(&body).unwrap(); + *captured.lock().unwrap() = Some((headers, body)); + let events = vec![chunk( + json!({ "role": "assistant", "content": "ok" }), + Some("stop"), + None, + )]; + Sse::new(stream::iter( + events.into_iter().map(Ok::<_, std::convert::Infallible>), + )) + } + }), + ); + let server = MockServer::spawn(app).await; + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let handle = SamplerActor::spawn( + test_config(server.base_url()), + RetryPolicy::default(), + event_tx, + ); + + // Multi-turn conversation exercising every request-side adaptation: + // reasoning folded onto the assistant, an empty-content tool-call turn, + // and an enum-only tool schema property. + let request = ConversationRequest { + items: vec![ + ConversationItem::User(UserItem { + content: vec![ContentPart::Text { + text: Arc::::from("read a.rs"), + }], + synthetic_reason: None, + ..Default::default() + }), + ConversationItem::Reasoning(synthesized_reasoning_item("planning the read")), + ConversationItem::Assistant(AssistantItem { + content: Arc::::from(""), + tool_calls: vec![ToolCall { + id: Arc::::from("call_9"), + name: "read_file".into(), + arguments: Arc::::from("{\"path\":\"a.rs\"}"), + }], + model_id: Some("kimi-for-coding".into()), + model_fingerprint: None, + reasoning_effort: None, + }), + ConversationItem::ToolResult(ToolResultItem { + tool_call_id: "call_9".into(), + content: Arc::::from("fn main() {}"), + images: vec![], + }), + ], + tools: vec![ToolSpec { + name: "read_file".into(), + description: Some("Read a file".into()), + parameters: json!({ + "type": "object", + "properties": { + // Enum-only property: Moonshot 400s without a `type`. + "mode": { "enum": ["full", "head"] }, + "path": { "type": "string" } + } + }), + }], + reasoning_effort: Some(ReasoningEffort::High), + ..Default::default() + }; + handle.submit(RequestId::from("req-wire"), request); + let events = drain_until_terminal(&mut event_rx, Duration::from_secs(30)).await; + server.shutdown(); + assert!(matches!( + events.last().unwrap(), + SamplingEvent::Completed { .. } + )); + + let (headers, body) = captured.lock().unwrap().take().expect("request captured"); + + // -- Auth: plain bearer, nothing else (PRD F3). + assert_eq!( + headers.get("authorization").unwrap().to_str().unwrap(), + "Bearer test-kimi-key" + ); + for gone in [ + "x-xai-token-auth", + "x-authenticateresponse", + "x-grok-conv-id", + "x-grok-req-id", + "x-grok-model-override", + "x-grok-session-id", + "x-grok-agent-id", + "x-grok-client-identifier", + "x-grok-client-version", + "x-grok-deployment-id", + "x-grok-user-id", + "x-grok-client-mode", + ] { + assert!( + headers.get(gone).is_none(), + "xAI proxy marker header must not be sent: {gone}" + ); + } + + // -- User-Agent: kigi/{version} (os; arch). + let ua = headers.get("user-agent").unwrap().to_str().unwrap(); + let expected_prefix = format!("kigi/{}", kigi_version::VERSION); + assert!( + ua.starts_with(&expected_prefix), + "UA must start with {expected_prefix}, got {ua}" + ); + + // -- Streaming fields exactly as kimi-cli sends them (kimi.py:174-181). + assert_eq!(body["stream"], json!(true)); + assert_eq!(body["stream_options"], json!({ "include_usage": true })); + + // -- Thinking mapping (kimi.py:214-223): effort → thinking, no + // reasoning_effort on the wire. + assert_eq!(body["thinking"], json!({ "type": "enabled" })); + assert_eq!(body.get("reasoning_effort"), None); + + // -- Message adaptations. + let messages = body["messages"].as_array().unwrap(); + let assistant = messages + .iter() + .find(|m| m["role"] == "assistant") + .expect("assistant turn present"); + assert_eq!( + assistant.get("content"), + None, + "empty tool-call content dropped (kimi.py:339-350)" + ); + assert_eq!(assistant.get("model_id"), None, "grok extension dropped"); + assert_eq!( + assistant["reasoning_content"], + json!("planning the read"), + "reasoning folded onto the assistant turn (kimi.py:351-352)" + ); + assert_eq!(assistant["tool_calls"][0]["id"], json!("call_9")); + + // -- Tool schema normalization (kosong jsonschema.py:88-142). + let props = &body["tools"][0]["function"]["parameters"]["properties"]; + assert_eq!(props["mode"]["type"], json!("string")); + assert_eq!(props["path"]["type"], json!("string")); +} + +// --------------------------------------------------------------------------- +// 429 with Retry-After +// --------------------------------------------------------------------------- + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rate_limit_honors_retry_after_then_succeeds() { + let counter = Arc::new(AtomicU32::new(0)); + let counter_handler = Arc::clone(&counter); + let app = Router::new().route( + "/v1/chat/completions", + post(move || { + let counter = Arc::clone(&counter_handler); + async move { + let n = counter.fetch_add(1, Ordering::SeqCst); + if n == 0 { + // Moonshot-shaped 429 body + standard Retry-After. + let mut headers = HeaderMap::new(); + headers.insert("retry-after", "1".parse().unwrap()); + Err::, (StatusCode, HeaderMap, String)>(( + StatusCode::TOO_MANY_REQUESTS, + headers, + json!({ "error": { + "message": "Your account is rate limited", + "type": "rate_limit_reached_error" + }}) + .to_string(), + )) + } else { + let events = vec![chunk( + json!({ "role": "assistant", "content": "after limit" }), + Some("stop"), + None, + )]; + Ok(Sse::new(stream::iter( + events.into_iter().map(Ok::<_, std::convert::Infallible>), + ))) + } + } + }), + ); + let server = MockServer::spawn(app).await; + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let handle = SamplerActor::spawn( + test_config(server.base_url()), + RetryPolicy::default(), + event_tx, + ); + + let started = std::time::Instant::now(); + handle.submit(RequestId::from("req-ra"), user_request("hi")); + let events = drain_until_terminal(&mut event_rx, Duration::from_secs(30)).await; + let elapsed = started.elapsed(); + server.shutdown(); + + // A Retrying event carried the classified rate-limit. + assert!(events.iter().any(|e| matches!( + e, + SamplingEvent::Retrying { kind, .. } + if *kind == kigi_sampler::SamplingErrorKind::RateLimited + ))); + match events.last().unwrap() { + SamplingEvent::Completed { response, .. } => { + assert_eq!( + response.assistant().unwrap().content.as_ref(), + "after limit" + ); + } + other => panic!("expected Completed after Retry-After wait, got {other:?}"), + } + assert_eq!(counter.load(Ordering::SeqCst), 2, "exactly one retry"); + // Retry-After: 1 replaces the ~2s jittered exponential backoff. The wait + // must be at least the advertised second (and clearly less than the + // exhaust-path 30s timeout). + assert!( + elapsed >= Duration::from_secs(1), + "waited less than Retry-After: {elapsed:?}" + ); +} + +// --------------------------------------------------------------------------- +// Mid-stream network drop → retry → recovery +// --------------------------------------------------------------------------- + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mid_stream_drop_recovers_via_retry() { + let counter = Arc::new(AtomicU32::new(0)); + let counter_handler = Arc::clone(&counter); + let app = Router::new().route( + "/v1/chat/completions", + post(move || { + let counter = Arc::clone(&counter_handler); + async move { + let n = counter.fetch_add(1, Ordering::SeqCst); + if n == 0 { + // First attempt: a partial chunk, then the connection + // dies mid-body (simulated network drop). + let events: Vec> = vec![ + Ok(chunk( + json!({ "role": "assistant", "content": "partial" }), + None, + None, + )), + Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "connection reset by peer", + )), + ]; + Sse::new(stream::iter(events)) + } else { + let events: Vec> = vec![Ok(chunk( + json!({ "role": "assistant", "content": "recovered" }), + Some("stop"), + None, + ))]; + Sse::new(stream::iter(events)) + } + } + }), + ); + let server = MockServer::spawn(app).await; + let (event_tx, mut event_rx) = mpsc::unbounded_channel(); + let handle = SamplerActor::spawn( + test_config(server.base_url()), + RetryPolicy::default(), + event_tx, + ); + + handle.submit(RequestId::from("req-drop"), user_request("hi")); + let events = drain_until_terminal(&mut event_rx, Duration::from_secs(60)).await; + server.shutdown(); + + assert!( + events + .iter() + .any(|e| matches!(e, SamplingEvent::Retrying { .. })), + "mid-stream drop must go through the retry loop" + ); + match events.last().unwrap() { + SamplingEvent::Completed { response, .. } => { + // The poisoned partial attempt is discarded; only the fresh + // attempt's content survives. + assert_eq!(response.assistant().unwrap().content.as_ref(), "recovered"); + } + other => panic!("expected Completed after recovery, got {other:?}"), + } + assert!( + counter.load(Ordering::SeqCst) >= 2, + "server hit at least twice" + ); +} diff --git a/crates/codegen/kigi-sampling-types/src/conversation.rs b/crates/codegen/kigi-sampling-types/src/conversation.rs index af60fe5..d82a352 100644 --- a/crates/codegen/kigi-sampling-types/src/conversation.rs +++ b/crates/codegen/kigi-sampling-types/src/conversation.rs @@ -666,10 +666,15 @@ impl TokenUsage { impl From for TokenUsage { fn from(u: Usage) -> Self { - let cached_prompt_tokens = u - .prompt_tokens_details - .as_ref() - .map_or(0, |d| d.cached_tokens); + // Kimi/Moonshot deviation: prefer the top-level `cached_tokens` field + // when present, falling back to the OpenAI-standard + // `prompt_tokens_details.cached_tokens`. Same precedence as kimi-cli + // (packages/kosong/src/kosong/chat_provider/kimi.py:427-437). + let cached_prompt_tokens = u.cached_tokens.unwrap_or_else(|| { + u.prompt_tokens_details + .as_ref() + .map_or(0, |d| d.cached_tokens) + }); Self { prompt_tokens: u.prompt_tokens, completion_tokens: u.completion_tokens, diff --git a/crates/codegen/kigi-sampling-types/src/error.rs b/crates/codegen/kigi-sampling-types/src/error.rs index a996cb8..8587267 100644 --- a/crates/codegen/kigi-sampling-types/src/error.rs +++ b/crates/codegen/kigi-sampling-types/src/error.rs @@ -96,13 +96,9 @@ pub enum SamplingError { status: StatusCode, message: String, model_metadata: Option, - /// Parsed from the `Retry-After` response header (seconds). + /// Parsed from the standard `Retry-After` response header (seconds). + /// The Kimi API only emits delta-seconds; HTTP-dates are ignored. retry_after_secs: Option, - /// Parsed from the `x-should-retry` response header. - /// `Some(true)` = transient, retry may help. - /// `Some(false)` = request-content error, don't retry. - /// `None` = header absent (old server or non-proxy origin). - should_retry: Option, }, #[error("reqwest error stream: {0}")] EventStreamError(String), @@ -271,14 +267,6 @@ impl SamplingError { } } - /// Server hint on whether this error is worth retrying. - pub fn should_retry_header(&self) -> Option { - match self { - SamplingError::Api { should_retry, .. } => *should_retry, - _ => None, - } - } - /// True when this error is a context-window/size overflow — deterministic, /// so retrying the same payload can't help. See [`is_context_length_error`]. pub fn is_context_length_error(&self) -> bool { @@ -304,7 +292,14 @@ impl From for SamplingError { } } -/// OpenAI-standard provider error format: `{"error": {"message": "...", "type": "..."}}`. +/// Kimi/Moonshot (OpenAI-compatible) error body: +/// `{"error": {"message": "...", "type": "..."}}`. +/// +/// This is the only error format the Kimi chat/completions endpoint emits — +/// the same shape the official client parses via the OpenAI SDK +/// (kimi-cli packages/kosong/src/kosong/chat_provider/openai_common.py:83-87 +/// maps `openai.APIStatusError` → status + message). The old xAI proxy's +/// flat `{"code": "...", "error": "..."}` format was removed with the proxy. #[derive(Debug, Deserialize)] struct ErrorResponse { error: ErrorBody, @@ -317,36 +312,20 @@ struct ErrorBody { kind: Option, } -/// Flat error from the Grok proxy/gateway: `{"code": "...", "error": "..."}`. -#[derive(Debug, Deserialize)] -struct FlatErrorResponse { - error: String, - #[serde(default)] - code: Option, -} - -/// Extract `(error_type, message)` from either error format. +/// Extract `(error_type, message)` from an OpenAI-compatible error body. fn try_parse_error(data: &str) -> Option<(String, String)> { - if let Ok(resp) = serde_json::from_str::(data) { - return Some(( - resp.error.kind.unwrap_or_else(|| "unknown".to_string()), - resp.error - .message - .unwrap_or_else(|| "unknown error".to_string()), - )); - } - if let Ok(flat) = serde_json::from_str::(data) { - return Some(( - flat.code.unwrap_or_else(|| "server_error".to_string()), - flat.error, - )); - } - None + let resp = serde_json::from_str::(data).ok()?; + Some(( + resp.error.kind.unwrap_or_else(|| "unknown".to_string()), + resp.error + .message + .unwrap_or_else(|| "unknown error".to_string()), + )) } pub fn parse_error_bytes(bytes: &[u8]) -> String { if let Some((error_type, message)) = std::str::from_utf8(bytes).ok().and_then(try_parse_error) { - if error_type == "unknown" || error_type == "server_error" { + if error_type == "unknown" { return message; } return format!("{error_type}: {message}"); @@ -420,7 +399,6 @@ mod tests { message: "none: The prompt is too long for this model's context window.".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!(api.is_context_length_error()); assert!( @@ -489,20 +467,20 @@ mod tests { ); } + /// Moonshot rate-limit body, OpenAI error shape (the format the Kimi + /// endpoints emit; see kimi-cli packages/kosong/src/kosong/chat_provider/chaos.py:88 + /// for the reference 429 body used by the official client's chaos tests). #[test] - fn try_parse_stream_error_flat_format() { - let data = r#"{"code":"The service is currently unavailable","error":"Service temporarily unavailable. The model did not respond to this request."}"#; - let err = try_parse_stream_error(data).expect("should parse flat error"); + fn try_parse_stream_error_openai_format() { + let data = r#"{"error":{"message":"Your account is rate limited","type":"rate_limit_reached_error"}}"#; + let err = try_parse_stream_error(data).expect("should parse OpenAI-shaped error"); match err { SamplingError::StreamError { error_type, message, } => { - assert_eq!(error_type, "The service is currently unavailable"); - assert_eq!( - message, - "Service temporarily unavailable. The model did not respond to this request." - ); + assert_eq!(error_type, "rate_limit_reached_error"); + assert_eq!(message, "Your account is rate limited"); } other => panic!("expected StreamError, got {other:?}"), } @@ -518,13 +496,19 @@ mod tests { } #[test] - fn parse_error_bytes_flat_format() { - let bytes = - br#"{"code":"The service is currently unavailable","error":"Service temporarily unavailable."}"#; - let msg = parse_error_bytes(bytes); + fn parse_error_bytes_openai_format_prefixes_type() { + let bytes = br#"{"error":{"message":"Your account is rate limited","type":"rate_limit_reached_error"}}"#; assert_eq!( - msg, - "The service is currently unavailable: Service temporarily unavailable." + parse_error_bytes(bytes), + "rate_limit_reached_error: Your account is rate limited" + ); + } + + #[test] + fn parse_error_bytes_non_json_falls_back_to_raw_text() { + assert_eq!( + parse_error_bytes(b" upstream exploded "), + "upstream exploded" ); } @@ -543,7 +527,6 @@ mod tests { message: "Content violates usage guidelines.".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!( !err.is_auth_error(), @@ -558,7 +541,6 @@ mod tests { message: "Invalid or expired credentials".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!( err.is_auth_error(), @@ -579,7 +561,6 @@ mod tests { message: "Rate limit exceeded".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!(err.is_rate_limited()); assert!(err.is_retryable(), "429 should be retryable"); @@ -594,7 +575,6 @@ mod tests { message: "internal".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!(!server_error.is_rate_limited()); @@ -612,7 +592,6 @@ mod tests { message: "slow down".into(), model_metadata: None, retry_after_secs: Some(42), - should_retry: None, }; assert_eq!(err.retry_after(), Some(42)); } @@ -624,7 +603,6 @@ mod tests { message: "slow down".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert_eq!(err.retry_after(), None); } @@ -645,7 +623,6 @@ mod tests { message: "Could not decrypt the provided encrypted_content. Ensure the value is the unmodified encrypted_content from a previous response.".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!(err.is_encrypted_content_error()); assert!( @@ -661,7 +638,6 @@ mod tests { message: "encrypted_content decryption failed".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!( !err.is_encrypted_content_error(), @@ -676,7 +652,6 @@ mod tests { message: "Invalid model parameter".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!( !err.is_encrypted_content_error(), @@ -691,7 +666,6 @@ mod tests { message: "Could not process image: unsupported format".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!(err.is_image_processing_error()); assert!(!err.is_encrypted_content_error()); @@ -704,7 +678,6 @@ mod tests { message: "upstream error: 400 Bad Request: Could not process image".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!(err.is_image_processing_error()); } @@ -716,7 +689,6 @@ mod tests { message: "Invalid model parameter".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!(!err.is_image_processing_error()); } @@ -728,7 +700,6 @@ mod tests { message: "internal server error".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!(!err.is_image_processing_error()); } @@ -740,7 +711,6 @@ mod tests { message: "Could not process image".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!( !err.is_image_processing_error(), @@ -755,7 +725,6 @@ mod tests { message: "Could not process image".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; assert!( !err.is_retryable(), diff --git a/crates/codegen/kigi-sampling-types/src/types.rs b/crates/codegen/kigi-sampling-types/src/types.rs index 47aaf7f..15a81a4 100644 --- a/crates/codegen/kigi-sampling-types/src/types.rs +++ b/crates/codegen/kigi-sampling-types/src/types.rs @@ -537,6 +537,15 @@ pub struct Usage { pub prompt_tokens: u32, pub completion_tokens: u32, pub total_tokens: u32, + /// Kimi/Moonshot deviation: the Moonshot chat/completions API reports + /// cache hits as a top-level `cached_tokens` field on `usage` instead of + /// the OpenAI-standard `prompt_tokens_details.cached_tokens`. Ported from + /// kimi-cli's parser (packages/kosong/src/kosong/chat_provider/kimi.py:427-437, + /// which checks the top-level field first and cites + /// platform.moonshot.cn/docs/api/chat). `From for TokenUsage` + /// applies the same precedence. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cached_tokens: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub prompt_tokens_details: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -592,6 +601,14 @@ pub struct ChatChunkChoice { pub delta: ChatChunkDelta, #[serde(skip_serializing_if = "Option::is_none")] pub finish_reason: Option, + /// Kimi/Moonshot deviation: some Kimi deployments attach the final + /// `usage` object to the last *choice* instead of (or in addition to) + /// the chunk's top-level `usage`. Ported from kimi-cli's + /// `extract_usage_from_chunk` + /// (packages/kosong/src/kosong/chat_provider/kimi.py:522-533), which + /// falls back to `choices[0].usage` when `chunk.usage` is absent. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub usage: Option, } /// Streaming delta for a tool call. diff --git a/crates/codegen/kigi-shared/Cargo.toml b/crates/codegen/kigi-shared/Cargo.toml index 9b6ec3d..22159a4 100644 --- a/crates/codegen/kigi-shared/Cargo.toml +++ b/crates/codegen/kigi-shared/Cargo.toml @@ -11,7 +11,6 @@ dirs = "6" dunce = { workspace = true } image = { workspace = true, features = ["png", "jpeg", "gif", "webp"] } parking_lot = { workspace = true } -prod-mc-cli-chat-proxy-types = { path = "../../../prod/mc/cli-chat-proxy-types" } regex = { workspace = true } serde = { workspace = true } serde_json = { workspace = true, features = ["preserve_order"] } diff --git a/crates/codegen/kigi-shared/src/session/mod.rs b/crates/codegen/kigi-shared/src/session/mod.rs index 15ec1fe..85d41d0 100644 --- a/crates/codegen/kigi-shared/src/session/mod.rs +++ b/crates/codegen/kigi-shared/src/session/mod.rs @@ -1,12 +1,44 @@ use std::path::PathBuf; +use serde::{Deserialize, Serialize}; + pub mod info; pub use info::Info; -// Re-export shared feedback wire types used by downstream crates -// (e.g. kigi-pager-render). -pub use prod_mc_cli_chat_proxy_types::feedback_types::FeedbackTerminalInfo; +/// Snapshot of the user's terminal environment at feedback time. +/// +/// Shared here (rather than in kigi-shell) because kigi-pager-render builds it +/// from its terminal probes and the shell attaches it to feedback records. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FeedbackTerminalInfo { + /// Terminal emulator brand (e.g. "Ghostty", "iTerm2", "Unknown"). + pub brand: String, + /// Multiplexer wrapping the session (e.g. "tmux", "Zellij", "None detected"). + pub multiplexer: String, + /// Whether the session is over SSH. + pub is_ssh: bool, + /// Whether Byobu is wrapping the session. + pub is_byobu: bool, + /// Raw `TERM` environment variable value. + pub term_var: String, + /// tmux server version if inside tmux, otherwise "n/a". + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tmux_version: Option, + /// Hyperlink (OSC 8) support level (e.g. "native", "hostile_parser"). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub hyperlink_osc8_support: Option, + /// Active clipboard legs, e.g. "native+osc52" or "native+tmux+osc52". + #[serde(default, skip_serializing_if = "Option::is_none")] + pub clipboard_route: Option, + /// Native clipboard tool: "pbcopy", "wl-copy", "xclip", "xsel", "arboard". + #[serde(default, skip_serializing_if = "Option::is_none")] + pub clipboard_native_tool: Option, + /// Display server: "wayland", "x11", "quartz", "win32", "unknown". + #[serde(default, skip_serializing_if = "Option::is_none")] + pub display_server: Option, +} pub fn session_dir(info: &Info) -> PathBuf { kigi_tools::util::kigi_home::sessions_cwd_dir(&info.cwd).join(info.id.to_string()) diff --git a/crates/codegen/kigi-shell-base/src/util/mod.rs b/crates/codegen/kigi-shell-base/src/util/mod.rs index 5d997a2..6084524 100644 --- a/crates/codegen/kigi-shell-base/src/util/mod.rs +++ b/crates/codegen/kigi-shell-base/src/util/mod.rs @@ -56,19 +56,19 @@ fn matches_trusted_base_url(candidate: &str, trusted_base: &str) -> bool { /// True for subscription coding-API URLs (the compiled production endpoint; /// deliberately NOT the env-overridable [`kigi_env::coding_api_base_url`] so a /// runtime override can't widen this trust set). -pub fn is_cli_chat_proxy_url(url: &str) -> bool { +pub fn is_production_coding_api_url(url: &str) -> bool { matches_trusted_base_url(url, kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url) } /// True for URLs the idle model-metadata refresh may re-fetch from: the /// *effective* subscription coding endpoint (the `KIGI_CODE_BASE_URL` /// override when set, else the compiled production endpoint), plus loopback -/// hosts (local dev proxies and test mocks). Unlike [`is_cli_chat_proxy_url`] +/// hosts (local dev proxies and test mocks). Unlike [`is_production_coding_api_url`] /// this honours the env override and loopback, so use it only to gate traffic /// that already flows to the session's configured base URL (the refresh /// re-fetches from the same host the session samples against); it must never /// widen a security trust set. pub fn is_effective_coding_endpoint_url(url: &str) -> bool { - if is_cli_chat_proxy_url(url) { + if is_production_coding_api_url(url) { return true; } if matches_trusted_base_url(url, &kigi_env::coding_api_base_url()) { @@ -83,18 +83,11 @@ pub fn is_effective_coding_endpoint_url(url: &str) -> bool { None => false, }) } -/// True for first-party xAI endpoints (`*.x.ai`, cli-chat-proxy, and optional -/// non-production first-party hosts when that feature is enabled). -/// `disable_api_key_auth` refuses keys only for these; other hosts are BYOK and -/// exempt. Safe against invalid URLs and suffix attacks (`evil-x.ai.example`). -pub fn is_first_party_xai_url(url: &str) -> bool { - if is_cli_chat_proxy_url(url) { - return true; - } - reqwest::Url::parse(url) - .ok() - .and_then(|u| u.host_str().map(|h| h.to_owned())) - .is_some_and(|host| host == "x.ai" || host.ends_with(".x.ai")) +/// True for first-party endpoints: the Kimi subscription coding API. The +/// session-token 401-refresh gate only refreshes against these; other hosts +/// are BYOK and exempt. Safe against invalid URLs and suffix attacks. +pub fn is_first_party_url(url: &str) -> bool { + is_production_coding_api_url(url) } /// Truncate a string to at most `max_chars` characters. /// Slices at char boundaries so multi-byte UTF-8 never panics. @@ -229,18 +222,18 @@ pub fn is_grok_process(pid: u32) -> bool { mod tests { use super::*; #[test] - fn test_is_cli_chat_proxy_url_accepts_proxy_subpath() { - assert!(is_cli_chat_proxy_url( + fn test_is_production_coding_api_url_accepts_proxy_subpath() { + assert!(is_production_coding_api_url( "https://api.kimi.com/coding/v1/chat/completions" )); } #[test] - fn test_is_cli_chat_proxy_url_rejects_public_api() { - assert!(!is_cli_chat_proxy_url("https://api.x.ai/v1")); + fn test_is_production_coding_api_url_rejects_public_api() { + assert!(!is_production_coding_api_url("https://api.x.ai/v1")); } #[test] - fn test_is_cli_chat_proxy_url_rejects_spoofed_hostname() { - assert!(!is_cli_chat_proxy_url( + fn test_is_production_coding_api_url_rejects_spoofed_hostname() { + assert!(!is_production_coding_api_url( "https://api.kimi.com.evil.example/coding/v1" )); } @@ -261,31 +254,22 @@ mod tests { )); } #[test] - fn test_is_cli_chat_proxy_url_rejects_v11_prefix_confusion() { - assert!(!is_cli_chat_proxy_url( + fn test_is_production_coding_api_url_rejects_v11_prefix_confusion() { + assert!(!is_production_coding_api_url( "https://api.kimi.com/coding/v11/chat/completions" )); } #[test] - fn test_is_first_party_xai_url() { - assert!(is_first_party_xai_url("https://api.x.ai/v1")); - assert!(is_first_party_xai_url( - "https://api.x.ai/v1/chat/completions" - )); - assert!(is_first_party_xai_url("https://x.ai")); - assert!(is_first_party_xai_url( + fn test_is_first_party_url() { + assert!(is_first_party_url( "https://api.kimi.com/coding/v1/chat/completions" )); - assert!(!is_first_party_xai_url("https://api.openai.com/v1")); - assert!(!is_first_party_xai_url("https://api.anthropic.com/v1")); - assert!(!is_first_party_xai_url( - "https://generativelanguage.googleapis.com" - )); - assert!(!is_first_party_xai_url("https://api.x.ai.evil.example/v1")); - assert!(!is_first_party_xai_url("https://evil-x.ai.attacker.com/v1")); - assert!(!is_first_party_xai_url("https://prefixx.ai/v1")); - assert!(!is_first_party_xai_url("not-a-url")); - assert!(!is_first_party_xai_url("")); + assert!(!is_first_party_url("https://api.x.ai/v1")); + assert!(!is_first_party_url("https://api.openai.com/v1")); + assert!(!is_first_party_url("https://api.anthropic.com/v1")); + assert!(!is_first_party_url("https://api.kimi.com.evil.example/v1")); + assert!(!is_first_party_url("not-a-url")); + assert!(!is_first_party_url("")); } #[test] fn test_truncate() { diff --git a/crates/codegen/kigi-shell/Cargo.toml b/crates/codegen/kigi-shell/Cargo.toml index 9673fa4..1069724 100644 --- a/crates/codegen/kigi-shell/Cargo.toml +++ b/crates/codegen/kigi-shell/Cargo.toml @@ -147,7 +147,6 @@ kigi-auth = { workspace = true, features = ["middleware"] } kigi-log = { workspace = true } kigi-http = { workspace = true } kigi-models = { workspace = true } -prod-mc-cli-chat-proxy-types = { path = "../../../prod/mc/cli-chat-proxy-types" } flate2 = { workspace = true } fs2 = { workspace = true } zstd = { workspace = true } diff --git a/crates/codegen/kigi-shell/benches/session_list.rs b/crates/codegen/kigi-shell/benches/session_list.rs index faace53..f478ee3 100644 --- a/crates/codegen/kigi-shell/benches/session_list.rs +++ b/crates/codegen/kigi-shell/benches/session_list.rs @@ -536,7 +536,6 @@ fn summary_ids(summaries: &[Summary]) -> Vec { async fn build_local_list_with_delayed_peer(cwd: String) -> UnifiedListResult { let local = build_unified_list( - None, None, ListReq { cwd: Some(cwd), @@ -642,7 +641,6 @@ fn bench_session_list(c: &mut Criterion) { |b| { b.iter_with_large_drop(|| { black_box(runtime.block_on(build_unified_list( - None, None, ListReq { cwd: Some(black_box(fixture.picker_cwd.clone())), diff --git a/crates/codegen/kigi-shell/src/agent/app.rs b/crates/codegen/kigi-shell/src/agent/app.rs index 5824a53..e9e9197 100644 --- a/crates/codegen/kigi-shell/src/agent/app.rs +++ b/crates/codegen/kigi-shell/src/agent/app.rs @@ -621,10 +621,8 @@ pub async fn run_leader( let fetch_auth_for_prefetch = ModelFetchAuth::resolve(&endpoints_for_prefetch); let platform_keys_for_prefetch = crate::agent::models::PlatformApiKeys::resolve(&agent_config.platforms); - // The shared pair helper owns the remote_fetch gate for both halves, so a - // disabled knob cannot block leader readiness on settings retries. - let (prefetched_models, remote_settings) = tokio::task::spawn_blocking(move || { - crate::agent::models::prefetch_models_and_settings_blocking( + let prefetched_models = tokio::task::spawn_blocking(move || { + crate::agent::models::prefetch_models_blocking( &endpoints_for_prefetch, auth_for_prefetch.as_ref(), fetch_auth_for_prefetch, @@ -632,20 +630,7 @@ pub async fn run_leader( ) }) .await - .unwrap_or((None, None)); - - // Process-wide image normalize cache: off by default, toggled here from - // `RemoteSettings.image_normalize_cache_enabled` once at startup. - let image_normalize_cache_enabled = remote_settings - .as_ref() - .and_then(|r| r.image_normalize_cache_enabled) - .unwrap_or(false); - crate::session::normalize_cache::NormalizeCache::global() - .set_enabled(image_normalize_cache_enabled); - tracing::debug!( - enabled = image_normalize_cache_enabled, - "image normalize cache toggle resolved from remote settings" - ); + .unwrap_or(None); // ── Phase 7: Signal readiness ───────────────────────────────────────────── // @@ -657,9 +642,7 @@ pub async fn run_leader( // ── Phase 8: LocalSet — agent, bridges, config watcher ─────────────────── let local_set = tokio::task::LocalSet::new(); - let remote_settings_for_reloader = remote_settings.clone(); let mut agent_config_for_spawn = agent_config.clone(); - agent_config_for_spawn.remote_settings = remote_settings; crate::util::config::sync_campaign_fields(&mut agent_config_for_spawn); let agent_to_ipc_tx_clone = agent_to_ipc_tx.clone(); let cancel_clone = cancel.clone(); @@ -891,7 +874,7 @@ pub async fn run_leader( initial_auth_key_hash, initial_config, auth_scope, - remote_settings_for_reloader, + None, config_update_tx, agent_config.cli_experimental_memory, agent_config.cli_no_memory, diff --git a/crates/codegen/kigi-shell/src/agent/chat_modes.rs b/crates/codegen/kigi-shell/src/agent/chat_modes.rs index bb026e2..b4c3df9 100644 --- a/crates/codegen/kigi-shell/src/agent/chat_modes.rs +++ b/crates/codegen/kigi-shell/src/agent/chat_modes.rs @@ -1,334 +1,16 @@ -//! grok.com chat-product model catalog: caches `/rest/modes` and maps modes to -//! the `SessionModelState` returned by `load_chat_session` (the chat analogue of -//! [`crate::agent::models::ModelsManager`]). NB: these "modes" populate the -//! desktop MODEL picker, not the ACP session plan-modes in `LoadSessionResponse.modes`. -use crate::auth::AuthManager; -use crate::remote::chat_models_client::{ - ChatModelsClient, ChatModelsError, ListModesResponse, Mode, -}; -use agent_client_protocol as acp; -use parking_lot::RwLock; -use std::sync::Arc; -use std::time::{Duration, Instant}; -/// ~54 min, matching grok-web's refetch cadence. -const CACHE_TTL: Duration = Duration::from_secs(54 * 60); -/// Cold-miss budget on the `session/load` critical path (warm/stale served instantly). -const COLD_FETCH_TIMEOUT: Duration = Duration::from_secs(2); -const DEFAULT_LOCALE: &str = "en"; -/// Process-wide flag set by the pager when started with `--chat` so initialize -/// and early UI seed the chat `/rest/modes` catalog instead of build models. +//! Legacy `--chat` gateway gate. +//! +//! The grok.com chat-product model picker (`/rest/modes`, `ChatModesManager`) +//! was removed with the xAI proxy: those "modes" came from a grok backend with +//! no Kimi counterpart. Only the process-mode gate survives so the `--chat` +//! frontend path stays a compile-time-off no-op across crates without a +//! cross-crate churn to delete every reference. + +/// Process-wide flag set by the pager when started with `--chat`. pub const KIGI_CHAT_MODE_ENV: &str = "KIGI_CHAT_MODE"; + /// True when the process is a gateway light-frontend (`--chat`) agent. -/// Hard-off in release builds so it can't be enabled via env. +/// Hard-off: the grok chat-modes backend is gone, so this is always `false`. pub fn process_chat_mode_enabled() -> bool { - if true { - return false; - } - match std::env::var(KIGI_CHAT_MODE_ENV) { - Ok(v) => { - let v = v.trim(); - !v.is_empty() && v != "0" && !v.eq_ignore_ascii_case("false") - } - Err(_) => false, - } -} -#[derive(Clone)] -struct CachedModes { - /// Keyed by identity; a mismatch is a miss so one user's modes never leak to another. - user_id: String, - locale: String, - fetched_at: Instant, - response: ListModesResponse, -} -/// Thread-safe, cheaply-cloneable manager. Cloning bumps the inner `Arc`. -#[derive(Clone)] -pub struct ChatModesManager { - inner: Arc, -} -struct Inner { - auth: Arc, - cache: RwLock>, - /// Single-flight guard so concurrent fetches coalesce. - fetch_lock: tokio::sync::Mutex<()>, -} -impl ChatModesManager { - pub fn new(auth: Arc) -> Self { - Self { - inner: Arc::new(Inner { - auth, - cache: RwLock::new(None), - fetch_lock: tokio::sync::Mutex::new(()), - }), - } - } - /// The active grok.com identity, or `None` when unauthenticated. Modes are - /// per-identity (tier/ACL), so every cache key and store is gated on it. - fn current_user_id(&self) -> Option { - self.inner.auth.current_or_expired().map(|a| a.user_id) - } - /// Chat model state for a `session/load` response. On missing auth or fetch - /// failure, serves last-good cache else empty — never the build catalog. - pub async fn model_state(&self) -> acp::SessionModelState { - let Some(user_id) = self.current_user_id() else { - return empty_state(); - }; - let locale = DEFAULT_LOCALE; - { - let guard = self.inner.cache.read(); - if let Some(c) = guard.as_ref() - && c.user_id == user_id - && c.locale == locale - { - if c.fetched_at.elapsed() < CACHE_TTL { - return modes_to_model_state(&c.response); - } - let stale = c.response.clone(); - drop(guard); - self.spawn_refresh(user_id, locale); - return modes_to_model_state(&stale); - } - } - let _flight = self.inner.fetch_lock.lock().await; - { - let guard = self.inner.cache.read(); - if let Some(c) = guard.as_ref() - && c.user_id == user_id - && c.locale == locale - && c.fetched_at.elapsed() < CACHE_TTL - { - return modes_to_model_state(&c.response); - } - } - match self.fetch(locale).await { - Ok(resp) if !resp.modes.is_empty() => { - if self.current_user_id().as_deref() != Some(user_id.as_str()) { - return empty_state(); - } - let mapped = modes_to_model_state(&resp); - if mapped.available_models.is_empty() { - tracing::warn!( - raw_modes = resp.modes.len(), - "chat modes: fetch returned modes but none selectable after availability filter" - ); - } - self.store(user_id, locale.to_owned(), resp); - mapped - } - Ok(_) => empty_state(), - Err(err) => { - tracing::warn!( - error = % err, "chat modes fetch failed; serving cache/empty" - ); - let guard = self.inner.cache.read(); - match guard.as_ref() { - Some(c) if c.user_id == user_id => modes_to_model_state(&c.response), - _ => empty_state(), - } - } - } - } - async fn fetch(&self, locale: &str) -> Result { - let client = ChatModelsClient::new(self.inner.auth.clone()); - match tokio::time::timeout(COLD_FETCH_TIMEOUT, client.list_modes(locale)).await { - Ok(result) => result, - Err(_elapsed) => Err(ChatModelsError::Timeout), - } - } - fn store(&self, user_id: String, locale: String, response: ListModesResponse) { - *self.inner.cache.write() = Some(CachedModes { - user_id, - locale, - fetched_at: Instant::now(), - response, - }); - } - /// Best-effort stale refresh; skips if a fetch is already in flight. - fn spawn_refresh(&self, user_id: String, locale: &'static str) { - let me = self.clone(); - tokio::spawn(async move { - let Ok(_flight) = me.inner.fetch_lock.try_lock() else { - return; - }; - if me.current_user_id().as_deref() != Some(user_id.as_str()) { - return; - } - if let Ok(resp) = me.fetch(locale).await - && !resp.modes.is_empty() - && me.current_user_id().as_deref() == Some(user_id.as_str()) - { - me.store(user_id, locale.to_owned(), resp); - } - }); - } - /// Kick a background `/rest/modes` fill when auth is already present so - /// `--chat` initialize / first `session/new` hit a warm cache. - pub fn warm_in_background(&self) { - let Some(user_id) = self.current_user_id() else { - return; - }; - self.spawn_refresh(user_id, DEFAULT_LOCALE); - } -} -fn empty_state() -> acp::SessionModelState { - acp::SessionModelState::new(acp::ModelId::from(String::new()), Vec::new()) -} -/// Maps grok.com modes → `SessionModelState`: keeps only `available` modes, -/// reconciles `current_model_id` (default → first available → empty, never -/// out-of-set), and stashes `badgeText`/`iconHint`/`tags` in `_meta`. -pub fn modes_to_model_state(resp: &ListModesResponse) -> acp::SessionModelState { - let available_models: Vec = resp - .modes - .iter() - .filter(|m| m.is_available()) - .map(mode_to_model_info) - .collect(); - let current_model_id = reconcile_current(&resp.default_mode_id, &available_models); - acp::SessionModelState::new(current_model_id, available_models) -} -fn mode_to_model_info(m: &Mode) -> acp::ModelInfo { - let name = if m.title.trim().is_empty() { - m.id.clone() - } else { - m.title.clone() - }; - acp::ModelInfo::new(acp::ModelId::from(m.id.clone()), name) - .description(if m.description.is_empty() { - None - } else { - Some(m.description.clone()) - }) - .meta(build_meta(m)) -} -fn build_meta(m: &Mode) -> Option { - let mut map = serde_json::Map::new(); - if let Some(badge) = m.badge_text.as_deref().filter(|s| !s.is_empty()) { - map.insert("badgeText".to_owned(), serde_json::json!(badge)); - } - if !m.icon_hint.is_empty() { - map.insert("iconHint".to_owned(), serde_json::json!(m.icon_hint)); - } - if !m.tags.is_empty() { - map.insert("tags".to_owned(), serde_json::json!(m.tags)); - } - if map.is_empty() { None } else { Some(map) } -} -fn reconcile_current(default_mode_id: &str, available: &[acp::ModelInfo]) -> acp::ModelId { - let in_set = |id: &str| available.iter().any(|m| m.model_id.0.as_ref() == id); - if !default_mode_id.is_empty() && in_set(default_mode_id) { - acp::ModelId::from(default_mode_id.to_owned()) - } else if let Some(first) = available.first() { - first.model_id.clone() - } else { - acp::ModelId::from(String::new()) - } -} -#[cfg(test)] -mod tests { - use super::*; - use crate::remote::chat_models_client::ModeAvailability; - fn available(id: &str, title: &str) -> Mode { - Mode { - id: id.to_owned(), - title: title.to_owned(), - availability: ModeAvailability { - available: Some(serde_json::json!({})), - ..Default::default() - }, - ..Default::default() - } - } - fn requires_upgrade(id: &str) -> Mode { - Mode { - id: id.to_owned(), - availability: ModeAvailability { - requires_upgrade: Some(serde_json::json!({ "message" : "Upgrade" })), - ..Default::default() - }, - ..Default::default() - } - } - #[test] - fn filters_to_available_modes() { - let resp = ListModesResponse { - modes: vec![ - available("auto", "Auto"), - requires_upgrade("heavy"), - available("fast", "Fast"), - ], - default_mode_id: "auto".to_owned(), - }; - let state = modes_to_model_state(&resp); - let ids: Vec = state - .available_models - .iter() - .map(|m| m.model_id.0.to_string()) - .collect(); - assert_eq!(ids, vec!["auto".to_string(), "fast".to_string()]); - assert_eq!(state.current_model_id.0.as_ref(), "auto"); - } - #[test] - fn default_outside_filtered_set_falls_back_to_first_available() { - let resp = ListModesResponse { - modes: vec![requires_upgrade("heavy"), available("fast", "Fast")], - default_mode_id: "heavy".to_owned(), - }; - let state = modes_to_model_state(&resp); - assert_eq!(state.current_model_id.0.as_ref(), "fast"); - assert!( - state - .available_models - .iter() - .any(|m| m.model_id == state.current_model_id) - ); - } - #[test] - fn empty_default_falls_back_to_first() { - let resp = ListModesResponse { - modes: vec![available("a", "A"), available("b", "B")], - default_mode_id: String::new(), - }; - let state = modes_to_model_state(&resp); - assert_eq!(state.current_model_id.0.as_ref(), "a"); - } - #[test] - fn no_available_modes_yields_empty_current() { - let resp = ListModesResponse { - modes: vec![requires_upgrade("heavy")], - default_mode_id: "heavy".to_owned(), - }; - let state = modes_to_model_state(&resp); - assert!(state.available_models.is_empty()); - assert_eq!(state.current_model_id.0.as_ref(), ""); - } - #[test] - fn maps_fields_and_meta() { - let mut m = available("auto", "Auto"); - m.description = "Picks the best model".to_owned(); - m.badge_text = Some("New".to_owned()); - m.icon_hint = "rocket".to_owned(); - m.tags = vec!["TAG_PRIMARY".to_owned()]; - let resp = ListModesResponse { - modes: vec![m], - default_mode_id: "auto".to_owned(), - }; - let state = modes_to_model_state(&resp); - let info = &state.available_models[0]; - assert_eq!(info.name, "Auto"); - assert_eq!(info.description.as_deref(), Some("Picks the best model")); - let meta = info.meta.as_ref().unwrap(); - assert_eq!(meta["badgeText"], serde_json::json!("New")); - assert_eq!(meta["iconHint"], serde_json::json!("rocket")); - assert_eq!(meta["tags"], serde_json::json!(["TAG_PRIMARY"])); - } - #[test] - fn name_falls_back_to_id_when_title_blank() { - let mut m = available("grok-4.5", ""); - m.title = " ".to_owned(); - let resp = ListModesResponse { - modes: vec![m], - default_mode_id: String::new(), - }; - let state = modes_to_model_state(&resp); - assert_eq!(state.available_models[0].name, "grok-4.5"); - } + false } diff --git a/crates/codegen/kigi-shell/src/agent/config.rs b/crates/codegen/kigi-shell/src/agent/config.rs index 8950982..47b05ad 100644 --- a/crates/codegen/kigi-shell/src/agent/config.rs +++ b/crates/codegen/kigi-shell/src/agent/config.rs @@ -1,6 +1,6 @@ use crate::agent::auth_method::ModelByok; +use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW; use crate::auth::{AuthManager, KimiCodeConfig}; -use crate::remote::DEFAULT_CONTEXT_WINDOW; use crate::{config::StorageMode, sampling::ApiBackend, tools::config::ShellToolsetConfig}; use agent_client_protocol as acp; use indexmap::IndexMap; @@ -141,7 +141,7 @@ pub struct EndpointsConfig { /// `Some` = explicitly configured. Tracking explicitness (vs comparing to the /// default value) lets an org pin the proxy to the default on purpose. #[serde(skip_serializing_if = "Option::is_none")] - pub cli_chat_proxy_base_url: Option, + pub coding_api_base_url: Option, /// Base URL for the public xAI API. pub xai_api_base_url: String, /// Optional extra access-header value (applied only with the optional @@ -271,11 +271,11 @@ impl EndpointsConfig { resolved } /// The subscription proxy base URL through which all auxiliary services (and - /// OAuth/session inference) resolve: explicit `cli_chat_proxy_base_url`, else + /// OAuth/session inference) resolve: explicit `coding_api_base_url`, else /// [`kigi_env::coding_api_base_url`]. NEVER falls back to `xai_api_base_url` — /// that is the inference endpoint (API-key auth) only. pub fn proxy_url(&self) -> String { - blank_as_unset(&self.cli_chat_proxy_base_url).unwrap_or_else(kigi_env::coding_api_base_url) + blank_as_unset(&self.coding_api_base_url).unwrap_or_else(kigi_env::coding_api_base_url) } pub fn resolve_inference_base_url(&self) -> String { self.models_base_url @@ -406,7 +406,7 @@ impl EndpointsConfig { impl Default for EndpointsConfig { fn default() -> Self { Self { - cli_chat_proxy_base_url: std::env::var("KIGI_CLI_CHAT_PROXY_BASE_URL").ok(), + coding_api_base_url: std::env::var("KIGI_CODE_BASE_URL").ok(), xai_api_base_url: std::env::var("KIGI_XAI_API_BASE_URL") .unwrap_or_else(|_| XAI_API_BASE_URL_DEFAULT.to_owned()), alpha_test_key: None, @@ -4169,19 +4169,11 @@ pub fn resolve_aux_model_sampling_config( endpoints: &EndpointsConfig, session_key: Option<&str>, alpha_test_key: Option, - client_version: Option, ) -> Option { let catalog_entry = find_model_by_id(models, model_id).cloned(); if let Some(entry) = &catalog_entry { let credentials = resolve_credentials(entry, session_key); - let sampler = sampling_config_for_model( - entry, - credentials, - alpha_test_key.clone(), - client_version.clone(), - None, - None, - ); + let sampler = sampling_config_for_model(entry, credentials, alpha_test_key.clone()); if sampler.api_key.is_some() { return Some(sampler); } @@ -4232,14 +4224,7 @@ pub fn resolve_aux_model_sampling_config( api_base_url: None, }; let credentials = resolve_credentials(&entry, session_key); - let sampler = sampling_config_for_model( - &entry, - credentials, - alpha_test_key, - client_version, - None, - None, - ); + let sampler = sampling_config_for_model(&entry, credentials, alpha_test_key); return Some(sampler); } tracing::warn!( @@ -4252,21 +4237,19 @@ pub fn resolve_aux_model_sampling_config( /// Shared so the aux resolve happy path and the /// `None` fallback cannot diverge between those entry points. /// -/// On aux resolve `Some`, stamp session-local fields (client id, attribution, bearer, +/// On aux resolve `Some`, stamp session-local fields (attribution, bearer, /// retries) onto the helper config. On `None`, fall back to the active session model and /// full config (not forcing `image_description_model` onto the agent endpoint, which 404s -/// on BYOK / non-proxy routes for internal slugs like `grok-build`). -/// Stamp the session-local fields (client id, attribution, bearer resolver, -/// retries) from the active session onto a routed aux `SamplerConfig` so a +/// on BYOK / non-proxy routes for internal slugs). +/// Stamp the session-local fields (attribution, bearer resolver, retries) +/// from the active session onto a routed aux `SamplerConfig` so a /// helper model keeps the session's auth/attribution. Shared by image-describe /// and the auto-mode classifier so the two can't drift. pub fn stamp_session_local_sampler_fields( cfg: &mut SamplerConfig, active_session_config: &SamplerConfig, - client_identifier: Option, max_retries: Option, ) { - cfg.client_identifier = client_identifier; cfg.attribution_callback = active_session_config.attribution_callback.clone(); cfg.bearer_resolver = active_session_config.bearer_resolver.clone(); cfg.max_retries = max_retries; @@ -4274,7 +4257,6 @@ pub fn stamp_session_local_sampler_fields( pub fn finalize_image_describe_sampler_config( resolved_aux: Option, active_session_config: &SamplerConfig, - client_identifier: Option, max_retries: Option, ) -> (String, SamplerConfig) { match resolved_aux { @@ -4282,7 +4264,6 @@ pub fn finalize_image_describe_sampler_config( stamp_session_local_sampler_fields( &mut describe_cfg, active_session_config, - client_identifier, max_retries, ); let model = describe_cfg.model.clone(); @@ -4310,9 +4291,6 @@ pub fn sampling_config_for_model( model: &ModelEntry, credentials: ResolvedCredentials, alpha_test_key: Option, - client_version: Option, - deployment_id: Option, - user_id: Option, ) -> SamplerConfig { let info = model.info(); let model_name = info.model.clone(); @@ -4337,15 +4315,11 @@ pub fn sampling_config_for_model( auth_scheme: credentials.auth_scheme, extra_headers, context_window: info.context_window.get(), - client_version, reasoning_effort: info.reasoning_effort, force_http1: false, max_retries: info.max_retries, stream_tool_calls: info.stream_tool_calls.unwrap_or(false), idle_timeout_secs: None, - client_identifier: None, - deployment_id, - user_id, origin_client: None, attribution_callback: None, bearer_resolver: None, @@ -4359,11 +4333,19 @@ pub fn sampling_config_for_model( /// Fold URL-derived headers into `extra_headers`. /// /// The sampler crate is intentionally URL-agnostic: it does not inspect -/// `base_url` to decide which auth or staging headers to add. Replicate the +/// `base_url` to decide which auth or identity headers to add. Replicate the /// URL-derived header logic at the shell boundary so callers downstream see a /// single homogenous header bag. /// -/// * First-party bases get the client-mode header. +/// * First-party (Kimi subscription) bases get the `X-Msh-Device-*` identity +/// headers, mirroring the official client, which sends its OAuth device +/// headers on every inference request (kimi-cli src/kimi_cli/llm.py:317-323 +/// `_kimi_default_headers` merges `oauth.common_headers()`). Third-party / +/// Moonshot-open-platform bases get none — only the bearer and User-Agent. +/// +/// A device-id failure only skips the headers (with a warning): inference +/// must not hard-fail because `~/.kigi/device_id` is unwritable — unlike +/// OAuth login, where the id is mandatory. /// /// Existing entries are never overwritten so callers can pre-set a value. pub fn inject_url_derived_headers( @@ -4371,10 +4353,20 @@ pub fn inject_url_derived_headers( alpha_test_key: Option<&str>, base_url: &str, ) { - if crate::util::is_cli_chat_proxy_url(base_url) { - headers - .entry(crate::http::CLIENT_MODE_HEADER.to_string()) - .or_insert_with(|| crate::http::process_client_mode().to_string()); + if crate::util::is_production_coding_api_url(base_url) { + match crate::auth::device_headers() { + Ok(device_headers) => { + for (name, value) in device_headers { + headers.entry(name.to_string()).or_insert(value); + } + } + Err(e) => { + tracing::warn!( + error = %e, + "device identity headers unavailable; sending inference request without them" + ); + } + } } let _ = (alpha_test_key, base_url); } @@ -4383,7 +4375,6 @@ pub fn resolve_model_to_sampling_config( models: &IndexMap, session_key: Option<&str>, alpha_test_key: Option, - client_version: Option, fallback_entry: Option, ) -> Option { let entry = find_model_by_id(models, model_id) @@ -4394,16 +4385,12 @@ pub fn resolve_model_to_sampling_config( &entry, credentials, alpha_test_key, - client_version, - None, - None, )) } fn resolve_hidden_default_web_search_sampling_config( model_id: &str, session_key: Option<&str>, alpha_test_key: Option, - client_version: Option, endpoints: &EndpointsConfig, ) -> SamplerConfig { let entry = ModelEntry { @@ -4445,21 +4432,13 @@ fn resolve_hidden_default_web_search_sampling_config( api_base_url: None, }; let credentials = resolve_credentials(&entry, session_key); - sampling_config_for_model( - &entry, - credentials, - alpha_test_key, - client_version, - None, - None, - ) + sampling_config_for_model(&entry, credentials, alpha_test_key) } pub fn resolve_web_search_sampling_config( model_id: &str, models: &IndexMap, session_key: Option<&str>, alpha_test_key: Option, - client_version: Option, endpoints: &EndpointsConfig, ) -> Option { let resolved = if let Some(entry) = find_model_by_id(models, model_id).cloned() { @@ -4468,16 +4447,12 @@ pub fn resolve_web_search_sampling_config( &entry, credentials, alpha_test_key, - client_version, - None, - None, )) } else if model_id == crate::models::default_web_search_model() { Some(resolve_hidden_default_web_search_sampling_config( model_id, session_key, alpha_test_key, - client_version, endpoints, )) } else { @@ -4757,21 +4732,23 @@ reasoning_effort = "low" } } #[test] - fn inject_url_derived_headers_adds_client_mode_for_first_party_url() { + fn inject_url_derived_headers_adds_device_identity_for_first_party_url() { let mut headers = IndexMap::new(); inject_url_derived_headers( &mut headers, None, kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url, ); - assert!(headers.get(crate::http::CLIENT_MODE_HEADER).is_some()); + assert!(headers.get("X-Msh-Device-Id").is_some()); + assert!(headers.get("X-Msh-Device-Name").is_some()); assert!(headers.get("X-XAI-Token-Auth").is_none()); } #[test] fn inject_url_derived_headers_skips_headers_for_external_url() { let mut headers = IndexMap::new(); inject_url_derived_headers(&mut headers, None, "https://api.example.com/v1"); - assert!(headers.get(crate::http::CLIENT_MODE_HEADER).is_none()); + assert!(headers.get("X-Msh-Device-Id").is_none()); + assert!(headers.get("X-Msh-Device-Name").is_none()); } #[test] fn inject_url_derived_headers_preserves_caller_extra_headers() { @@ -4924,7 +4901,6 @@ reasoning_effort = "low" &IndexMap::new(), Some("session-token"), None, - None, &endpoints, ) .expect("hidden default web search model should resolve"); @@ -4943,7 +4919,7 @@ reasoning_effort = "low" model: "composer-session-model".into(), ..Default::default() }; - let (model, cfg) = finalize_image_describe_sampler_config(None, &active, None, Some(3)); + let (model, cfg) = finalize_image_describe_sampler_config(None, &active, Some(3)); assert_eq!(model, "composer-session-model"); assert_eq!(cfg.model, "composer-session-model"); assert_ne!(cfg.model, "grok-build"); @@ -4958,11 +4934,9 @@ reasoning_effort = "low" model: "grok-build".into(), ..Default::default() }; - let (model, cfg) = - finalize_image_describe_sampler_config(Some(aux), &active, Some("cli".into()), Some(7)); + let (model, cfg) = finalize_image_describe_sampler_config(Some(aux), &active, Some(7)); assert_eq!(model, "grok-build"); assert_eq!(cfg.model, "grok-build"); - assert_eq!(cfg.client_identifier.as_deref(), Some("cli")); assert_eq!(cfg.max_retries, Some(7)); } #[test] @@ -4980,7 +4954,7 @@ reasoning_effort = "low" ), ); let resolved = - resolve_aux_model_sampling_config("grok-build", &catalog, &endpoints, None, None, None) + resolve_aux_model_sampling_config("grok-build", &catalog, &endpoints, None, None) .expect("override entry has an API key, so resolution succeeds"); assert_eq!(resolved.model, "v9m-rl-learnability-tp8"); assert_eq!(resolved.base_url, "https://vendor.example/v1"); @@ -5085,14 +5059,8 @@ reasoning_effort = "low" None, None, ); - let sampling_config = sampling_config_for_model( - &model, - resolve_credentials(&model, None), - None, - None, - None, - None, - ); + let sampling_config = + sampling_config_for_model(&model, resolve_credentials(&model, None), None); assert_eq!( sampling_config.api_key, Some("model-specific-key".to_string()) @@ -5111,9 +5079,6 @@ reasoning_effort = "low" auth_scheme: AuthScheme::Bearer, }, None, - None, - None, - None, ); assert_eq!(sampling_config.api_key, Some("fallback-key".to_string())); } @@ -5388,14 +5353,8 @@ reasoning_effort = "low" None, ); model.info.api_backend = ApiBackend::Messages; - let config = sampling_config_for_model( - &model, - resolve_credentials(&model, Some("tok")), - None, - None, - None, - None, - ); + let config = + sampling_config_for_model(&model, resolve_credentials(&model, Some("tok")), None); assert_eq!(config.api_backend, ApiBackend::Messages); assert_eq!(config.auth_scheme, AuthScheme::Bearer); assert_eq!(config.api_key, Some("tok".to_string())); @@ -5440,7 +5399,7 @@ reasoning_effort = "low" assert_eq!(creds.auth_scheme, AuthScheme::XApiKey); assert_eq!(creds.auth_type, kigi_chat_state::AuthType::ApiKey); assert_eq!(creds.api_key, Some("sk-ant-test-key".to_string())); - let config = sampling_config_for_model(&model, creds, None, None, None, None); + let config = sampling_config_for_model(&model, creds, None); assert_eq!(config.auth_scheme, AuthScheme::XApiKey); assert_eq!(config.api_backend, ApiBackend::Messages); let client = kigi_sampler::SamplingClient::new(config).expect("client should build"); @@ -5459,7 +5418,7 @@ reasoning_effort = "low" assert_eq!(model.info.auth_scheme, AuthScheme::Bearer); let creds = resolve_credentials(&model, None); assert_eq!(creds.auth_scheme, AuthScheme::Bearer); - let config = sampling_config_for_model(&model, creds, None, None, None, None); + let config = sampling_config_for_model(&model, creds, None); assert_eq!(config.auth_scheme, AuthScheme::Bearer); let client = kigi_sampler::SamplingClient::new(config).expect("client should build"); let info = client.auth_info(); @@ -5771,25 +5730,11 @@ reasoning_effort = "low" #[test] fn sampling_config_context_window_from_entry_or_default() { let model = test_model_entry("any-model", "https://api.x.ai/v1", None, None, None); - let config = sampling_config_for_model( - &model, - resolve_credentials(&model, None), - None, - None, - None, - None, - ); + let config = sampling_config_for_model(&model, resolve_credentials(&model, None), None); assert_eq!(config.context_window, 200_000); let mut model = test_model_entry("any-model", "https://api.x.ai/v1", None, None, None); model.info.context_window = NonZeroU64::new(256_000).unwrap(); - let config = sampling_config_for_model( - &model, - resolve_credentials(&model, None), - None, - None, - None, - None, - ); + let config = sampling_config_for_model(&model, resolve_credentials(&model, None), None); assert_eq!(config.context_window, 256_000); } #[test] @@ -5917,14 +5862,8 @@ reasoning_effort = "low" let mut model = test_model_entry("test-model", "https://api.example.com/v1", None, None, None); model.info.api_backend = ApiBackend::Responses; - let sampling_config = sampling_config_for_model( - &model, - resolve_credentials(&model, None), - None, - None, - None, - None, - ); + let sampling_config = + sampling_config_for_model(&model, resolve_credentials(&model, None), None); assert_eq!(sampling_config.api_backend, ApiBackend::Responses); } #[test] @@ -6636,7 +6575,7 @@ reasoning_effort = "low" } fn resolve_sampling(model: &ModelEntry, session_key: Option<&str>) -> SamplerConfig { let credentials = resolve_credentials(model, session_key); - sampling_config_for_model(model, credentials, None, None, None, None) + sampling_config_for_model(model, credentials, None) } #[test] #[serial] @@ -7022,7 +6961,7 @@ reasoning_effort = "low" &format!( r#" [endpoints] - cli_chat_proxy_base_url = "https://enterprise-proxy.acme.com/v1" + coding_api_base_url = "https://enterprise-proxy.acme.com/v1" [model."{BUNDLED_DEFAULT_KEY}"] api_key = "acme-api-key" @@ -7052,14 +6991,14 @@ reasoning_effort = "low" let (_, models) = resolve_models_from_toml( r#" [endpoints] - cli_chat_proxy_base_url = "https://enterprise-proxy.acme.com/v1" + coding_api_base_url = "https://enterprise-proxy.acme.com/v1" "#, None, ); let model = models.get(BUNDLED_DEFAULT_KEY).expect("model should exist"); assert_eq!( model.info.base_url, "https://enterprise-proxy.acme.com/v1", - "default model should use enterprise cli_chat_proxy_base_url" + "default model should use enterprise coding_api_base_url" ); // The open-platform fallback entries keep their fixed moonshot bases; // only the subscription entry follows the proxy override. @@ -7073,7 +7012,7 @@ reasoning_effort = "low" /// the ambient environment. Gated behind `#[serial]`. fn unset_endpoint_env_vars() { for k in [ - "KIGI_CLI_CHAT_PROXY_BASE_URL", + "KIGI_CODE_BASE_URL", kigi_env::CODE_BASE_URL_ENV, "KIGI_XAI_API_BASE_URL", "KIGI_FEEDBACK_BASE_URL", @@ -7101,7 +7040,7 @@ reasoning_effort = "low" let inference = "https://inference.acme-corp.example/xai/v1"; let cfg = EndpointsConfig { xai_api_base_url: inference.to_string(), - cli_chat_proxy_base_url: None, + coding_api_base_url: None, ..Default::default() }; let proxy = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url; @@ -7119,7 +7058,7 @@ reasoning_effort = "low" ); assert_eq!(cfg.xai_api_base_url, inference); let overridden = EndpointsConfig { - cli_chat_proxy_base_url: Some("https://proxy.enterprise.example/v1".to_string()), + coding_api_base_url: Some("https://proxy.enterprise.example/v1".to_string()), managed_config_url: Some( "https://control.enterprise.example/deployment/config".to_string(), ), @@ -7159,7 +7098,7 @@ reasoning_effort = "low" .unwrap(), ) .expect("config should parse"); - assert!(cfg.endpoints.cli_chat_proxy_base_url.is_none()); + assert!(cfg.endpoints.coding_api_base_url.is_none()); assert_eq!( cfg.endpoints.resolve_managed_config_url(), format!( @@ -7181,7 +7120,7 @@ reasoning_effort = "low" &format!( r#" [endpoints] - cli_chat_proxy_base_url = "https://enterprise-proxy.acme.com/v1" + coding_api_base_url = "https://enterprise-proxy.acme.com/v1" [model."{dm}"] base_url = "https://my-special-proxy.example.com/v1" @@ -8656,7 +8595,7 @@ agent_type = "cursor" fn otlp_traces_endpoint_precedence() { let proxy = "https://inference.acme.com/v1".to_string(); let derived = EndpointsConfig { - cli_chat_proxy_base_url: Some(proxy.clone()), + coding_api_base_url: Some(proxy.clone()), ..Default::default() }; assert_eq!( @@ -8664,7 +8603,7 @@ agent_type = "cursor" "https://inference.acme.com/v1/traces" ); let base = EndpointsConfig { - cli_chat_proxy_base_url: Some(proxy.clone()), + coding_api_base_url: Some(proxy.clone()), otel_exporter_otlp_endpoint: Some("https://otel.acme.com".to_string()), ..Default::default() }; @@ -8673,7 +8612,7 @@ agent_type = "cursor" "https://otel.acme.com/v1/traces" ); let full = EndpointsConfig { - cli_chat_proxy_base_url: Some(proxy), + coding_api_base_url: Some(proxy), otel_exporter_otlp_endpoint: Some("https://ignored.example".to_string()), otel_exporter_otlp_traces_endpoint: Some("https://otel.acme.com/v1/traces".to_string()), ..Default::default() @@ -8702,7 +8641,7 @@ agent_type = "cursor" /// explicitly unset so ambient env (via `Default`) can't leak in. fn internal_otlp_test_config() -> EndpointsConfig { EndpointsConfig { - cli_chat_proxy_base_url: Some("https://proxy.example/v1".to_string()), + coding_api_base_url: Some("https://proxy.example/v1".to_string()), otel_exporter_otlp_endpoint: None, otel_exporter_otlp_traces_endpoint: None, otel_exporter_otlp_headers: None, diff --git a/crates/codegen/kigi-shell/src/agent/feedback_client.rs b/crates/codegen/kigi-shell/src/agent/feedback_client.rs index 9f9a41c..a2d26cb 100644 --- a/crates/codegen/kigi-shell/src/agent/feedback_client.rs +++ b/crates/codegen/kigi-shell/src/agent/feedback_client.rs @@ -1,304 +1,25 @@ -//! REST client for feedback collection via cli-chat-proxy. +//! Feedback client for the Kimi Code platform. //! -//! This client handles: -//! - Syncing session signals to cli-chat-proxy -//! - Submitting feedback responses -//! - Completing/dismissing feedback requests -//! - Creating feedback requests (when triggered by heuristics) +//! Port of kimi-cli's `/feedback` slash command (kimi-cli +//! `src/kimi_cli/ui/shell/slash.py`, `feedback()`): subscription (OAuth) +//! sessions POST the user's feedback text to `{coding_api_base_url}/feedback` +//! with a Bearer token; everyone else is pointed at the GitHub issue tracker. +//! The request body carries exactly the fields kimi-cli sends: +//! `session_id`, `content`, `version`, `os`, `model`. use std::sync::Arc; -use anyhow::{Context, Result}; -use reqwest::RequestBuilder; -use serde::de::DeserializeOwned; +use serde::Serialize; -// Import feedback wire types from cli-chat-proxy -use prod_mc_cli_chat_proxy_types::feedback_types::{ - ClientType, CreateFeedbackRequestInput, CreateFeedbackRequestResponse, - FeedbackHeuristicsConfig, FeedbackRequestUpdateResponse, FeedbackResponse, FeedbackSubmission, - SessionEventRequest, SessionEventResponse, SessionSignalsUpdate, SessionSignalsUpdateResponse, -}; +/// Where non-subscription users (no OAuth session) submit feedback instead. +pub const FEEDBACK_ISSUES_URL: &str = "https://github.com/ZacharyZhang-NY/Kigi-CLI/issues"; -/// Client version header sent on every request to cli-chat-proxy for version gating. -const CLIENT_VERSION_HEADER: &str = "x-grok-client-version"; - -// ============================================================================ -// Turn delta wire types (local to kigi-shell until cli-chat-proxy catches up) -// ============================================================================ - -/// Per-turn delta sent at the end of every turn via -/// `POST /v1/sessions/{session_id}/turn-deltas`. -/// -/// Each field falls into one of four categories: -/// -/// - **Delta** — the *change* since the previous turn end (computed as -/// `current_cumulative - previous_turn_snapshot`). For the first turn, -/// the previous snapshot is zero. -/// - **Turn-level** — an absolute value measured only for *this* turn, -/// reset between turns. `None` when the event did not occur this turn. -/// - **Accumulated** — a cumulative total since session start, monotonically -/// increasing across turns. -/// - **Context** — session/turn metadata that is neither a counter nor a -/// measurement (e.g. IDs, timestamps, client type). -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SessionTurnDelta { - // ── Context fields ────────────────────────────────────────────────── - /// **[context]** Which client surface produced this record (e.g. CLI, TUI). - pub client_type: ClientType, - - /// **[context]** 1-based turn number at the time of this snapshot. Equals - /// the cumulative `turn_count` from `SessionSignals`. - pub turn_number: i64, - - // ── Delta counters ────────────────────────────────────────────────── - // Each is `current_cumulative - previous_turn_snapshot`. - /// **[delta]** Number of tool calls made during this turn. - pub delta_tool_calls: i64, - - /// **[delta]** Number of tool calls that failed during this turn. - pub delta_tool_failures: i64, - - /// **[delta]** Number of errors (including sampling errors) during this turn. - pub delta_errors: i64, - - /// **[delta]** Number of user cancellations (Ctrl+C) during this turn. - pub delta_cancellations: i64, - - /// **[delta]** Number of regeneration requests during this turn. - pub delta_regenerations: i64, - - /// **[delta]** Number of conversation compactions during this turn. - pub delta_compactions: i64, - - /// **[delta]** Number of edit-and-retry actions (user rewinds prompt) - /// during this turn. - pub delta_edit_and_retries: i64, - - /// **[delta]** Number of positive ratings (thumbs-up) during this turn. - pub delta_positive_ratings: i64, - - /// **[delta]** Number of negative ratings (thumbs-down) during this turn. - pub delta_negative_ratings: i64, - - /// **[delta]** Number of assistant messages produced during this turn - /// (may be >1 when tool-call rounds generate intermediate messages). - pub delta_assistant_messages: i64, - - /// **[delta]** Number of long idle pauses (>60 s) that occurred during - /// this turn. - pub delta_long_pauses: i64, - - /// **[delta]** Number of successful tool uses during this turn. Derived - /// as `delta_tool_calls − delta_tool_failures`. - pub delta_successful_tool_uses: i64, - - // ── Turn-level snapshot values ────────────────────────────────────── - /// **[turn-level]** Consecutive cancellation streak at turn end. This is - /// a point-in-time snapshot (not a diff) — it resets to 0 when a turn - /// completes normally. - pub consecutive_cancellations: i64, - - // ── Turn-level latency ────────────────────────────────────────────── - // Absolute measurements for this turn's inference request only. - // `None` when no inference occurred during the turn. - /// **[turn-level]** Time-to-first-token for this turn's model response - /// (milliseconds). `None` when no inference occurred. - #[serde(skip_serializing_if = "Option::is_none")] - pub time_to_first_token_ms: Option, - - /// **[turn-level]** Total wall-clock response time for this turn's model - /// response (milliseconds). `None` when no inference occurred. - #[serde(skip_serializing_if = "Option::is_none")] - pub total_response_time_ms: Option, - - /// **[turn-level]** Inter-token latency p50 for this turn (ms). - /// Computed from the token intervals collected during this turn only. - #[serde(skip_serializing_if = "Option::is_none")] - pub itl_p50_ms: Option, - - /// **[turn-level]** Inter-token latency p99 for this turn (ms). - #[serde(skip_serializing_if = "Option::is_none")] - pub itl_p99_ms: Option, - - /// **[turn-level]** Inter-token latency maximum for this turn (ms). - #[serde(skip_serializing_if = "Option::is_none")] - pub itl_max_ms: Option, - - /// **[turn-level]** Inter-token latency mean for this turn (ms). - #[serde(skip_serializing_if = "Option::is_none")] - pub itl_mean_ms: Option, - - // ── Accumulated / snapshot session-level values ───────────────────── - /// **[accumulated]** Current context window usage as a percentage (0–100) - /// at turn end. Read from cumulative `SessionSignals.context_window_usage`. - pub context_window_usage: i64, - - /// **[accumulated]** Primary model ID (most recently used model). Read - /// from cumulative `SessionSignals.primary_model_id`. - #[serde(skip_serializing_if = "Option::is_none")] - pub model_id: Option, - - // ── Turn-level outcome / served checkpoint ────────────────────────── - /// Whole-turn wall-clock duration (prompt→final response), ms. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub turn_duration_ms: Option, - - /// Terminal outcome: `"completed"` | `"cancelled"` | `"error"`. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub turn_outcome: Option, - - /// Served model fingerprint (upstream `system_fingerprint`). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub model_fingerprint: Option, - - // ── Turn-level tool / error detail ────────────────────────────────── - /// **[turn-level]** Distinct tool names invoked during this turn - /// (deduplicated, sorted, capped at 100 entries). Reset each turn. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub tools_used_this_turn: Vec, - - /// **[turn-level]** Error type strings that occurred during this turn - /// (e.g. `"timeout"`, `"rate_limit"`, `"tool_error"`). Reset each turn. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub error_types_this_turn: Vec, - - /// **[turn-level]** Per-tool success/failure breakdown for this turn, - /// JSON-serialized array of `{ tool_name, successes, failures }`. - /// Empty string when no tool calls occurred. Reset each turn. - #[serde(default, skip_serializing_if = "String::is_empty")] - pub tool_outcomes: String, - - // ── Accumulated totals ────────────────────────────────────────────── - /// **[accumulated]** Total tool calls since session start. - /// Read from cumulative `SessionSignals.tool_call_count`. - pub cumulative_tool_calls: i64, - - /// **[accumulated]** Total errors since session start. - /// Read from cumulative `SessionSignals.error_count`. - pub cumulative_errors: i64, - - /// **[accumulated]** Wall-clock seconds elapsed since session start. - /// Read from cumulative `SessionSignals.session_duration_seconds`. - pub session_duration_seconds: i64, - - /// **[accumulated]** Sum of token counts across all compactions since - /// session start. Read from `SessionSignals.total_tokens_before_compaction`. - #[serde(default)] - pub total_tokens_before_compaction: i64, - - /// **[context]** Arbitrary JSON metadata blob. - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, - - /// **[context]** Prompt/request ID that initiated this turn. - #[serde(skip_serializing_if = "Option::is_none")] - pub request_id: Option, - - /// **[context]** Wall-clock time when the session was created. Used for - /// BQ partitioning on the backend. - #[serde(skip_serializing_if = "Option::is_none")] - pub session_start_at: Option>, - - // ── Feedback state ────────────────────────────────────────────────── - /// **[accumulated]** Total number of feedback requests sent this session. - /// Supplied by `FeedbackHeuristics`, not the signals actor. - #[serde(default)] - pub feedback_requests_sent: i64, - - /// **[accumulated]** Wall-clock timestamp of the most recent feedback - /// request sent this session. Supplied by `FeedbackHeuristics`. - #[serde(skip_serializing_if = "Option::is_none")] - pub last_feedback_request_at: Option>, - - // ── Turn-level token counts ───────────────────────────────────────── - /// **[turn-level]** Number of response (completion minus reasoning) - /// tokens generated during this turn. `None` when no inference occurred. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub response_tokens: Option, - - /// **[turn-level]** Number of thinking/reasoning tokens generated during - /// this turn. `None` when no inference occurred. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub thinking_tokens: Option, - - // ── LOC Attribution Deltas ────────────────────────────────────────── - // Each is `current_cumulative - previous_turn_snapshot`, same as the - // counter deltas above. Tracks lines-of-code changes attributed to - // the agent vs. the human during this turn. - /// **[delta]** Lines added by the agent during this turn. - #[serde(default)] - pub delta_agent_lines_added: i64, - - /// **[delta]** Lines removed by the agent during this turn. - #[serde(default)] - pub delta_agent_lines_removed: i64, - - /// **[delta]** Agent-added lines that were reverted during this turn. - #[serde(default)] - pub delta_agent_lines_added_reverted: i64, - - /// **[delta]** Agent-removed lines that were reverted during this turn. - #[serde(default)] - pub delta_agent_lines_removed_reverted: i64, - - /// **[delta]** Lines added by the human during this turn. - #[serde(default)] - pub delta_human_lines_added: i64, - - /// **[delta]** Lines removed by the human during this turn. - #[serde(default)] - pub delta_human_lines_removed: i64, - - /// **[delta]** Human-added lines that were reverted during this turn. - #[serde(default)] - pub delta_human_lines_added_reverted: i64, - - /// **[delta]** Human-removed lines that were reverted during this turn. - #[serde(default)] - pub delta_human_lines_removed_reverted: i64, - - /// **[delta]** New distinct files touched by the agent during this turn. - #[serde(default)] - pub delta_agent_files_touched: i64, - - /// **[delta]** New distinct files touched by the human during this turn. - #[serde(default)] - pub delta_human_files_touched: i64, - - /// **[delta]** New distinct files touched (union of agent + human) - /// during this turn. - #[serde(default)] - pub delta_total_files_touched: i64, - - /// **[context]** Whether LOC (lines-of-code) attribution tracking was - /// enabled for this session. When `false`, all `delta_*` LOC fields - /// above are meaningless zeros — the hunk tracker was never spawned. - /// When `true`, zeros mean "tracking was active but no code changed." - /// Defaults to `false` for backwards-compat with old clients that - /// don't send this field. - #[serde(default)] - pub loc_tracking_enabled: bool, -} - -/// Response from the turn-deltas endpoint. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SessionTurnDeltaResponse { - pub session_id: String, - pub turn_number: i64, - pub recorded_at: chrono::DateTime, -} - -/// HTTP error from the feedback/signals API with a preserved status code. -/// -/// Used to let callers distinguish auth failures (401) from transient errors -/// without fragile string matching on error messages. +/// HTTP error from the feedback endpoint with a preserved status code, so +/// callers can distinguish auth failures (401) without string matching. #[derive(Debug, thiserror::Error)] -#[error("{context} failed with status {status}: {body}")] +#[error("feedback submission failed with status {status}: {body}")] pub struct FeedbackApiError { pub status: reqwest::StatusCode, - pub context: &'static str, pub body: String, } @@ -307,1041 +28,224 @@ impl FeedbackApiError { pub fn is_unauthorized(&self) -> bool { self.status == reqwest::StatusCode::UNAUTHORIZED } - - /// Returns `true` if this is a 403 Forbidden response. - pub fn is_forbidden(&self) -> bool { - self.status == reqwest::StatusCode::FORBIDDEN - } } -/// Client for the feedback collection API via cli-chat-proxy. +/// JSON body for `POST {base}/feedback` — field names exactly as kimi-cli +/// sends them (slash.py `payload = {...}`). +#[derive(Debug, Clone, Serialize)] +struct FeedbackPayload<'a> { + session_id: &'a str, + content: &'a str, + version: &'a str, + os: &'a str, + model: Option<&'a str>, +} + +/// Bearer-token source for the feedback POST: the live OAuth session by +/// default, or a fixed token in tests. +#[derive(Clone)] +enum BearerSource { + AuthManager(Arc), + Static(String), +} + +/// Client for the Kimi Code feedback endpoint. #[derive(Clone)] pub struct FeedbackClient { http: reqwest::Client, - client: reqwest_middleware::ClientWithMiddleware, base_url: String, - credentials: crate::util::kigi_auth_credentials::KigiAuthCredentials, + bearer: BearerSource, session_id: Option, } impl FeedbackClient { - pub fn new(base_url: impl Into, user_token: Option) -> Self { - let http = crate::http::shared_client(); - let credentials = crate::util::kigi_auth_credentials::KigiAuthCredentials::new(user_token); - let client = Self::build_middleware_client(&http, &credentials); + /// Client bound to the live OAuth session. Callers gate construction on + /// an existing session auth (see `MvpAgent::feedback_client`). + pub(crate) fn new( + base_url: impl Into, + auth_manager: Arc, + ) -> Self { Self { - http, - client, + http: crate::http::shared_client(), base_url: base_url.into(), - credentials, + bearer: BearerSource::AuthManager(auth_manager), session_id: None, } } + /// Client with a fixed Bearer token (tests). + #[cfg(test)] + pub(crate) fn with_static_token(base_url: impl Into, token: impl Into) -> Self { + Self { + http: crate::http::shared_client(), + base_url: base_url.into(), + bearer: BearerSource::Static(token.into()), + session_id: None, + } + } + + /// Default session id used when a submission doesn't carry one. pub fn with_session_id(mut self, session_id: impl Into) -> Self { self.session_id = Some(session_id.into()); self } - pub fn with_alpha_test_key(mut self, key: Option) -> Self { - self.credentials.alpha_test_key = key; - self.rebuild_middleware(); - self - } - - pub fn with_deployment_key(mut self, key: Option) -> Self { - self.credentials.deployment_key = key; - self.rebuild_middleware(); - self - } - - /// Create a FeedbackClient with a custom reqwest Client. - pub fn with_client( - http: reqwest::Client, - base_url: impl Into, - user_token: Option, - ) -> Self { - let credentials = crate::util::kigi_auth_credentials::KigiAuthCredentials::new(user_token); - let client = Self::build_middleware_client(&http, &credentials); - Self { - http, - client, - base_url: base_url.into(), - credentials, - session_id: None, + fn bearer_token(&self) -> Option { + match &self.bearer { + BearerSource::AuthManager(am) => am + .current_or_expired() + .filter(|a| a.is_session_auth()) + .map(|a| a.key.clone()), + BearerSource::Static(token) => Some(token.clone()), } } - pub(crate) fn with_auth_manager( - mut self, - auth_manager: std::sync::Arc, - ) -> Self { - self.credentials = self.credentials.with_auth_manager(auth_manager); - self.rebuild_middleware(); - self - } - - /// Whether this client can refresh credentials on a 401: requires both an - /// attached `AuthManager` and a wired `TokenRefresher` (e.g. static - /// deployment-key sessions return false). - pub fn has_token_refresher(&self) -> bool { - self.credentials - .auth_manager() - .is_some_and(|am| am.has_refresher_attached()) - } - - /// Rebuild the middleware-wrapped client from the current credentials. - /// Called by each builder method so the middleware sees the final state. - fn rebuild_middleware(&mut self) { - self.client = Self::build_middleware_client(&self.http, &self.credentials); - } - - fn build_middleware_client( - http: &reqwest::Client, - credentials: &crate::util::kigi_auth_credentials::KigiAuthCredentials, - ) -> reqwest_middleware::ClientWithMiddleware { - let provider = Self::make_auth_provider(credentials); - // max_retries=0: the middleware stamps the auth header but does NOT - // drive its own ServerRejected recovery on 401. Background consumers - // (signals sync, turn deltas) handle retry at the application level - // via with_one_shot_auth_retry / try_refresh_and_retry_sync, which - // first wait for the proactive refresh to complete before falling back - // to active recovery. This prevents the 401-amplification pattern - // where the middleware's eager ServerRejected refresh races with every - // other auth consumer during token-expiry windows. - reqwest_middleware::ClientBuilder::new(http.clone()) - .with(kigi_auth::AuthRetryMiddleware::new(provider, 0)) - .build() - } - - fn make_auth_provider( - credentials: &crate::util::kigi_auth_credentials::KigiAuthCredentials, - ) -> Arc { - if let Some(am) = credentials.auth_manager() { - Arc::new( - crate::auth::credential_provider::ShellAuthCredentialProvider::new( - am.clone(), - credentials.deployment_key.clone(), - credentials.alpha_test_key.clone(), - ), - ) - } else { - let wire_bearer = credentials - .deployment_key - .clone() - .or(credentials.user_token.clone()); - Arc::new(kigi_auth::StaticAuthCredentialProvider::new( - Box::new(credentials.clone()), - wire_bearer, - )) - } - } - - fn record_401_attribution_if_needed(&self, response: &reqwest::Response, op: &str) { - if response.status() == reqwest::StatusCode::UNAUTHORIZED - && let Some(am) = self.credentials.auth_manager() - { - let bearer_prefix = self - .credentials - .deployment_key - .as_deref() - .or(self.credentials.user_token.as_deref()); - crate::auth::attribution::record_consumer_401( - am.as_ref(), - self.session_id.as_deref(), - crate::auth::attribution::ConsumerKind::FeedbackClient, - op, - bearer_prefix, - ); - } - } - - pub async fn try_refresh_credentials(&self) -> bool { - let Some(manager) = self.credentials.auth_manager() else { - return false; - }; - manager.try_recover_unauthorized().await - } - - /// Wait for another consumer (proactive refresh, main request path) to - /// refresh the token. Returns `true` if the token changed within the - /// timeout. Background consumers call this before driving their own - /// `ServerRejected` recovery to avoid amplifying 401 bursts. - pub(crate) async fn wait_for_token_refresh(&self, timeout: std::time::Duration) -> bool { - let Some(manager) = self.credentials.auth_manager() else { - return false; - }; - manager.wait_for_token_refresh(timeout).await - } - - /// `true` iff the attached `AuthManager` has a non-aged-out - /// permanent-failure verdict from the IdP. - pub(crate) fn is_auth_permanently_failed(&self) -> bool { - self.credentials - .auth_manager() - .is_some_and(|am| am.has_permanent_failure()) - } - - /// Create a POST request builder with common headers. - fn post(&self, url: &str) -> RequestBuilder { - self.add_common_headers(self.http.post(url)) - } - - /// Create a GET request builder with common headers. - fn get(&self, url: &str) -> RequestBuilder { - self.add_common_headers(self.http.get(url)) - } - - fn add_common_headers(&self, builder: RequestBuilder) -> RequestBuilder { - let builder = builder - .header(CLIENT_VERSION_HEADER, kigi_version::VERSION) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ); - // User-token auth requires the companion marker header for proxy - // routing. Deployment keys do not need it. - if self.credentials.deployment_key.is_none() { - builder.header("X-XAI-Token-Auth", "xai-grok-cli") - } else { - builder - } - } - - async fn send_json( - &self, - request: RequestBuilder, - context: &'static str, - ) -> Result { - let request = kigi_file_utils::trace_context::inject_trace_context_into_request(request); - let req = request.build().context(context)?; - let response = self.client.execute(req).await.context(context)?; - - self.record_401_attribution_if_needed(&response, context); - - if response.status() == reqwest::StatusCode::FORBIDDEN { - tracing::debug!("{context} rejected (403), skipping"); - anyhow::bail!("{context} rejected (403)"); - } - - if !response.status().is_success() { - let status = response.status(); - let body = response - .text() - .await - .unwrap_or_else(|_| "Unknown error".to_string()); - return Err(FeedbackApiError { - status, - context, - body, - } - .into()); - } - - response - .json() - .await - .with_context(|| format!("Failed to parse {} response", context)) - } - - async fn send_empty(&self, request: RequestBuilder, context: &'static str) -> Result<()> { - let request = kigi_file_utils::trace_context::inject_trace_context_into_request(request); - let req = request.build().context(context)?; - let response = self.client.execute(req).await.context(context)?; - - self.record_401_attribution_if_needed(&response, context); - - if response.status() == reqwest::StatusCode::FORBIDDEN { - tracing::debug!("{context} rejected (403), skipping"); - return Ok(()); - } - - if !response.status().is_success() { - let status = response.status(); - let body = response - .text() - .await - .unwrap_or_else(|_| "Unknown error".to_string()); - return Err(FeedbackApiError { - status, - context, - body, - } - .into()); - } - - Ok(()) - } - - /// Update session signals. - /// POST /v1/sessions/{session_id}/signals - pub async fn update_signals( - &self, - session_id: &str, - update: &SessionSignalsUpdate, - ) -> Result { - let url = format!("{}/sessions/{}/signals", self.base_url, session_id); - let request = self.post(&url).json(update); - self.send_json(request, "Signals update").await - } - - /// Record a session event. - /// POST /v1/sessions/{session_id}/events - pub async fn record_event( - &self, - session_id: &str, - event: &SessionEventRequest, - ) -> Result { - let url = format!("{}/sessions/{}/events", self.base_url, session_id); - let request = self.post(&url).json(event); - self.send_json(request, "Event recording").await - } - - /// Submit feedback. - /// POST /v1/feedback + /// `POST {base}/feedback` (kimi-cli slash.py parity). `model` is the + /// active model key, when known. On 401 the OAuth session is refreshed + /// once and the request retried. pub async fn submit_feedback( &self, - submission: &FeedbackSubmission, - ) -> Result { - let url = format!("{}/feedback", self.base_url); - let request = self.post(&url).json(submission); - self.send_json(request, "Feedback submission").await + session_id: &str, + content: &str, + model: Option<&str>, + ) -> Result<(), anyhow::Error> { + let sid = if session_id.is_empty() { + self.session_id.as_deref().unwrap_or_default() + } else { + session_id + }; + match self.post_feedback(sid, content, model).await { + Err(e) + if e.downcast_ref::() + .is_some_and(FeedbackApiError::is_unauthorized) + && self.try_refresh_credentials().await => + { + self.post_feedback(sid, content, model).await + } + other => other, + } } - /// Complete a feedback request. - /// POST /v1/feedback/requests/{request_id}/complete - pub async fn complete_request( - &self, - request_id: &str, - submission: &FeedbackSubmission, - ) -> Result<()> { - let url = format!( - "{}/feedback/requests/{}/complete", - self.base_url, request_id - ); - let request = self.post(&url).json(submission); - self.send_empty(request, "Completing feedback request") - .await - } - - /// Dismiss a feedback request. - /// POST /v1/feedback/requests/{request_id}/dismiss - pub async fn dismiss_request(&self, request_id: &str) -> Result { - let url = format!("{}/feedback/requests/{}/dismiss", self.base_url, request_id); - let request = self.post(&url); - self.send_json(request, "Dismissing feedback request").await - } - - /// Create a new feedback request. - /// POST /v1/feedback/requests - /// - /// Called when the agent decides to request feedback (based on heuristics). - /// This creates a record in BigQuery before the FeedbackRequest notification - /// is sent to the client. - pub async fn create_feedback_request( - &self, - input: &CreateFeedbackRequestInput, - ) -> Result { - let url = format!("{}/feedback/requests", self.base_url); - let request = self.post(&url).json(input); - self.send_json(request, "Creating feedback request").await - } - - /// Get the active feedback heuristics configuration. - /// GET /v1/feedback/config - /// - /// This fetches the current feedback configuration from the server, - /// including tier thresholds, sample rates, and feedback modes. - pub async fn get_feedback_config(&self) -> Result { - let url = format!("{}/feedback/config", self.base_url); - let request = self.get(&url); - self.send_json(request, "Fetching feedback config").await - } - - /// Send a per-turn delta to the backend. - /// POST /v1/sessions/{session_id}/turn-deltas - /// - /// Called at the end of every turn to stream time-series data for - /// regression tracking and session analytics. - pub async fn send_turn_delta( + async fn post_feedback( &self, session_id: &str, - delta: &SessionTurnDelta, - ) -> Result { - let url = format!("{}/sessions/{}/turn-deltas", self.base_url, session_id); - let request = self.post(&url).json(delta); - self.send_json(request, "Sending turn delta").await - } -} - -/// Helper to create a SessionSignalsUpdate from local session signals. -pub fn signals_to_update( - signals: &crate::session::signals::SessionSignals, - client_type: ClientType, -) -> SessionSignalsUpdate { - SessionSignalsUpdate { - client_type, - total_turns: Some(signals.turn_count as i64), - user_message_count: Some(signals.user_message_count as i64), - assistant_message_count: Some(signals.assistant_message_count as i64), - cancellation_count: Some(signals.cancellation_count as i64), - consecutive_cancellations: Some(signals.consecutive_cancellations as i64), - error_count: Some(signals.error_count as i64), - tool_failure_count: Some(signals.tool_failure_count as i64), - tool_call_count: Some(signals.tool_call_count as i64), - compaction_count: Some(signals.compaction_count as i64), - regeneration_count: Some(signals.regeneration_count as i64), - edit_and_retry_count: Some(signals.edit_and_retry_count as i64), - positive_ratings: Some(signals.positive_ratings as i64), - negative_ratings: Some(signals.negative_ratings as i64), - long_pauses_count: Some(signals.long_pauses_count as i64), - session_duration_seconds: Some(signals.session_duration_seconds as i64), - tools_used: signals.tools_used.clone(), - models_used: signals.models_used.clone(), - primary_model_id: signals.primary_model_id.clone(), - // Latency metrics - avg_time_to_first_token_ms: Some(signals.avg_time_to_first_token_ms as i64), - avg_response_time_ms: Some(signals.avg_response_time_ms as i64), - min_time_to_first_token_ms: Some(signals.min_time_to_first_token_ms as i64), - max_time_to_first_token_ms: Some(signals.max_time_to_first_token_ms as i64), - latency_sample_count: Some(signals.latency_sample_count as i64), - // ITL metrics (session-level aggregates) - // Guard p50/p99 with itl_sample_count > 0 so that fresh sessions - // (no ITL measured) send None → SQL NULL, preserving the "not yet - // reported" semantic in the nullable PG columns. - last_itl_p50_ms: signals.itl_p50_ms.map(|v| v as i64), - last_itl_p99_ms: signals.itl_p99_ms.map(|v| v as i64), - worst_itl_max_ms: signals.itl_max_ms.map(|v| v as i64), - avg_itl_mean_ms: signals.itl_mean_ms.map(|v| v as i64), - total_chunk_count: Some(signals.total_chunk_count as i64), - itl_sample_count: Some(signals.itl_sample_count as i64), - // Inference idle timeout tracing - inference_idle_timeouts: Some(signals.inference_idle_timeouts as i64), - inference_idle_timeout_configured_secs: signals - .inference_idle_timeout_configured_secs - .map(|v| v as i64), - // Legacy client-side doom-loop detection removed; keep its columns null. - doom_loop_warnings: None, - doom_loop_terminations: None, - doom_loop_threshold: None, - doom_loop_ro_threshold: None, - // Doom-loop recovery (server-detected, client-resampled) tracing - doom_loop_recovery_fired: Some( - signals.doom_loop_recovery_attempts > 0 - || signals.doom_loop_recovery_accepted_after_budget > 0, - ), - doom_loop_recovery_attempts: Some(signals.doom_loop_recovery_attempts as i64), - doom_loop_recovery_accepted_after_budget: Some( - signals.doom_loop_recovery_accepted_after_budget as i64, - ), - doom_loop_recovery_top_trigger: signals.doom_loop_recovery_top_trigger.clone(), - doom_loop_recovery_aborted_chunks: Some(signals.doom_loop_recovery_aborted_chunks as i64), - // GCS upload queue removed (zero-egress build); keep its columns null. - gcs_queue_enqueued: None, - gcs_queue_uploaded: None, - gcs_queue_failed: None, - gcs_queue_fallbacks: None, - gcs_queue_circuit_breaker_trips: None, - gcs_queue_pending: None, - gcs_queue_pending_bytes: None, - gcs_queue_orphans_cleaned: None, - // LOC Attribution - agent_lines_added: Some(signals.agent_lines_added), - agent_lines_removed: Some(signals.agent_lines_removed), - agent_lines_added_reverted: Some(signals.agent_lines_added_reverted), - agent_lines_removed_reverted: Some(signals.agent_lines_removed_reverted), - human_lines_added: Some(signals.human_lines_added), - human_lines_removed: Some(signals.human_lines_removed), - human_lines_added_reverted: Some(signals.human_lines_added_reverted), - human_lines_removed_reverted: Some(signals.human_lines_removed_reverted), - agent_files_touched: Some(signals.agent_files_touched as i64), - human_files_touched: Some(signals.human_files_touched as i64), - total_files_touched: Some(signals.total_files_touched as i64), - metadata: None, - } -} - -/// Build a `SessionTurnDelta` from a `TurnDeltaSnapshot` produced by the signals actor. -/// -/// `feedback_requests_sent` and `last_feedback_request_at` are supplied by the -/// caller (from `FeedbackHeuristics`) because the signals actor does not track -/// feedback state. -/// `request_id` is the prompt/request identifier for this turn. -/// `loc_tracking_enabled` indicates whether the LOC attribution hunk tracker -/// was active for this session. When `false`, LOC delta fields are zeros -/// because the tracker was never spawned — not because no code changed. -pub fn snapshot_to_turn_delta( - snapshot: &crate::session::signals::TurnDeltaSnapshot, - client_type: ClientType, - request_id: Option, - feedback_requests_sent: u32, - last_feedback_request_at: Option>, - loc_tracking_enabled: bool, - turn_duration_ms: Option, - turn_outcome: Option, - model_fingerprint: Option, -) -> SessionTurnDelta { - let metadata = { - let mut metadata = serde_json::Map::new(); - if let Some(mode) = snapshot.start_prompt_mode.as_ref() { - metadata.insert("startPromptMode".to_owned(), serde_json::json!(mode)); + content: &str, + model: Option<&str>, + ) -> Result<(), anyhow::Error> { + let token = self + .bearer_token() + .ok_or_else(|| anyhow::anyhow!("no OAuth session token for feedback submission"))?; + let url = format!("{}/feedback", self.base_url.trim_end_matches('/')); + let payload = FeedbackPayload { + session_id, + content, + version: kigi_version::VERSION, + os: crate::auth::device::device_model(), + model, + }; + let response = self + .http + .post(&url) + .bearer_auth(&token) + .json(&payload) + .timeout(std::time::Duration::from_secs(15)) + .send() + .await?; + let status = response.status(); + if status.is_success() { + tracing::info!(session_id = %session_id, "feedback submitted"); + return Ok(()); } - if let Some(mode) = snapshot.end_prompt_mode.as_ref() { - metadata.insert("endPromptMode".to_owned(), serde_json::json!(mode)); + let body = response.text().await.unwrap_or_default(); + tracing::warn!(status = status.as_u16(), "feedback submission rejected"); + Err(FeedbackApiError { status, body }.into()) + } + + async fn try_refresh_credentials(&self) -> bool { + match &self.bearer { + BearerSource::AuthManager(am) => am.try_recover_unauthorized().await, + BearerSource::Static(_) => false, } - (!metadata.is_empty()).then_some(serde_json::Value::Object(metadata)) - }; - let d = &snapshot.delta; - let c = &snapshot.current; - SessionTurnDelta { - client_type, - turn_number: d.turn_number as i64, - // Deltas - delta_tool_calls: d.delta_tool_calls, - delta_tool_failures: d.delta_tool_failures, - delta_errors: d.delta_errors, - delta_cancellations: d.delta_cancellations, - delta_regenerations: d.delta_regenerations, - delta_compactions: d.delta_compactions, - delta_edit_and_retries: d.delta_edit_and_retries, - delta_positive_ratings: d.delta_positive_ratings, - delta_negative_ratings: d.delta_negative_ratings, - delta_assistant_messages: d.delta_assistant_messages, - delta_long_pauses: d.delta_long_pauses, - delta_successful_tool_uses: d.delta_successful_tool_uses, - // Turn-level snapshot values - consecutive_cancellations: d.consecutive_cancellations as i64, - // Turn-level absolute values - time_to_first_token_ms: d.last_time_to_first_token_ms.map(|v| v as i64), - total_response_time_ms: d.last_total_response_time_ms.map(|v| v as i64), - // Per-turn ITL (delta uses u64, wire type uses i64) - itl_p50_ms: d.last_itl_p50_ms.map(|v| v as i64), - itl_p99_ms: d.last_itl_p99_ms.map(|v| v as i64), - itl_max_ms: d.last_itl_max_ms.map(|v| v as i64), - itl_mean_ms: d.last_itl_mean_ms.map(|v| v as i64), - context_window_usage: c.context_window_usage as i64, - model_id: c.primary_model_id.clone(), - turn_duration_ms, - turn_outcome, - model_fingerprint, - tools_used_this_turn: d.tools_this_turn.clone(), - error_types_this_turn: d.error_types_this_turn.clone(), - tool_outcomes: if d.tool_outcomes_this_turn.is_empty() { - String::new() - } else { - serde_json::to_string(&d.tool_outcomes_this_turn).unwrap_or_default() - }, - // Cumulative totals - cumulative_tool_calls: c.tool_call_count as i64, - cumulative_errors: c.error_count as i64, - session_duration_seconds: c.session_duration_seconds as i64, - total_tokens_before_compaction: c.total_tokens_before_compaction as i64, - metadata, - request_id, - session_start_at: None, // set by caller if available - feedback_requests_sent: feedback_requests_sent as i64, - last_feedback_request_at, - response_tokens: d.response_tokens.map(|v| v as i64), - thinking_tokens: d.thinking_tokens.map(|v| v as i64), - // LOC Attribution - delta_agent_lines_added: d.delta_agent_lines_added, - delta_agent_lines_removed: d.delta_agent_lines_removed, - delta_agent_lines_added_reverted: d.delta_agent_lines_added_reverted, - delta_agent_lines_removed_reverted: d.delta_agent_lines_removed_reverted, - delta_human_lines_added: d.delta_human_lines_added, - delta_human_lines_removed: d.delta_human_lines_removed, - delta_human_lines_added_reverted: d.delta_human_lines_added_reverted, - delta_human_lines_removed_reverted: d.delta_human_lines_removed_reverted, - delta_agent_files_touched: d.delta_agent_files_touched, - delta_human_files_touched: d.delta_human_files_touched, - delta_total_files_touched: d.delta_total_files_touched, - loc_tracking_enabled, } } #[cfg(test)] mod tests { use super::*; + use wiremock::matchers::{bearer_token, body_partial_json, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; - #[test] - fn test_signals_to_update() { - let signals = crate::session::signals::SessionSignals { - turn_count: 10, - user_message_count: 10, - assistant_message_count: 10, - error_count: 1, - tool_failure_count: 0, - cancellation_count: 0, - consecutive_cancellations: 0, - regeneration_count: 1, - edit_and_retry_count: 2, - positive_ratings: 3, - negative_ratings: 1, - long_pauses_count: 4, - has_reverted: false, - compaction_count: 2, - total_tokens_before_compaction: 20_000, - context_window_usage: 50, - tool_call_count: 5, - tools_used: vec!["read_file".to_string(), "search_replace".to_string()], - models_used: vec!["grok-3".to_string()], - primary_model_id: Some("grok-3".to_string()), - session_duration_seconds: 120, - // Latency metrics - avg_time_to_first_token_ms: 150, - avg_response_time_ms: 2500, - min_time_to_first_token_ms: 100, - max_time_to_first_token_ms: 300, - latency_sample_count: 5, - // ITL metrics - itl_p50_ms: Some(45), - itl_p99_ms: Some(180), - itl_max_ms: Some(350), - itl_mean_ms: Some(62), - total_chunk_count: 1200, - itl_sample_count: 15, - itl_digest: None, - itl_sum_ms: 0, - itl_interval_count: 0, - doom_loop_recovery_attempts: 2, - doom_loop_recovery_accepted_after_budget: 1, - doom_loop_recovery_top_trigger: Some("tail_repetition:4@thinking".to_string()), - doom_loop_recovery_aborted_chunks: 421, - ..Default::default() - }; - - let update = signals_to_update(&signals, ClientType::Agent); - - assert_eq!(update.total_turns, Some(10)); - assert_eq!(update.error_count, Some(1)); - assert_eq!(update.compaction_count, Some(2)); - assert_eq!(update.cancellation_count, Some(0)); - assert_eq!(update.consecutive_cancellations, Some(0)); - assert_eq!(update.tool_failure_count, Some(0)); - assert_eq!(update.tool_call_count, Some(5)); - assert_eq!(update.session_duration_seconds, Some(120)); - assert_eq!(update.tools_used.len(), 2); - assert_eq!(update.models_used.len(), 1); - assert_eq!(update.primary_model_id, Some("grok-3".to_string())); - // New counter assertions - assert_eq!(update.edit_and_retry_count, Some(2)); - assert_eq!(update.positive_ratings, Some(3)); - assert_eq!(update.negative_ratings, Some(1)); - assert_eq!(update.long_pauses_count, Some(4)); - // Latency assertions - assert_eq!(update.avg_time_to_first_token_ms, Some(150)); - assert_eq!(update.avg_response_time_ms, Some(2500)); - assert_eq!(update.min_time_to_first_token_ms, Some(100)); - assert_eq!(update.max_time_to_first_token_ms, Some(300)); - assert_eq!(update.latency_sample_count, Some(5)); - // ITL assertions - assert_eq!(update.last_itl_p50_ms, Some(45)); - assert_eq!(update.last_itl_p99_ms, Some(180)); - assert_eq!(update.worst_itl_max_ms, Some(350)); - assert_eq!(update.avg_itl_mean_ms, Some(62)); - assert_eq!(update.total_chunk_count, Some(1200)); - assert_eq!(update.itl_sample_count, Some(15)); - // Doom-loop recovery assertions (legacy detection columns stay null) - assert_eq!(update.doom_loop_recovery_fired, Some(true)); - assert_eq!(update.doom_loop_recovery_attempts, Some(2)); - assert_eq!(update.doom_loop_recovery_accepted_after_budget, Some(1)); - assert_eq!( - update.doom_loop_recovery_top_trigger.as_deref(), - Some("tail_repetition:4@thinking") - ); - assert_eq!(update.doom_loop_recovery_aborted_chunks, Some(421)); - assert_eq!(update.doom_loop_warnings, None); - } - - #[test] - fn test_snapshot_to_turn_delta_includes_prompt_mode_metadata() { - let snapshot = crate::session::signals::TurnDeltaSnapshot { - current: crate::session::signals::SessionSignals::default(), - delta: crate::session::signals::SessionSignalsDelta { - turn_number: 1, - ..Default::default() - }, - start_prompt_mode: Some("plan".to_string()), - end_prompt_mode: Some("agent".to_string()), - turn_input_tokens: 0, - turn_output_tokens: 0, - turn_cached_input_tokens: 0, - }; - - let delta = snapshot_to_turn_delta( - &snapshot, - ClientType::Agent, - Some("request-1".to_string()), - 0, - None, - false, - Some(1500), - Some("completed".to_string()), - Some("fp_test_123".to_string()), - ); - - assert_eq!( - delta.metadata, - Some(serde_json::json!({ - "startPromptMode": "plan", - "endPromptMode": "agent" - })) - ); - assert_eq!(delta.turn_duration_ms, Some(1500)); - assert_eq!(delta.turn_outcome.as_deref(), Some("completed")); - assert_eq!(delta.model_fingerprint.as_deref(), Some("fp_test_123")); - } - - #[test] - fn test_signals_to_update_fresh_session_itl_none() { - // When itl_sample_count == 0, p50/p99 must be None (SQL NULL) - // to preserve the "not yet reported" semantic in the nullable PG columns. - let signals = crate::session::signals::SessionSignals { - turn_count: 1, - user_message_count: 1, - assistant_message_count: 1, - error_count: 0, - tool_failure_count: 0, - cancellation_count: 0, - consecutive_cancellations: 0, - regeneration_count: 0, - edit_and_retry_count: 0, - positive_ratings: 0, - negative_ratings: 0, - long_pauses_count: 0, - has_reverted: false, - compaction_count: 0, - total_tokens_before_compaction: 0, - context_window_usage: 0, - tool_call_count: 0, - tools_used: vec![], - models_used: vec![], - primary_model_id: None, - session_duration_seconds: 10, - avg_time_to_first_token_ms: 0, - avg_response_time_ms: 0, - min_time_to_first_token_ms: 0, - max_time_to_first_token_ms: 0, - latency_sample_count: 0, - // No ITL measured yet - itl_p50_ms: None, - itl_p99_ms: None, - itl_max_ms: None, - itl_mean_ms: None, - total_chunk_count: 0, - itl_sample_count: 0, - itl_digest: None, - itl_sum_ms: 0, - itl_interval_count: 0, - ..Default::default() - }; - - let update = signals_to_update(&signals, ClientType::Agent); - - // p50/p99 must be None when no ITL has been measured - assert_eq!(update.last_itl_p50_ms, None); - assert_eq!(update.last_itl_p99_ms, None); - // max and mean are also None when no ITL has been measured - assert_eq!(update.worst_itl_max_ms, None); - assert_eq!(update.avg_itl_mean_ms, None); - assert_eq!(update.total_chunk_count, Some(0)); - assert_eq!(update.itl_sample_count, Some(0)); - } -} - -#[cfg(test)] -mod forbidden_tests { - use super::*; - use axum::{Router, response::IntoResponse, routing::post}; - use std::net::SocketAddr; - use tokio::net::TcpListener; - - async fn start_server(router: Router) -> (SocketAddr, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let handle = tokio::spawn(async move { - axum::serve(listener, router).await.unwrap(); - }); - (addr, handle) - } - - fn forbidden_handler() -> impl IntoResponse { - (axum::http::StatusCode::FORBIDDEN, "ZDR team") - } - + /// Happy path: POST /feedback with Bearer and the exact kimi-cli body + /// fields succeeds. #[tokio::test] - async fn send_empty_returns_ok_on_403() { - let router = Router::new().route( - "/v1/feedback/requests/{id}/complete", - post(|| async { forbidden_handler() }), - ); - let (addr, _) = start_server(router).await; + async fn submit_feedback_posts_kimi_body_with_bearer() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/feedback")) + .and(bearer_token("tok-123")) + .and(body_partial_json(serde_json::json!({ + "session_id": "sess-1", + "content": "love it", + "version": kigi_version::VERSION, + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .expect(1) + .mount(&server) + .await; - let client = FeedbackClient::with_client( - reqwest::Client::new(), - format!("http://{addr}/v1"), - Some("tok".into()), - ); - let submission: FeedbackSubmission = serde_json::from_value(serde_json::json!({ - "sessionId": "s1", - "clientType": "agent", - "feedbackType": "rating", - })) - .unwrap(); - let result = client.complete_request("req-1", &submission).await; - assert!(result.is_ok(), "403 on send_empty must return Ok"); - } - - #[tokio::test] - async fn send_json_bails_on_403_with_clear_message() { - let router = Router::new().route( - "/v1/feedback/config", - axum::routing::get(|| async { forbidden_handler() }), - ); - let (addr, _) = start_server(router).await; - - let client = FeedbackClient::with_client( - reqwest::Client::new(), - format!("http://{addr}/v1"), - Some("tok".into()), - ); - let result = client.get_feedback_config().await; - let err = result.unwrap_err(); - assert!( - err.to_string().contains("(403)"), - "error must mention 403, got: {err}" - ); - } -} - -/// Auth resolve + 401 recovery tests. -#[cfg(test)] -mod auth_refresh_tests { - use super::*; - use crate::auth::{AuthManager, AuthMode, KimiAuth, KimiCodeConfig}; - use axum::{Router, routing::get}; - use chrono::{Duration, Utc}; - use std::net::SocketAddr; - use std::sync::Arc; - use std::sync::atomic::{AtomicU32, Ordering}; - use tokio::net::TcpListener; - - async fn start_server(router: Router) -> (SocketAddr, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let handle = tokio::spawn(async move { - axum::serve(listener, router).await.unwrap(); - }); - (addr, handle) - } - - /// `send_empty` must sign with the AuthManager bearer. Pre-fix - /// it skipped the resolve and went out unauthenticated. - #[tokio::test] - async fn send_empty_signs_request_with_active_auth() { - let captured = Arc::new(parking_lot::Mutex::new(None::)); - let captured_for_handler = captured.clone(); - let router = Router::new().route( - "/v1/feedback/requests/{id}/complete", - axum::routing::post(move |headers: axum::http::HeaderMap| { - let captured = captured_for_handler.clone(); - async move { - if let Some(auth) = headers.get(axum::http::header::AUTHORIZATION) { - *captured.lock() = Some(auth.to_str().unwrap_or("").to_owned()); - } - axum::http::StatusCode::OK - } - }), - ); - let (addr, _server) = start_server(router).await; - - let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); - am.hot_swap(KimiAuth { - key: "fresh-from-auth-manager".into(), - auth_mode: AuthMode::ApiKey, - create_time: Utc::now(), - user_id: "user-42".into(), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..KimiAuth::test_default() - }); - - let client = FeedbackClient::new( - format!("http://{addr}/v1"), - Some("STALE-build-time-token".into()), - ) - .with_auth_manager(am.clone()); - - let submission: FeedbackSubmission = serde_json::from_value(serde_json::json!({ - "sessionId": "s1", - "clientType": "agent", - "feedbackType": "rating", - })) - .unwrap(); + let client = FeedbackClient::with_static_token(server.uri(), "tok-123"); client - .complete_request("req-1", &submission) + .submit_feedback("sess-1", "love it", Some("kimi-k2")) .await - .expect("send_empty must succeed when authenticated"); - - let sent = captured.lock().clone().expect("server saw the request"); - assert_eq!(sent, "Bearer fresh-from-auth-manager"); + .expect("feedback should succeed"); } - /// Outgoing bearer must match `AuthManager.current()`, not the - /// FeedbackClient's build-time snapshot. + /// The body carries the `os` and `model` fields kimi-cli sends. #[tokio::test] - async fn feedback_client_uses_active_auth_for_each_request() { - let captured = Arc::new(parking_lot::Mutex::new(None::)); - let captured_for_handler = captured.clone(); - let router = Router::new().route( - "/v1/feedback/config", - get(move |headers: axum::http::HeaderMap| { - let captured = captured_for_handler.clone(); - async move { - if let Some(auth) = headers.get(axum::http::header::AUTHORIZATION) { - *captured.lock() = Some(auth.to_str().unwrap_or("").to_owned()); - } - axum::Json(serde_json::json!({})) - } - }), - ); - let (addr, _server) = start_server(router).await; + async fn submit_feedback_includes_os_and_model() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/feedback")) + .and(body_partial_json(serde_json::json!({ + "model": "kimi-k2", + "os": crate::auth::device::device_model(), + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .expect(1) + .mount(&server) + .await; - let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); - let fresh = KimiAuth { - key: "fresh-from-auth-manager".into(), - auth_mode: AuthMode::ApiKey, - create_time: Utc::now(), - user_id: "user-42".into(), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..KimiAuth::test_default() - }; - am.hot_swap(fresh); - - let client = FeedbackClient::new( - format!("http://{addr}/v1"), - Some("STALE-build-time-token".into()), - ) - .with_auth_manager(am.clone()); - - let _ = client.get_feedback_config().await; - - let sent = captured.lock().clone().expect("server saw the request"); - assert_eq!( - sent, "Bearer fresh-from-auth-manager", - "outgoing bearer must come from AuthManager (not the build-time snapshot)" - ); + let client = FeedbackClient::with_static_token(server.uri(), "tok"); + client + .submit_feedback("sess-2", "hi", Some("kimi-k2")) + .await + .expect("feedback should succeed"); } - /// Counts refresh() calls -- proves disk-reload short-circuits - /// before the IdP is hit. - struct CountingRefresher { - calls: Arc, - } - - #[async_trait::async_trait] - impl crate::auth::refresh::TokenRefresher for CountingRefresher { - async fn refresh( - &self, - _reason: crate::auth::refresh::RefreshReason, - ) -> crate::auth::refresh::RefreshOutcome { - self.calls.fetch_add(1, Ordering::SeqCst); - crate::auth::refresh::RefreshOutcome::Success(Box::new(KimiAuth { - key: "fresh-from-refresher".into(), - auth_mode: AuthMode::OAuth, - create_time: Utc::now(), - user_id: "user-42".into(), - refresh_token: Some("rt-fresh".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..KimiAuth::test_default() - })) - } - } - - /// Disk has a fresher RT than memory; recovery must succeed via - /// reload without calling the refresher. + /// Auth failure: a 401 surfaces as `FeedbackApiError::is_unauthorized` + /// (static-token clients cannot refresh, so no retry). #[tokio::test] - async fn try_refresh_credentials_picks_up_disk_rotation_without_hitting_idp() { - let dir = tempfile::tempdir().unwrap(); - let cfg = KimiCodeConfig::default(); - let scope = cfg.auth_scope(); - let am = Arc::new(AuthManager::new(dir.path(), cfg)); + async fn submit_feedback_auth_failure_is_typed_401() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/feedback")) + .respond_with( + ResponseTemplate::new(401) + .set_body_json(serde_json::json!({"error": "invalid token"})), + ) + .expect(1) + .mount(&server) + .await; - // In-memory: stale token (the one the server rejected). - am.hot_swap(KimiAuth { - key: "stale-rejected".into(), - auth_mode: AuthMode::OAuth, - create_time: Utc::now() - Duration::hours(2), - user_id: "user-42".into(), - refresh_token: Some("rt-stale".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..KimiAuth::test_default() - }); - - // Disk: a sibling already rotated to a fresh token. - let disk_auth = KimiAuth { - key: "fresh-from-sibling-on-disk".into(), - auth_mode: AuthMode::OAuth, - create_time: Utc::now(), - user_id: "user-42".into(), - refresh_token: Some("rt-fresh".into()), - expires_at: Some(Utc::now() + Duration::hours(1)), - ..KimiAuth::test_default() - }; - let mut store = std::collections::BTreeMap::new(); - store.insert(scope, disk_auth); - let json = serde_json::to_string_pretty(&store).unwrap(); - std::fs::write(dir.path().join("auth.json"), json).unwrap(); - - let calls = Arc::new(AtomicU32::new(0)); - let refresher = Arc::new(CountingRefresher { - calls: calls.clone(), - }); - am.set_refresher(refresher); - - let client = FeedbackClient::new("http://example/v1", Some("stale-rejected".into())) - .with_auth_manager(am.clone()); - - let recovered = client.try_refresh_credentials().await; - assert!(recovered, "recovery must succeed via disk reload"); - assert_eq!( - calls.load(Ordering::SeqCst), - 0, - "refresher must NOT be called when disk already holds a fresh token \ - (proves we routed through unauthorized_recovery, not direct refresh)" - ); - assert_eq!( - am.current().unwrap().key, - "fresh-from-sibling-on-disk", - "AuthManager's current token must be the disk-loaded one after recovery" - ); - } - - /// LegacySession -> `ServerRejectedNoRecovery` -> `false` - /// (caller stops retrying, doesn't loop on a no-op refresher). - #[tokio::test] - async fn try_refresh_credentials_returns_false_on_terminal_failure() { - let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); - - // LegacySession: no refresh_token, no recovery possible. - am.hot_swap(KimiAuth { - key: "legacy-rejected".into(), - auth_mode: AuthMode::OAuth, - create_time: Utc::now() - Duration::days(60), - user_id: "user-42".into(), - ..KimiAuth::test_default() - }); - - let client = FeedbackClient::new("http://example/v1", Some("legacy-rejected".into())) - .with_auth_manager(am); - - let recovered = client.try_refresh_credentials().await; - assert!( - !recovered, - "LegacySession must surface ServerRejectedNoRecovery as `false`, \ - not loop on a refresher that can't help" - ); + let client = FeedbackClient::with_static_token(server.uri(), "bad-token"); + let err = client + .submit_feedback("sess-3", "hello", None) + .await + .expect_err("401 must fail"); + let api_err = err + .downcast_ref::() + .expect("typed FeedbackApiError"); + assert!(api_err.is_unauthorized()); + assert!(api_err.body.contains("invalid token")); } } diff --git a/crates/codegen/kigi-shell/src/agent/handlers/mod.rs b/crates/codegen/kigi-shell/src/agent/handlers/mod.rs index 0c4b799..4c225e6 100644 --- a/crates/codegen/kigi-shell/src/agent/handlers/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/handlers/mod.rs @@ -1,3 +1,2 @@ pub(crate) mod model_switch; pub(crate) mod session; -pub(crate) mod workspaces; diff --git a/crates/codegen/kigi-shell/src/agent/handlers/session.rs b/crates/codegen/kigi-shell/src/agent/handlers/session.rs index a7a9f14..8fd60ca 100644 --- a/crates/codegen/kigi-shell/src/agent/handlers/session.rs +++ b/crates/codegen/kigi-shell/src/agent/handlers/session.rs @@ -271,19 +271,10 @@ async fn handle_session_list( // (never union) so every list surface is conversations-only. let req = unified_list::parse_list_req(args.params.get()) .map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?; - tracing::debug!( - chat_mode_forced_kind = crate::agent::chat_modes::process_chat_mode_enabled(), - "session/list" - ); + tracing::debug!("session/list"); let registry_client = agent.session_registry_client(); - let conversations_client = agent.conversations_client(); - let result = unified_list::build_unified_list( - registry_client.as_ref(), - conversations_client.as_ref(), - req, - ) - .await; + let result = unified_list::build_unified_list(registry_client.as_ref(), req).await; ExtMethodResult::success(unified_list::ext_list_response(result)) .to_ext_response() diff --git a/crates/codegen/kigi-shell/src/agent/handlers/workspaces.rs b/crates/codegen/kigi-shell/src/agent/handlers/workspaces.rs deleted file mode 100644 index 93a1411..0000000 --- a/crates/codegen/kigi-shell/src/agent/handlers/workspaces.rs +++ /dev/null @@ -1,174 +0,0 @@ -use agent_client_protocol::{self as acp}; -use serde::{Deserialize, Serialize}; - -use super::super::mvp_agent::MvpAgent; -use crate::remote::{ListWorkspacesPage, WsError, WsQuery}; -use crate::session::ExtMethodResult; - -const DEFAULT_PAGE_SIZE: i64 = 50; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -struct WorkspacesListRequest { - #[serde(default)] - page_size: Option, - #[serde(default)] - page_token: Option, - #[serde(default)] - query: Option, - #[serde(default)] - kind: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct WorkspaceRow { - id: String, - name: String, - #[serde(skip_serializing_if = "Option::is_none")] - kind: Option, - #[serde(skip_serializing_if = "Option::is_none")] - create_time: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct WorkspacesListResponse { - workspaces: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - next_page_token: Option, - #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] - meta: Option, -} - -#[derive(Debug, Serialize)] -struct WorkspacesMeta { - #[serde(rename = "x.ai/partial")] - partial: PartialInfo, -} - -#[derive(Debug, Serialize)] -struct PartialInfo { - workspaces: bool, - reason: &'static str, -} - -pub async fn handle( - agent: &MvpAgent, - args: &acp::ExtRequest, -) -> Result { - let req: WorkspacesListRequest = serde_json::from_str(args.params.get()) - .map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?; - - let q = WsQuery { - // Clamp to a sane positive page size: a missing, zero, or negative - // `pageSize` falls back to the default rather than being forwarded - // verbatim to `/rest/workspaces`. - page_size: match req.page_size { - Some(n) if n > 0 => n, - _ => DEFAULT_PAGE_SIZE, - }, - page_token: req.page_token, - query: req.query, - kind: req.kind, - }; - - let response = match agent.workspaces_client().list_workspaces(&q).await { - Ok(page) => success_response(page), - Err(WsError::NoOauth) => degraded_response("no_oauth"), - Err(e) => { - // Degrade to a partial result, but don't silently swallow the - // cause — log it so field failures are diagnosable. - tracing::warn!("workspaces/list fetch failed: {e}"); - degraded_response("error") - } - }; - - ExtMethodResult::success(response) - .to_ext_response() - .map_err(|e| acp::Error::internal_error().data(e.to_string())) -} - -fn success_response(page: ListWorkspacesPage) -> WorkspacesListResponse { - WorkspacesListResponse { - workspaces: page - .workspaces - .into_iter() - .map(|w| WorkspaceRow { - id: w.workspace_id, - name: w.name, - kind: w.kind, - create_time: w.create_time, - }) - .collect(), - next_page_token: page.next_page_token, - meta: None, - } -} - -fn degraded_response(reason: &'static str) -> WorkspacesListResponse { - WorkspacesListResponse { - workspaces: Vec::new(), - next_page_token: None, - meta: Some(WorkspacesMeta { - partial: PartialInfo { - workspaces: true, - reason, - }, - }), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::remote::Workspace; - - #[test] - fn request_parses_camelcase_and_defaults_page_size() { - let req: WorkspacesListRequest = - serde_json::from_value(serde_json::json!({})).expect("empty params parse"); - assert!(req.page_size.is_none()); - - let req: WorkspacesListRequest = serde_json::from_value(serde_json::json!({ - "pageSize": 10, - "pageToken": "tok", - "query": "gpu", - "kind": "WORKSPACE_KIND_IMAGINE" - })) - .expect("full params parse"); - assert_eq!(req.page_size, Some(10)); - assert_eq!(req.page_token.as_deref(), Some("tok")); - assert_eq!(req.query.as_deref(), Some("gpu")); - assert_eq!(req.kind.as_deref(), Some("WORKSPACE_KIND_IMAGINE")); - } - - #[test] - fn success_response_projects_grok_workspace_fields() { - let page = ListWorkspacesPage { - workspaces: vec![Workspace { - workspace_id: "ws_1".into(), - name: "Research".into(), - create_time: Some("2026-06-18T17:30:00Z".into()), - kind: Some("WORKSPACE_KIND_IMAGINE".into()), - }], - next_page_token: Some("tok2".into()), - }; - let value = serde_json::to_value(success_response(page)).unwrap(); - assert_eq!(value["workspaces"][0]["id"], "ws_1"); - assert_eq!(value["workspaces"][0]["name"], "Research"); - assert_eq!(value["workspaces"][0]["kind"], "WORKSPACE_KIND_IMAGINE"); - assert_eq!(value["workspaces"][0]["createTime"], "2026-06-18T17:30:00Z"); - assert_eq!(value["nextPageToken"], "tok2"); - assert!(value.get("_meta").is_none()); - } - - #[test] - fn degraded_response_carries_partial_reason() { - let value = serde_json::to_value(degraded_response("no_oauth")).unwrap(); - assert_eq!(value["workspaces"].as_array().unwrap().len(), 0); - assert!(value.get("nextPageToken").is_none()); - assert_eq!(value["_meta"]["x.ai/partial"]["workspaces"], true); - assert_eq!(value["_meta"]["x.ai/partial"]["reason"], "no_oauth"); - } -} diff --git a/crates/codegen/kigi-shell/src/agent/init.rs b/crates/codegen/kigi-shell/src/agent/init.rs index e96a1c8..2deed31 100644 --- a/crates/codegen/kigi-shell/src/agent/init.rs +++ b/crates/codegen/kigi-shell/src/agent/init.rs @@ -73,27 +73,6 @@ fn resolve_config(cfg: &AgentConfig, auth_manager: &AuthManager) -> AgentConfig tracing::info!(field = %e.path, value = %e.value, source = %e.source, "policy override"); } - // Fallback: if the client didn't pre-supply remote settings, fetch them - // now so remote-settings-gated features work regardless of which client - // spawned us. Clients that already call `start_early_prefetch()` and - // thread the result into `cfg.remote_settings` skip this entirely. - if cfg.remote_settings.is_none() - && let Some(handle) = - crate::agent::models::start_early_prefetch(Some(cfg.kimi_code_config.clone())) - { - match handle.join() { - Ok(result) => { - cfg.remote_settings = result.settings; - crate::util::config::set_remote_campaigns_from_settings( - cfg.remote_settings.as_ref(), - ); - tracing::info!("remote_settings fetched as shell-level fallback"); - } - Err(_) => { - tracing::warn!("remote_settings fallback prefetch thread panicked"); - } - } - } crate::util::config::sync_campaign_fields(&mut cfg); crate::agent::config::apply_remote_settings_side_effects(cfg.remote_settings.as_ref()); diff --git a/crates/codegen/kigi-shell/src/agent/mod.rs b/crates/codegen/kigi-shell/src/agent/mod.rs index f3ee54a..3cd6a96 100644 --- a/crates/codegen/kigi-shell/src/agent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/mod.rs @@ -5,11 +5,12 @@ pub mod chat_modes; pub mod config; pub mod config_model_override_parse; mod ext_parsers; -pub mod feedback_client; +pub(crate) mod feedback_client; pub mod folder_trust; pub(crate) mod handlers; pub mod init; pub mod models; +pub(crate) mod models_fetch; pub mod mvp_agent; pub(crate) mod proxy; pub(crate) mod restore_code; diff --git a/crates/codegen/kigi-shell/src/agent/models.rs b/crates/codegen/kigi-shell/src/agent/models.rs index e35b6a6..343a7e5 100644 --- a/crates/codegen/kigi-shell/src/agent/models.rs +++ b/crates/codegen/kigi-shell/src/agent/models.rs @@ -10,8 +10,8 @@ use chrono::{DateTime, Duration as ChronoDuration, Utc}; use indexmap::IndexMap; use crate::agent::config::{self, ModelEntry, resolve_credentials, sampling_config_for_model}; +use crate::agent::models_fetch::{FetchModelsResult, fetch_models_blocking}; use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig}; -use crate::remote::{FetchModelsResult, fetch_models_blocking}; use crate::sampling::SamplerConfig as SamplingConfig; use globset::{Glob, GlobSet, GlobSetBuilder}; use kigi_sampling_types::{ReasoningEffort, ReasoningEffortOption}; @@ -290,7 +290,7 @@ impl ModelsManager { cache .load_fresh( &fetch_auth.cache_auth_method(), - &crate::remote::models_fetch_origin( + &crate::agent::models_fetch::models_fetch_origin( &cfg.endpoints, fetch_auth, has_session, @@ -1035,11 +1035,6 @@ impl ModelsManager { current_model, credentials, config.endpoints.alpha_test_key.clone(), - config.client_version.clone(), - crate::managed_config::resolve_deployment_id( - config.endpoints.deployment_key.as_deref(), - ), - None, ) } @@ -1053,7 +1048,12 @@ impl ModelsManager { let fetch_auth = *self.inner.fetch_auth.read(); let has_oauth = self.inner.auth_manager.current_or_expired().is_some(); let platform_keys = PlatformApiKeys::resolve(&platforms); - crate::remote::models_fetch_origin(&endpoints, fetch_auth, has_oauth, &platform_keys) + crate::agent::models_fetch::models_fetch_origin( + &endpoints, + fetch_auth, + has_oauth, + &platform_keys, + ) } fn try_load_cache(&self) -> bool { @@ -1327,7 +1327,7 @@ struct ModelsCache { #[serde(default, skip_serializing_if = "Option::is_none")] auth_method: Option, /// Models-list URL this catalog was fetched from - /// ([`crate::remote::models_list_url`]). Compared on load so a cache + /// ([`crate::agent::models_fetch::models_fetch_origin`]). Compared on load so a cache /// written against one backend is a miss for another: entries embed /// absolute `base_url`s, so adopting a foreign-origin cache silently /// re-points inference (the windows lifecycle e2e failed exactly this @@ -1575,43 +1575,6 @@ pub(crate) fn prefetch_models_blocking( .models } -/// Blocking models + `/v1/settings` prefetch pair, shared by the early -/// prefetch thread and the leader's startup phase so the settings gate lives -/// once. The remote_fetch knob is resolved a single time so the two fetch -/// decisions cannot disagree mid-startup. -pub(crate) fn prefetch_models_and_settings_blocking( - endpoints: &config::EndpointsConfig, - auth: Option<&KimiAuth>, - fetch_auth: ModelFetchAuth, - platform_keys: &PlatformApiKeys, -) -> ( - Option>, - Option, -) { - let remote_fetch_enabled = crate::util::config::resolve_remote_fetch_enabled(); - let models = prefetch_models_blocking_gated( - endpoints, - auth, - fetch_auth, - platform_keys, - remote_fetch_enabled, - ) - .models; - // Settings need a subscription session; skip for API-key-only setups. - let settings = match auth { - Some(auth) if remote_fetch_enabled => { - let _timer = crate::instrumentation_timer!("startup.early_settings_fetch"); - crate::remote::fetch_settings_blocking( - &endpoints.proxy_url(), - auth, - endpoints.alpha_test_key.as_deref(), - ) - } - _ => None, - }; - (models, settings) -} - /// `remote_fetch_enabled` is a parameter so the pair helper above resolves the /// knob once for both halves. fn prefetch_models_blocking_gated( @@ -1624,8 +1587,12 @@ fn prefetch_models_blocking_gated( let cache_auth = fetch_auth.cache_auth_method(); // Same fetch plan the network path below executes — the cache is only // valid for it. - let cache_origin = - crate::remote::models_fetch_origin(endpoints, fetch_auth, auth.is_some(), platform_keys); + let cache_origin = crate::agent::models_fetch::models_fetch_origin( + endpoints, + fetch_auth, + auth.is_some(), + platform_keys, + ); let cache = ModelsCacheManager::new(); if let Some(cached) = cache.load_fresh(&cache_auth, &cache_origin) { tracing::info!( @@ -1703,10 +1670,9 @@ fn stale_cache_or_failure( ModelsFetchOutcome::failed(oauth_unauthorized) } -/// Startup prefetch result: models + remote settings. +/// Startup prefetch result: the model catalog, when a fetch plan existed. pub struct EarlyPrefetchResult { pub models: Option>, - pub settings: Option, } /// Handle for a startup prefetch thread. @@ -1743,7 +1709,7 @@ fn resolve_prefetch_env_with_auth(auth: Option) -> Option /// `has_custom_endpoint()` (which otherwise forces the prefetch to run): the /// explicit off switch must hold even when a stray login, a platform API key, /// or a `deployment_key` would re-arm the prefetch — and with it the -/// `/v1/settings` fetch and the deployment-config sync on the prefetch thread. +/// deployment-config sync on the prefetch thread. /// /// PRD F2 acceptance: a moonshot API key alone (no subscription login) must /// arm the prefetch so the catalog syncs on startup. @@ -1754,7 +1720,7 @@ fn resolve_prefetch_env_from_parts( remote_fetch_enabled: bool, ) -> Option { if !remote_fetch_enabled { - tracing::info!("startup model/settings prefetch skipped: remote_fetch disabled"); + tracing::info!("startup model prefetch skipped: remote_fetch disabled"); return None; } @@ -1779,7 +1745,7 @@ fn resolve_prefetch_env(kimi_code_config: Option) -> Option) -> Option EarlyPrefetchHandle { let mut timer = crate::instrumentation_timer!("startup.early_prefetch"); let proxy_endpoint = env.endpoints.proxy_url(); timer.with_field("endpoint", proxy_endpoint.as_str()); - let (models, settings) = prefetch_models_and_settings_blocking( + let models = prefetch_models_blocking( &env.endpoints, env.auth.as_ref(), env.model_fetch_auth, @@ -1824,7 +1790,7 @@ fn spawn_prefetch_thread(env: PrefetchEnv) -> EarlyPrefetchHandle { let _ = rt.block_on(crate::managed_config::sync()); } - EarlyPrefetchResult { models, settings } + EarlyPrefetchResult { models } }) } @@ -3868,7 +3834,7 @@ mod tests { fn proxied_endpoints(server_uri: &str) -> config::EndpointsConfig { config::EndpointsConfig { - cli_chat_proxy_base_url: Some(server_uri.to_string()), + coding_api_base_url: Some(server_uri.to_string()), models_base_url: None, models_list_url: None, ..config::EndpointsConfig::default() @@ -3900,7 +3866,7 @@ mod tests { ..KimiAuth::test_default() }; let result = tokio::task::spawn_blocking(move || { - crate::remote::fetch_models_blocking( + crate::agent::models_fetch::fetch_models_blocking( &endpoints, Some(&auth), ModelFetchAuth::Platforms, @@ -3969,7 +3935,12 @@ mod tests { let endpoints = config::EndpointsConfig::default(); let keys = PlatformApiKeys::test_keys(Some("sk-cn-secret"), None); let result = tokio::task::spawn_blocking(move || { - crate::remote::fetch_models_blocking(&endpoints, None, ModelFetchAuth::Platforms, &keys) + crate::agent::models_fetch::fetch_models_blocking( + &endpoints, + None, + ModelFetchAuth::Platforms, + &keys, + ) }) .await .unwrap() @@ -4062,7 +4033,7 @@ mod tests { auth_manager.set_refresher(Arc::new(SwapRefresher)); let mut cfg = config::Config::default(); - cfg.endpoints.cli_chat_proxy_base_url = Some(server.uri()); + cfg.endpoints.coding_api_base_url = Some(server.uri()); let mgr = ModelsManager::new( None, IndexMap::new(), @@ -4120,8 +4091,12 @@ mod tests { assert!(bundled.contains_key("moonshot-ai/kimi-k2-turbo-preview")); // 2. A STALE cache for the same fetch plan is served on sync failure. - let origin = - crate::remote::models_fetch_origin(&endpoints, ModelFetchAuth::Platforms, true, &keys); + let origin = crate::agent::models_fetch::models_fetch_origin( + &endpoints, + ModelFetchAuth::Platforms, + true, + &keys, + ); let cache = ModelsCacheManager::new(); let stale = ModelsCache { fetched_at: Utc::now() - ChronoDuration::seconds(86_400), @@ -4203,7 +4178,7 @@ mod tests { let _cache = EnvGuard::set("KIGI_MODELS_CACHE_DIR", cache_dir.path().to_str().unwrap()); let endpoints = proxied_endpoints("http://127.0.0.1:9"); // Cache written when a moonshot key was ALSO configured... - let with_key_origin = crate::remote::models_fetch_origin( + let with_key_origin = crate::agent::models_fetch::models_fetch_origin( &endpoints, ModelFetchAuth::Platforms, true, diff --git a/crates/codegen/kigi-shell/src/agent/models_fetch.rs b/crates/codegen/kigi-shell/src/agent/models_fetch.rs new file mode 100644 index 0000000..95fe5a5 --- /dev/null +++ b/crates/codegen/kigi-shell/src/agent/models_fetch.rs @@ -0,0 +1,1084 @@ +//! Model catalog fetch (PRD F4). +//! +//! Fetches the model catalog with `GET {base}/models` per enabled platform +//! (the subscription platform via the OAuth session, the open platforms via +//! their API keys), plus the custom-endpoint OpenAI-compatible listing path. +//! +//! This is the sole surviving network surface relocated out of the deleted +//! xAI-proxy backend client (`remote/`); it talks only to the configured +//! Kimi/Moonshot model endpoints, never to a proxy backend. +use crate::auth::KimiAuth; +use indexmap::IndexMap; +use serde::Deserialize; + +/// Errors from a model-catalog fetch. +#[derive(Debug, thiserror::Error)] +pub(crate) enum BackendError { + #[error("Network error: {0}")] + Network(#[from] reqwest::Error), + #[error("Request failed: {status} - {body}")] + RequestFailed { status: u16, body: String }, + #[error("Auth error: {0}")] + Auth(String), +} +pub(crate) const DEFAULT_CONTEXT_WINDOW: u64 = 256_000; +#[derive(Debug, Deserialize)] +struct ModelsResponse { + data: Vec, +} +/// The models-fetch origin key for this endpoints/auth shape. Used as the +/// models disk-cache origin: cached entries embed absolute `base_url`s from +/// the backend(s) that served them, so a catalog fetched against one fetch +/// plan (env override, different set of platform credentials, a test's mock +/// server) must be a cache miss for any other. Encodes URLs and enabled +/// platform NAMES only — never credential values. +pub(crate) fn models_fetch_origin( + endpoints: &crate::agent::config::EndpointsConfig, + fetch_auth: crate::agent::models::ModelFetchAuth, + has_oauth: bool, + platform_keys: &crate::agent::models::PlatformApiKeys, +) -> String { + match fetch_auth { + crate::agent::models::ModelFetchAuth::CustomEndpoint => endpoints.resolve_models_list_url(), + crate::agent::models::ModelFetchAuth::Platforms => { + let parts: Vec = enabled_platforms(has_oauth, platform_keys) + .into_iter() + .map(|p| format!("{}={}", p.as_str(), platform_models_url(p, endpoints))) + .collect(); + format!("platforms[{}]", parts.join(";")) + } + } +} +/// The platforms with usable credentials, in registry order (kimi-code first +/// so "default model = first list item" favors the subscription). +fn enabled_platforms( + has_oauth: bool, + platform_keys: &crate::agent::models::PlatformApiKeys, +) -> Vec { + kigi_models::PlatformId::ALL + .into_iter() + .filter(|p| { + if p.uses_oauth() { + has_oauth + } else { + platform_keys.key_for(*p).is_some() + } + }) + .collect() +} +/// `{base}/models` for one platform. The subscription platform resolves its +/// base through the endpoints config (`coding_api_base_url` override, +/// else `KIGI_CODE_BASE_URL` / production default via kigi-env); the open +/// platforms use their fixed bases. +fn platform_models_url( + platform: kigi_models::PlatformId, + endpoints: &crate::agent::config::EndpointsConfig, +) -> String { + let base = if platform.uses_oauth() { + endpoints.proxy_url() + } else { + platform.base_url() + }; + format!("{}/models", base.trim_end_matches('/')) +} +/// Fetch result: model entries + optional etag from the subscription platform. +pub struct FetchModelsResult { + pub models: Vec, + pub etag: Option, + /// The OAuth platform answered 401. The async layer forces a token + /// refresh and retries once (port of kimi-cli `refresh_managed_models`). + pub oauth_unauthorized: bool, +} +/// Fetch the model catalog (PRD F4). +/// +/// - Custom endpoint mode (`KIGI_MODELS_BASE_URL` / `models_list_url`): a +/// single OpenAI-compatible listing fetched with the BYOK key or session +/// bearer, parsed leniently ([`parse_remote_model_value`]). +/// - Otherwise, the fixed platform registry: `GET {base}/models` with +/// `Authorization: Bearer ` per enabled platform, +/// parsed per the F4 wire contract with capability derivation and the +/// `kimi-k` prefix filter for the open platforms. +/// +/// Succeeds when at least one platform delivers; per-platform failures are +/// logged (status codes only, never credentials). +pub(crate) fn fetch_models_blocking( + endpoints: &crate::agent::config::EndpointsConfig, + auth: Option<&KimiAuth>, + fetch_auth: crate::agent::models::ModelFetchAuth, + platform_keys: &crate::agent::models::PlatformApiKeys, +) -> Result { + match fetch_auth { + crate::agent::models::ModelFetchAuth::CustomEndpoint => { + fetch_custom_endpoint_models_blocking(endpoints, auth) + } + crate::agent::models::ModelFetchAuth::Platforms => { + fetch_platform_models_blocking(endpoints, auth, platform_keys) + } + } +} +fn fetch_custom_endpoint_models_blocking( + endpoints: &crate::agent::config::EndpointsConfig, + auth: Option<&KimiAuth>, +) -> Result { + let client = crate::http::shared_blocking_client(); + let url = endpoints.resolve_models_list_url(); + let inference_base_url = endpoints.resolve_inference_base_url(); + tracing::info!("Fetching models from custom endpoint {}", url); + let api_key = crate::agent::auth_method::read_xai_api_key_env() + .or_else(|_| { + auth.map(|a| a.key.clone()) + .ok_or(std::env::VarError::NotPresent) + }) + .map_err(|_| { + BackendError::Auth("No API key for custom models endpoint. Set XAI_API_KEY.".into()) + })?; + let request = client + .get(&url) + .header("Authorization", format!("Bearer {}", api_key)); + let response = request.send()?; + if !response.status().is_success() { + let status = response.status().as_u16(); + let body = response.text().unwrap_or_default(); + tracing::warn!("Failed to fetch models: {} - {}", status, body); + return Err(BackendError::RequestFailed { status, body }); + } + let etag = response + .headers() + .get("etag") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()); + let models_response: ModelsResponse = response.json()?; + tracing::info!("Fetched {} models from {}", models_response.data.len(), url); + let mut models = Vec::with_capacity(models_response.data.len()); + for (idx, value) in models_response.data.into_iter().enumerate() { + match parse_remote_model_value(&value, &inference_base_url) { + Some(model) => models.push(model), + None => { + tracing::warn!( + "Skipping model at index {}: missing required field ('model' or 'context_window') or invalid types", + idx + ) + } + } + } + Ok(FetchModelsResult { + models, + etag, + oauth_unauthorized: false, + }) +} +/// Registry fetch across all platforms with usable credentials. +fn fetch_platform_models_blocking( + endpoints: &crate::agent::config::EndpointsConfig, + auth: Option<&KimiAuth>, + platform_keys: &crate::agent::models::PlatformApiKeys, +) -> Result { + let enabled = enabled_platforms(auth.is_some(), platform_keys); + if enabled.is_empty() { + return Err(BackendError::Auth( + "No platform credentials: log in with `kigi login` or configure a moonshot API key \ + (KIGI_MOONSHOT_API_KEY or [platforms.*] in ~/.kigi/config.toml)." + .into(), + )); + } + + let mut models = Vec::new(); + let mut etag = None; + let mut oauth_unauthorized = false; + let mut successes = 0usize; + let mut last_error: Option = None; + for platform in &enabled { + let bearer = if platform.uses_oauth() { + auth.map(|a| a.key.clone()) + .expect("enabled_platforms gated on auth presence") + } else { + platform_keys + .key_for(*platform) + .expect("enabled_platforms gated on key presence") + .to_owned() + }; + match fetch_one_platform_models(*platform, endpoints, &bearer) { + Ok((platform_models, platform_etag)) => { + tracing::info!( + platform = platform.as_str(), + count = platform_models.len(), + "platform models fetch succeeded" + ); + successes += 1; + if platform.uses_oauth() { + etag = platform_etag; + } + models.extend(platform_models); + } + Err(e) => { + if platform.uses_oauth() + && matches!(&e, BackendError::RequestFailed { status: 401, .. }) + { + oauth_unauthorized = true; + } + tracing::warn!( + platform = platform.as_str(), + error = %e, + "platform models fetch failed" + ); + last_error = Some(e); + } + } + } + + if successes == 0 { + // All enabled platforms failed. When the failure includes an OAuth + // 401, return `Ok` with the flag set (and no models) so the async + // layer can force a token refresh and retry — an `Err` would drop + // the signal. Non-401 failures propagate as the last error. + if oauth_unauthorized { + return Ok(FetchModelsResult { + models: Vec::new(), + etag: None, + oauth_unauthorized: true, + }); + } + return Err(last_error.unwrap_or_else(|| { + BackendError::Auth("no platform models fetch was attempted".into()) + })); + } + Ok(FetchModelsResult { + models, + etag, + oauth_unauthorized, + }) +} +/// `GET {base}/models` for one platform (PRD F4 wire contract): +/// `Authorization: Bearer ` → `{data:[{id, context_length, +/// supports_reasoning, supports_image_in, supports_video_in, display_name?}]}`. +/// Applies the platform's `kimi-k` prefix filter and capability derivation, +/// and keys each entry `{platform_id}/{model_id}`. +fn fetch_one_platform_models( + platform: kigi_models::PlatformId, + endpoints: &crate::agent::config::EndpointsConfig, + bearer: &str, +) -> Result<(Vec, Option), BackendError> { + let client = crate::http::shared_blocking_client(); + let url = platform_models_url(platform, endpoints); + tracing::info!(platform = platform.as_str(), url = %url, "fetching platform models"); + let response = client + .get(&url) + .header("Authorization", format!("Bearer {}", bearer)) + .send()?; + if !response.status().is_success() { + let status = response.status().as_u16(); + let body = response.text().unwrap_or_default(); + return Err(BackendError::RequestFailed { status, body }); + } + let etag = response + .headers() + .get("etag") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()); + let listing: kigi_models::WireModelsResponse = response.json()?; + let total = listing.data.len(); + let filtered = kigi_models::filter_allowed_models(platform, listing.data); + if filtered.len() != total { + tracing::info!( + platform = platform.as_str(), + total, + kept = filtered.len(), + "applied platform model-prefix filter" + ); + } + let base_url = if platform.uses_oauth() { + endpoints.proxy_url() + } else { + platform.base_url() + }; + let models = filtered + .into_iter() + .map(|wire| platform_wire_model_to_entry(platform, wire, &base_url)) + .collect(); + Ok((models, etag)) +} +/// Map one F4 wire model to a catalog entry config. +/// +/// SECURITY: the entry carries only env-var NAMES (`env_key`) for the open +/// platforms — never key values — because raw fetched entries are persisted +/// to the models disk cache. Config-file keys are stamped in-memory later by +/// `resolve_model_list`'s platform-credentials layer. +fn platform_wire_model_to_entry( + platform: kigi_models::PlatformId, + wire: kigi_models::WireModel, + base_url: &str, +) -> crate::agent::config::ModelEntryConfig { + let capabilities = wire.capabilities(); + let context_window = std::num::NonZeroU64::new(wire.context_length).unwrap_or_else(|| { + tracing::debug!( + model = %wire.id, + default = DEFAULT_CONTEXT_WINDOW, + "platform model missing context_length; using default" + ); + std::num::NonZeroU64::new(DEFAULT_CONTEXT_WINDOW).expect("non-zero") + }); + let env_key = (!platform.uses_oauth()) + .then(|| crate::agent::config::EnvKeys::new(platform.api_key_env_names().iter().copied())); + crate::agent::config::ModelEntryConfig { + id: Some(platform.managed_model_key(&wire.id)), + name: Some(wire.display_name.clone().unwrap_or_else(|| wire.id.clone())), + model: wire.id, + base_url: base_url.to_owned(), + description: None, + max_completion_tokens: None, + temperature: None, + top_p: None, + api_key: None, + env_key, + api_backend: Default::default(), + auth_scheme: None, + reasoning_effort: None, + supports_reasoning_effort: false, + reasoning_efforts: Vec::new(), + capabilities, + extra_headers: IndexMap::new(), + context_window, + auto_compact_threshold_percent: None, + system_prompt_label: None, + api_base_url: None, + use_concise: false, + agent_type: crate::agent::config::default_agent_type(), + inference_idle_timeout_secs: None, + max_retries: None, + hidden: false, + // Subscription models require the OAuth session; open-platform + // models are usable by API-key users. + supported_in_api: !platform.uses_oauth(), + supports_backend_search: false, + compactions_remaining: None, + compaction_at_tokens: None, + show_model_fingerprint: false, + stream_tool_calls: None, + laziness_detector: Default::default(), + } +} +/// Parse a single model entry from the /models response. +/// Used by both initial model fetch and session-resume metadata refresh. +pub fn parse_remote_model_value( + value: &serde_json::Value, + default_base_url: &str, +) -> Option { + let obj = value.as_object()?; + let meta = obj.get("_meta").and_then(|v| v.as_object()); + let id = get_string(obj, "id"); + let model = get_string(obj, "model") + .or_else(|| get_string(obj, "modelId")) + .or_else(|| id.clone()) + .or_else(|| meta.and_then(|m| get_string(m, "model"))) + .or_else(|| meta.and_then(|m| get_string(m, "modelId")))?; + let base_url = get_string(obj, "baseUrl") + .or_else(|| get_string(obj, "base_url")) + .unwrap_or_else(|| default_base_url.to_owned()); + let name = get_string(obj, "name").or_else(|| Some(model.clone())); + let context_window = get_u64(obj, "contextWindow") + .or_else(|| get_u64(obj, "context_window")) + .or_else(|| meta.and_then(|m| get_u64(m, "contextWindow"))) + .or_else(|| meta.and_then(|m| get_u64(m, "totalContextTokens"))) + .unwrap_or(DEFAULT_CONTEXT_WINDOW); + let context_window = std::num::NonZeroU64::new(context_window)?; + let agent_type = get_string(obj, "systemPromptType") + .or_else(|| get_string(obj, "system_prompt_type")) + .or_else(|| get_string(obj, "agent_type")) + .or_else(|| get_string(obj, "agentType")) + .or_else(|| meta.and_then(|m| get_string(m, "agentType"))) + .or_else(|| meta.and_then(|m| get_string(m, "agent_type"))) + .unwrap_or_else(crate::agent::config::default_agent_type); + let api_backend = get_string(obj, "apiBackend") + .or_else(|| get_string(obj, "api_backend")) + .and_then(|s| match s.as_str() { + "responses" => Some(crate::sampling::ApiBackend::Responses), + "chat_completions" => Some(crate::sampling::ApiBackend::ChatCompletions), + "messages" => Some(crate::sampling::ApiBackend::Messages), + _ => None, + }) + .unwrap_or_default(); + Some(crate::agent::config::ModelEntryConfig { + id, + model, + base_url, + name, + description: get_string(obj, "description"), + max_completion_tokens: get_u64(obj, "maxCompletionTokens") + .or_else(|| get_u64(obj, "max_completion_tokens")) + .and_then(|v| u32::try_from(v).ok()), + temperature: get_f64(obj, "temperature").map(|v| v as f32), + top_p: get_f64(obj, "topP").or_else(|| get_f64(obj, "top_p")).map(|v| v as f32), + api_key: get_string(obj, "apiKey").or_else(|| get_string(obj, "api_key")), + env_key: get_env_keys(obj, "envKey").or_else(|| get_env_keys(obj, "env_key")), + api_backend, + context_window, + auto_compact_threshold_percent: get_u64(obj, "autoCompactThresholdPercent") + .or_else(|| get_u64(obj, "auto_compact_threshold_percent")) + .and_then(|v| u8::try_from(v).ok()), + system_prompt_label: get_string(obj, "systemPromptLabel") + .or_else(|| get_string(obj, "system_prompt_label")) + .filter(|s| !s.trim().is_empty()), + extra_headers: get_string_map(obj, "extraHeaders"), + api_base_url: get_string(obj, "apiBaseUrl") + .or_else(|| get_string(obj, "api_base_url")), + use_concise: obj + .get("useConcise") + .or_else(|| obj.get("use_concise")) + .and_then(|v| v.as_bool()) + .unwrap_or(false), + agent_type, + inference_idle_timeout_secs: get_u64(obj, "inferenceIdleTimeoutSecs") + .or_else(|| get_u64(obj, "inference_idle_timeout_secs")), + max_retries: get_u64(obj, "maxRetries") + .or_else(|| get_u64(obj, "max_retries")) + .and_then(|v| u32::try_from(v).ok()), + hidden: obj + .get("hidden") + .or_else(|| meta.and_then(|m| m.get("hidden"))) + .and_then(|v| v.as_bool()) + .unwrap_or(false), + supported_in_api: obj + .get("supportedInApi") + .or_else(|| obj.get("supported_in_api")) + .or_else(|| meta.and_then(|m| m.get("supportedInApi"))) + .and_then(|v| v.as_bool()) + .unwrap_or(true), + auth_scheme: None, + reasoning_effort: get_string(obj, "reasoningEffort") + .or_else(|| get_string(obj, "reasoning_effort")) + .or_else(|| meta.and_then(|m| get_string(m, "reasoningEffort"))) + .and_then(|s| s.parse().ok()), + supports_reasoning_effort: obj + .get("supportsReasoningEffort") + .or_else(|| obj.get("supports_reasoning_effort")) + .or_else(|| meta.and_then(|m| m.get("supportsReasoningEffort"))) + .and_then(|v| v.as_bool()) + .unwrap_or(false), + reasoning_efforts: obj + .get("reasoningEfforts") + .or_else(|| obj.get("reasoning_efforts")) + .or_else(|| meta.and_then(|m| m.get("reasoningEfforts"))) + .and_then(|v| v.as_array()) + .map(|arr| kigi_sampling_types::parse_reasoning_effort_options(arr)) + .unwrap_or_default(), + capabilities: obj + .get("capabilities") + .and_then(|v| { + serde_json::from_value::>(v.clone()).ok() + }) + .unwrap_or_default(), + supports_backend_search: obj + .get("supportsBackendSearch") + .or_else(|| obj.get("supports_backend_search")) + .or_else(|| meta.and_then(|m| m.get("supportsBackendSearch"))) + .and_then(|v| v.as_bool()) + .unwrap_or(false), + compactions_remaining: obj + .get("compactionsRemaining") + .or_else(|| obj.get("compactions_remaining")) + .or_else(|| meta.and_then(|m| m.get("compactionsRemaining"))) + .and_then(parse_compactions_remaining) + .or_else(|| { + obj + .get("sendCompactionsRemaining") + .or_else(|| obj.get("send_compactions_remaining")) + .or_else(|| meta.and_then(|m| m.get("sendCompactionsRemaining"))) + .and_then(|v| v.as_bool()) + .map(kigi_sampling_types::CompactionsRemaining::Dynamic) + }), + compaction_at_tokens: obj + .get("compactionAtTokens") + .or_else(|| obj.get("compaction_at_tokens")) + .or_else(|| meta.and_then(|m| m.get("compactionAtTokens"))) + .and_then(parse_compaction_at_tokens), + show_model_fingerprint: obj + .get("showModelFingerprint") + .or_else(|| obj.get("show_model_fingerprint")) + .or_else(|| meta.and_then(|m| m.get("showModelFingerprint"))) + .and_then(|v| v.as_bool()) + .unwrap_or(false), + stream_tool_calls: obj + .get("streamToolCalls") + .or_else(|| obj.get("stream_tool_calls")) + .and_then(|v| v.as_bool()), + laziness_detector: get_object(obj, "lazinessDetector") + .or_else(|| get_object(obj, "laziness_detector")) + .or_else(|| meta.and_then(|m| get_object(m, "lazinessDetector"))) + .and_then(|v| match serde_json::from_value::< + crate::agent::config::LazinessDetectorPerModelConfig, + >(v.clone()) { + Ok(cfg) => Some(cfg), + Err(e) => { + tracing::warn!( + error = % e, + "Failed to deserialize laziness_detector block from remote model; falling back to default" + ); + None + } + }) + .unwrap_or_default(), + }) +} +fn get_string(obj: &serde_json::Map, key: &str) -> Option { + obj.get(key).and_then(|v| v.as_str()).map(|s| s.to_string()) +} +/// Parse `env_key` / `envKey` as a single string or a string array. +fn get_env_keys( + obj: &serde_json::Map, + key: &str, +) -> Option { + let v = obj.get(key)?; + if let Some(s) = v.as_str() { + return Some(crate::agent::config::EnvKeys::single(s)); + } + if let Some(arr) = v.as_array() { + let mut names = Vec::with_capacity(arr.len()); + for item in arr { + let Some(s) = item.as_str() else { + tracing::warn!( + key, + "env_key array has a non-string element; ignoring env_key" + ); + return None; + }; + if !s.is_empty() { + names.push(s.to_owned()); + } + } + if names.is_empty() { + return None; + } + return Some(crate::agent::config::EnvKeys::new(names)); + } + None +} +fn parse_compaction_at_tokens( + v: &serde_json::Value, +) -> Option { + use kigi_sampling_types::CompactionAtTokens; + v.as_bool() + .map(CompactionAtTokens::Enabled) + .or_else(|| v.as_u64().map(CompactionAtTokens::Fixed)) +} +fn parse_compactions_remaining( + v: &serde_json::Value, +) -> Option { + use kigi_sampling_types::CompactionsRemaining; + v.as_bool().map(CompactionsRemaining::Dynamic).or_else(|| { + v.as_u64() + .and_then(|n| u8::try_from(n).ok()) + .map(CompactionsRemaining::Fixed) + }) +} +fn get_u64(obj: &serde_json::Map, key: &str) -> Option { + obj.get(key).and_then(|v| v.as_u64()) +} +fn get_f64(obj: &serde_json::Map, key: &str) -> Option { + obj.get(key).and_then(|v| v.as_f64()) +} +fn get_object<'a>( + obj: &'a serde_json::Map, + key: &str, +) -> Option<&'a serde_json::Value> { + obj.get(key).filter(|v| v.is_object()) +} +fn get_string_map( + obj: &serde_json::Map, + key: &str, +) -> IndexMap { + obj.get(key) + .and_then(|v| v.as_object()) + .map(|map| { + map.iter() + .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string()))) + .collect() + }) + .unwrap_or_default() +} +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn get_env_keys_parses_strings_and_rejects_non_strings() { + use crate::agent::config::EnvKeys; + let parse = |v: serde_json::Value| { + let obj = serde_json::json!({ "env_key" : v }); + get_env_keys(obj.as_object().unwrap(), "env_key") + }; + assert_eq!(parse(serde_json::json!("A")), Some(EnvKeys::single("A"))); + assert_eq!( + parse(serde_json::json!(["A", "B"])), + Some(EnvKeys::new(["A", "B"])) + ); + assert_eq!(parse(serde_json::json!(["A", 123])), None); + assert_eq!(parse(serde_json::json!([])), None); + } + #[test] + fn parse_openai_format_uses_id_field() { + let value = serde_json::json!( + { "id" : "grok-3", "object" : "model", "owned_by" : "xai", "context_window" : + 131072 } + ); + let result = parse_remote_model_value(&value, "https://api.x.ai/v1").unwrap(); + assert_eq!(result.model, "grok-3"); + assert_eq!(result.base_url, "https://api.x.ai/v1"); + assert_eq!(result.name.as_deref(), Some("grok-3")); + } + #[test] + fn parse_model_field_takes_priority_over_id() { + let value = serde_json::json!( + { "id" : "display-key", "model" : "actual-model-id", "name" : "Display Name", + "context_window" : 131072 } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert_eq!(result.model, "actual-model-id"); + assert_eq!(result.name.as_deref(), Some("Display Name")); + } + #[test] + fn parse_reads_reasoning_effort_fields() { + use kigi_sampling_types::ReasoningEffort; + let value = serde_json::json!( + { "model" : "grok-4.5", "context_window" : 1_000_000, + "supports_reasoning_effort" : true, "reasoning_effort" : "high" } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert!(result.supports_reasoning_effort); + assert_eq!(result.reasoning_effort, Some(ReasoningEffort::High)); + let value = serde_json::json!( + { "model" : "grok-4.5", "contextWindow" : 1_000_000, + "supportsReasoningEffort" : true, "reasoningEffort" : "xhigh" } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert!(result.supports_reasoning_effort); + assert_eq!(result.reasoning_effort, Some(ReasoningEffort::Xhigh)); + let value = serde_json::json!({ "model" : "x", "context_window" : 256_000 }); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert!(!result.supports_reasoning_effort); + assert!(result.reasoning_effort.is_none()); + } + #[test] + fn parse_reads_reasoning_efforts_list() { + use kigi_sampling_types::ReasoningEffort; + let value = serde_json::json!( + { "model" : "grok-4.5", "context_window" : 1_000_000, "reasoning_efforts" : + [{ "id" : "deep", "value" : "xhigh", "label" : "Deep" }, { "value" : + "quantum" }, "low",] } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert_eq!(result.reasoning_efforts.len(), 2); + assert_eq!(result.reasoning_efforts[0].id, "deep"); + assert_eq!(result.reasoning_efforts[0].value, ReasoningEffort::Xhigh); + assert_eq!(result.reasoning_efforts[1].value, ReasoningEffort::Low); + for value in [ + serde_json::json!( + { "model" : "m", "context_window" : 256_000, "reasoningEfforts" : [{ + "value" : "high" }] } + ), + serde_json::json!( + { "model" : "m", "context_window" : 256_000, "_meta" : { + "reasoningEfforts" : [{ "value" : "high" }] } } + ), + ] { + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert_eq!(result.reasoning_efforts.len(), 1); + assert_eq!(result.reasoning_efforts[0].value, ReasoningEffort::High); + } + let value = serde_json::json!({ "model" : "x", "context_window" : 256_000 }); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert!(result.reasoning_efforts.is_empty()); + } + #[test] + fn parse_reads_meta_fallback_fields() { + let value = serde_json::json!( + { "_meta" : { "model" : "meta-model-id", "contextWindow" : 131072, + "agentType" : "concise" } } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert_eq!(result.model, "meta-model-id"); + assert_eq!( + result.context_window, + std::num::NonZeroU64::new(131072).unwrap() + ); + assert_eq!(result.agent_type, "concise"); + } + #[test] + fn parse_remote_model_value_no_laziness_detector_block_yields_default() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert_eq!( + result.laziness_detector, + crate::agent::config::LazinessDetectorPerModelConfig::default() + ); + } + #[test] + fn parse_remote_model_value_parses_camelcase_key() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { + "enabled" : true, "max_nudges_per_session" : 2, "idle_threshold_ms" : 12_000, + "min_confidence" : 0.75, }, } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + let expected = crate::agent::config::LazinessDetectorPerModelConfig { + enabled: true, + max_nudges_per_session: 2, + idle_threshold_ms: Some(12_000), + min_confidence: Some(0.75), + include_reasoning: None, + }; + assert_eq!(result.laziness_detector, expected); + } + #[test] + fn parse_remote_model_value_parses_snake_case_laziness_detector() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, "laziness_detector" : { + "enabled" : true, "max_nudges_per_session" : 3, "idle_threshold_ms" : 8_000, + "min_confidence" : 0.6, }, } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + let expected = crate::agent::config::LazinessDetectorPerModelConfig { + enabled: true, + max_nudges_per_session: 3, + idle_threshold_ms: Some(8_000), + min_confidence: Some(0.6), + include_reasoning: None, + }; + assert_eq!(result.laziness_detector, expected); + } + #[test] + fn parse_remote_model_value_parses_meta_laziness_detector() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, "_meta" : { + "lazinessDetector" : { "enabled" : true, "max_nudges_per_session" : 1, + "idle_threshold_ms" : 15_000, "min_confidence" : 0.9, }, }, } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + let expected = crate::agent::config::LazinessDetectorPerModelConfig { + enabled: true, + max_nudges_per_session: 1, + idle_threshold_ms: Some(15_000), + min_confidence: Some(0.9), + include_reasoning: None, + }; + assert_eq!(result.laziness_detector, expected); + } + #[test] + fn parse_remote_model_value_partial_block_uses_field_defaults() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { + "enabled" : true, }, } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + let expected = crate::agent::config::LazinessDetectorPerModelConfig { + enabled: true, + max_nudges_per_session: 0, + idle_threshold_ms: None, + min_confidence: None, + include_reasoning: None, + }; + assert_eq!(result.laziness_detector, expected); + } + #[test] + fn parse_remote_model_value_malformed_block_falls_back_to_default() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { + "enabled" : true, "max_nudges_per_session" : "abc", }, } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert_eq!( + result.laziness_detector, + crate::agent::config::LazinessDetectorPerModelConfig::default() + ); + } + #[test] + fn parse_remote_model_value_non_object_value_falls_back_to_default() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : + "not-an-object", } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert_eq!( + result.laziness_detector, + crate::agent::config::LazinessDetectorPerModelConfig::default() + ); + } + #[test] + fn parse_remote_model_value_top_level_camelcase_wins_over_snake_case() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { + "enabled" : true, "max_nudges_per_session" : 7, }, "laziness_detector" : { + "enabled" : false, "max_nudges_per_session" : 99, }, } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + let expected = crate::agent::config::LazinessDetectorPerModelConfig { + enabled: true, + max_nudges_per_session: 7, + idle_threshold_ms: None, + min_confidence: None, + include_reasoning: None, + }; + assert_eq!(result.laziness_detector, expected); + } + /// `include_reasoning: false` parses cleanly under the per-model + /// `lazinessDetector` block (camelCase wrapper, snake_case inner — + /// matching the existing field-naming convention used for the + /// sibling `min_confidence`, `idle_threshold_ms`, etc.). + #[test] + fn parse_remote_model_value_parses_include_reasoning_under_camelcase_wrapper() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { + "enabled" : true, "include_reasoning" : false, }, } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert_eq!(result.laziness_detector.include_reasoning, Some(false)); + } + #[test] + fn parse_remote_model_value_parses_include_reasoning_under_snake_case_wrapper() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, "laziness_detector" : { + "enabled" : true, "include_reasoning" : true, }, } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert_eq!(result.laziness_detector.include_reasoning, Some(true)); + } + #[test] + fn parse_remote_model_value_omitted_include_reasoning_defaults_to_none() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { + "enabled" : true, "max_nudges_per_session" : 2, }, } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert_eq!( + result.laziness_detector.include_reasoning, None, + "absent include_reasoning defers to harness default via None", + ); + } + #[test] + fn parse_remote_model_value_top_level_wins_over_meta() { + let value = serde_json::json!( + { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { + "enabled" : true, "max_nudges_per_session" : 5, }, "_meta" : { + "lazinessDetector" : { "enabled" : false, "max_nudges_per_session" : 99, }, + }, } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + let expected = crate::agent::config::LazinessDetectorPerModelConfig { + enabled: true, + max_nudges_per_session: 5, + idle_threshold_ms: None, + min_confidence: None, + include_reasoning: None, + }; + assert_eq!(result.laziness_detector, expected); + } + #[test] + fn parse_reads_show_model_fingerprint_field() { + let value = serde_json::json!( + { "model" : "grok-build", "context_window" : 256_000, + "show_model_fingerprint" : true } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert!(result.show_model_fingerprint); + let value = serde_json::json!( + { "model" : "grok-build", "contextWindow" : 256_000, "showModelFingerprint" : + true } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert!(result.show_model_fingerprint); + let value = serde_json::json!( + { "model" : "grok-build", "context_window" : 256_000, "_meta" : { + "showModelFingerprint" : true } } + ); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert!(result.show_model_fingerprint); + let value = serde_json::json!({ "model" : "x", "context_window" : 256_000 }); + let result = parse_remote_model_value(&value, "https://default.url").unwrap(); + assert!(!result.show_model_fingerprint); + } + #[test] + fn get_object_returns_none_for_non_object_values() { + let value = serde_json::json!( + { "string" : "hello", "number" : 42, "bool" : true, "array" : [1, 2, 3], + "null" : null, } + ); + let obj = value.as_object().unwrap(); + assert!(get_object(obj, "string").is_none()); + assert!(get_object(obj, "number").is_none()); + assert!(get_object(obj, "bool").is_none()); + assert!(get_object(obj, "array").is_none()); + assert!(get_object(obj, "null").is_none()); + assert!(get_object(obj, "missing").is_none()); + } + #[test] + fn get_object_returns_some_for_actual_object() { + let value = serde_json::json!({ "nested" : { "a" : 1, "b" : "two" }, }); + let obj = value.as_object().unwrap(); + let nested = get_object(obj, "nested").expect("nested key should resolve to object"); + assert!(nested.is_object()); + assert_eq!(nested["a"], serde_json::json!(1)); + assert_eq!(nested["b"], serde_json::json!("two")); + } + fn endpoints( + proxy: &str, + models_base_url: Option<&str>, + models_list_url: Option<&str>, + ) -> crate::agent::config::EndpointsConfig { + crate::agent::config::EndpointsConfig { + coding_api_base_url: Some(proxy.to_owned()), + models_base_url: models_base_url.map(|s| s.to_owned()), + models_list_url: models_list_url.map(|s| s.to_owned()), + ..Default::default() + } + } + #[test] + fn inference_url_defaults_to_proxy() { + let ep = endpoints("https://proxy.kigi.com/v1", None, None); + assert_eq!(ep.resolve_inference_base_url(), "https://proxy.kigi.com/v1"); + } + #[test] + fn inference_url_uses_models_base_url() { + let ep = endpoints( + "https://proxy.kigi.com/v1", + Some("https://enterprise.acme.com/v1"), + None, + ); + assert_eq!( + ep.resolve_inference_base_url(), + "https://enterprise.acme.com/v1" + ); + } + #[test] + fn inference_url_base_url_wins_over_proxy() { + let ep = endpoints( + "https://proxy.kigi.com/v1", + Some("https://inference.acme.com/v1"), + Some("https://registry.acme.com/api/models"), + ); + assert_eq!( + ep.resolve_inference_base_url(), + "https://inference.acme.com/v1" + ); + } + #[test] + fn list_url_defaults_to_proxy_models() { + let ep = endpoints("https://proxy.kigi.com/v1", None, None); + assert_eq!( + ep.resolve_models_list_url(), + "https://proxy.kigi.com/v1/models" + ); + } + #[test] + fn list_url_derived_from_base_url() { + let ep = endpoints( + "https://proxy.kigi.com/v1", + Some("https://api.x.ai/v1"), + None, + ); + assert_eq!(ep.resolve_models_list_url(), "https://api.x.ai/v1/models"); + } + #[test] + fn list_url_explicit_overrides_derivation() { + let ep = endpoints( + "https://proxy.kigi.com/v1", + Some("https://inference.acme.com/v1"), + Some("https://registry.acme.com/api/list-models"), + ); + assert_eq!( + ep.resolve_models_list_url(), + "https://registry.acme.com/api/list-models" + ); + } + /// INVARIANT: each platform's `/models` URL matches its registry base — + /// kimi-code → the subscription proxy (config override respected, else the + /// kigi-env default), moonshot platforms → their fixed bases — and the + /// cache-origin key encodes the enabled fetch plan without any secrets. + #[test] + #[serial_test::serial] + fn platform_models_urls_and_fetch_origin() { + use crate::agent::config::EndpointsConfig; + use crate::agent::models::{ModelFetchAuth, PlatformApiKeys}; + for k in [ + "KIGI_CODE_BASE_URL", + "KIGI_CODE_BASE_URL", + "KIGI_MODELS_LIST_URL", + ] { + unsafe { std::env::remove_var(k) }; + } + let cfg = EndpointsConfig::from_config_value(&toml::Value::Table(Default::default())); + assert_eq!( + platform_models_url(kigi_models::PlatformId::KimiCode, &cfg), + "https://api.kimi.com/coding/v1/models" + ); + assert_eq!( + platform_models_url(kigi_models::PlatformId::MoonshotCn, &cfg), + "https://api.moonshot.cn/v1/models" + ); + assert_eq!( + platform_models_url(kigi_models::PlatformId::MoonshotAi, &cfg), + "https://api.moonshot.ai/v1/models" + ); + // Proxy override re-points the subscription platform only. + let proxied = EndpointsConfig::from_config_value( + &toml::from_str( + r#"[endpoints] + coding_api_base_url = "https://proxy.acme.example/v1""#, + ) + .unwrap(), + ); + assert_eq!( + platform_models_url(kigi_models::PlatformId::KimiCode, &proxied), + "https://proxy.acme.example/v1/models" + ); + assert_eq!( + platform_models_url(kigi_models::PlatformId::MoonshotCn, &proxied), + "https://api.moonshot.cn/v1/models" + ); + + // Origin key: OAuth-only plan lists kimi-code only; adding a moonshot + // key changes the plan (→ cache miss); the key VALUE never appears. + let oauth_only = models_fetch_origin( + &cfg, + ModelFetchAuth::Platforms, + true, + &PlatformApiKeys::default(), + ); + assert_eq!( + oauth_only, + "platforms[kimi-code=https://api.kimi.com/coding/v1/models]" + ); + let with_cn = models_fetch_origin( + &cfg, + ModelFetchAuth::Platforms, + true, + &crate::agent::models::PlatformApiKeys::test_keys(Some("sk-secret-cn"), None), + ); + assert_ne!( + oauth_only, with_cn, + "enabling a platform must change the origin" + ); + assert!(with_cn.contains("moonshot-cn=https://api.moonshot.cn/v1/models")); + assert!( + !with_cn.contains("sk-secret-cn"), + "origin key must never embed credential values" + ); + + // Custom endpoint mode → the explicit list URL verbatim. + let custom = EndpointsConfig::from_config_value( + &toml::from_str( + r#"[endpoints] + models_base_url = "https://models.acme.com/v1""#, + ) + .unwrap(), + ); + assert_eq!( + models_fetch_origin( + &custom, + ModelFetchAuth::CustomEndpoint, + false, + &PlatformApiKeys::default(), + ), + "https://models.acme.com/v1/models" + ); + } +} diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs index e7bbb1c..8f62131 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs @@ -48,7 +48,6 @@ impl acp::Agent for MvpAgent { ); }); kigi_workspace::trust::migrate_legacy_hook_trust(); - self.maybe_sync_bundle_in_background(false); let mut client_type = arguments .meta .as_ref() @@ -278,11 +277,7 @@ impl acp::Agent for MvpAgent { } self.spawn_initialize_launch_mcp_setup(fetch_managed_mcps); self.spawn_managed_gateway_tool_catalog_fetch(); - let init_model_state = if crate::agent::chat_modes::process_chat_mode_enabled() { - self.chat_modes.model_state().await - } else { - self.model_state(None) - }; + let init_model_state = self.model_state(None); Ok( acp::InitializeResponse::new(acp::ProtocolVersion::V1) .agent_capabilities( @@ -374,9 +369,6 @@ impl acp::Agent for MvpAgent { } } self.set_auth_method(arguments.method_id.clone()); - if crate::agent::chat_modes::process_chat_mode_enabled() { - self.chat_modes.warm_in_background(); - } emit_login_span(true, "api_key", None, None); Ok(Default::default()) } @@ -421,10 +413,8 @@ impl acp::Agent for MvpAgent { .authenticate_after_cached_token_unavailable(arguments) .await; }; - self.refresh_remote_settings(&auth).await; self.emit_settings_update_notification(); - self.maybe_sync_bundle_in_background(false); - { + { let mut sampling_config = self.sampling_config.borrow_mut(); sampling_config.api_key = Some(auth.key); tracing::debug!( @@ -437,12 +427,8 @@ impl acp::Agent for MvpAgent { ); } self.set_auth_method(arguments.method_id.clone()); - if crate::agent::chat_modes::process_chat_mode_enabled() { - self.chat_modes.warm_in_background(); - } let uid = self.auth_manager.current().map(|a| a.user_id); emit_login_span(true, "cached_token", uid.as_deref(), None); - self.maybe_fetch_post_auth_settings().await; Ok(self.auth_response_with_meta()) } auth_method::KIGI_COM_METHOD_ID => { @@ -517,21 +503,15 @@ impl acp::Agent for MvpAgent { ); } self.auth_manager.hot_swap(auth.clone()); - self.refresh_remote_settings(&auth).await; self.emit_settings_update_notification(); - self.maybe_sync_bundle_in_background(false); - self.set_auth_method(arguments.method_id.clone()); + self.set_auth_method(arguments.method_id.clone()); self.models_manager.on_auth_changed().await; - if crate::agent::chat_modes::process_chat_mode_enabled() { - self.chat_modes.warm_in_background(); - } emit_login_span( true, arguments.method_id.0.as_ref(), Some(auth.user_id.as_str()), None, ); - self.maybe_fetch_post_auth_settings().await; Ok(self.auth_response_with_meta()) } _ => { @@ -559,9 +539,7 @@ impl acp::Agent for MvpAgent { .data("initialize must be called before new_session") })?; self.seed_client_config_auth_if_available(); - if let Ok(auth) = self.auth_manager.auth().await { - self.refresh_settings_and_reapply(&auth).await; - } + self.refresh_settings_and_reapply().await; let cwd = AbsPathBuf::new(arguments.cwd.clone()) .map_err(|e| acp::Error::invalid_params().data(e.to_string()))?; let remote_settings = self.cfg.borrow().remote_settings.clone(); @@ -858,8 +836,10 @@ impl acp::Agent for MvpAgent { Some(serde_json::json!({ "cwd" : cwd.as_str() })), ); let models = if is_chat_kind { + // The grok.com chat-mode model picker was removed with the xAI + // proxy; a chat-kind session has no managed catalog to offer. chat_new_session_model_state( - self.chat_modes.model_state().await, + acp::SessionModelState::new(acp::ModelId::from(String::new()), Vec::new()), session_initial_model .filter(|_| matches!(bridge_attach, BridgeAttach::Spawned)), ) @@ -982,14 +962,6 @@ impl acp::Agent for MvpAgent { .build_summary_client(&load_session_sampling)?; let mut persistence_timer = crate::instrumentation_timer!("session.load_light"); persistence_timer.with_field("session_id", session_id.0.as_ref()); - let backend = if self.build_registry_config().is_some() { - Some( - crate::remote::BackendClient::new() - .with_auth_manager(self.auth_manager.clone()), - ) - } else { - None - }; let registry_title_sync = self .session_registry_client() .map(|client| crate::session::persistence::RegistryGeneratedTitleSync { @@ -999,9 +971,6 @@ impl acp::Agent for MvpAgent { let (persistence_info, persistence) = crate::session::persistence::load_light( &session_info, summary_client, - self.storage_mode, - Some(self.auth_manager.clone()), - backend.as_ref(), Some(self.gateway.clone()), summary_model, registry_title_sync, @@ -2095,9 +2064,6 @@ impl acp::Agent for MvpAgent { | "x.ai/sessions/list" => { crate::agent::handlers::session::handle(self, &args).await } - "x.ai/workspaces/list" => { - crate::agent::handlers::workspaces::handle(self, &args).await - } "x.ai/session/updates" => { crate::extensions::session_updates::handle(&args, &self.gateway).await } @@ -2122,6 +2088,7 @@ impl acp::Agent for MvpAgent { crate::extensions::session_admin::handle(self, &args).await } "x.ai/session/repair" => crate::extensions::repair::handle(self, &args).await, + "x.ai/billing" => crate::extensions::billing::handle(self, &args).await, "x.ai/memory/flush" | "x.ai/memory/rewrite" => { crate::extensions::memory::handle(self, &args).await } @@ -2136,206 +2103,6 @@ impl acp::Agent for MvpAgent { crate::extensions::feedback::handle(self, &args).await } "x.ai/recap" => crate::extensions::recap::handle(self, &args).await, - "x.ai/cloud/terminate" => { - crate::extensions::auth_gate::require_xai_auth( - &self.auth_manager, - "Authentication required", - "Run `grok login` to authenticate.", - )?; - let params: serde_json::Value = serde_json::from_str(args.params.get()) - .map_err(|e| acp::Error::invalid_params().data(e.to_string()))?; - let sandbox_id = params - .get("sandbox_id") - .and_then(|v| v.as_str()) - .ok_or_else(|| { - acp::Error::invalid_params().data("missing sandbox_id") - })?; - let sandbox_client = crate::remote::SandboxClient::new( - self.cli_chat_proxy_base_url(), - self.auth_manager.clone(), - ); - sandbox_client - .terminate_session( - sandbox_id, - &crate::remote::SandboxTerminateRequest { - environment_id: None, - }, - ) - .await - .map_err(|e| { - acp::Error::internal_error() - .data(format!("Failed to terminate sandbox: {e}")) - })?; - crate::extensions::to_raw_response(&serde_json::json!({ "ok" : true })) - } - "x.ai/cloud/env/list" => { - crate::extensions::auth_gate::require_xai_auth( - &self.auth_manager, - "Authentication required", - "Run `grok login` to authenticate.", - )?; - let sandbox_client = crate::remote::SandboxClient::new( - self.cli_chat_proxy_base_url(), - self.auth_manager.clone(), - ); - let resp = sandbox_client - .list_environments( - &crate::remote::SandboxListEnvironmentsRequest::default(), - ) - .await - .map_err(|e| { - acp::Error::internal_error() - .data(format!("Failed to list environments: {e}")) - })?; - crate::extensions::to_raw_response( - &serde_json::json!({ "environments" : resp.environments, }), - ) - } - "x.ai/cloud/env/create" => { - crate::extensions::auth_gate::require_xai_auth( - &self.auth_manager, - "Authentication required", - "Run `grok login` to authenticate.", - )?; - let params: serde_json::Value = serde_json::from_str(args.params.get()) - .map_err(|e| acp::Error::invalid_params().data(e.to_string()))?; - let sandbox_client = crate::remote::SandboxClient::new( - self.cli_chat_proxy_base_url(), - self.auth_manager.clone(), - ); - let resp = sandbox_client - .create_environment( - &crate::remote::SandboxCreateEnvironmentRequest { - name: params - .get("name") - .and_then(|v| v.as_str()) - .map(String::from), - description: params - .get("description") - .and_then(|v| v.as_str()) - .map(String::from), - repository: params - .get("repository") - .and_then(|v| v.as_str()) - .map(String::from), - default_branch: params - .get("default_branch") - .and_then(|v| v.as_str()) - .map(String::from), - container_image: params - .get("container_image") - .and_then(|v| v.as_str()) - .map(String::from), - setup_script: params - .get("setup_script") - .and_then(|v| v.as_str()) - .map(String::from), - workspace_directory: Some("/workspace".to_string()), - internet_enabled: Some(true), - domain_allowlist_preset: Some("common".to_string()), - allowed_http_methods: Some("all".to_string()), - ..Default::default() - }, - ) - .await - .map_err(|e| { - acp::Error::internal_error() - .data(format!("Failed to create environment: {e}")) - })?; - crate::extensions::to_raw_response( - &serde_json::json!({ "environment" : resp.environment, }), - ) - } - "x.ai/cloud/env/update" => { - crate::extensions::auth_gate::require_xai_auth( - &self.auth_manager, - "Authentication required", - "Run `grok login` to authenticate.", - )?; - let params: serde_json::Value = serde_json::from_str(args.params.get()) - .map_err(|e| acp::Error::invalid_params().data(e.to_string()))?; - let environment_id = params - .get("environment_id") - .and_then(|v| v.as_str()) - .ok_or_else(|| { - acp::Error::invalid_params().data("missing environment_id") - })?; - let sandbox_client = crate::remote::SandboxClient::new( - self.cli_chat_proxy_base_url(), - self.auth_manager.clone(), - ); - let resp = sandbox_client - .update_environment( - environment_id, - &crate::remote::SandboxUpdateEnvironmentRequest { - name: params - .get("name") - .and_then(|v| v.as_str()) - .map(String::from), - description: params - .get("description") - .and_then(|v| v.as_str()) - .map(String::from), - repository: params - .get("repository") - .and_then(|v| v.as_str()) - .map(String::from), - default_branch: params - .get("default_branch") - .and_then(|v| v.as_str()) - .map(String::from), - container_image: params - .get("container_image") - .and_then(|v| v.as_str()) - .map(String::from), - setup_script: params - .get("setup_script") - .and_then(|v| v.as_str()) - .map(String::from), - ..Default::default() - }, - ) - .await - .map_err(|e| { - acp::Error::internal_error() - .data(format!("Failed to update environment: {e}")) - })?; - crate::extensions::to_raw_response( - &serde_json::json!({ "environment" : resp.environment, }), - ) - } - "x.ai/cloud/env/delete" => { - crate::extensions::auth_gate::require_xai_auth( - &self.auth_manager, - "Authentication required", - "Run `grok login` to authenticate.", - )?; - let params: serde_json::Value = serde_json::from_str(args.params.get()) - .map_err(|e| acp::Error::invalid_params().data(e.to_string()))?; - let environment_id = params - .get("environment_id") - .and_then(|v| v.as_str()) - .ok_or_else(|| { - acp::Error::invalid_params().data("missing environment_id") - })?; - let sandbox_client = crate::remote::SandboxClient::new( - self.cli_chat_proxy_base_url(), - self.auth_manager.clone(), - ); - sandbox_client - .delete_environment(environment_id) - .await - .map_err(|e| { - acp::Error::internal_error() - .data(format!("Failed to delete environment: {e}")) - })?; - crate::extensions::to_raw_response(&serde_json::json!({ "ok" : true })) - } - "x.ai/billing" => crate::extensions::billing::handle(self, &args).await, - "x.ai/auto-topup-rule" => { - crate::extensions::billing::handle(self, &args).await - } - "x.ai/share_session" => crate::extensions::share::handle(self, &args).await, "x.ai/rollout/survey" => { crate::extensions::rollout::handle(self, &args).await } @@ -2395,9 +2162,6 @@ impl acp::Agent for MvpAgent { s if s.starts_with("x.ai/search/") => { crate::extensions::search::handle(self, &args).await } - s if s.starts_with("x.ai/bundle/") => { - crate::extensions::bundle::handle(self, &args).await - } s if s.starts_with("x.ai/code/") => { let ops = self.resolve_workspace_ops()?; crate::extensions::code_nav::handle(self, &ops, &args).await diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs index 431ea17..e02bd52 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs @@ -28,23 +28,15 @@ impl MvpAgent { let session_key = self.auth_manager.current_or_expired().map(|a| a.key.clone()); let models = self.models_manager.models(); let endpoints = self.models_manager.endpoints(); - let (alpha_test_key, client_version) = { - let cfg = self.cfg.borrow(); - ( - cfg.endpoints.alpha_test_key.clone(), - cfg.client_version.clone(), - ) - }; + let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone(); let config = match crate::agent::config::resolve_aux_model_sampling_config( &slug, &models, &endpoints, session_key.as_deref(), alpha_test_key, - client_version, ) { Some(mut cfg) => { - cfg.client_identifier = primary.client_identifier.clone(); cfg.attribution_callback = primary.attribution_callback.clone(); cfg.bearer_resolver = primary.bearer_resolver.clone(); cfg.max_retries = primary.max_retries; @@ -60,10 +52,6 @@ impl MvpAgent { let client = OaiCompatClient::new(config).map_err(map_sampling_err_to_acp)?; Ok((client, model)) } - fn has_proxy_credentials(&self) -> bool { - self.cfg.borrow().endpoints.deployment_key.is_some() - || self.auth_manager.current_or_expired().is_some_and(|a| a.is_session_auth()) - } /// `true` for session-based ACP auth methods. fn is_session_based_auth(&self) -> bool { self.auth_method_id @@ -384,39 +372,23 @@ impl MvpAgent { ); } } - /// Extract feedback credentials when proxy credentials are available. - /// - /// Returns `(base_url, user_token, optional_extra_access_key, deployment_key)`. - /// Used by both [`feedback_client`] and session spawning to avoid - /// duplicating the credential assembly logic. - #[allow(clippy::type_complexity)] - fn feedback_credentials( - &self, - ) -> Option<(String, Option, Option, Option)> { - if !self.has_proxy_credentials() { - return None; - } - let user_token = self + /// Feedback endpoint base when this is a subscription (OAuth) session — + /// the Kimi Code feedback endpoint only takes the OAuth Bearer, so + /// API-key-only setups get `None` (they are pointed at the issue + /// tracker instead; kimi-cli slash.py parity). + fn feedback_base_url(&self) -> Option { + let has_session = self .auth_manager .current_or_expired() - .filter(|a| a.is_session_auth()) - .map(|a| a.key.clone()); - let cfg = self.cfg.borrow(); - let base_url = cfg.endpoints.resolve_feedback_base_url(); - let alpha_test_key = cfg.endpoints.alpha_test_key.clone(); - let deployment_key = cfg.endpoints.deployment_key.clone(); - Some((base_url, user_token, alpha_test_key, deployment_key)) + .is_some_and(|a| a.is_session_auth()); + has_session.then(|| self.cfg.borrow().endpoints.resolve_feedback_base_url()) } - /// Build a `FeedbackClient` with resolved feedback URL and credentials. + /// Build a `FeedbackClient` for subscription sessions. pub(crate) fn feedback_client(&self) -> Option { - let (base_url, user_token, alpha_test_key, deployment_key) = self - .feedback_credentials()?; - Some( - FeedbackClient::new(base_url, user_token) - .with_alpha_test_key(alpha_test_key) - .with_deployment_key(deployment_key) - .with_auth_manager(self.auth_manager.clone()), - ) + Some(FeedbackClient::new( + self.feedback_base_url()?, + self.auth_manager.clone(), + )) } /// Build a `RegistryConfig` if the feature is enabled (for passing to persistence actor). pub(super) fn build_registry_config( @@ -460,17 +432,6 @@ impl MvpAgent { .with_auth(self.auth_manager.clone()), ) } - pub(crate) fn conversations_client( - &self, - ) -> Option { - if !crate::session::unified_list::conversations_lane_active() { - return None; - } - Some(crate::remote::ConversationsClient::new(self.auth_manager.clone())) - } - pub(crate) fn workspaces_client(&self) -> crate::remote::WorkspacesClient { - crate::remote::WorkspacesClient::new(self.auth_manager.clone()) - } /// Pre-session command availability snapshot. /// /// Used by the `x.ai/commands/list` ext method and the @@ -515,13 +476,9 @@ impl MvpAgent { ) -> &kigi_agent::plugins::SharedPluginRegistryHandle { &self.plugin_registry_handle } - /// `true` when the agent runs in writeback storage mode. - pub(crate) fn is_writeback_storage(&self) -> bool { - matches!(self.storage_mode, StorageMode::Writeback) - } /// Resolved cli-chat-proxy base for session features (via /// `proxy_url`). Not for the deployment-config fetch. - pub(crate) fn cli_chat_proxy_base_url(&self) -> String { + pub(crate) fn coding_api_base_url(&self) -> String { self.cfg.borrow().endpoints.proxy_url() } pub(crate) fn alpha_test_key(&self) -> Option { @@ -635,54 +592,14 @@ impl MvpAgent { pub(crate) fn deployment_key(&self) -> Option { self.cfg.borrow().endpoints.deployment_key.clone() } - /// Re-fetch remote settings and re-init the telemetry client. - /// - /// Called unconditionally from both auth handlers so that: - /// - First install / expired OIDC token: settings are fetched for - /// the first time (the early prefetch had no auth to use). - /// - Reauth / account switch: settings are refreshed to reflect - /// the new user's remote settings targeting attributes. - /// - /// This only refreshes `cfg.remote_settings` and re-inits the - /// telemetry client (the only global static). Other settings - /// derived from `remote_settings` (`web_fetch_enabled`, etc.) are - /// resolved lazily per-turn from `cfg` and pick up the new values - /// automatically. - /// Agent-level fields materialised at startup (`worktree_type`, - /// `restore_code`) are NOT re-resolved here; that requires a - /// broader refactor of the init path. - pub(super) async fn refresh_remote_settings(&self, auth: &crate::auth::KimiAuth) { - if !crate::util::config::resolve_remote_fetch_enabled() { - tracing::debug!("post-auth settings refresh skipped: remote_fetch disabled"); - return; - } - let Some(settings) = self.fetch_remote_settings(auth.clone()).await else { - tracing::warn!("post-auth settings refresh failed (HTTP or parse error)"); - return; - }; - tracing::info!("post-auth settings refreshed"); - { - let mut cfg = self.cfg.borrow_mut(); - cfg.remote_settings = Some(settings); - crate::util::config::sync_campaign_fields(&mut cfg); - crate::agent::config::apply_remote_settings_side_effects( - cfg.remote_settings.as_ref(), - ); - } - } - /// Refresh remote settings settings and re-resolve eagerly-resolved config fields. + /// Re-resolve eagerly-resolved config fields from the local config. /// /// Called on `/new` session creation so feature flags reflect the latest - /// remote settings state without requiring a TUI restart. Extends - /// [`refresh_remote_settings`] by also re-running [`resolve_runtime_fields`] - /// with the fresh settings. + /// on-disk config without requiring a TUI restart. (Formerly this also + /// re-fetched the xAI proxy's remote settings; that endpoint is gone.) /// /// In-flight sessions are unaffected — they snapshot config at creation. - pub(super) async fn refresh_settings_and_reapply( - &self, - auth: &crate::auth::KimiAuth, - ) { - self.refresh_remote_settings(auth).await; + pub(super) async fn refresh_settings_and_reapply(&self) { let cwd = std::env::current_dir().ok(); { let mut cfg = self.cfg.borrow_mut(); @@ -698,36 +615,6 @@ impl MvpAgent { } self.emit_settings_update_notification(); } - /// Shared fetch half of every settings refresh: endpoint fields from a - /// scoped `cfg` borrow, `fetch_settings_blocking` off-executor (it already - /// retries transient errors internally), failures normalized to `None`. - /// Callers own their miss logging. - pub(super) async fn fetch_remote_settings( - &self, - auth: crate::auth::KimiAuth, - ) -> Option { - if !crate::util::config::resolve_remote_fetch_enabled() { - tracing::debug!("settings fetch skipped: remote_fetch disabled"); - return None; - } - let (base_url, alpha_test_key) = { - let cfg = self.cfg.borrow(); - (cfg.endpoints.proxy_url(), cfg.endpoints.alpha_test_key.clone()) - }; - match tokio::task::spawn_blocking(move || crate::remote::fetch_settings_blocking( - &base_url, - &auth, - alpha_test_key.as_deref(), - )) - .await - { - Ok(settings) => settings, - Err(e) => { - tracing::warn!(error = % e, "settings fetch task panicked"); - None - } - } - } pub(super) async fn send_model_auto_switched( &self, session_id: &acp::SessionId, @@ -828,26 +715,9 @@ impl MvpAgent { ), ); } - let cfg = self.cfg.borrow(); - let alpha_test_key = cfg.endpoints.alpha_test_key.clone(); - let client_version = cfg.client_version.clone(); - let deployment_id = crate::managed_config::resolve_deployment_id( - cfg.endpoints.deployment_key.as_deref(), - ); - drop(cfg); - let user_id = self - .auth_manager - .current_or_expired() - .filter(|a| a.is_session_auth()) - .map(|a| a.user_id); - let mut config = crate::agent::config::sampling_config_for_model( - model, - credentials, - alpha_test_key, - client_version, - deployment_id, - user_id, - ); + let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone(); + let mut config = + crate::agent::config::sampling_config_for_model(model, credentials, alpha_test_key); config.origin_client = origin_client; config } @@ -912,13 +782,7 @@ impl MvpAgent { .unwrap_or_else(|| kigi_version::VERSION.to_string()); let alpha_test_key = cfg.endpoints.alpha_test_key.clone(); let mut headers = indexmap::IndexMap::new(); - headers.insert("user-agent".to_string(), format!("xai-grok-build/{version}")); - inject_proxy_headers( - &mut headers, - cfg.client_version.as_deref(), - alpha_test_key.as_deref(), - &base_url, - ); + headers.insert("user-agent".to_string(), format!("kigi/{version}")); ImageGenConfig::Enabled { api_key: api_key.clone(), base_url, @@ -961,13 +825,7 @@ impl MvpAgent { .unwrap_or_else(|| kigi_version::VERSION.to_string()); let alpha_test_key = cfg.endpoints.alpha_test_key.clone(); let mut headers = indexmap::IndexMap::new(); - headers.insert("user-agent".to_string(), format!("xai-grok-build/{version}")); - inject_proxy_headers( - &mut headers, - cfg.client_version.as_deref(), - alpha_test_key.as_deref(), - &base_url, - ); + headers.insert("user-agent".to_string(), format!("kigi/{version}")); VideoGenConfig::Enabled { api_key, base_url, @@ -987,15 +845,8 @@ impl MvpAgent { &models, session.as_ref().map(|a| a.key.as_str()), alpha_test_key.clone(), - client_version, &self.cfg.borrow().endpoints, )?; - inject_proxy_headers( - &mut cfg.extra_headers, - cfg.client_version.as_deref(), - alpha_test_key.as_deref(), - &cfg.base_url, - ); Some(cfg) } /// Returns `Err` with a user-facing message on invalid config; the caller at @@ -1112,15 +963,6 @@ impl MvpAgent { .map(|(name, p)| p.render_io_summary(name)) .collect(), models_manager, - chat_modes: { - let chat_modes = crate::agent::chat_modes::ChatModesManager::new( - auth_manager.clone(), - ); - if crate::agent::chat_modes::process_chat_mode_enabled() { - chat_modes.warm_in_background(); - } - chat_modes - }, cfg: RefCell::new(cfg.clone()), auth_method_id: crate::agent::auth_method::new_shared_auth_method_id(None), sampling_config: RefCell::new(sampling_config), @@ -1159,7 +1001,6 @@ impl MvpAgent { subagent_event_rx: RefCell::new(Some(subagent_event_rx)), subagent_coordinator: RefCell::new(subagent_coordinator), monitor_event_buffer: kigi_tools::implementations::grok_build::task::types::MonitorEventBuffer::default(), - bundle_sync_in_flight: Arc::new(std::sync::atomic::AtomicBool::new(false)), workspace_ops: RefCell::new(None), require_gateway_sessions: Rc::new( RefCell::new(std::collections::HashSet::new()), @@ -2221,19 +2062,9 @@ impl MvpAgent { let auto_update = self.cfg.borrow().cli.auto_update; let client_type = *self.client_type.borrow(); let buffering_settings = self.buffering_settings.borrow().clone(); - let ( - feedback_proxy_url, - feedback_user_token, - feedback_alpha_test_key, - deployment_key, - ) = if let Some((url, token, alpha, deploy)) = self.feedback_credentials() { - (Some(url), token, alpha, deploy) - } else { - (None, None, None, None) - }; + let feedback_base_url = self.feedback_base_url(); tracing::info!( - session_id = % session_info.id.0, feedback_url = ? feedback_proxy_url, - authenticated = feedback_user_token.is_some(), + session_id = % session_info.id.0, feedback_url = ? feedback_base_url, "Initializing feedback manager for session" ); let skills = self.cfg.borrow().skills.clone(); @@ -2489,7 +2320,6 @@ impl MvpAgent { self.auth_type(), ), alpha_test_key: self.alpha_test_key(), - client_version: sampling_config.client_version.clone(), }; let attribution_callback: Option< kigi_sampler::SharedAttributionCallback, @@ -2599,10 +2429,7 @@ impl MvpAgent { self.codebase_indexes.clone(), client_code_nav_enabled, fs_watch_caps, - feedback_proxy_url, - feedback_user_token, - feedback_alpha_test_key, - deployment_key, + feedback_base_url, client_terminal, client_fs_read && client_fs_write, gateway_enabled, @@ -2617,7 +2444,6 @@ impl MvpAgent { persisted_goal_mode, persisted_announcement_state, self.memory_config.clone(), - loc_tracking_enabled, feedback_flags, self.managed_mcp_cache.clone(), managed_mcp_expires_at, diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs index e2c7b01..c71098f 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs @@ -406,17 +406,11 @@ struct SettingsUpdateNotification { sharing_enabled: Option, session_picker_grouped: Option, tips: Option>, - gate_message: Option, - gate_url: Option, - gate_label: Option, - allow_access: Option, - subscription_tier_display: Option, auto_permission_mode_enabled: Option, /// Soft-default permission mode for the pager (post-auth / `/new` refresh). permission_mode: Option, group_tool_verbs: Option, collapsed_edit_blocks: Option, - subscription_watch_interval_secs: Option, } /// Reason why a client is not eligible to use codebase indexing. /// @@ -509,9 +503,6 @@ pub struct MvpAgent { pub(crate) sampling_config: RefCell, pub(crate) auth_manager: Arc, pub(crate) models_manager: crate::agent::models::ModelsManager, - /// grok.com chat-product catalog (`/rest/modes`) for chat sessions; distinct - /// from `models_manager` (the build `/v1/models` catalog). - pub(crate) chat_modes: crate::agent::chat_modes::ChatModesManager, /// Forwards pasted codes from `handle_auth_submit_code` to the auth flow. pub(crate) auth_code_tx: RefCell>>, /// Receives the auth URL from the auth flow; read by `handle_auth_get_url`. @@ -672,20 +663,6 @@ pub struct MvpAgent { /// this flag keeps that to a single discovery walk. plugin_registry_initialized: std::cell::Cell, persona_io_summaries: Vec, - /// Single-flight guard for the proactive bundle sync background task. - /// - /// `maybe_sync_bundle_in_background` is invoked from each post-auth path - /// (initialize, cached-token reauth, oidc) and a rapid reconnect can fire - /// all three within the TTL window, giving us multiple concurrent - /// `tokio::task::spawn_local` tasks racing to extract the tar archive, - /// rewrite `manifest.json`, and prune stale files. The non-atomic - /// per-file write/prune semantics in `bundle::extract_bundle_archive` - /// make that race observable as a partially-written cache. - /// - /// We use an `Arc` so the spawned task can clear the flag - /// on completion without re-borrowing `&self`. `Send` is required - /// because the inner `sync_bundle_to_root` now uses `spawn_blocking`. - bundle_sync_in_flight: Arc, /// Local workspace ops, built lazily via [`Self::ensure_local_workspace_ops`]. /// The agent never opens Computer Hub as a harness/client; remote cloud /// sandboxes are gateway-owned (`gateway_bridge` / `computer_sessions`). @@ -944,50 +921,6 @@ impl AuthRequestMeta { .unwrap_or_default() } } -/// Inject standard proxy headers into an `extra_headers` map. -/// -/// Every authenticated request to cli-chat-proxy (web search, image gen, and -/// any future tools that go through the proxy) must carry these headers. -/// Centralising them here means new tool code paths only need one call instead -/// of remembering which headers the proxy expects. -/// -/// Headers injected: -/// - `x-grok-client-version` -- required by the proxy's version-gate check. -/// Uses `client_version` when provided, otherwise falls back to cli-chat-proxy -/// compile-time `CARGO_PKG_VERSION`. -/// - `X-XAI-Token-Auth` / `x-authenticateresponse` -- required by the -/// cli-chat-proxy auth middleware when the `base_url` is a known proxy URL. -/// - optional extra access header -- only set when the corresponding key is -/// `Some` *and* the `base_url` points at a matching non-production host -/// (requires the optional non-production feature). -/// -/// Existing entries are never overwritten so callers can pre-set a value. -fn inject_proxy_headers( - headers: &mut indexmap::IndexMap, - client_version: Option<&str>, - alpha_test_key: Option<&str>, - base_url: &str, -) { - headers - .entry("x-grok-client-version".to_string()) - .or_insert_with(|| { - client_version - .map(String::from) - .unwrap_or_else(|| kigi_version::VERSION.to_string()) - }); - if crate::util::is_cli_chat_proxy_url(base_url) { - headers - .entry("X-XAI-Token-Auth".to_string()) - .or_insert_with(|| "xai-grok-cli".to_string()); - headers - .entry("x-authenticateresponse".to_string()) - .or_insert_with(|| "authenticate-response".to_string()); - headers - .entry(crate::http::CLIENT_MODE_HEADER.to_string()) - .or_insert_with(|| crate::http::process_client_mode().to_string()); - } - let _ = (alpha_test_key, base_url); -} fn resolve_inference_idle_timeout_secs( models: &indexmap::IndexMap, model: &str, @@ -1580,47 +1513,6 @@ impl MvpAgent { }); AuthenticateResponse::new().meta(meta) } - /// Fetch remote settings after authentication when early prefetch had none. - /// Notifies the pager so soft-default permission_mode applies post-login. - pub(super) async fn maybe_fetch_post_auth_settings(&self) { - if self.cfg.borrow().remote_settings.is_some() { - return; - } - let Some(auth) = self.auth_manager.current() else { - return; - }; - let is_session_auth = auth.is_session_auth(); - let Some(settings) = self.fetch_remote_settings(auth).await else { - return; - }; - tracing::info!("post-auth remote_settings fetch succeeded"); - { - let mut cfg = self.cfg.borrow_mut(); - cfg.remote_settings = Some(settings); - crate::agent::config::apply_remote_settings_side_effects( - cfg.remote_settings.as_ref(), - ); - if cfg.storage_mode == StorageMode::Local - && cfg.mode != crate::agent::config::AgentMode::Generic - { - cfg.storage_mode = StorageMode::resolve( - None, - cfg.remote_settings.as_ref(), - ); - if cfg.storage_mode == StorageMode::Writeback && !is_session_auth { - cfg.storage_mode = StorageMode::Local; - } - } - if let Some(v) = cfg - .remote_settings - .as_ref() - .and_then(|s| s.path_not_found_hints) - { - cfg.path_not_found_hints = v; - } - } - self.emit_settings_update_notification(); - } /// Fire-and-forget `x.ai/settings/update` from the current remote snapshot. pub(super) fn emit_settings_update_notification(&self) { let payload = { @@ -1631,20 +1523,12 @@ impl MvpAgent { sharing_enabled: rs.and_then(|s| s.sharing_enabled), session_picker_grouped: rs.and_then(|s| s.session_picker_grouped), tips: rs.and_then(|s| s.tips.clone()), - gate_message: rs.and_then(|s| s.gate_message.clone()), - gate_url: rs.and_then(|s| s.gate_url.clone()), - gate_label: rs.and_then(|s| s.gate_label.clone()), - allow_access: rs.and_then(|s| s.allow_access), - subscription_tier_display: rs - .and_then(|s| s.subscription_tier_display.clone()), auto_permission_mode_enabled: crate::util::config::remote_auto_mode_enabled( rs, ), permission_mode: rs.and_then(|s| s.permission_mode.clone()), group_tool_verbs: rs.and_then(|s| s.group_tool_verbs), collapsed_edit_blocks: rs.and_then(|s| s.collapsed_edit_blocks), - subscription_watch_interval_secs: rs - .and_then(|s| s.subscription_watch_interval_secs), } }; if let Ok(params) = serde_json::value::to_raw_value(&payload) { @@ -1719,78 +1603,6 @@ impl MvpAgent { }); } } - /// Spawn a best-effort bundle sync. Re-fires on every call site (init, - /// cached_token, grok.com/oidc); the cheap pre-checks below absorb repeats - /// so reconnects are cheap. - /// - /// Pre-spawn gating order (cheapest first, all synchronous): - /// 1. Auth gate — avoid spawning a no-op task on every init. - /// 2. Freshness check — skip the sender snapshot + spawn entirely on - /// cache hits, which is the steady-state on every reconnect. - /// 3. Single-flight guard — if a previous sync is still in flight (e.g., - /// initialize + cached_token + oidc fired in quick succession before - /// the first sync's tar extract finished), drop this call to avoid - /// racing concurrent extracts that would interleave per-file writes - /// against `~/.kigi/bundled/` and the manifest. - pub(crate) fn maybe_sync_bundle_in_background(&self, force: bool) { - use crate::extensions::bundle::{ - BUNDLE_SYNC_TTL, bundle_cache_is_fresh, has_bundle_credentials, - maybe_sync_bundle_to_root, - }; - use std::sync::atomic::Ordering; - let am = self.auth_manager.clone(); - let deployment_key = self.deployment_key(); - if !has_bundle_credentials(Some(&am), deployment_key.as_deref()) { - return; - } - let root = crate::bundle::bundled_root(); - if !force && bundle_cache_is_fresh(&root, BUNDLE_SYNC_TTL) { - tracing::debug!("proactive bundle sync skipped pre-spawn: cache is fresh"); - return; - } - let in_flight = self.bundle_sync_in_flight.clone(); - if in_flight - .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) - .is_err() - { - tracing::debug!( - "proactive bundle sync skipped: another sync is already in flight" - ); - return; - } - let proxy_base_url = self.cli_chat_proxy_base_url(); - let alpha_test_key = self.alpha_test_key(); - let senders: Vec< - tokio::sync::mpsc::UnboundedSender, - > = self.sessions.borrow().values().map(|h| h.cmd_tx.clone()).collect(); - tokio::task::spawn_local(async move { - let result = maybe_sync_bundle_to_root( - &root, - &proxy_base_url, - Some(&am), - deployment_key.as_deref(), - alpha_test_key.as_deref(), - force, - BUNDLE_SYNC_TTL, - ) - .await; - in_flight.store(false, Ordering::Release); - match result { - Ok(Some(res)) => { - tracing::info!( - version = % res.version, personas = res.personas_count, roles = - res.roles_count, agents = res.agents_count, skills = res - .skills_count, "proactive bundle sync complete" - ); - Self::broadcast_refresh_skill_baseline(senders); - } - Ok(None) => {} - Err(err) => { - tracing::warn!(error = % err, "proactive bundle sync failed"); - } - } - }); - } } /// Parse `_meta.agentProfile` as a JSON object or string name. /// Returns `None` if absent or invalid. diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/subagent_coordinator.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/subagent_coordinator.rs index 96b0633..a872b0f 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/subagent_coordinator.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/subagent_coordinator.rs @@ -402,7 +402,7 @@ impl MvpAgent { client_hooks: Default::default(), sampling_config: self.sampling_config.borrow().clone(), managed_mcp_proxy_base_url: parent_managed_mcp_proxy_base_url - .unwrap_or_else(|| self.cli_chat_proxy_base_url()), + .unwrap_or_else(|| self.coding_api_base_url()), alpha_test_key: self.alpha_test_key(), auth_method_id: self .auth_method_id diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs index 23bc959..c71f413 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs @@ -1822,18 +1822,6 @@ fn orphaned_tasks_filters_rewind_dead_branches() { ); } #[test] -fn allow_access_from_remote_settings() { - let json = serde_json::json!({ "allow_access" : true }); - let rs: crate::util::config::RemoteSettings = serde_json::from_value(json).unwrap(); - assert_eq!(rs.allow_access, Some(true)); - let json = serde_json::json!({ "allow_access" : false }); - let rs: crate::util::config::RemoteSettings = serde_json::from_value(json).unwrap(); - assert_eq!(rs.allow_access, Some(false)); - let json = serde_json::json!({}); - let rs: crate::util::config::RemoteSettings = serde_json::from_value(json).unwrap(); - assert_eq!(rs.allow_access, None); -} -#[test] fn on_demand_enabled_from_remote_settings() { let json = serde_json::json!({ "on_demand_enabled" : false }); let rs: crate::util::config::RemoteSettings = serde_json::from_value(json).unwrap(); diff --git a/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs b/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs index bf5d91d..f4dd857 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs @@ -708,7 +708,6 @@ pub(crate) async fn handle_subagent_request( api_key: effective_sampling_config.api_key.clone(), auth_type: inherited_auth_type, alpha_test_key: ctx.alpha_test_key.clone(), - client_version: effective_sampling_config.client_version.clone(), }; kigi_log::unified_log::info( "subagent spawn credentials", @@ -1020,9 +1019,6 @@ pub(crate) async fn handle_subagent_request( false, subagent_fs_watch, None, - None, - None, - None, false, false, std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true)), @@ -1057,7 +1053,6 @@ pub(crate) async fn handle_subagent_request( } else { ctx.memory_config.clone() }, - false, Default::default(), ctx.managed_mcp_state.clone(), None, diff --git a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs index 754913b..32d58cd 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/mod.rs @@ -898,15 +898,11 @@ async fn read_parent_sampling_config( auth_scheme, extra_headers, context_window: cfg.context_window.get(), - client_version: creds.client_version, reasoning_effort: cfg.reasoning_effort, force_http1: false, max_retries: None, stream_tool_calls: cfg.stream_tool_calls.unwrap_or(false), idle_timeout_secs: None, - client_identifier: ctx.sampling_config.client_identifier.clone(), - deployment_id: ctx.sampling_config.deployment_id.clone(), - user_id: ctx.sampling_config.user_id.clone(), origin_client: ctx.sampling_config.origin_client.clone(), attribution_callback: ctx.attribution_callback.clone(), bearer_resolver: None, @@ -997,14 +993,7 @@ fn resolve_model_override_to_config( let mut credentials = resolve_credentials(&entry, session_key); credentials.auth_type = subagent_auth_type(Some(&entry), &ctx.auth_method_id); let resolved_auth_type = credentials.auth_type; - let config = sampling_config_for_model( - &entry, - credentials, - ctx.alpha_test_key.clone(), - ctx.sampling_config.client_version.clone(), - ctx.sampling_config.deployment_id.clone(), - ctx.sampling_config.user_id.clone(), - ); + let config = sampling_config_for_model(&entry, credentials, ctx.alpha_test_key.clone()); kigi_log::unified_log::debug( "subagent resolve_model_override_to_config", None, diff --git a/crates/codegen/kigi-shell/src/auth/device.rs b/crates/codegen/kigi-shell/src/auth/device.rs index 9cb86f9..112dce0 100644 --- a/crates/codegen/kigi-shell/src/auth/device.rs +++ b/crates/codegen/kigi-shell/src/auth/device.rs @@ -29,12 +29,14 @@ pub(crate) fn ascii_header_value(value: &str) -> String { } } -/// The three device-identity headers sent on every OAuth call. +/// The three device-identity headers sent on every OAuth call and, via +/// `agent::config::inject_url_derived_headers`, on every first-party +/// inference request (mirroring kimi-cli src/kimi_cli/llm.py:317-323). /// /// Errors when the persistent device id cannot be created (e.g. read-only /// `~/.kigi`): the OAuth endpoints require `X-Msh-Device-Id`, so login cannot -/// proceed without it. -pub(crate) fn device_headers() -> anyhow::Result<[(&'static str, String); 3]> { +/// proceed without it. Inference callers treat the error as skip-with-warning. +pub fn device_headers() -> anyhow::Result<[(&'static str, String); 3]> { Ok([ ("X-Msh-Device-Name", ascii_header_value(&device_name())), ("X-Msh-Device-Model", ascii_header_value(device_model())), diff --git a/crates/codegen/kigi-shell/src/auth/device_code.rs b/crates/codegen/kigi-shell/src/auth/device_code.rs index a69a69b..e2c1f97 100644 --- a/crates/codegen/kigi-shell/src/auth/device_code.rs +++ b/crates/codegen/kigi-shell/src/auth/device_code.rs @@ -150,6 +150,11 @@ async fn complete_device_code_login( /// caller can decide how to notify the user (eprintln on CLI, nothing on TUI /// where the URL is already rendered in the widget). async fn open_browser_detached(url: &str) -> bool { + // Unit tests drive the full login flow against mock servers — their + // fixture URLs must never reach a real browser. + if cfg!(test) { + return false; + } let url = url.to_owned(); match tokio::task::spawn_blocking(move || webbrowser::open(&url)).await { Ok(Ok(())) => true, @@ -171,13 +176,16 @@ mod tests { use wiremock::matchers::{body_string_contains, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; + /// Fixture mirroring the live `device_authorization` payload (verified + /// against auth.kimi.com): verification URLs are passed through verbatim + /// by the login flow, so they use the real shape. fn device_auth_json(code: &str) -> serde_json::Value { serde_json::json!({ - "user_code": "ABCD-1234", + "user_code": "WXYZ-6789", "device_code": code, - "verification_uri": "https://auth.kimi.com/device", - "verification_uri_complete": "https://auth.kimi.com/device?code=ABCD-1234", - "expires_in": 600, + "verification_uri": "https://www.kimi.com/code/authorize_device", + "verification_uri_complete": "https://www.kimi.com/code/authorize_device?user_code=WXYZ-6789", + "expires_in": 1800, "interval": 0, // floored to 1s by the poll loop }) } diff --git a/crates/codegen/kigi-shell/src/auth/kimi_oauth.rs b/crates/codegen/kigi-shell/src/auth/kimi_oauth.rs index d1d4659..cb0dd33 100644 --- a/crates/codegen/kigi-shell/src/auth/kimi_oauth.rs +++ b/crates/codegen/kigi-shell/src/auth/kimi_oauth.rs @@ -359,11 +359,11 @@ mod tests { "client_id={KIMI_CODE_CLIENT_ID}" ))) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "user_code": "ABCD-1234", + "user_code": "WXYZ-6789", "device_code": "dev-code-1", - "verification_uri": "https://auth.kimi.com/device", - "verification_uri_complete": "https://auth.kimi.com/device?code=ABCD-1234", - "expires_in": 600, + "verification_uri": "https://www.kimi.com/code/authorize_device", + "verification_uri_complete": "https://www.kimi.com/code/authorize_device?user_code=WXYZ-6789", + "expires_in": 1800, "interval": 7, }))) .expect(1) @@ -371,13 +371,13 @@ mod tests { .await; let auth = request_device_authorization(&server.uri()).await.unwrap(); - assert_eq!(auth.user_code, "ABCD-1234"); + assert_eq!(auth.user_code, "WXYZ-6789"); assert_eq!(auth.device_code, "dev-code-1"); assert_eq!(auth.interval, 7); - assert_eq!(auth.expires_in, Some(600)); + assert_eq!(auth.expires_in, Some(1800)); assert_eq!( auth.verification_uri_complete, - "https://auth.kimi.com/device?code=ABCD-1234" + "https://www.kimi.com/code/authorize_device?user_code=WXYZ-6789" ); } @@ -392,7 +392,7 @@ mod tests { .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "user_code": "AAAA", "device_code": "d", - "verification_uri_complete": "https://auth.kimi.com/device?code=AAAA", + "verification_uri_complete": "https://www.kimi.com/code/authorize_device?user_code=AAAA", "interval": 5, }))) .expect(1) @@ -409,7 +409,7 @@ mod tests { .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "user_code": "AAAA", "device_code": "d", - "verification_uri_complete": "https://auth.kimi.com/device?code=AAAA", + "verification_uri_complete": "https://www.kimi.com/code/authorize_device?user_code=AAAA", }))) .mount(&server) .await; diff --git a/crates/codegen/kigi-shell/src/auth/meta.rs b/crates/codegen/kigi-shell/src/auth/meta.rs index 30bfaa4..5e318ee 100644 --- a/crates/codegen/kigi-shell/src/auth/meta.rs +++ b/crates/codegen/kigi-shell/src/auth/meta.rs @@ -1,17 +1,5 @@ use serde::{Deserialize, Serialize}; -/// Access-gate copy resolved from remote settings (message + optional CTA). -/// Auth no longer produces gates (tier gating was an xAI concept); the pager -/// still renders one when remote settings carry a gate message. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct GateInfo { - pub message: String, - #[serde(default)] - pub url: Option, - #[serde(default)] - pub label: Option, -} - /// Typed auth metadata passed from the shell to the pager via ACP. #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct AuthMeta { diff --git a/crates/codegen/kigi-shell/src/auth/mod.rs b/crates/codegen/kigi-shell/src/auth/mod.rs index 6c60d36..2342621 100644 --- a/crates/codegen/kigi-shell/src/auth/mod.rs +++ b/crates/codegen/kigi-shell/src/auth/mod.rs @@ -20,9 +20,10 @@ pub use flow::{ run_auth_flow_with_stderr_bridge, run_cli_login, run_cli_logout, try_ensure_fresh_auth, }; mod meta; +pub use device::device_headers; pub use error::{AuthError, RefreshTokenError, RefreshTokenFailedReason}; pub use manager::{AuthManager, shared_api_key_provider}; -pub use meta::{AuthMeta, GateInfo}; +pub use meta::AuthMeta; pub use model::{AuthMode, KimiAuth, lookup_auth}; pub(crate) use model::{TOKEN_TTL, is_expired, token_suffix}; pub use storage::{ diff --git a/crates/codegen/kigi-shell/src/bundle.rs b/crates/codegen/kigi-shell/src/bundle.rs index c66336b..03fc43b 100644 --- a/crates/codegen/kigi-shell/src/bundle.rs +++ b/crates/codegen/kigi-shell/src/bundle.rs @@ -1,1511 +1,19 @@ -use anyhow::{Context, Result, bail}; -use prod_mc_cli_chat_proxy_types::SubagentBundle; -use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; -use std::collections::HashMap; -use std::io::{ErrorKind, Read}; -use std::path::{Path, PathBuf}; +//! Location of the local subagent-content cache (`~/.kigi/bundled/`). +//! +//! Formerly this module managed a synced bundle of personas/roles/agents/ +//! skills fetched from the xAI cli-chat-proxy (`GET /v1/subagents/bundle`). +//! That backend is gone; the directory remains a passive, locally-populated +//! content root that role/persona discovery scans (see +//! `config::resolve_*` discovery in `config/mod.rs`). + +use std::path::PathBuf; const BUNDLED_DIR_NAME: &str = "bundled"; -const MANIFEST_FILE_NAME: &str = "manifest.json"; - -const ARCHIVE_MAX_DECOMPRESSED_SIZE: usize = 50 * 1024 * 1024; -const ARCHIVE_MAX_ENTRIES: usize = 1000; -const ARCHIVE_MAX_ENTRY_SIZE: u64 = 1024 * 1024; - -#[derive(Deserialize)] -struct ArchiveBundleMetadata { - version: String, -} - -#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] -pub struct BundleManifest { - pub version: String, - pub checksums: HashMap, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum BundleFileKind { - Persona, - Role, - Agent, - Skill, -} - -impl BundleFileKind { - fn dir_name(self) -> &'static str { - match self { - Self::Persona => "personas", - Self::Role => "roles", - Self::Agent => "agents", - Self::Skill => "skills", - } - } - - fn extension(self) -> &'static str { - match self { - Self::Agent | Self::Skill => "md", - Self::Persona | Self::Role => "toml", - } - } - - fn label(self) -> &'static str { - match self { - Self::Persona => "persona", - Self::Role => "role", - Self::Agent => "agent", - Self::Skill => "skill", - } - } - - fn from_dir_name(dir_name: &str) -> Option { - match dir_name { - "personas" => Some(Self::Persona), - "roles" => Some(Self::Role), - "agents" => Some(Self::Agent), - "skills" => Some(Self::Skill), - _ => None, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum BundleFileState { - Absent, - MatchesManaged, - ModifiedOrUnmanaged, -} - -#[derive(Debug)] -struct BundleFile<'a> { - relative_path: String, - checksum: String, - content: &'a str, -} +/// `~/.kigi/bundled/` — the on-disk root for bundled subagent content. pub fn bundled_root() -> PathBuf { dirs::home_dir() .unwrap_or_else(|| PathBuf::from(".")) .join(".kigi") .join(BUNDLED_DIR_NAME) } - -pub fn read_cached_manifest(root: &Path) -> Result> { - let manifest_path = manifest_path(root); - let bytes = match std::fs::read(&manifest_path) { - Ok(bytes) => bytes, - Err(error) if error.kind() == ErrorKind::NotFound => return Ok(None), - Err(error) => { - return Err(error) - .with_context(|| format!("failed to read {}", manifest_path.display())); - } - }; - - serde_json::from_slice(&bytes) - .with_context(|| format!("failed to parse {}", manifest_path.display())) - .map(Some) -} - -pub fn write_bundle_to_cache(root: &Path, bundle: &SubagentBundle) -> Result { - let old_manifest = read_cached_manifest(root)?.map(sanitize_manifest); - ensure_bundle_dirs(root)?; - - let bundle_files = bundle_files(bundle)?; - let mut next_checksums = HashMap::new(); - - for bundle_file in &bundle_files { - let previous_checksum = old_manifest - .as_ref() - .and_then(|manifest| manifest.checksums.get(&bundle_file.relative_path)); - let absolute_path = root.join(&bundle_file.relative_path); - - match bundle_file_state(&absolute_path, previous_checksum.map(String::as_str))? { - BundleFileState::Absent | BundleFileState::MatchesManaged => { - write_bundle_file(&absolute_path, bundle_file.content.as_bytes())?; - next_checksums.insert( - bundle_file.relative_path.clone(), - bundle_file.checksum.clone(), - ); - } - BundleFileState::ModifiedOrUnmanaged => { - if let Some(previous_checksum) = previous_checksum { - next_checksums - .insert(bundle_file.relative_path.clone(), previous_checksum.clone()); - } - } - } - } - - if let Some(old_manifest) = old_manifest.as_ref() { - prune_removed_files(root, old_manifest, &mut next_checksums)?; - } - - let next_manifest = BundleManifest { - version: bundle.version.clone(), - checksums: next_checksums, - }; - let manifest_json = - serde_json::to_vec_pretty(&next_manifest).context("failed to serialize bundle manifest")?; - std::fs::write(manifest_path(root), manifest_json) - .with_context(|| format!("failed to write {}", manifest_path(root).display()))?; - - Ok(next_manifest) -} - -pub fn extract_bundle_archive(root: &Path, archive_bytes: &[u8]) -> Result { - let decoder = flate2::read::GzDecoder::new(archive_bytes); - let mut archive = tar::Archive::new(decoder); - - let old_manifest = read_cached_manifest(root)?.map(sanitize_manifest); - ensure_bundle_dirs(root)?; - - let mut next_checksums = HashMap::new(); - let mut version = String::new(); - let mut total_decompressed: usize = 0; - let mut entry_count: usize = 0; - - for entry_result in archive - .entries() - .context("failed to read archive entries")? - { - let mut entry = entry_result.context("failed to read archive entry")?; - - if entry.header().entry_type() != tar::EntryType::Regular { - continue; - } - - entry_count += 1; - if entry_count > ARCHIVE_MAX_ENTRIES { - bail!("archive exceeds maximum entry count ({ARCHIVE_MAX_ENTRIES})"); - } - - let entry_size = entry.header().size().context("failed to read entry size")?; - if entry_size > ARCHIVE_MAX_ENTRY_SIZE { - bail!("archive entry exceeds maximum size ({ARCHIVE_MAX_ENTRY_SIZE} bytes)"); - } - - total_decompressed = total_decompressed - .checked_add(entry_size as usize) - .context("decompressed size overflow")?; - if total_decompressed > ARCHIVE_MAX_DECOMPRESSED_SIZE { - bail!( - "archive exceeds maximum decompressed size ({ARCHIVE_MAX_DECOMPRESSED_SIZE} bytes)" - ); - } - - let raw_path = entry - .path() - .context("failed to read entry path")? - .to_string_lossy() - .into_owned(); - let path = raw_path.strip_prefix("./").unwrap_or(&raw_path); - - if path == "bundle.json" { - let mut content = String::new(); - entry - .read_to_string(&mut content) - .context("failed to read bundle.json")?; - let meta: ArchiveBundleMetadata = - serde_json::from_str(&content).context("failed to parse bundle.json")?; - version = meta.version; - continue; - } - - let cache_relative_path = match map_archive_path_to_cache_path(path) { - Some(p) => p, - None => continue, - }; - - let mut content = Vec::with_capacity(entry_size as usize); - entry - .read_to_end(&mut content) - .with_context(|| format!("failed to read archive entry: {path}"))?; - let checksum = checksum_bytes(&content); - - let absolute_path = root.join(&cache_relative_path); - let previous_checksum = old_manifest - .as_ref() - .and_then(|m| m.checksums.get(&cache_relative_path)); - - match bundle_file_state(&absolute_path, previous_checksum.map(String::as_str))? { - BundleFileState::Absent | BundleFileState::MatchesManaged => { - write_bundle_file(&absolute_path, &content)?; - next_checksums.insert(cache_relative_path, checksum); - } - BundleFileState::ModifiedOrUnmanaged => { - if let Some(prev) = previous_checksum { - next_checksums.insert(cache_relative_path, prev.clone()); - } - } - } - } - - if version.is_empty() { - bail!("archive missing bundle.json with version field"); - } - - if let Some(old_manifest) = old_manifest.as_ref() { - prune_removed_files(root, old_manifest, &mut next_checksums)?; - } - - let next_manifest = BundleManifest { - version, - checksums: next_checksums, - }; - let manifest_json = - serde_json::to_vec_pretty(&next_manifest).context("failed to serialize bundle manifest")?; - std::fs::write(manifest_path(root), manifest_json) - .with_context(|| format!("failed to write {}", manifest_path(root).display()))?; - - Ok(next_manifest) -} - -pub fn checksum_bytes(bytes: &[u8]) -> String { - let mut hasher = Sha256::new(); - hasher.update(bytes); - format!("{:x}", hasher.finalize()) -} - -pub fn checksum_file(path: &Path) -> Result { - let bytes = std::fs::read(path) - .with_context(|| format!("failed to read {} for checksum", path.display()))?; - Ok(checksum_bytes(&bytes)) -} - -pub fn prune_removed_files( - root: &Path, - old_manifest: &BundleManifest, - retained_checksums: &mut HashMap, -) -> Result<()> { - for (relative_path, previous_checksum) in sanitize_manifest(old_manifest.clone()).checksums { - if retained_checksums.contains_key(&relative_path) { - continue; - } - - let absolute_path = root.join(&relative_path); - match bundle_file_state(&absolute_path, Some(previous_checksum.as_str()))? { - BundleFileState::Absent => {} - BundleFileState::MatchesManaged => { - std::fs::remove_file(&absolute_path) - .with_context(|| format!("failed to remove {}", absolute_path.display()))?; - } - BundleFileState::ModifiedOrUnmanaged => { - retained_checksums.insert(relative_path, previous_checksum); - } - } - } - - Ok(()) -} - -fn ensure_bundle_dirs(root: &Path) -> Result<()> { - std::fs::create_dir_all(root) - .with_context(|| format!("failed to create {}", root.display()))?; - - for dir_name in ["personas", "roles", "agents", "skills"] { - let dir = root.join(dir_name); - std::fs::create_dir_all(&dir) - .with_context(|| format!("failed to create {}", dir.display()))?; - } - - Ok(()) -} - -fn manifest_path(root: &Path) -> PathBuf { - root.join(MANIFEST_FILE_NAME) -} - -fn checksum_file_if_exists(path: &Path) -> Result> { - match std::fs::read(path) { - Ok(bytes) => Ok(Some(checksum_bytes(&bytes))), - Err(error) if error.kind() == ErrorKind::NotFound => Ok(None), - Err(error) => { - Err(error).with_context(|| format!("failed to read {} for checksum", path.display())) - } - } -} - -fn write_bundle_file(absolute_path: &Path, content: &[u8]) -> Result<()> { - if let Some(parent) = absolute_path.parent() { - std::fs::create_dir_all(parent) - .with_context(|| format!("failed to create {}", parent.display()))?; - } - std::fs::write(absolute_path, content) - .with_context(|| format!("failed to write {}", absolute_path.display()))?; - Ok(()) -} - -fn bundle_file_state(path: &Path, old_checksum: Option<&str>) -> Result { - let current_checksum = match checksum_file_if_exists(path)? { - Some(checksum) => checksum, - None => return Ok(BundleFileState::Absent), - }; - - Ok(match old_checksum { - Some(old_checksum) if current_checksum == old_checksum => BundleFileState::MatchesManaged, - Some(_) | None => BundleFileState::ModifiedOrUnmanaged, - }) -} - -fn sanitize_manifest(manifest: BundleManifest) -> BundleManifest { - let checksums = manifest - .checksums - .into_iter() - .filter_map(|(relative_path, checksum)| { - sanitize_relative_path(&relative_path).map(|relative_path| (relative_path, checksum)) - }) - .collect(); - - BundleManifest { - version: manifest.version, - checksums, - } -} - -fn sanitize_relative_path(relative_path: &str) -> Option { - if relative_path.is_empty() || relative_path.starts_with('/') || relative_path.contains('\\') { - return None; - } - - let mut parts = relative_path.split('/'); - let dir_name = parts.next()?; - let second = parts.next()?; - - match parts.next() { - None => { - if second.is_empty() { - return None; - } - let kind = BundleFileKind::from_dir_name(dir_name)?; - if kind == BundleFileKind::Skill { - return None; - } - let file_stem = second.strip_suffix(&format!(".{}", kind.extension()))?; - validate_bundle_name(kind, file_stem).ok()?; - Some(relative_path_for(kind, file_stem)) - } - Some(third) => { - if dir_name != "skills" { - return None; - } - validate_bundle_name(BundleFileKind::Skill, second).ok()?; - // Reject components that would let extraction escape the per-skill directory. - for component in std::iter::once(third).chain(parts) { - if component.is_empty() - || component == "." - || component == ".." - || component.chars().any(char::is_control) - { - return None; - } - } - Some(relative_path.to_string()) - } - } -} - -fn map_archive_path_to_cache_path(archive_path: &str) -> Option { - if let Some(rest) = archive_path.strip_prefix("subagents/") { - return sanitize_relative_path(rest); - } - if archive_path.starts_with("skills/") { - return sanitize_relative_path(archive_path); - } - None -} - -pub fn count_entries_by_prefix(manifest: &BundleManifest, prefix: &str) -> usize { - manifest - .checksums - .keys() - .filter(|k| k.starts_with(prefix)) - .count() -} - -fn bundle_files(bundle: &SubagentBundle) -> Result>> { - let mut files = Vec::new(); - extend_bundle_files(&mut files, BundleFileKind::Persona, &bundle.personas)?; - extend_bundle_files(&mut files, BundleFileKind::Role, &bundle.roles)?; - extend_bundle_files(&mut files, BundleFileKind::Agent, &bundle.agents)?; - extend_bundle_files(&mut files, BundleFileKind::Skill, &bundle.skills)?; - Ok(files) -} - -fn extend_bundle_files<'a>( - files: &mut Vec>, - kind: BundleFileKind, - entries: &'a HashMap, -) -> Result<()> { - for (name, content) in entries { - validate_bundle_name(kind, name)?; - files.push(BundleFile { - relative_path: relative_path_for(kind, name), - checksum: checksum_bytes(content.as_bytes()), - content, - }); - } - Ok(()) -} - -fn relative_path_for(kind: BundleFileKind, name: &str) -> String { - match kind { - BundleFileKind::Skill => format!("{}/{name}/SKILL.md", kind.dir_name()), - _ => format!("{}/{name}.{}", kind.dir_name(), kind.extension()), - } -} - -fn validate_bundle_name(kind: BundleFileKind, name: &str) -> Result<()> { - if name.is_empty() - || name == "." - || name == ".." - || name.contains('/') - || name.contains('\\') - || name.chars().any(char::is_control) - { - bail!("invalid bundled {} name: {name:?}", kind.label()); - } - - Ok(()) -} - -#[cfg(test)] -pub(crate) mod test_helpers { - pub fn make_test_archive(entries: &[(&str, &[u8])]) -> Vec { - let encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); - let mut builder = tar::Builder::new(encoder); - for &(path, content) in entries { - let mut header = tar::Header::new_gnu(); - header.set_size(content.len() as u64); - header.set_mode(0o644); - header.set_cksum(); - builder.append_data(&mut header, path, content).unwrap(); - } - let encoder = builder.into_inner().unwrap(); - encoder.finish().unwrap() - } - - pub fn bundle_json(version: &str) -> String { - format!(r#"{{"version":"{version}"}}"#) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use tempfile::TempDir; - - fn bundle_with_persona(version: &str, name: &str, content: &str) -> SubagentBundle { - let mut bundle = SubagentBundle::empty(version); - bundle - .personas - .insert(name.to_string(), content.to_string()); - bundle - } - - fn bundle_with_skill(version: &str, name: &str, content: &str) -> SubagentBundle { - let mut bundle = SubagentBundle::empty(version); - bundle.skills.insert(name.to_string(), content.to_string()); - bundle - } - - fn cache_root(tmp: &TempDir) -> PathBuf { - tmp.path().join("bundled") - } - - #[test] - fn write_new_file() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let bundle = bundle_with_persona("v1", "researcher", "instructions = \"hello\""); - - let manifest = write_bundle_to_cache(&root, &bundle).unwrap(); - - assert_eq!(manifest.version, "v1"); - assert_eq!( - std::fs::read_to_string(root.join("personas/researcher.toml")).unwrap(), - "instructions = \"hello\"" - ); - assert_eq!(read_cached_manifest(&root).unwrap(), Some(manifest)); - } - - #[test] - fn overwrite_unchanged_file() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - write_bundle_to_cache( - &root, - &bundle_with_persona("v1", "researcher", "instructions = \"old\""), - ) - .unwrap(); - - write_bundle_to_cache( - &root, - &bundle_with_persona("v2", "researcher", "instructions = \"new\""), - ) - .unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("personas/researcher.toml")).unwrap(), - "instructions = \"new\"" - ); - } - - #[test] - fn skip_user_modified_file() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let manifest_v1 = write_bundle_to_cache( - &root, - &bundle_with_persona("v1", "researcher", "instructions = \"old\""), - ) - .unwrap(); - std::fs::write( - root.join("personas/researcher.toml"), - "instructions = \"user edit\"", - ) - .unwrap(); - - let manifest_v2 = write_bundle_to_cache( - &root, - &bundle_with_persona("v2", "researcher", "instructions = \"new\""), - ) - .unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("personas/researcher.toml")).unwrap(), - "instructions = \"user edit\"" - ); - assert_eq!( - manifest_v2.checksums.get("personas/researcher.toml"), - manifest_v1.checksums.get("personas/researcher.toml") - ); - } - - #[test] - fn prune_removed_unmodified_file() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - write_bundle_to_cache( - &root, - &bundle_with_persona("v1", "researcher", "instructions = \"old\""), - ) - .unwrap(); - - let manifest = write_bundle_to_cache(&root, &SubagentBundle::empty("v2")).unwrap(); - - assert!(!root.join("personas/researcher.toml").exists()); - assert!(!manifest.checksums.contains_key("personas/researcher.toml")); - } - - #[test] - fn keep_removed_modified_file() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let manifest_v1 = write_bundle_to_cache( - &root, - &bundle_with_persona("v1", "researcher", "instructions = \"old\""), - ) - .unwrap(); - std::fs::write( - root.join("personas/researcher.toml"), - "instructions = \"user edit\"", - ) - .unwrap(); - - let manifest_v2 = write_bundle_to_cache(&root, &SubagentBundle::empty("v2")).unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("personas/researcher.toml")).unwrap(), - "instructions = \"user edit\"" - ); - assert_eq!( - manifest_v2.checksums.get("personas/researcher.toml"), - manifest_v1.checksums.get("personas/researcher.toml") - ); - } - - #[test] - fn write_rejects_path_traversal_names() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let outside = tmp.path().join("outside.toml"); - let bundle = bundle_with_persona("v1", "../../outside", "instructions = \"evil\""); - - let error = write_bundle_to_cache(&root, &bundle).unwrap_err(); - - assert!(error.to_string().contains("invalid bundled persona name")); - assert!(!outside.exists()); - } - - #[test] - fn prune_skips_unsafe_manifest_paths() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - std::fs::create_dir_all(&root).unwrap(); - - let outside = tmp.path().join("outside.toml"); - std::fs::write(&outside, "keep me").unwrap(); - - let old_manifest = BundleManifest { - version: "v1".to_string(), - checksums: HashMap::from([( - "personas/../../outside.toml".to_string(), - checksum_file(&outside).unwrap(), - )]), - }; - let mut retained = HashMap::new(); - - prune_removed_files(&root, &old_manifest, &mut retained).unwrap(); - - assert_eq!(std::fs::read_to_string(&outside).unwrap(), "keep me"); - assert!(retained.is_empty()); - } - - #[test] - fn user_revert_after_skipped_update_allows_future_update() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - write_bundle_to_cache( - &root, - &bundle_with_persona("v1", "researcher", "instructions = \"old\""), - ) - .unwrap(); - std::fs::write( - root.join("personas/researcher.toml"), - "instructions = \"user edit\"", - ) - .unwrap(); - write_bundle_to_cache( - &root, - &bundle_with_persona("v2", "researcher", "instructions = \"new\""), - ) - .unwrap(); - - std::fs::write( - root.join("personas/researcher.toml"), - "instructions = \"old\"", - ) - .unwrap(); - - write_bundle_to_cache( - &root, - &bundle_with_persona("v3", "researcher", "instructions = \"latest\""), - ) - .unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("personas/researcher.toml")).unwrap(), - "instructions = \"latest\"" - ); - } - - #[test] - fn user_revert_after_preserved_remove_allows_future_prune() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - write_bundle_to_cache( - &root, - &bundle_with_persona("v1", "researcher", "instructions = \"old\""), - ) - .unwrap(); - std::fs::write( - root.join("personas/researcher.toml"), - "instructions = \"user edit\"", - ) - .unwrap(); - write_bundle_to_cache(&root, &SubagentBundle::empty("v2")).unwrap(); - - std::fs::write( - root.join("personas/researcher.toml"), - "instructions = \"old\"", - ) - .unwrap(); - - let manifest = write_bundle_to_cache(&root, &SubagentBundle::empty("v3")).unwrap(); - - assert!(!root.join("personas/researcher.toml").exists()); - assert!(!manifest.checksums.contains_key("personas/researcher.toml")); - } - - #[test] - fn same_version_retry_repairs_missing_managed_file() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let bundle = bundle_with_persona("v1", "researcher", "instructions = \"hello\""); - - write_bundle_to_cache(&root, &bundle).unwrap(); - std::fs::remove_file(root.join("personas/researcher.toml")).unwrap(); - - write_bundle_to_cache(&root, &bundle).unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("personas/researcher.toml")).unwrap(), - "instructions = \"hello\"" - ); - } - - #[test] - fn write_new_skill_file() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let bundle = bundle_with_skill("v1", "commit", "# Commit Skill\nRun git commit."); - - let manifest = write_bundle_to_cache(&root, &bundle).unwrap(); - - assert_eq!(manifest.version, "v1"); - assert_eq!( - std::fs::read_to_string(root.join("skills/commit/SKILL.md")).unwrap(), - "# Commit Skill\nRun git commit." - ); - assert!(manifest.checksums.contains_key("skills/commit/SKILL.md")); - } - - #[test] - fn skip_user_modified_skill() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let manifest_v1 = - write_bundle_to_cache(&root, &bundle_with_skill("v1", "commit", "# Original")).unwrap(); - std::fs::write(root.join("skills/commit/SKILL.md"), "# User custom").unwrap(); - - let manifest_v2 = - write_bundle_to_cache(&root, &bundle_with_skill("v2", "commit", "# Updated")).unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("skills/commit/SKILL.md")).unwrap(), - "# User custom" - ); - assert_eq!( - manifest_v2.checksums.get("skills/commit/SKILL.md"), - manifest_v1.checksums.get("skills/commit/SKILL.md") - ); - } - - #[test] - fn prune_removed_unmodified_skill() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - write_bundle_to_cache(&root, &bundle_with_skill("v1", "commit", "# Original")).unwrap(); - - let manifest = write_bundle_to_cache(&root, &SubagentBundle::empty("v2")).unwrap(); - - assert!(!root.join("skills/commit/SKILL.md").exists()); - assert!(!manifest.checksums.contains_key("skills/commit/SKILL.md")); - } - - #[test] - fn keep_removed_modified_skill() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let manifest_v1 = - write_bundle_to_cache(&root, &bundle_with_skill("v1", "commit", "# Original")).unwrap(); - std::fs::write(root.join("skills/commit/SKILL.md"), "# User custom").unwrap(); - - let manifest_v2 = write_bundle_to_cache(&root, &SubagentBundle::empty("v2")).unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("skills/commit/SKILL.md")).unwrap(), - "# User custom" - ); - assert_eq!( - manifest_v2.checksums.get("skills/commit/SKILL.md"), - manifest_v1.checksums.get("skills/commit/SKILL.md") - ); - } - - #[test] - fn skill_rejects_path_traversal() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let outside = tmp.path().join("outside.md"); - let bundle = bundle_with_skill("v1", "../../outside", "# evil"); - - let error = write_bundle_to_cache(&root, &bundle).unwrap_err(); - - assert!(error.to_string().contains("invalid bundled skill name")); - assert!(!outside.exists()); - } - - #[test] - fn sanitize_accepts_valid_skill_path() { - assert_eq!( - sanitize_relative_path("skills/commit/SKILL.md"), - Some("skills/commit/SKILL.md".to_string()) - ); - } - - #[test] - fn sanitize_accepts_nested_skill_paths() { - assert_eq!( - sanitize_relative_path("skills/implement/scripts/memory.py"), - Some("skills/implement/scripts/memory.py".to_string()) - ); - assert_eq!( - sanitize_relative_path("skills/implement/tests/test_memory.py"), - Some("skills/implement/tests/test_memory.py".to_string()) - ); - assert_eq!( - sanitize_relative_path("skills/commit/README.md"), - Some("skills/commit/README.md".to_string()) - ); - assert_eq!( - sanitize_relative_path("skills/foo/a/b/c/d.txt"), - Some("skills/foo/a/b/c/d.txt".to_string()) - ); - } - - #[test] - fn sanitize_rejects_invalid_skill_paths() { - // Wrong top-level directory. - assert_eq!(sanitize_relative_path("personas/commit/SKILL.md"), None); - // Two-component skill path (must be at least 3). - assert_eq!(sanitize_relative_path("skills/commit.md"), None); - // Empty skill name. - assert_eq!(sanitize_relative_path("skills//SKILL.md"), None); - // Path traversal in skill name. - assert_eq!(sanitize_relative_path("skills/../SKILL.md"), None); - assert_eq!(sanitize_relative_path("skills/../etc/SKILL.md"), None); - // Path traversal in nested components. - assert_eq!(sanitize_relative_path("skills/foo/../bar/SKILL.md"), None); - assert_eq!( - sanitize_relative_path("skills/foo/scripts/../../etc/passwd"), - None - ); - // `.` and empty components in the nested portion are rejected. - assert_eq!(sanitize_relative_path("skills/foo/./SKILL.md"), None); - assert_eq!(sanitize_relative_path("skills/foo//SKILL.md"), None); - } - - #[test] - fn sanitize_accepts_valid_two_component_paths() { - assert_eq!( - sanitize_relative_path("personas/researcher.toml"), - Some("personas/researcher.toml".to_string()) - ); - assert_eq!( - sanitize_relative_path("roles/reviewer.toml"), - Some("roles/reviewer.toml".to_string()) - ); - assert_eq!( - sanitize_relative_path("agents/coder.md"), - Some("agents/coder.md".to_string()) - ); - } - - // --- archive extraction tests --- - - use super::test_helpers::{bundle_json, make_test_archive}; - - #[test] - fn extract_archive_writes_personas_roles_agents_and_skills() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let v = bundle_json("v1"); - let archive = make_test_archive(&[ - ("bundle.json", v.as_bytes()), - ( - "subagents/personas/researcher.toml", - b"instructions = \"hello\"", - ), - ("subagents/roles/reviewer.toml", b"description = \"review\""), - ("subagents/agents/default.md", b"# agent"), - ("skills/commit/SKILL.md", b"# Commit skill"), - ]); - - let manifest = extract_bundle_archive(&root, &archive).unwrap(); - - assert_eq!(manifest.version, "v1"); - assert_eq!( - std::fs::read_to_string(root.join("personas/researcher.toml")).unwrap(), - "instructions = \"hello\"" - ); - assert_eq!( - std::fs::read_to_string(root.join("roles/reviewer.toml")).unwrap(), - "description = \"review\"" - ); - assert_eq!( - std::fs::read_to_string(root.join("agents/default.md")).unwrap(), - "# agent" - ); - assert_eq!( - std::fs::read_to_string(root.join("skills/commit/SKILL.md")).unwrap(), - "# Commit skill" - ); - assert!(manifest.checksums.contains_key("personas/researcher.toml")); - assert!(manifest.checksums.contains_key("roles/reviewer.toml")); - assert!(manifest.checksums.contains_key("agents/default.md")); - assert!(manifest.checksums.contains_key("skills/commit/SKILL.md")); - } - - #[test] - fn extract_archive_writes_nested_skill_files() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let v = bundle_json("v1"); - let archive = make_test_archive(&[ - ("bundle.json", v.as_bytes()), - ("skills/implement/SKILL.md", b"# Implement skill"), - ("skills/implement/scripts/memory.py", b"print('memory')\n"), - ( - "skills/implement/tests/test_memory.py", - b"def test_memory():\n pass\n", - ), - ]); - - let manifest = extract_bundle_archive(&root, &archive).unwrap(); - - assert_eq!(manifest.version, "v1"); - assert_eq!( - std::fs::read_to_string(root.join("skills/implement/SKILL.md")).unwrap(), - "# Implement skill" - ); - assert_eq!( - std::fs::read_to_string(root.join("skills/implement/scripts/memory.py")).unwrap(), - "print('memory')\n" - ); - assert_eq!( - std::fs::read_to_string(root.join("skills/implement/tests/test_memory.py")).unwrap(), - "def test_memory():\n pass\n" - ); - assert!(manifest.checksums.contains_key("skills/implement/SKILL.md")); - assert!( - manifest - .checksums - .contains_key("skills/implement/scripts/memory.py") - ); - assert!( - manifest - .checksums - .contains_key("skills/implement/tests/test_memory.py") - ); - } - - #[test] - fn extract_archive_prunes_removed_nested_skill_files() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let v1 = bundle_json("v1"); - let v1_archive = make_test_archive(&[ - ("bundle.json", v1.as_bytes()), - ("skills/implement/SKILL.md", b"# Implement"), - ("skills/implement/scripts/memory.py", b"# v1 helper\n"), - ]); - extract_bundle_archive(&root, &v1_archive).unwrap(); - assert!(root.join("skills/implement/scripts/memory.py").exists()); - - let v2 = bundle_json("v2"); - let v2_archive = make_test_archive(&[ - ("bundle.json", v2.as_bytes()), - ("skills/implement/SKILL.md", b"# Implement"), - ]); - let manifest = extract_bundle_archive(&root, &v2_archive).unwrap(); - - assert!(!root.join("skills/implement/scripts/memory.py").exists()); - assert!( - !manifest - .checksums - .contains_key("skills/implement/scripts/memory.py") - ); - assert!(manifest.checksums.contains_key("skills/implement/SKILL.md")); - } - - #[test] - fn extract_archive_keeps_user_modified_nested_skill_file() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let v1 = bundle_json("v1"); - let v1_archive = make_test_archive(&[ - ("bundle.json", v1.as_bytes()), - ("skills/implement/SKILL.md", b"# v1"), - ("skills/implement/scripts/memory.py", b"# v1 helper\n"), - ]); - let manifest_v1 = extract_bundle_archive(&root, &v1_archive).unwrap(); - - std::fs::write( - root.join("skills/implement/scripts/memory.py"), - b"# user edit\n", - ) - .unwrap(); - - let v2 = bundle_json("v2"); - let v2_archive = make_test_archive(&[ - ("bundle.json", v2.as_bytes()), - ("skills/implement/SKILL.md", b"# v2"), - ("skills/implement/scripts/memory.py", b"# v2 helper\n"), - ]); - let manifest_v2 = extract_bundle_archive(&root, &v2_archive).unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("skills/implement/scripts/memory.py")).unwrap(), - "# user edit\n" - ); - assert_eq!( - manifest_v2 - .checksums - .get("skills/implement/scripts/memory.py"), - manifest_v1 - .checksums - .get("skills/implement/scripts/memory.py") - ); - } - - #[test] - fn extract_archive_overwrites_unchanged_nested_skill_file() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let v1 = bundle_json("v1"); - let v1_archive = make_test_archive(&[ - ("bundle.json", v1.as_bytes()), - ("skills/implement/SKILL.md", b"# v1"), - ("skills/implement/scripts/memory.py", b"# v1 helper\n"), - ]); - extract_bundle_archive(&root, &v1_archive).unwrap(); - - let v2 = bundle_json("v2"); - let v2_archive = make_test_archive(&[ - ("bundle.json", v2.as_bytes()), - ("skills/implement/SKILL.md", b"# v2"), - ("skills/implement/scripts/memory.py", b"# v2 helper\n"), - ]); - extract_bundle_archive(&root, &v2_archive).unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("skills/implement/scripts/memory.py")).unwrap(), - "# v2 helper\n" - ); - } - - #[test] - fn extract_archive_skips_user_modified_files() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let v1 = bundle_json("v1"); - let v1_archive = make_test_archive(&[ - ("bundle.json", v1.as_bytes()), - ( - "subagents/personas/researcher.toml", - b"instructions = \"old\"", - ), - ]); - let manifest_v1 = extract_bundle_archive(&root, &v1_archive).unwrap(); - - std::fs::write( - root.join("personas/researcher.toml"), - "instructions = \"user edit\"", - ) - .unwrap(); - - let v2 = bundle_json("v2"); - let v2_archive = make_test_archive(&[ - ("bundle.json", v2.as_bytes()), - ( - "subagents/personas/researcher.toml", - b"instructions = \"new\"", - ), - ]); - let manifest_v2 = extract_bundle_archive(&root, &v2_archive).unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("personas/researcher.toml")).unwrap(), - "instructions = \"user edit\"" - ); - assert_eq!( - manifest_v2.checksums.get("personas/researcher.toml"), - manifest_v1.checksums.get("personas/researcher.toml") - ); - } - - #[test] - fn extract_archive_prunes_removed_files() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let v1 = bundle_json("v1"); - let v1_archive = make_test_archive(&[ - ("bundle.json", v1.as_bytes()), - ( - "subagents/personas/researcher.toml", - b"instructions = \"hello\"", - ), - ]); - extract_bundle_archive(&root, &v1_archive).unwrap(); - - let v2 = bundle_json("v2"); - let v2_archive = make_test_archive(&[("bundle.json", v2.as_bytes())]); - let manifest = extract_bundle_archive(&root, &v2_archive).unwrap(); - - assert!(!root.join("personas/researcher.toml").exists()); - assert!(!manifest.checksums.contains_key("personas/researcher.toml")); - } - - #[test] - fn extract_archive_keeps_removed_modified_file() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let v1 = bundle_json("v1"); - let v1_archive = make_test_archive(&[ - ("bundle.json", v1.as_bytes()), - ( - "subagents/personas/researcher.toml", - b"instructions = \"old\"", - ), - ]); - let manifest_v1 = extract_bundle_archive(&root, &v1_archive).unwrap(); - - std::fs::write( - root.join("personas/researcher.toml"), - "instructions = \"user edit\"", - ) - .unwrap(); - - let v2 = bundle_json("v2"); - let v2_archive = make_test_archive(&[("bundle.json", v2.as_bytes())]); - let manifest_v2 = extract_bundle_archive(&root, &v2_archive).unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("personas/researcher.toml")).unwrap(), - "instructions = \"user edit\"" - ); - assert_eq!( - manifest_v2.checksums.get("personas/researcher.toml"), - manifest_v1.checksums.get("personas/researcher.toml") - ); - } - - #[test] - fn extract_archive_rejects_oversized_entry() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let v = bundle_json("v1"); - let big_content = vec![0u8; ARCHIVE_MAX_ENTRY_SIZE as usize + 1]; - let archive = make_test_archive(&[ - ("bundle.json", v.as_bytes()), - ("subagents/personas/big.toml", &big_content), - ]); - - let err = extract_bundle_archive(&root, &archive).unwrap_err(); - assert!(err.to_string().contains("exceeds maximum size")); - } - - #[test] - fn extract_archive_rejects_too_many_entries() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); - let mut builder = tar::Builder::new(encoder); - - let v = r#"{"version":"v1"}"#; - let mut header = tar::Header::new_gnu(); - header.set_size(v.len() as u64); - header.set_mode(0o644); - header.set_cksum(); - builder - .append_data(&mut header, "bundle.json", v.as_bytes()) - .unwrap(); - - let small = b"x"; - for i in 0..ARCHIVE_MAX_ENTRIES { - let path = format!("unknown/f{i}"); - let mut h = tar::Header::new_gnu(); - h.set_size(small.len() as u64); - h.set_mode(0o644); - h.set_cksum(); - builder.append_data(&mut h, &path, &small[..]).unwrap(); - } - - let encoder = builder.into_inner().unwrap(); - let archive = encoder.finish().unwrap(); - - let err = extract_bundle_archive(&root, &archive).unwrap_err(); - assert!(err.to_string().contains("exceeds maximum entry count")); - } - - #[test] - fn extract_archive_skips_directory_entries() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); - let mut builder = tar::Builder::new(encoder); - - let v = r#"{"version":"v1"}"#; - let mut h = tar::Header::new_gnu(); - h.set_size(v.len() as u64); - h.set_mode(0o644); - h.set_cksum(); - builder - .append_data(&mut h, "bundle.json", v.as_bytes()) - .unwrap(); - - let mut dir_h = tar::Header::new_gnu(); - dir_h.set_size(0); - dir_h.set_mode(0o755); - dir_h.set_entry_type(tar::EntryType::Directory); - dir_h.set_cksum(); - builder - .append_data(&mut dir_h, "subagents/personas/", &[] as &[u8]) - .unwrap(); - - let content = b"instructions = \"hello\""; - let mut fh = tar::Header::new_gnu(); - fh.set_size(content.len() as u64); - fh.set_mode(0o644); - fh.set_cksum(); - builder - .append_data(&mut fh, "subagents/personas/researcher.toml", &content[..]) - .unwrap(); - - let encoder = builder.into_inner().unwrap(); - let archive = encoder.finish().unwrap(); - - let manifest = extract_bundle_archive(&root, &archive).unwrap(); - - assert_eq!(manifest.checksums.len(), 1); - assert!(manifest.checksums.contains_key("personas/researcher.toml")); - } - - #[test] - fn extract_archive_missing_bundle_json_fails() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let archive = make_test_archive(&[( - "subagents/personas/researcher.toml", - b"instructions = \"hello\"" as &[u8], - )]); - - let err = extract_bundle_archive(&root, &archive).unwrap_err(); - assert!(err.to_string().contains("bundle.json")); - } - - #[test] - fn extract_archive_handles_dot_slash_prefix() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let v = bundle_json("v1"); - let archive = make_test_archive(&[ - ("./bundle.json", v.as_bytes()), - ( - "./subagents/personas/researcher.toml", - b"instructions = \"hello\"", - ), - ("./skills/commit/SKILL.md", b"# Commit"), - ]); - - let manifest = extract_bundle_archive(&root, &archive).unwrap(); - assert_eq!(manifest.version, "v1"); - assert!(manifest.checksums.contains_key("personas/researcher.toml")); - assert!(manifest.checksums.contains_key("skills/commit/SKILL.md")); - } - - #[test] - fn extract_archive_overwrites_unchanged_file() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let v1 = bundle_json("v1"); - let v1_archive = make_test_archive(&[ - ("bundle.json", v1.as_bytes()), - ( - "subagents/personas/researcher.toml", - b"instructions = \"old\"", - ), - ]); - extract_bundle_archive(&root, &v1_archive).unwrap(); - - let v2 = bundle_json("v2"); - let v2_archive = make_test_archive(&[ - ("bundle.json", v2.as_bytes()), - ( - "subagents/personas/researcher.toml", - b"instructions = \"new\"", - ), - ]); - extract_bundle_archive(&root, &v2_archive).unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("personas/researcher.toml")).unwrap(), - "instructions = \"new\"" - ); - } - - #[test] - fn extract_archive_skips_unknown_top_level_paths() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let v = bundle_json("v1"); - let archive = make_test_archive(&[ - ("bundle.json", v.as_bytes()), - ("README.md", b"# readme"), - ("unknown/file.txt", b"data"), - ("subagents/personas/valid.toml", b"instructions = \"ok\""), - ]); - - let manifest = extract_bundle_archive(&root, &archive).unwrap(); - assert_eq!(manifest.checksums.len(), 1); - assert!(manifest.checksums.contains_key("personas/valid.toml")); - } - - #[test] - fn extract_archive_rejects_excessive_total_decompressed_size() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - - let encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); - let mut builder = tar::Builder::new(encoder); - - let v = r#"{"version":"v1"}"#; - let mut h = tar::Header::new_gnu(); - h.set_size(v.len() as u64); - h.set_mode(0o644); - h.set_cksum(); - builder - .append_data(&mut h, "bundle.json", v.as_bytes()) - .unwrap(); - - // 51 entries of 1 MB each = 51 MB > 50 MB limit. - // Each entry is at the per-entry limit (not over), so only the aggregate check fires. - let big = vec![0u8; ARCHIVE_MAX_ENTRY_SIZE as usize]; - for i in 0..51 { - let path = format!("subagents/personas/p{i}.toml"); - let mut fh = tar::Header::new_gnu(); - fh.set_size(big.len() as u64); - fh.set_mode(0o644); - fh.set_cksum(); - builder.append_data(&mut fh, &path, big.as_slice()).unwrap(); - } - - let encoder = builder.into_inner().unwrap(); - let archive = encoder.finish().unwrap(); - - let err = extract_bundle_archive(&root, &archive).unwrap_err(); - assert!( - err.to_string() - .contains("exceeds maximum decompressed size") - ); - } - - // --- map_archive_path_to_cache_path tests --- - - #[test] - fn map_archive_strips_subagents_prefix() { - assert_eq!( - map_archive_path_to_cache_path("subagents/personas/researcher.toml"), - Some("personas/researcher.toml".to_string()) - ); - assert_eq!( - map_archive_path_to_cache_path("subagents/roles/reviewer.toml"), - Some("roles/reviewer.toml".to_string()) - ); - assert_eq!( - map_archive_path_to_cache_path("subagents/agents/default.md"), - Some("agents/default.md".to_string()) - ); - } - - #[test] - fn map_archive_preserves_skills_path() { - assert_eq!( - map_archive_path_to_cache_path("skills/commit/SKILL.md"), - Some("skills/commit/SKILL.md".to_string()) - ); - } - - #[test] - fn map_archive_preserves_nested_skill_paths() { - assert_eq!( - map_archive_path_to_cache_path("skills/implement/scripts/memory.py"), - Some("skills/implement/scripts/memory.py".to_string()) - ); - assert_eq!( - map_archive_path_to_cache_path("skills/implement/tests/test_memory.py"), - Some("skills/implement/tests/test_memory.py".to_string()) - ); - } - - #[test] - fn sanitize_accepts_shared_data_under_skills() { - // Non-skill directories under skills/ (e.g., shared/personas/) are - // valid archive entries -- they carry data that skills read at runtime. - assert_eq!( - sanitize_relative_path("skills/shared/personas/reviewer.md"), - Some("skills/shared/personas/reviewer.md".to_string()) - ); - assert_eq!( - sanitize_relative_path("skills/shared/personas/implementer.md"), - Some("skills/shared/personas/implementer.md".to_string()) - ); - } - - #[test] - fn extract_archive_writes_shared_data_under_skills() { - let tmp = TempDir::new().unwrap(); - let root = cache_root(&tmp); - let v = bundle_json("v1"); - let archive = make_test_archive(&[ - ("bundle.json", v.as_bytes()), - ("skills/review/SKILL.md", b"# Review skill"), - ("skills/shared/personas/reviewer.md", b"You are a reviewer."), - ]); - - let manifest = extract_bundle_archive(&root, &archive).unwrap(); - - assert_eq!( - std::fs::read_to_string(root.join("skills/shared/personas/reviewer.md")).unwrap(), - "You are a reviewer." - ); - assert!( - manifest - .checksums - .contains_key("skills/shared/personas/reviewer.md") - ); - } - - #[test] - fn map_archive_skips_unknown_paths() { - assert_eq!(map_archive_path_to_cache_path("unknown/file.txt"), None); - assert_eq!(map_archive_path_to_cache_path("README.md"), None); - assert_eq!(map_archive_path_to_cache_path(""), None); - } - - #[test] - fn map_archive_rejects_traversal_under_subagents() { - assert_eq!( - map_archive_path_to_cache_path("subagents/personas/../../etc/passwd"), - None - ); - } - - // --- count_entries_by_prefix tests --- - - #[test] - fn count_entries_by_prefix_counts_correctly() { - let manifest = BundleManifest { - version: "v1".to_string(), - checksums: HashMap::from([ - ("personas/a.toml".to_string(), "abc".to_string()), - ("personas/b.toml".to_string(), "def".to_string()), - ("roles/r.toml".to_string(), "ghi".to_string()), - ("skills/commit/SKILL.md".to_string(), "jkl".to_string()), - ]), - }; - assert_eq!(count_entries_by_prefix(&manifest, "personas/"), 2); - assert_eq!(count_entries_by_prefix(&manifest, "roles/"), 1); - assert_eq!(count_entries_by_prefix(&manifest, "skills/"), 1); - assert_eq!(count_entries_by_prefix(&manifest, "agents/"), 0); - } -} diff --git a/crates/codegen/kigi-shell/src/config/mod.rs b/crates/codegen/kigi-shell/src/config/mod.rs index 8c31f45..5c262f6 100644 --- a/crates/codegen/kigi-shell/src/config/mod.rs +++ b/crates/codegen/kigi-shell/src/config/mod.rs @@ -1094,11 +1094,11 @@ fn apply_requirements_inner( config.endpoints.xai_api_base_url = val.to_owned(); push("endpoints.xai_api_base_url", val.to_owned()); } - if let Some(val) = req_str(req, "endpoints", "cli_chat_proxy_base_url") - && config.endpoints.cli_chat_proxy_base_url.as_deref() != Some(val) + if let Some(val) = req_str(req, "endpoints", "coding_api_base_url") + && config.endpoints.coding_api_base_url.as_deref() != Some(val) { - config.endpoints.cli_chat_proxy_base_url = Some(val.to_owned()); - push("endpoints.cli_chat_proxy_base_url", val.to_owned()); + config.endpoints.coding_api_base_url = Some(val.to_owned()); + push("endpoints.coding_api_base_url", val.to_owned()); } enforce_str!( "endpoints", diff --git a/crates/codegen/kigi-shell/src/config/tests.rs b/crates/codegen/kigi-shell/src/config/tests.rs index cb16736..77919e8 100644 --- a/crates/codegen/kigi-shell/src/config/tests.rs +++ b/crates/codegen/kigi-shell/src/config/tests.rs @@ -2564,7 +2564,7 @@ fn config_layers_user_overrides_managed() { fn enterprise_two_file_merge_routes_deployment_key_to_proxy() { for k in [ "KIGI_MANAGED_CONFIG_URL", - "KIGI_CLI_CHAT_PROXY_BASE_URL", + "KIGI_CODE_BASE_URL", "KIGI_TRACE_UPLOAD_ENDPOINT_URL", ] { unsafe { std::env::remove_var(k) }; @@ -2573,7 +2573,7 @@ fn enterprise_two_file_merge_routes_deployment_key_to_proxy() { r#" [endpoints] xai_api_base_url = "https://inference.acme-corp.example/xai/v1" -cli_chat_proxy_base_url = "https://cli-chat-proxy.kigi.com/v1" +coding_api_base_url = "https://cli-chat-proxy.kigi.com/v1" [model.kigi-build] base_url = "https://inference.acme-corp.example/xai/v1" @@ -2668,13 +2668,13 @@ fn config_layers_system_managed_lowest_priority() { #[test] fn apply_requirements_value_overrides_user_settings() { let raw_config: toml::Value = toml::from_str( - "[cli]\nauto_update = true\nchannel = \"beta\"\n\n[features]\nfeedback = true\nlsp_tools = true\nweb_fetch = true\nwrite_file = true\n\n[ui]\nyolo = true\n\n[models]\ndefault = \"user-model\"\nweb_search = \"user-ws-model\"\n\n[endpoints]\ncli_chat_proxy_base_url = \"https://user-proxy.example/v1\"\nxai_api_base_url = \"https://user-api.example/v1\"\nmodels_base_url = \"https://user-models.example/v1\"\nmodels_list_url = \"https://user-models.example/v1/models\"\n", + "[cli]\nauto_update = true\nchannel = \"beta\"\n\n[features]\nfeedback = true\nlsp_tools = true\nweb_fetch = true\nwrite_file = true\n\n[ui]\nyolo = true\n\n[models]\ndefault = \"user-model\"\nweb_search = \"user-ws-model\"\n\n[endpoints]\ncoding_api_base_url = \"https://user-proxy.example/v1\"\nxai_api_base_url = \"https://user-api.example/v1\"\nmodels_base_url = \"https://user-models.example/v1\"\nmodels_list_url = \"https://user-models.example/v1/models\"\n", ) .unwrap(); let mut cfg = crate::agent::config::Config::new_from_toml_cfg(&raw_config).unwrap(); cfg.default_yolo_mode = true; let requirements: toml::Value = toml::from_str( - "[cli]\nauto_update = false\nchannel = \"stable\"\n\n[features]\nfeedback = false\nlsp_tools = false\nweb_fetch = false\nwrite_file = false\nremote_fetch = false\n\n[ui]\nyolo = false\n\n[models]\ndefault = \"managed-model\"\nweb_search = \"managed-ws-model\"\n\n[endpoints]\ncli_chat_proxy_base_url = \"https://managed-proxy.example/v1\"\nxai_api_base_url = \"https://managed-api.example/v1\"\nmodels_base_url = \"https://managed-models.example/v1\"\nmodels_list_url = \"https://managed-models.example/v1/models\"\ndeployment_key = \"enterprise-deploy-key-should-not-log\"\n", + "[cli]\nauto_update = false\nchannel = \"stable\"\n\n[features]\nfeedback = false\nlsp_tools = false\nweb_fetch = false\nwrite_file = false\nremote_fetch = false\n\n[ui]\nyolo = false\n\n[models]\ndefault = \"managed-model\"\nweb_search = \"managed-ws-model\"\n\n[endpoints]\ncoding_api_base_url = \"https://managed-proxy.example/v1\"\nxai_api_base_url = \"https://managed-api.example/v1\"\nmodels_base_url = \"https://managed-models.example/v1\"\nmodels_list_url = \"https://managed-models.example/v1/models\"\ndeployment_key = \"enterprise-deploy-key-should-not-log\"\n", ) .unwrap(); let source = RequirementSource::Requirements { @@ -2697,7 +2697,7 @@ fn apply_requirements_value_overrides_user_settings() { assert_eq!(Some("managed-ws-model"), cfg.models.web_search.as_deref()); assert_eq!(Some("stable"), cfg.cli.channel.as_deref()); assert_eq!( - Some("https://managed-proxy.example/v1"), cfg.endpoints.cli_chat_proxy_base_url + Some("https://managed-proxy.example/v1"), cfg.endpoints.coding_api_base_url .as_deref() ); assert_eq!("https://managed-api.example/v1", cfg.endpoints.xai_api_base_url); diff --git a/crates/codegen/kigi-shell/src/extensions/billing.rs b/crates/codegen/kigi-shell/src/extensions/billing.rs index 2f4b1f2..d9848eb 100644 --- a/crates/codegen/kigi-shell/src/extensions/billing.rs +++ b/crates/codegen/kigi-shell/src/extensions/billing.rs @@ -1,8 +1,11 @@ -//! `x.ai/billing` extension handler. +//! `x.ai/billing` extension handler — Kimi Code usage/quota. //! -//! Fetches the authenticated user's Grok Build billing configuration -//! (credit limit, usage, on-demand cap, billing period, history) from -//! the backend. Used by the pager/desktop to display credits and usage. +//! Port of kimi-cli's `/usage` command (kimi-cli `src/kimi_cli/ui/shell/usage.py`): +//! `GET {coding_api_base_url}/usages` with the OAuth Bearer token, parsed into +//! display rows (`{usage: {...}, limits: [{detail, window, ...}]}` payload +//! shape). The TUI renders the rows as label + remaining-quota bar + +//! reset hint. The xAI credits/auto-topup surface this file used to serve is +//! gone with the xAI proxy. use agent_client_protocol as acp; use serde::{Deserialize, Serialize}; @@ -10,593 +13,396 @@ use serde::{Deserialize, Serialize}; use super::{ExtResult, to_raw_response}; use crate::agent::MvpAgent; -/// Billing period cycle identifier. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct BillingCycle { - pub year: i32, - pub month: i32, -} - -/// Cent value from the billing API (USD cents). -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct Cent { - /// proto3 JSON omits zero-valued scalars, so a `$0` Cent arrives as `{}`; - /// default to 0 rather than failing the whole parse. - #[serde(default)] - pub val: i64, -} - -/// A usage period (weekly or monthly) from the newer credits config. +/// One usage row: a named quota with `used`/`limit` and an optional +/// human-readable reset hint (e.g. "resets in 2h 5m"). /// -/// `start`/`end` are RFC 3339 timestamps. `period_type` is the proto enum name -/// (e.g. `USAGE_PERIOD_TYPE_WEEKLY`); kept so callers can distinguish weekly -/// vs monthly cycles. +/// `Deserialize` because the TUI parses this back out of the +/// `x.ai/billing` ext response. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UsageRow { + pub label: String, + pub used: i64, + pub limit: i64, + #[serde(skip_serializing_if = "Option::is_none")] + pub reset_hint: Option, +} + +/// Response for `x.ai/billing`: the parsed usage rows, in display order +/// (summary row first when the payload carries one). #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] -pub struct UsagePeriod { - #[serde(rename = "type", default, skip_serializing_if = "Option::is_none")] - pub period_type: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub start: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub end: Option, +pub struct UsageResponse { + pub rows: Vec, } -/// Usage summary for one past billing period. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct BillingPeriodUsage { - #[serde(skip_serializing_if = "Option::is_none")] - pub billing_cycle: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub included_used: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub on_demand_used: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub total_used: Option, -} - -/// Current billing configuration for Grok Build coding credits. -/// -/// Carries both the newer credits-config fields (`credit_usage_percent`, -/// `current_period`) and the deprecated `GrokBuildBillingConfig` fields -/// (`monthly_limit`, `used`, `billing_period_*`). Consumers should prefer the -/// new fields and fall back to the deprecated ones, so the same struct works -/// against both the new `GetGrokCreditsConfig` and the legacy -/// `GetGrokBuildBillingConfig` backend responses. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct BillingConfig { - /// Included credit usage as a percentage of the allowance (0.0–100.0). - /// Preferred over deriving from `monthly_limit`/`used`. - #[serde(skip_serializing_if = "Option::is_none")] - pub credit_usage_percent: Option, - /// Current usage period (weekly or monthly). Preferred over - /// `billing_period_start`/`billing_period_end`. - #[serde(skip_serializing_if = "Option::is_none")] - pub current_period: Option, - /// Deprecated: included monthly credit budget. Use `credit_usage_percent`. - #[serde(skip_serializing_if = "Option::is_none")] - pub monthly_limit: Option, - /// Deprecated: credits used this period. Use `credit_usage_percent`. - #[serde(skip_serializing_if = "Option::is_none")] - pub used: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub on_demand_cap: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub on_demand_used: Option, - /// Remaining prepaid (purchased) credit balance, positive — the "bought - /// credits" the user has topped up. Populated from the credits config - /// (`GetGrokCreditsConfig.prepaid_balance`); absent in the legacy billing - /// shape. - #[serde(skip_serializing_if = "Option::is_none")] - pub prepaid_balance: Option, - /// Whether this user is on unified usage billing (shared weekly/monthly - /// pool). From `GrokCreditsConfig.is_unified_billing_user`, which billing - /// sets from remote settings `unified_consumer_billing_enabled`. `None` when - /// absent (legacy `GetGrokBuildBillingConfig` shape or older servers). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub is_unified_billing_user: Option, - /// Deprecated: use `current_period.start`. - #[serde(skip_serializing_if = "Option::is_none")] - pub billing_period_start: Option, - /// Deprecated: use `current_period.end`. - #[serde(skip_serializing_if = "Option::is_none")] - pub billing_period_end: Option, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub history: Vec, -} - -/// Top-level response (primarily from `GET /rest/grok/credits` + auto-topup-rule). -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct BillingConfigResponse { - pub config: Option, - /// Whether on-demand credit usage is enabled. When `false`, the pager - /// should hide on-demand controls. Populated from `RemoteSettings`. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub on_demand_enabled: Option, - /// User-friendly subscription tier name (e.g. "SuperGrok Heavy"). - /// Populated from `RemoteSettings` so the pager can update its cached - /// tier on every billing fetch without an extra request. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub subscription_tier: Option, -} - -/// Auto top-up configuration (from GetAutoTopupRule). -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct AutoTopupRule { - /// proto3 JSON omits `false`, so a disabled rule arrives without this field; - /// default to `false` rather than failing the parse (which would otherwise - /// keep a stale cached rule in the pager). - #[serde(default)] - pub enabled: bool, - pub min_before_hitting_sl: Option, - pub topup_amount: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub max_amount_per_month: Option, -} - -/// Wrapper for the auto top-up rule response. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct GetAutoTopupRuleResponse { - #[serde(default)] - pub rule: Option, +/// Error from the usages fetch, mapped to the same user-facing messages +/// kimi-cli shows (usage.py error handling). +#[derive(Debug, thiserror::Error)] +pub enum UsageError { + #[error("Authorization failed. Please check your credentials.")] + Unauthorized, + #[error("Usage endpoint not available. Try Kimi for Coding.")] + NotFound, + #[error("Failed to fetch usage (HTTP {status}).")] + Http { status: u16 }, + #[error("Failed to fetch usage: {0}")] + Network(#[from] reqwest::Error), + #[error("Failed to parse usage response: {0}")] + Parse(#[from] serde_json::Error), } #[tracing::instrument(skip_all, fields(method = %args.method))] pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { match args.method.as_ref() { "x.ai/billing" => { - tracing::info!("handling billing config request"); - handle_get_billing(agent).await - } - "x.ai/auto-topup-rule" => { - tracing::info!("handling auto top-up rule request"); - handle_get_auto_topup_rule(agent).await + tracing::info!("handling usage request"); + handle_get_usage(agent).await } _ => Err(acp::Error::method_not_found()), } } -/// Structured context for unified-log entries from a successful billing fetch. -/// -/// Keeps history to a count + the most recent period so `~/.kigi/logs/unified.jsonl` -/// stays useful without dumping unbounded period arrays. -fn billing_unified_log_ctx(billing: &BillingConfigResponse) -> serde_json::Value { - let history_len = billing - .config - .as_ref() - .map(|c| c.history.len()) - .unwrap_or(0); - let latest_history = billing - .config - .as_ref() - .and_then(|c| c.history.last()) - .and_then(|p| serde_json::to_value(p).ok()); - - let mut config_value = billing - .config - .as_ref() - .and_then(|c| serde_json::to_value(c).ok()) - .unwrap_or(serde_json::Value::Null); - if let Some(obj) = config_value.as_object_mut() { - // Drop full history array; surface length + latest entry instead. - obj.remove("history"); - obj.insert("historyLen".into(), serde_json::json!(history_len)); - if let Some(latest) = latest_history { - obj.insert("latestHistory".into(), latest); - } - } - - serde_json::json!({ - "config": config_value, - "onDemandEnabled": billing.on_demand_enabled, - "subscriptionTier": billing.subscription_tier, - }) -} - -async fn handle_get_billing(agent: &MvpAgent) -> ExtResult { +async fn handle_get_usage(agent: &MvpAgent) -> ExtResult { let auth = super::auth_gate::require_xai_auth( &agent.auth_manager, - "Authentication required to fetch billing data", - "Billing data requires auth with grok.com. Run `grok login` to authenticate.", + "Authentication required to fetch usage data", + "Usage data requires a Kimi Code subscription session. Run `kigi login` to authenticate.", )?; - let proxy_base = agent.cli_chat_proxy_base_url(); - let base = proxy_base.trim_end_matches('/'); - - // Credits balance / usage (new billing system) via the CLI proxy, which - // forwards to the backend `GetGrokCreditsConfig`. - let credits_url = format!("{}/billing?format=credits", base); - let credits_resp = crate::http::shared_client() - .get(&credits_url) - .header("Authorization", format!("Bearer {}", auth.key)) - .header("x-userid", &auth.user_id) - .header("x-grok-client-version", kigi_version::VERSION) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ) - .timeout(std::time::Duration::from_secs(15)) - .send() + let base = agent.cfg.borrow().endpoints.proxy_url(); + let usage = fetch_usage(&crate::http::shared_client(), &base, &auth.key) .await .map_err(|e| { - tracing::error!(error = %e, "billing: upstream request failed"); + tracing::warn!(error = %e, "usage fetch failed"); kigi_log::unified_log::warn( - "billing: upstream request failed", + "usage: fetch failed", None, Some(serde_json::json!({ "error": e.to_string() })), ); - acp::Error::internal_error().data(format!("Failed to fetch billing data: {e}")) + acp::Error::internal_error().data(e.to_string()) })?; - if !credits_resp.status().is_success() { - let status = credits_resp.status().as_u16(); - let body = credits_resp.text().await.unwrap_or_default(); - tracing::warn!(status, url = %credits_url, "billing: upstream error"); - - let detail = serde_json::from_str::(&body) - .ok() - .and_then(|v| v.get("error").and_then(|e| e.as_str()).map(String::from)) - .unwrap_or_else(|| format!("HTTP {status}")); - - kigi_log::unified_log::warn( - "billing: upstream error", - None, - Some(serde_json::json!({ - "status": status, - "detail": detail, - })), - ); - - return Err(acp::Error::internal_error().data(format!("Billing service error: {detail}"))); - } - - let mut billing: BillingConfigResponse = credits_resp.json().await.map_err(|e| { - tracing::error!(error = %e, "billing: failed to parse response"); - kigi_log::unified_log::warn( - "billing: failed to parse response", - None, - Some(serde_json::json!({ "error": e.to_string() })), - ); - acp::Error::internal_error().data(format!("Failed to parse billing data: {e}")) - })?; - - // Enrich with fields from remote settings. - let rs = agent.cfg.borrow().remote_settings.clone(); - billing.on_demand_enabled = rs.as_ref().and_then(|rs| rs.on_demand_enabled); - billing.subscription_tier = rs.as_ref().and_then(|rs| { - rs.subscription_tier_display - .clone() - .or_else(|| rs.subscription_tier.clone()) - }); - - // Every prompt / /usage / poll path hits `x.ai/billing`; log the fetched - // credits snapshot so support can correlate limit UX with real balances. kigi_log::unified_log::info( - "billing: fetched credits config", + "usage: fetched quota rows", None, - Some(billing_unified_log_ctx(&billing)), + serde_json::to_value(&usage).ok(), ); - to_raw_response(&billing) + to_raw_response(&usage) } -async fn handle_get_auto_topup_rule(agent: &MvpAgent) -> ExtResult { - let auth = super::auth_gate::require_xai_auth( - &agent.auth_manager, - "Authentication required to fetch auto top-up rule", - "Auto top-up data requires auth with grok.com. Run `grok login` to authenticate.", - )?; - - let proxy_base = agent.cli_chat_proxy_base_url(); - let base = proxy_base.trim_end_matches('/'); - - // Auto top-up rule via the CLI proxy, which forwards to the backend - // `GetAutoTopupRule`. - let url = format!("{}/auto-topup-rule", base); - let response = crate::http::shared_client() +/// `GET {base}/usages` with a Bearer token, parsed per kimi-cli usage.py. +pub(crate) async fn fetch_usage( + http: &reqwest::Client, + base_url: &str, + token: &str, +) -> Result { + let url = format!("{}/usages", base_url.trim_end_matches('/')); + let response = http .get(&url) - .header("Authorization", format!("Bearer {}", auth.key)) - .header("x-userid", &auth.user_id) - .header("x-grok-client-version", kigi_version::VERSION) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ) - .timeout(std::time::Duration::from_secs(10)) + .bearer_auth(token) + .timeout(std::time::Duration::from_secs(15)) .send() - .await - .map_err(|e| { - tracing::error!(error = %e, "auto-topup: upstream request failed"); - acp::Error::internal_error().data(format!("Failed to fetch auto top-up rule: {e}")) - })?; + .await?; + match response.status().as_u16() { + 200..=299 => {} + 401 => return Err(UsageError::Unauthorized), + 404 => return Err(UsageError::NotFound), + status => return Err(UsageError::Http { status }), + } + let payload: serde_json::Value = serde_json::from_str(&response.text().await?)?; + Ok(parse_usage_payload(&payload)) +} - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); - tracing::warn!(status, url = %url, "auto-topup: upstream error"); +/// Port of usage.py `_parse_usage_payload`: `usage` (summary) + `limits[]`. +fn parse_usage_payload(payload: &serde_json::Value) -> UsageResponse { + let mut rows = Vec::new(); - let detail = serde_json::from_str::(&body) - .ok() - .and_then(|v| v.get("error").and_then(|e| e.as_str()).map(String::from)) - .unwrap_or_else(|| format!("HTTP {status}")); - - return Err( - acp::Error::internal_error().data(format!("Auto top-up service error: {detail}")) - ); + if let Some(usage) = payload.get("usage").filter(|v| v.is_object()) + && let Some(row) = to_usage_row(usage, "Weekly limit") + { + rows.push(row); } - // Return the upstream response body verbatim (as a JSON value) so /usage - // can print the exact data from this request unformatted. - let body_text = response.text().await.unwrap_or_default(); - let value: serde_json::Value = - serde_json::from_str(&body_text).unwrap_or(serde_json::json!({"raw": body_text})); - to_raw_response(&value) + if let Some(limits) = payload.get("limits").and_then(|v| v.as_array()) { + for (idx, item) in limits.iter().enumerate() { + if !item.is_object() { + continue; + } + let detail = match item.get("detail") { + Some(d) if d.is_object() => d, + _ => item, + }; + let empty = serde_json::json!({}); + let window = match item.get("window") { + Some(w) if w.is_object() => w, + _ => &empty, + }; + let label = limit_label(item, detail, window, idx); + if let Some(row) = to_usage_row(detail, &label) { + rows.push(row); + } + } + } + + UsageResponse { rows } +} + +/// Port of usage.py `_to_usage_row`: `used`/`limit`, with +/// `used = limit - remaining` fallback; row dropped when both absent. +fn to_usage_row(data: &serde_json::Value, default_label: &str) -> Option { + let limit = to_int(data.get("limit")); + let used = to_int(data.get("used")).or_else(|| match (to_int(data.get("remaining")), limit) { + (Some(remaining), Some(limit)) => Some(limit - remaining), + _ => None, + }); + if used.is_none() && limit.is_none() { + return None; + } + let label = data + .get("name") + .and_then(non_empty_str) + .or_else(|| data.get("title").and_then(non_empty_str)) + .map(str::to_owned) + .unwrap_or_else(|| default_label.to_owned()); + Some(UsageRow { + label, + used: used.unwrap_or(0), + limit: limit.unwrap_or(0), + reset_hint: reset_hint(data), + }) +} + +/// Port of usage.py `_limit_label`: name/title/scope, else the window +/// duration ("5h limit"), else "Limit #N". +fn limit_label( + item: &serde_json::Value, + detail: &serde_json::Value, + window: &serde_json::Value, + idx: usize, +) -> String { + for key in ["name", "title", "scope"] { + if let Some(val) = item + .get(key) + .and_then(non_empty_str) + .or_else(|| detail.get(key).and_then(non_empty_str)) + { + return val.to_owned(); + } + } + + let duration = to_int(window.get("duration")) + .or_else(|| to_int(item.get("duration"))) + .or_else(|| to_int(detail.get("duration"))); + let time_unit = window + .get("timeUnit") + .and_then(non_empty_str) + .or_else(|| item.get("timeUnit").and_then(non_empty_str)) + .or_else(|| detail.get("timeUnit").and_then(non_empty_str)) + .unwrap_or(""); + if let Some(duration) = duration.filter(|&d| d != 0) { + if time_unit.contains("MINUTE") { + if duration >= 60 && duration % 60 == 0 { + return format!("{}h limit", duration / 60); + } + return format!("{duration}m limit"); + } + if time_unit.contains("HOUR") { + return format!("{duration}h limit"); + } + if time_unit.contains("DAY") { + return format!("{duration}d limit"); + } + return format!("{duration}s limit"); + } + + format!("Limit #{}", idx + 1) +} + +/// Port of usage.py `_reset_hint`: absolute reset keys first, then +/// seconds-until keys. +fn reset_hint(data: &serde_json::Value) -> Option { + for key in ["reset_at", "resetAt", "reset_time", "resetTime"] { + if let Some(val) = data.get(key).and_then(non_empty_str) { + return Some(format_reset_time(val)); + } + } + for key in ["reset_in", "resetIn", "ttl", "window"] { + if let Some(seconds) = to_int(data.get(key)).filter(|&s| s != 0) { + return Some(format!( + "resets in {}", + format_duration(seconds.max(0) as u64) + )); + } + } + None +} + +/// Port of usage.py `_format_reset_time`: ISO timestamp → "resets in …" / +/// "reset" (already past) / "resets at " when unparseable. +fn format_reset_time(val: &str) -> String { + match chrono::DateTime::parse_from_rfc3339(val) { + Ok(dt) => { + let delta = dt.with_timezone(&chrono::Utc) - chrono::Utc::now(); + let seconds = delta.num_seconds(); + if seconds <= 0 { + "reset".to_owned() + } else { + format!("resets in {}", format_duration(seconds as u64)) + } + } + Err(_) => format!("resets at {val}"), + } +} + +/// Port of kimi-cli `utils/datetime.py` `format_duration`: short units, +/// seconds shown only for sub-minute durations. +fn format_duration(seconds: u64) -> String { + let days = seconds / 86_400; + let hours = (seconds % 86_400) / 3_600; + let minutes = (seconds % 3_600) / 60; + let secs = seconds % 60; + let mut parts = Vec::new(); + if days > 0 { + parts.push(format!("{days}d")); + } + if hours > 0 { + parts.push(format!("{hours}h")); + } + if minutes > 0 { + parts.push(format!("{minutes}m")); + } + if secs > 0 && parts.is_empty() { + parts.push(format!("{secs}s")); + } + if parts.is_empty() { + "0s".to_owned() + } else { + parts.join(" ") + } +} + +fn non_empty_str(v: &serde_json::Value) -> Option<&str> { + v.as_str().filter(|s| !s.is_empty()) +} + +/// Port of usage.py `_to_int`: ints and int-shaped floats/strings; anything +/// else is `None`. +fn to_int(value: Option<&serde_json::Value>) -> Option { + let value = value?; + if let Some(i) = value.as_i64() { + return Some(i); + } + if let Some(f) = value.as_f64() { + return Some(f as i64); + } + value.as_str()?.trim().parse::().ok() } #[cfg(test)] mod tests { use super::*; + use wiremock::matchers::{bearer_token, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; - #[test] - fn auto_topup_disabled_rule_omits_enabled_field() { - // proto3 JSON omits `false` / `0`, so a disabled rule arrives without - // `enabled` (and zero Cents as `{}`). It must still deserialize (as - // disabled) rather than erroring — otherwise the pager keeps a stale - // cached rule. - let json = serde_json::json!({ - "rule": { "topupAmount": {"val": 500}, "minBeforeHittingSl": {} } - }); - let resp: GetAutoTopupRuleResponse = serde_json::from_value(json).unwrap(); - let rule = resp.rule.expect("rule present"); - assert!(!rule.enabled); - assert_eq!(rule.topup_amount.unwrap().val, 500); - assert_eq!(rule.min_before_hitting_sl.unwrap().val, 0); - } - - #[test] - fn billing_config_response_deserializes_from_backend_json() { - let json = serde_json::json!({ - "config": { - "monthlyLimit": {"val": 2000}, - "used": {"val": 1234}, - "onDemandCap": {"val": 500}, - "billingPeriodStart": "2025-04-01T00:00:00Z", - "billingPeriodEnd": "2025-05-01T00:00:00Z", - "history": [ + /// Happy path: GET /usages with Bearer, kimi payload shape → rows with + /// summary first, remaining-derived `used`, and window-derived labels. + #[tokio::test] + async fn fetch_usage_parses_kimi_payload() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/usages")) + .and(bearer_token("tok-42")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "usage": { "limit": 1000, "used": 250, "reset_at": "2099-01-01T00:00:00Z" }, + "limits": [ { - "billingCycle": {"year": 2025, "month": 3}, - "includedUsed": {"val": 1800}, - "onDemandUsed": {"val": 0}, - "totalUsed": {"val": 1800} - } + "window": { "duration": 300, "timeUnit": "TIME_UNIT_MINUTE" }, + "detail": { "limit": 50, "remaining": 30, "resetIn": 1800 } + }, + { "name": "RPM", "limit": 60, "used": 12 } ] - } - }); - let resp: BillingConfigResponse = serde_json::from_value(json).unwrap(); - let config = resp.config.unwrap(); - assert_eq!(config.monthly_limit.unwrap().val, 2000); - assert_eq!(config.used.unwrap().val, 1234); - assert_eq!(config.on_demand_cap.unwrap().val, 500); - assert_eq!( - config.billing_period_start.as_deref(), - Some("2025-04-01T00:00:00Z") - ); - assert_eq!(config.history.len(), 1); - let period = &config.history[0]; - let cycle = period.billing_cycle.as_ref().unwrap(); - assert_eq!(cycle.year, 2025); - assert_eq!(cycle.month, 3); - assert_eq!(period.included_used.as_ref().unwrap().val, 1800); - assert_eq!(period.total_used.as_ref().unwrap().val, 1800); - } + }))) + .expect(1) + .mount(&server) + .await; - #[test] - fn billing_unified_log_ctx_includes_credits_and_collapses_history() { - let resp = BillingConfigResponse { - config: Some(BillingConfig { - credit_usage_percent: Some(42.5), - current_period: Some(UsagePeriod { - period_type: Some("USAGE_PERIOD_TYPE_WEEKLY".into()), - start: Some("2025-04-01T00:00:00Z".into()), - end: Some("2025-04-08T00:00:00Z".into()), - }), - monthly_limit: Some(Cent { val: 2000 }), - used: Some(Cent { val: 850 }), - on_demand_cap: Some(Cent { val: 500 }), - on_demand_used: Some(Cent { val: 0 }), - prepaid_balance: Some(Cent { val: 100 }), - is_unified_billing_user: Some(true), - billing_period_start: None, - billing_period_end: None, - history: vec![ - BillingPeriodUsage { - billing_cycle: Some(BillingCycle { - year: 2025, - month: 2, - }), - included_used: Some(Cent { val: 1000 }), - on_demand_used: Some(Cent { val: 0 }), - total_used: Some(Cent { val: 1000 }), - }, - BillingPeriodUsage { - billing_cycle: Some(BillingCycle { - year: 2025, - month: 3, - }), - included_used: Some(Cent { val: 1800 }), - on_demand_used: Some(Cent { val: 0 }), - total_used: Some(Cent { val: 1800 }), - }, - ], - }), - on_demand_enabled: Some(true), - subscription_tier: Some("SuperGrok".into()), - }; - let ctx = billing_unified_log_ctx(&resp); - assert_eq!(ctx["onDemandEnabled"], true); - assert_eq!(ctx["subscriptionTier"], "SuperGrok"); - let config = ctx["config"].as_object().expect("config object"); + let usage = fetch_usage(&reqwest::Client::new(), &server.uri(), "tok-42") + .await + .expect("usage fetch should succeed"); + + assert_eq!(usage.rows.len(), 3); + assert_eq!(usage.rows[0].label, "Weekly limit"); + assert_eq!(usage.rows[0].used, 250); + assert_eq!(usage.rows[0].limit, 1000); assert!( - config.get("history").is_none(), - "full history must be collapsed" + usage.rows[0] + .reset_hint + .as_deref() + .is_some_and(|h| h.starts_with("resets in")), + "absolute reset_at renders a relative hint: {:?}", + usage.rows[0].reset_hint ); - assert_eq!(config["historyLen"], 2); - assert_eq!( - config["latestHistory"]["billingCycle"]["month"], 3, - "latest history period retained" - ); - assert_eq!(config["creditUsagePercent"], 42.5); - assert_eq!(config["prepaidBalance"]["val"], 100); + // 300 minutes → "5h limit"; used derived from remaining (50-30=20). + assert_eq!(usage.rows[1].label, "5h limit"); + assert_eq!(usage.rows[1].used, 20); + assert_eq!(usage.rows[1].limit, 50); + assert_eq!(usage.rows[1].reset_hint.as_deref(), Some("resets in 30m")); + // Item-level fields when there is no `detail` object. + assert_eq!(usage.rows[2].label, "RPM"); + assert_eq!(usage.rows[2].used, 12); + assert_eq!(usage.rows[2].limit, 60); + } + + /// Auth failure: 401 maps to the typed `Unauthorized` error (kimi-cli's + /// "Authorization failed" path). + #[tokio::test] + async fn fetch_usage_maps_401_to_unauthorized() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/usages")) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + + let err = fetch_usage(&reqwest::Client::new(), &server.uri(), "bad") + .await + .expect_err("401 must fail"); + assert!(matches!(err, UsageError::Unauthorized)); + } + + /// 404 maps to the "endpoint not available" error (kimi-cli parity). + #[tokio::test] + async fn fetch_usage_maps_404_to_not_found() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/usages")) + .respond_with(ResponseTemplate::new(404)) + .expect(1) + .mount(&server) + .await; + + let err = fetch_usage(&reqwest::Client::new(), &server.uri(), "tok") + .await + .expect_err("404 must fail"); + assert!(matches!(err, UsageError::NotFound)); + } + + /// Empty payload parses to zero rows (TUI shows "No usage data"). + #[test] + fn parse_usage_payload_empty_object_yields_no_rows() { + let usage = parse_usage_payload(&serde_json::json!({})); + assert!(usage.rows.is_empty()); } #[test] - fn billing_config_response_roundtrips_through_json() { - let config = BillingConfig { - credit_usage_percent: None, - current_period: None, - monthly_limit: Some(Cent { val: 5000 }), - used: Some(Cent { val: 123 }), - on_demand_cap: Some(Cent { val: 0 }), - on_demand_used: Some(Cent { val: 50 }), - prepaid_balance: Some(Cent { val: 750 }), - is_unified_billing_user: None, - billing_period_start: Some("2025-04-01T00:00:00Z".to_string()), - billing_period_end: Some("2025-05-01T00:00:00Z".to_string()), - history: vec![BillingPeriodUsage { - billing_cycle: Some(BillingCycle { - year: 2025, - month: 3, - }), - included_used: Some(Cent { val: 4500 }), - on_demand_used: Some(Cent { val: 100 }), - total_used: Some(Cent { val: 4600 }), - }], - }; - let resp = BillingConfigResponse { - config: Some(config), - on_demand_enabled: None, - subscription_tier: None, - }; - let json = serde_json::to_value(&resp).unwrap(); - let roundtripped: BillingConfigResponse = serde_json::from_value(json).unwrap(); - let rt_config = roundtripped.config.unwrap(); - assert_eq!(rt_config.monthly_limit.unwrap().val, 5000); - assert_eq!(rt_config.used.unwrap().val, 123); - assert_eq!(rt_config.prepaid_balance.unwrap().val, 750); - assert_eq!(rt_config.history.len(), 1); - } - - #[test] - fn billing_config_response_handles_null_config() { - let json = serde_json::json!({"config": null}); - let resp: BillingConfigResponse = serde_json::from_value(json).unwrap(); - assert!(resp.config.is_none()); - } - - #[test] - fn billing_config_response_handles_empty_history() { - let json = serde_json::json!({ - "config": { - "monthlyLimit": {"val": 1000}, - "used": {"val": 0} - } - }); - let resp: BillingConfigResponse = serde_json::from_value(json).unwrap(); - let config = resp.config.unwrap(); - assert_eq!(config.monthly_limit.unwrap().val, 1000); - assert!(config.history.is_empty()); - } - - #[test] - fn billing_config_serializes_camel_case() { - let config = BillingConfig { - credit_usage_percent: None, - current_period: None, - monthly_limit: Some(Cent { val: 100 }), - used: None, - on_demand_cap: None, - on_demand_used: None, - prepaid_balance: None, - is_unified_billing_user: None, - billing_period_start: None, - billing_period_end: None, - history: vec![], - }; - let json = serde_json::to_value(&config).unwrap(); - assert!(json.get("monthlyLimit").is_some()); - // Fields with None are skipped - assert!(json.get("creditUsagePercent").is_none()); - assert!(json.get("currentPeriod").is_none()); - assert!(json.get("used").is_none()); - assert!(json.get("onDemandCap").is_none()); - assert!(json.get("onDemandUsed").is_none()); - assert!(json.get("prepaidBalance").is_none()); - assert!(json.get("billingPeriodStart").is_none()); - // Empty history is skipped - assert!(json.get("history").is_none()); - } - - #[test] - fn billing_config_deserializes_credits_config_shape() { - // Newer `GetGrokCreditsConfig` response: percentage-based usage, - // a typed current period, and history keyed by `period`. - let json = serde_json::json!({ - "config": { - "creditUsagePercent": 42.5, - "currentPeriod": { - "type": "USAGE_PERIOD_TYPE_WEEKLY", - "start": "2026-06-01T00:00:00Z", - "end": "2026-06-08T00:00:00Z" - }, - "onDemandCap": {"val": 5000}, - "onDemandUsed": {"val": 300}, - "prepaidBalance": {"val": 1250}, - "isUnifiedBillingUser": true, - "productUsage": [ - {"product": "PRODUCT_GROK_BUILD", "usagePercent": 61.2} - ], - "history": [ - { - "period": { - "type": "USAGE_PERIOD_TYPE_WEEKLY", - "start": "2026-05-25T00:00:00Z", - "end": "2026-06-01T00:00:00Z" - }, - "onDemandUsed": {"val": 120} - } - ] - } - }); - let resp: BillingConfigResponse = serde_json::from_value(json).unwrap(); - let config = resp.config.unwrap(); - assert_eq!(config.credit_usage_percent, Some(42.5)); - let period = config.current_period.as_ref().unwrap(); - assert_eq!( - period.period_type.as_deref(), - Some("USAGE_PERIOD_TYPE_WEEKLY") - ); - assert_eq!(period.end.as_deref(), Some("2026-06-08T00:00:00Z")); - // Deprecated fields are absent in the credits shape. - assert!(config.monthly_limit.is_none()); - assert!(config.billing_period_end.is_none()); - assert_eq!(config.on_demand_cap.unwrap().val, 5000); - assert_eq!(config.on_demand_used.unwrap().val, 300); - // Bought (prepaid) credit balance is parsed from the credits config. - assert_eq!(config.prepaid_balance.unwrap().val, 1250); - assert_eq!(config.is_unified_billing_user, Some(true)); - // productUsage is still unused by the CLI billing surface. - assert_eq!(config.history.len(), 1); - assert_eq!(config.history[0].on_demand_used.as_ref().unwrap().val, 120); - } - - #[test] - fn cent_serializes_as_val_field() { - let c = Cent { val: 4299 }; - let json = serde_json::to_value(&c).unwrap(); - assert_eq!(json, serde_json::json!({"val": 4299})); + fn format_duration_matches_kimi_semantics() { + assert_eq!(format_duration(0), "0s"); + assert_eq!(format_duration(45), "45s"); + assert_eq!(format_duration(90), "1m"); + assert_eq!(format_duration(3_661), "1h 1m"); + assert_eq!(format_duration(90_000), "1d 1h"); } } diff --git a/crates/codegen/kigi-shell/src/extensions/bundle.rs b/crates/codegen/kigi-shell/src/extensions/bundle.rs deleted file mode 100644 index 0a0267c..0000000 --- a/crates/codegen/kigi-shell/src/extensions/bundle.rs +++ /dev/null @@ -1,1204 +0,0 @@ -//! ACP extension handlers for bundled subagent cache sync and status. -//! -//! These endpoints operate on the on-disk bundled cache only. Sync updates the -//! cache for future agent construction / future conversations; it does not live -//! reload the currently running `MvpAgent` instance. -use super::{ExtResult, parse_params, to_ext_response}; -use crate::agent::MvpAgent; -use crate::bundle::{self, BundleManifest}; -use crate::remote::{FetchedBundle, fetch_bundle}; -use agent_client_protocol as acp; -use anyhow::Context; -use kigi_tools::implementations::skills::discovery::extract_first_paragraph; -use serde::{Deserialize, Serialize}; -use std::path::Path; -use std::time::Duration; -/// Default freshness window for the proactive bundle sync. Bypassed by `force`. -pub(crate) const BUNDLE_SYNC_TTL: Duration = Duration::from_secs(60 * 60); -/// Error message returned when no auth source is available for a bundle sync. -/// -/// Hoisted to a constant so the user-facing wording stays in lockstep -/// across `sync_bundle`, `sync_bundle_to_root`, and any future call sites. -pub(crate) const NO_BUNDLE_CREDENTIALS_ERROR: &str = - "bundle sync requires either an authenticated cli-chat-proxy session or a deployment key"; -/// Whether the caller has any source of authentication that the -/// cli-chat-proxy `/v1/subagents/bundle` endpoint will accept. -/// -/// Centralised so the auth gate predicate stays consistent across: -/// - `sync_bundle` (user-triggered ACP entrypoint) -/// - `sync_bundle_to_root` (defense-in-depth on the public function) -/// - `maybe_sync_bundle_to_root` (proactive wrapper, silent skip on miss) -/// - `MvpAgent::maybe_sync_bundle_in_background` (post-auth pre-spawn gate) -/// -/// All four call sites previously inlined the same predicate; a future -/// auth-source addition (e.g., service-account token) only needs to land -/// here. -#[inline] -pub(crate) fn has_bundle_credentials( - auth_manager: Option<&std::sync::Arc>, - deployment_key: Option<&str>, -) -> bool { - auth_manager - .as_ref() - .is_some_and(|am| am.current_or_expired().is_some()) - || deployment_key.is_some() -} -#[derive(Debug, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -struct BundleSyncRequest { - #[serde(default)] - force: bool, -} -#[derive(Debug, Serialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct BundleSyncResult { - pub updated: bool, - pub version: String, - pub personas_count: usize, - pub roles_count: usize, - pub agents_count: usize, - pub skills_count: usize, -} -#[derive(Debug, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -struct BundleStatusRequest {} -#[derive(Debug, Serialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct BundleStatusResult { - pub has_cache: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub version: Option, - pub personas: Vec, - pub roles: Vec, - pub agents: Vec, - pub skills: Vec, - #[serde(skip_serializing_if = "Vec::is_empty")] - pub persona_details: Vec, - #[serde(skip_serializing_if = "Vec::is_empty")] - pub role_details: Vec, -} -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct PersonaDetail { - pub name: String, - pub description: Option, - pub has_inputs: bool, - pub has_outputs: bool, -} -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct RoleDetail { - pub name: String, - pub description: String, -} -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -struct EntryGetRequest { - kind: String, - name: String, -} -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct EntryGetResult { - pub kind: String, - pub name: String, - pub content: String, -} -pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { - match args.method.as_ref() { - "x.ai/bundle/sync" => { - let req: BundleSyncRequest = parse_params(args)?; - to_ext_response(sync_bundle(agent, req).await) - } - "x.ai/bundle/status" => { - let _req: BundleStatusRequest = parse_params(args)?; - to_ext_response(status_bundle()) - } - "x.ai/bundle/entry/get" => { - let req: EntryGetRequest = parse_params(args)?; - to_ext_response(get_entry(&req.kind, &req.name)) - } - _ => Err(acp::Error::method_not_found()), - } -} -async fn sync_bundle(agent: &MvpAgent, req: BundleSyncRequest) -> anyhow::Result { - let deployment_key = agent.deployment_key(); - if !has_bundle_credentials(Some(&agent.auth_manager), deployment_key.as_deref()) { - anyhow::bail!(NO_BUNDLE_CREDENTIALS_ERROR); - } - sync_bundle_to_root( - &bundle::bundled_root(), - &agent.cli_chat_proxy_base_url(), - Some(&agent.auth_manager), - deployment_key.as_deref(), - agent.alpha_test_key().as_deref(), - req.force, - ) - .await -} -/// `true` when `/manifest.json` exists, was written within `ttl`, and -/// is parseable as a [`BundleManifest`]. -/// -/// The parse check guards against the silent-skip failure mode where the -/// mtime is recent (e.g., a partial/aborted write) but the manifest is -/// truncated or otherwise corrupt. A bare mtime check would let -/// `maybe_sync_bundle_to_root` proactively skip a re-sync, leaving callers -/// (`status_bundle_at`, `SubagentsConfig::resolve`) to fail later with an -/// empty or stale catalog. Treating an unparseable manifest as "not fresh" -/// forces a re-sync on the next post-auth event. -pub(crate) fn bundle_cache_is_fresh(root: &Path, ttl: Duration) -> bool { - let manifest = root.join("manifest.json"); - let Ok(meta) = std::fs::metadata(&manifest) else { - return false; - }; - let Ok(modified) = meta.modified() else { - return false; - }; - let within_ttl = modified - .elapsed() - .map(|elapsed| elapsed < ttl) - .unwrap_or(false); - if !within_ttl { - return false; - } - matches!(bundle::read_cached_manifest(root), Ok(Some(_))) -} -/// Proactive variant of [`sync_bundle_to_root`] that respects an auth gate -/// and a TTL guard. -/// -/// Returns: -/// - `Ok(Some(result))` when a sync was performed. -/// - `Ok(None)` when the call was skipped (no credentials or cache fresh). -/// - `Err(_)` when sync was attempted but the network call or extract failed. -pub(crate) async fn maybe_sync_bundle_to_root( - root: &Path, - proxy_base_url: &str, - auth_manager: Option<&std::sync::Arc>, - deployment_key: Option<&str>, - alpha_test_key: Option<&str>, - force: bool, - ttl: Duration, -) -> anyhow::Result> { - if !has_bundle_credentials(auth_manager, deployment_key) { - tracing::debug!("proactive bundle sync skipped: no auth and no deployment key"); - return Ok(None); - } - if !force && bundle_cache_is_fresh(root, ttl) { - tracing::debug!( - ttl_secs = ttl.as_secs(), - "proactive bundle sync skipped: cache is fresh" - ); - return Ok(None); - } - sync_bundle_to_root( - root, - proxy_base_url, - auth_manager, - deployment_key, - alpha_test_key, - force, - ) - .await - .map(Some) -} -pub(crate) async fn sync_bundle_to_root( - root: &Path, - proxy_base_url: &str, - auth_manager: Option<&std::sync::Arc>, - deployment_key: Option<&str>, - alpha_test_key: Option<&str>, - _force: bool, -) -> anyhow::Result { - if !has_bundle_credentials(auth_manager, deployment_key) { - anyhow::bail!(NO_BUNDLE_CREDENTIALS_ERROR); - } - let fetched = - fetch_bundle(proxy_base_url, auth_manager, deployment_key, alpha_test_key).await?; - match fetched { - FetchedBundle::Archive(bytes) => { - let root_owned = root.to_path_buf(); - let manifest = tokio::task::spawn_blocking(move || { - bundle::extract_bundle_archive(&root_owned, &bytes) - }) - .await - .context("bundle extract task panicked")??; - let personas_count = bundle::count_entries_by_prefix(&manifest, "personas/"); - let roles_count = bundle::count_entries_by_prefix(&manifest, "roles/"); - let agents_count = bundle::count_entries_by_prefix(&manifest, "agents/"); - let skills_count = bundle::count_entries_by_prefix(&manifest, "skills/"); - Ok(BundleSyncResult { - updated: true, - version: manifest.version, - personas_count, - roles_count, - agents_count, - skills_count, - }) - } - FetchedBundle::Legacy(legacy_bundle) => { - let version = legacy_bundle.version.clone(); - let personas_count = legacy_bundle.personas.len(); - let roles_count = legacy_bundle.roles.len(); - let agents_count = legacy_bundle.agents.len(); - let skills_count = legacy_bundle.skills.len(); - let root_owned = root.to_path_buf(); - tokio::task::spawn_blocking(move || { - bundle::write_bundle_to_cache(&root_owned, &legacy_bundle) - }) - .await - .context("bundle write task panicked")??; - Ok(BundleSyncResult { - updated: true, - version, - personas_count, - roles_count, - agents_count, - skills_count, - }) - } - } -} -fn get_entry(kind: &str, name: &str) -> anyhow::Result { - get_entry_at(&bundle::bundled_root(), kind, name) -} -fn validate_entry_name(name: &str) -> anyhow::Result<()> { - if name.is_empty() - || name.contains('/') - || name.contains('\\') - || name.contains("..") - || name == "." - { - anyhow::bail!("invalid entry name: {name}"); - } - Ok(()) -} -fn get_entry_at(root: &Path, kind: &str, name: &str) -> anyhow::Result { - validate_entry_name(name)?; - let (dir_name, ext) = match kind { - "persona" => ("personas", "toml"), - "role" => ("roles", "toml"), - "agent" => ("agents", "md"), - _ => anyhow::bail!("unknown entry kind: {kind}"), - }; - let path = root.join(dir_name).join(format!("{name}.{ext}")); - let content = std::fs::read_to_string(&path) - .with_context(|| format!("{kind} '{name}' not found in bundle cache"))?; - Ok(EntryGetResult { - kind: kind.to_owned(), - name: name.to_owned(), - content, - }) -} -fn status_bundle() -> anyhow::Result { - status_bundle_at(&bundle::bundled_root()) -} -fn status_bundle_at(root: &Path) -> anyhow::Result { - let Some(manifest) = bundle::read_cached_manifest(root)? else { - return Ok(BundleStatusResult { - has_cache: false, - version: None, - personas: Vec::new(), - roles: Vec::new(), - agents: Vec::new(), - skills: Vec::new(), - persona_details: Vec::new(), - role_details: Vec::new(), - }); - }; - let personas = list_cached_entries(root, &manifest, "personas", "toml"); - let roles = list_cached_entries(root, &manifest, "roles", "toml"); - let agents = list_cached_entries(root, &manifest, "agents", "md"); - let skills = list_cached_skill_entries(root, &manifest); - let persona_details = personas - .iter() - .filter_map(|name| persona_detail_from_toml(name, root)) - .collect(); - let role_details = roles - .iter() - .filter_map(|name| role_detail_from_toml(name, root)) - .collect(); - Ok(BundleStatusResult { - has_cache: true, - version: Some(manifest.version.clone()), - personas, - roles, - agents, - skills, - persona_details, - role_details, - }) -} -fn persona_detail_from_toml(name: &str, root: &Path) -> Option { - let path = root.join("personas").join(format!("{name}.toml")); - let content = std::fs::read_to_string(&path).ok()?; - let table: toml::Value = toml::from_str(&content).ok()?; - let desc = table - .get("description") - .and_then(|v| v.as_str()) - .filter(|s| !s.trim().is_empty()) - .map(str::to_owned) - .or_else(|| { - table - .get("instructions") - .and_then(|v| v.as_str()) - .and_then(extract_first_paragraph) - }); - let has_inputs = table - .get("inputs") - .and_then(|v| v.as_array()) - .is_some_and(|a| !a.is_empty()); - let has_outputs = table - .get("outputs") - .and_then(|v| v.as_array()) - .is_some_and(|a| !a.is_empty()); - Some(PersonaDetail { - name: name.to_owned(), - description: desc, - has_inputs, - has_outputs, - }) -} -fn role_detail_from_toml(name: &str, root: &Path) -> Option { - let path = root.join("roles").join(format!("{name}.toml")); - let content = std::fs::read_to_string(&path).ok()?; - let table: toml::Value = toml::from_str(&content).ok()?; - let desc = table - .get("description") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_owned(); - Some(RoleDetail { - name: name.to_owned(), - description: desc, - }) -} -fn list_cached_entries( - root: &Path, - manifest: &BundleManifest, - dir_name: &str, - extension: &str, -) -> Vec { - let dir = root.join(dir_name); - let Ok(entries) = std::fs::read_dir(&dir) else { - return Vec::new(); - }; - let mut names: Vec = entries - .flatten() - .filter_map(|entry| { - let path = entry.path(); - if !path.is_file() { - return None; - } - let file_name = path.file_name()?.to_str()?; - let relative_path = format!("{dir_name}/{file_name}"); - if !manifest.checksums.contains_key(&relative_path) { - return None; - } - match path.extension().and_then(|ext| ext.to_str()) { - Some(ext) if ext == extension => path - .file_stem() - .and_then(|stem| stem.to_str()) - .map(ToOwned::to_owned), - _ => None, - } - }) - .collect(); - names.sort(); - names -} -fn list_cached_skill_entries(root: &Path, manifest: &BundleManifest) -> Vec { - let prefix = "skills/"; - let mut names: Vec = manifest - .checksums - .keys() - .filter_map(|k| { - let name = k.strip_prefix(prefix)?.strip_suffix("/SKILL.md")?; - root.join(k).is_file().then(|| name.to_owned()) - }) - .collect(); - names.sort(); - names -} -#[cfg(test)] -mod tests { - use super::*; - use axum::{ - Router, - extract::State, - http::{HeaderMap, StatusCode}, - routing::get, - }; - use prod_mc_cli_chat_proxy_types::SubagentBundle; - use serial_test::serial; - use std::sync::{Arc, Mutex}; - use tempfile::TempDir; - struct HomeGuard { - previous: Option, - } - impl Drop for HomeGuard { - fn drop(&mut self) { - match self.previous.take() { - Some(previous) => unsafe { - std::env::set_var("HOME", previous); - }, - None => unsafe { - std::env::remove_var("HOME"); - }, - } - } - } - fn with_bundled_home(tmp: &TempDir) -> HomeGuard { - let previous = std::env::var_os("HOME"); - unsafe { - std::env::set_var("HOME", tmp.path()); - } - HomeGuard { previous } - } - fn sample_bundle() -> SubagentBundle { - let mut bundle = SubagentBundle::empty("bundle-v1"); - bundle.personas.insert( - "researcher".to_string(), - concat!( - "instructions = \"You are a thorough researcher.\\nDig deep.\"\n", - "[[inputs]]\nname = \"topic\"\n", - "[[outputs]]\nname = \"report\"\n", - ) - .to_string(), - ); - bundle.roles.insert( - "reviewer".to_string(), - "description = \"Meticulous code reviewer\"\n".to_string(), - ); - bundle - .agents - .insert("default".to_string(), "# agent\n".to_string()); - bundle - } - fn sample_bundle_with_skills() -> SubagentBundle { - let mut bundle = sample_bundle(); - bundle - .skills - .insert("commit".to_string(), "# Commit skill\n".to_string()); - bundle - .skills - .insert("review".to_string(), "# Review skill\n".to_string()); - bundle - } - fn test_auth() -> crate::auth::KimiAuth { - crate::auth::KimiAuth { - key: "token".to_string(), - auth_mode: crate::auth::AuthMode::OAuth, - create_time: chrono::Utc::now(), - user_id: "user-1".to_string(), - email: Some("test@example.com".to_string()), - refresh_token: None, - expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - expires_in: Some(3600), - scope: None, - token_type: None, - } - } - fn test_auth_manager() -> Arc { - let dir = tempfile::tempdir().unwrap(); - let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::KimiCodeConfig::default()); - mgr.hot_swap(test_auth()); - std::mem::forget(dir); - Arc::new(mgr) - } - #[derive(Debug, Default, Clone)] - struct SeenHeaders { - authorization: Option, - token_auth: Option, - user_id: Option, - email: Option, - alpha_test_key: Option, - } - #[derive(Clone)] - struct BundleServerState { - body: serde_json::Value, - status_code: StatusCode, - seen_headers: Arc>>, - } - async fn start_bundle_server( - status_code: StatusCode, - body: serde_json::Value, - ) -> ( - String, - Arc>>, - tokio::task::JoinHandle<()>, - ) { - let seen_headers = Arc::new(Mutex::new(Vec::new())); - let state = BundleServerState { - body, - status_code, - seen_headers: seen_headers.clone(), - }; - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let app = Router::new() - .route( - "/v1/subagents/bundle", - get( - |State(state): State, headers: HeaderMap| async move { - state.seen_headers.lock().unwrap().push(SeenHeaders { - authorization: headers - .get("authorization") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - token_auth: headers - .get("x-xai-token-auth") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - user_id: headers - .get("x-userid") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - email: headers - .get("x-email") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - alpha_test_key: { - let _ = &headers; - None - }, - }); - (state.status_code, axum::Json(state.body)) - }, - ), - ) - .with_state(state); - let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (format!("{base}/v1"), seen_headers, handle) - } - #[test] - #[serial] - fn status_reports_no_cache_when_manifest_missing() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let status = status_bundle_at(&bundle::bundled_root()).unwrap(); - assert_eq!( - status, - BundleStatusResult { - has_cache: false, - version: None, - personas: vec![], - roles: vec![], - agents: vec![], - skills: vec![], - persona_details: vec![], - role_details: vec![], - } - ); - } - #[test] - #[serial] - fn status_reports_cached_entries_from_manifest_and_disk() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let root = bundle::bundled_root(); - bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap(); - std::fs::write( - root.join("personas/local-only.toml"), - "instructions = \"ignore\"", - ) - .unwrap(); - let status = status_bundle_at(&bundle::bundled_root()).unwrap(); - assert!(status.has_cache); - assert_eq!(status.version.as_deref(), Some("bundle-v1")); - assert_eq!(status.personas, vec!["researcher"]); - assert_eq!(status.roles, vec!["reviewer"]); - assert_eq!(status.agents, vec!["default"]); - assert_eq!(status.skills, Vec::::new()); - } - #[tokio::test(flavor = "current_thread")] - #[serial] - async fn sync_success_writes_cache_and_returns_counts() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let bundle = sample_bundle(); - let (proxy_base_url, _seen_headers, server) = start_bundle_server( - axum::http::StatusCode::OK, - serde_json::to_value(&bundle).unwrap(), - ) - .await; - let root = bundle::bundled_root(); - let am = test_auth_manager(); - let result = sync_bundle_to_root(&root, &proxy_base_url, Some(&am), None, None, false) - .await - .unwrap(); - assert_eq!(result.version, "bundle-v1"); - assert_eq!(result.personas_count, 1); - assert_eq!(result.roles_count, 1); - assert_eq!(result.agents_count, 1); - assert_eq!(result.skills_count, 0); - assert!(root.join("personas/researcher.toml").exists()); - assert!(root.join("roles/reviewer.toml").exists()); - assert!(root.join("agents/default.md").exists()); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - #[serial] - async fn sync_force_true_has_same_write_semantics() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let bundle = sample_bundle(); - let (proxy_base_url, _seen_headers, server) = - start_bundle_server(StatusCode::OK, serde_json::to_value(&bundle).unwrap()).await; - let root = bundle::bundled_root(); - let am = test_auth_manager(); - let normal = sync_bundle_to_root(&root, &proxy_base_url, Some(&am), None, None, false) - .await - .unwrap(); - let forced = sync_bundle_to_root(&root, &proxy_base_url, Some(&am), None, None, true) - .await - .unwrap(); - assert_eq!(forced, normal); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - #[serial] - async fn sync_http_failure_surfaces_error() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let (proxy_base_url, _seen_headers, server) = start_bundle_server( - StatusCode::UNAUTHORIZED, - serde_json::json!({ "error" : "unauthorized" }), - ) - .await; - let am = test_auth_manager(); - let error = sync_bundle_to_root( - &bundle::bundled_root(), - &proxy_base_url, - Some(&am), - None, - None, - false, - ) - .await - .unwrap_err(); - assert!(error.to_string().contains("401")); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - #[serial] - async fn sync_uses_deployment_key_auth_mode() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let bundle = sample_bundle(); - let (proxy_base_url, seen_headers, server) = - start_bundle_server(StatusCode::OK, serde_json::to_value(&bundle).unwrap()).await; - let am = test_auth_manager(); - let result = sync_bundle_to_root( - &bundle::bundled_root(), - &proxy_base_url, - Some(&am), - Some("deploy-key"), - None, - false, - ) - .await - .unwrap(); - assert_eq!(result.version, "bundle-v1"); - let headers = seen_headers.lock().unwrap(); - let headers = headers.last().unwrap(); - assert_eq!(headers.authorization.as_deref(), Some("Bearer deploy-key")); - assert_eq!(headers.token_auth, None); - assert_eq!(headers.user_id, None); - assert_eq!(headers.email, None); - assert_eq!(headers.alpha_test_key, None); - server.abort(); - } - #[test] - #[serial] - fn status_only_reports_bundled_cache_not_higher_priority_sources() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let root = bundle::bundled_root(); - bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap(); - let project_root = tmp.path().join("workspace"); - std::fs::create_dir_all(project_root.join(".kigi/personas")).unwrap(); - std::fs::create_dir_all(project_root.join(".kigi/roles")).unwrap(); - std::fs::write( - project_root.join(".kigi/personas/researcher.toml"), - "instructions = \"project persona\"\n", - ) - .unwrap(); - std::fs::write( - project_root.join(".kigi/roles/reviewer.toml"), - "description = \"project role\"\n", - ) - .unwrap(); - let config = crate::config::SubagentsConfig::resolve( - false, - &toml::Value::Table(Default::default()), - Some(&project_root), - ); - assert_eq!( - config - .personas - .get("researcher") - .and_then(|persona| persona.instructions.as_deref()), - Some("project persona") - ); - assert_eq!( - config - .roles - .get("reviewer") - .map(|role| role.description.as_str()), - Some("project role") - ); - let status = status_bundle_at(&root).unwrap(); - assert_eq!(status.personas, vec!["researcher"]); - assert_eq!(status.roles, vec!["reviewer"]); - assert_eq!(status.agents, vec!["default"]); - assert_eq!(status.skills, Vec::::new()); - } - #[test] - #[serial] - fn sync_requires_auth_or_deployment_key() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let error = futures::executor::block_on(sync_bundle_to_root( - &bundle::bundled_root(), - "http://127.0.0.1:1/v1", - None, - None, - None, - false, - )) - .unwrap_err(); - assert!( - error.to_string() - .contains("bundle sync requires either an authenticated cli-chat-proxy session or a deployment key") - ); - } - #[test] - #[serial] - fn get_entry_reads_persona_file() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let root = bundle::bundled_root(); - bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap(); - let result = get_entry_at(&root, "persona", "researcher").unwrap(); - assert_eq!(result.kind, "persona"); - assert_eq!(result.name, "researcher"); - assert!(result.content.contains("instructions")); - } - #[test] - #[serial] - fn get_entry_unknown_kind_returns_error() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let root = bundle::bundled_root(); - let err = get_entry_at(&root, "widget", "foo").unwrap_err(); - assert!(err.to_string().contains("unknown entry kind: widget")); - } - #[test] - #[serial] - fn get_entry_missing_file_returns_error() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let root = bundle::bundled_root(); - bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap(); - let err = get_entry_at(&root, "persona", "nonexistent").unwrap_err(); - assert!(err.to_string().contains("not found in bundle cache")); - } - #[test] - fn get_entry_rejects_path_traversal() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("bundled"); - for bad_name in ["../../../etc/passwd", "foo/bar", "a\\b", "..", "."] { - let err = get_entry_at(&root, "persona", bad_name).unwrap_err(); - assert!( - err.to_string().contains("invalid entry name"), - "expected rejection for {bad_name:?}, got: {err}" - ); - } - } - #[test] - fn get_entry_rejects_empty_name() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("bundled"); - let err = get_entry_at(&root, "persona", "").unwrap_err(); - assert!(err.to_string().contains("invalid entry name")); - } - #[test] - #[serial] - fn status_includes_persona_and_role_details() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let root = bundle::bundled_root(); - bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap(); - let status = status_bundle_at(&root).unwrap(); - assert_eq!(status.persona_details.len(), 1); - let pd = &status.persona_details[0]; - assert_eq!(pd.name, "researcher"); - assert_eq!( - pd.description.as_deref(), - Some("You are a thorough researcher. Dig deep.") - ); - assert!(pd.has_inputs); - assert!(pd.has_outputs); - assert_eq!(status.role_details.len(), 1); - let rd = &status.role_details[0]; - assert_eq!(rd.name, "reviewer"); - assert_eq!(rd.description, "Meticulous code reviewer"); - } - #[test] - #[serial] - fn status_without_toml_files_returns_empty_details() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let root = bundle::bundled_root(); - let mut bundle = SubagentBundle::empty("v1"); - bundle - .personas - .insert("ghost".to_string(), "instructions = \"x\"".to_string()); - bundle::write_bundle_to_cache(&root, &bundle).unwrap(); - std::fs::remove_file(root.join("personas/ghost.toml")).unwrap(); - let status = status_bundle_at(&root).unwrap(); - assert!(status.personas.is_empty()); - assert!(status.persona_details.is_empty()); - } - #[test] - fn malformed_toml_skipped_gracefully() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("cache"); - std::fs::create_dir_all(root.join("personas")).unwrap(); - std::fs::write(root.join("personas/bad.toml"), "{{{{not toml").unwrap(); - assert!(persona_detail_from_toml("bad", &root).is_none()); - } - #[test] - fn persona_detail_without_inputs_outputs() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path(); - std::fs::create_dir_all(root.join("personas")).unwrap(); - std::fs::write( - root.join("personas/simple.toml"), - "instructions = \"Just a simple persona\"", - ) - .unwrap(); - let detail = persona_detail_from_toml("simple", root).unwrap(); - assert_eq!(detail.name, "simple"); - assert_eq!(detail.description.as_deref(), Some("Just a simple persona")); - assert!(!detail.has_inputs); - assert!(!detail.has_outputs); - } - #[test] - fn role_detail_missing_description_defaults_empty() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path(); - std::fs::create_dir_all(root.join("roles")).unwrap(); - std::fs::write(root.join("roles/bare.toml"), "instructions = \"hi\"").unwrap(); - let detail = role_detail_from_toml("bare", root).unwrap(); - assert_eq!(detail.name, "bare"); - assert_eq!(detail.description, ""); - } - #[tokio::test(flavor = "current_thread")] - #[serial] - async fn sync_with_skills_reports_skills_count() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let bundle = sample_bundle_with_skills(); - let (proxy_base_url, _seen_headers, server) = - start_bundle_server(StatusCode::OK, serde_json::to_value(&bundle).unwrap()).await; - let result = sync_bundle_to_root( - &bundle::bundled_root(), - &proxy_base_url, - Some(&test_auth_manager()), - None, - None, - false, - ) - .await - .unwrap(); - assert_eq!(result.skills_count, 2); - assert_eq!(result.personas_count, 1); - assert_eq!(result.roles_count, 1); - assert_eq!(result.agents_count, 1); - server.abort(); - } - #[test] - #[serial] - fn status_lists_skill_names_from_manifest() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let root = bundle::bundled_root(); - bundle::write_bundle_to_cache(&root, &sample_bundle_with_skills()).unwrap(); - let status = status_bundle_at(&root).unwrap(); - assert!(status.has_cache); - assert_eq!(status.skills, vec!["commit", "review"]); - assert_eq!(status.personas, vec!["researcher"]); - } - #[test] - #[serial] - fn status_skills_only_lists_files_present_on_disk() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let root = bundle::bundled_root(); - bundle::write_bundle_to_cache(&root, &sample_bundle_with_skills()).unwrap(); - std::fs::remove_file(root.join("skills/commit/SKILL.md")).unwrap(); - let status = status_bundle_at(&root).unwrap(); - assert_eq!(status.skills, vec!["review"]); - } - use crate::bundle::test_helpers::make_test_archive; - #[derive(Clone)] - struct ArchiveServerState { - archive_bytes: Vec, - } - async fn start_archive_bundle_server( - archive_bytes: Vec, - ) -> (String, tokio::task::JoinHandle<()>) { - let state = ArchiveServerState { archive_bytes }; - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let app = Router::new() - .route( - "/v1/bundle/archive", - get(|State(state): State| async move { - (StatusCode::OK, state.archive_bytes) - }), - ) - .with_state(state); - let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (format!("{base}/v1"), handle) - } - #[tokio::test(flavor = "current_thread")] - #[serial] - async fn sync_with_archive_endpoint_extracts_and_reports_counts() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let archive = make_test_archive(&[ - ("bundle.json", br#"{"version":"archive-v1"}"#), - ( - "subagents/personas/researcher.toml", - b"instructions = \"hello\"", - ), - ("subagents/roles/reviewer.toml", b"description = \"review\""), - ("skills/commit/SKILL.md", b"# Commit skill"), - ]); - let (proxy_base_url, server) = start_archive_bundle_server(archive).await; - let result = sync_bundle_to_root( - &bundle::bundled_root(), - &proxy_base_url, - Some(&test_auth_manager()), - None, - None, - false, - ) - .await - .unwrap(); - assert_eq!(result.version, "archive-v1"); - assert_eq!(result.personas_count, 1); - assert_eq!(result.roles_count, 1); - assert_eq!(result.agents_count, 0); - assert_eq!(result.skills_count, 1); - assert!( - bundle::bundled_root() - .join("personas/researcher.toml") - .exists() - ); - assert!( - bundle::bundled_root() - .join("skills/commit/SKILL.md") - .exists() - ); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - #[serial] - async fn sync_falls_back_to_legacy_when_archive_unavailable() { - let tmp = TempDir::new().unwrap(); - let _home = with_bundled_home(&tmp); - let bundle = sample_bundle_with_skills(); - let (proxy_base_url, _seen_headers, server) = - start_bundle_server(StatusCode::OK, serde_json::to_value(&bundle).unwrap()).await; - let result = sync_bundle_to_root( - &bundle::bundled_root(), - &proxy_base_url, - Some(&test_auth_manager()), - None, - None, - false, - ) - .await - .unwrap(); - assert_eq!(result.version, "bundle-v1"); - assert_eq!(result.personas_count, 1); - assert_eq!(result.skills_count, 2); - server.abort(); - } - fn backdate_manifest(root: &std::path::Path, age: Duration) { - let path = root.join("manifest.json"); - let stale_time = std::time::SystemTime::now() - age; - let times = std::fs::FileTimes::new().set_modified(stale_time); - std::fs::File::options() - .write(true) - .open(&path) - .expect("open manifest for backdate") - .set_times(times) - .expect("set manifest mtime"); - } - #[tokio::test(flavor = "current_thread")] - async fn maybe_sync_skips_without_auth_or_deployment_key() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("bundled"); - let (proxy_base_url, seen_headers, server) = start_bundle_server( - StatusCode::OK, - serde_json::to_value(sample_bundle()).unwrap(), - ) - .await; - let result = maybe_sync_bundle_to_root( - &root, - &proxy_base_url, - None, - None, - None, - false, - BUNDLE_SYNC_TTL, - ) - .await - .unwrap(); - assert!(result.is_none(), "expected sync skipped"); - assert!( - seen_headers.lock().unwrap().is_empty(), - "auth-gated sync must not hit the network" - ); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - async fn maybe_sync_skips_when_cache_is_fresh() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("bundled"); - bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap(); - let (proxy_base_url, seen_headers, server) = start_bundle_server( - StatusCode::OK, - serde_json::to_value(sample_bundle()).unwrap(), - ) - .await; - let result = maybe_sync_bundle_to_root( - &root, - &proxy_base_url, - Some(&test_auth_manager()), - None, - None, - false, - BUNDLE_SYNC_TTL, - ) - .await - .unwrap(); - assert!(result.is_none(), "fresh cache should skip sync"); - assert!( - seen_headers.lock().unwrap().is_empty(), - "fresh cache must not hit the network" - ); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - async fn maybe_sync_runs_when_cache_is_stale() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("bundled"); - bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap(); - backdate_manifest(&root, BUNDLE_SYNC_TTL + Duration::from_secs(60)); - let (proxy_base_url, seen_headers, server) = start_bundle_server( - StatusCode::OK, - serde_json::to_value(sample_bundle()).unwrap(), - ) - .await; - let outcome = maybe_sync_bundle_to_root( - &root, - &proxy_base_url, - Some(&test_auth_manager()), - None, - None, - false, - BUNDLE_SYNC_TTL, - ) - .await - .unwrap() - .expect("stale cache should trigger a sync"); - assert_eq!(outcome.version, "bundle-v1"); - assert_eq!(seen_headers.lock().unwrap().len(), 1); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - async fn maybe_sync_force_bypasses_ttl() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("bundled"); - bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap(); - let (proxy_base_url, seen_headers, server) = start_bundle_server( - StatusCode::OK, - serde_json::to_value(sample_bundle()).unwrap(), - ) - .await; - let outcome = maybe_sync_bundle_to_root( - &root, - &proxy_base_url, - Some(&test_auth_manager()), - None, - None, - true, - BUNDLE_SYNC_TTL, - ) - .await - .unwrap() - .expect("force=true should always sync"); - assert_eq!(outcome.version, "bundle-v1"); - assert_eq!(seen_headers.lock().unwrap().len(), 1); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - async fn maybe_sync_runs_when_no_manifest_exists() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("bundled"); - let (proxy_base_url, seen_headers, server) = start_bundle_server( - StatusCode::OK, - serde_json::to_value(sample_bundle()).unwrap(), - ) - .await; - let outcome = maybe_sync_bundle_to_root( - &root, - &proxy_base_url, - Some(&test_auth_manager()), - None, - None, - false, - BUNDLE_SYNC_TTL, - ) - .await - .unwrap() - .expect("missing manifest should trigger a sync"); - assert_eq!(outcome.personas_count, 1); - assert_eq!(seen_headers.lock().unwrap().len(), 1); - server.abort(); - } - #[test] - fn bundle_cache_is_fresh_returns_false_without_manifest() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("bundled"); - assert!(!bundle_cache_is_fresh(&root, BUNDLE_SYNC_TTL)); - } - #[test] - fn bundle_cache_is_fresh_true_for_recent_manifest() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("bundled"); - bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap(); - assert!(bundle_cache_is_fresh(&root, BUNDLE_SYNC_TTL)); - } - #[test] - fn bundle_cache_is_fresh_false_for_stale_manifest() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("bundled"); - bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap(); - backdate_manifest(&root, BUNDLE_SYNC_TTL + Duration::from_secs(60)); - assert!(!bundle_cache_is_fresh(&root, BUNDLE_SYNC_TTL)); - } - #[test] - fn bundle_cache_is_fresh_false_for_corrupted_manifest() { - let tmp = TempDir::new().unwrap(); - let root = tmp.path().join("bundled"); - bundle::write_bundle_to_cache(&root, &sample_bundle()).unwrap(); - std::fs::write(root.join("manifest.json"), "{not json}").unwrap(); - assert!(!bundle_cache_is_fresh(&root, BUNDLE_SYNC_TTL)); - } -} diff --git a/crates/codegen/kigi-shell/src/extensions/feedback.rs b/crates/codegen/kigi-shell/src/extensions/feedback.rs index 711c709..f277373 100644 --- a/crates/codegen/kigi-shell/src/extensions/feedback.rs +++ b/crates/codegen/kigi-shell/src/extensions/feedback.rs @@ -1,12 +1,15 @@ //! `x.ai/feedback`, `x.ai/feedback/dismiss`, `x.ai/btw`, and `x.ai/review/*` //! extension handlers. //! -//! - `feedback`/`feedback/dismiss`: persist user ratings/text locally and -//! forward to cli-chat-proxy. +//! - `feedback`/`feedback/dismiss`: persist user ratings/text locally; text +//! feedback from subscription (OAuth) sessions is forwarded to the Kimi +//! Code feedback endpoint (`POST {base}/feedback`, kimi-cli slash.py +//! parity). Without a subscription session the record stays local and the +//! response points at the GitHub issue tracker. //! - `btw`: dispatch a side question to the active session via //! `SessionCommand::SideQuestion` and return the answer. //! - `review/comment` and `review/comment/delete`: record inline code review -//! events to cloud storage. +//! events locally. use std::sync::Arc; @@ -15,6 +18,7 @@ use tokio::sync::oneshot; use super::{ExtResult, parse_params}; use crate::agent::MvpAgent; +use crate::agent::feedback_client::FEEDBACK_ISSUES_URL; use crate::session::persistence::{LocalFeedbackEntry, UserFeedbackEntry}; use crate::session::{ ClientFeedbackInput, CommentDeleteRequest, CommentDeleteResponse, CommentRequest, @@ -34,7 +38,7 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { } m if m.starts_with("x.ai/review") => { tracing::info!("handling review comment"); - handle_review(agent, args).await + handle_review(args).await } _ => Err(acp::Error::method_not_found()), } @@ -97,8 +101,7 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult let simple: crate::session::FeedbackRequest = parse_params(args)?; ClientFeedbackInput { session_id: simple.session_id, - client_type: - prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Tui, + client_type: crate::session::feedback_types::ClientType::Tui, rating_type: None, rating_value: None, feedback_text: Some(simple.feedback_text), @@ -151,7 +154,7 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult submission.merge_metadata(user_meta); } - // Enrich with session context for Slack notifications (best-effort). + // Enrich with session context (persisted alongside the record). if let Some(ref session_handle) = session_handle { let (tx, rx) = tokio::sync::oneshot::channel(); let _ = session_handle @@ -174,7 +177,7 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult if let (Some(session_handle), Some(rating_value)) = (&session_handle, feedback_input.rating_value) { - use prod_mc_cli_chat_proxy_types::feedback_types::RatingType; + use crate::session::feedback_types::RatingType; let (is_positive, is_negative) = match feedback_input.rating_type { // Thumbs: -1 = down, 0 = neutral, 1 = up Some(RatingType::Thumbs) | None => (rating_value > 0, rating_value < 0), @@ -208,8 +211,8 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult let client = agent.feedback_client(); if client.is_none() { - tracing::warn!( - "no feedback client available (missing proxy credentials); feedback saved locally only" + tracing::info!( + "no subscription session; feedback saved locally — submit at {FEEDBACK_ISSUES_URL}" ); } let outcome = crate::session::feedback_manager::submit_feedback_workflow( @@ -222,15 +225,17 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult match &outcome { crate::session::feedback_manager::SubmitOutcome::Submitted => { - tracing::info!("feedback submitted to proxy successfully"); + tracing::info!("feedback submitted to the Kimi Code endpoint"); } crate::session::feedback_manager::SubmitOutcome::LocalOnly => { - tracing::warn!("feedback saved locally only (no proxy client)"); + tracing::info!("feedback saved locally only"); } crate::session::feedback_manager::SubmitOutcome::Failed(e) => { - tracing::error!(error = %e, "feedback submission to proxy failed"); - return Err(acp::Error::internal_error() - .data(format!("Feedback submission failed: {e}"))); + tracing::error!(error = %e, "feedback submission failed"); + return Err(acp::Error::internal_error().data(format!( + "Feedback submission failed: {e}. \ + Please submit feedback at {FEEDBACK_ISSUES_URL}" + ))); } } @@ -281,36 +286,14 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult } } - let request_id = dismiss_input.request_id.clone(); - let client = agent - .feedback_client() - .ok_or_else(|| acp::Error::internal_error().data("No credentials for feedback"))?; - let feedback_base_url = agent.cfg.borrow().endpoints.resolve_feedback_base_url(); - match client.dismiss_request(&request_id).await { - Ok(response) => { - tracing::info!( - request_id = %response.request_id, - status = %response.status, - feedback_url = %feedback_base_url, - "Feedback request dismissed" - ); - let value = serde_json::to_value(&response) - .map(|value| serde_json::value::to_raw_value(&value).map(Arc::from)) - .expect("to work") - .expect("to work"); - Ok(acp::ExtResponse::new(value)) - } - Err(e) => { - tracing::warn!( - error = %e, - request_id = %request_id, - feedback_url = %feedback_base_url, - "Failed to dismiss feedback request" - ); - Err(acp::Error::internal_error() - .data(format!("Failed to dismiss feedback request: {e}"))) - } - } + let value = serde_json::to_value(serde_json::json!({ + "requestId": dismiss_input.request_id, + "status": "dismissed", + })) + .map(|value| serde_json::value::to_raw_value(&value).map(Arc::from)) + .expect("to work") + .expect("to work"); + Ok(acp::ExtResponse::new(value)) } _ => Err(acp::Error::method_not_found()), } @@ -319,9 +302,9 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult /// Record inline code review events. /// /// Methods: -/// - `x.ai/review/comment`: record a new inline code comment to cloud storage +/// - `x.ai/review/comment`: record a new inline code comment /// - `x.ai/review/comment/delete`: record a tombstone event for a deleted comment -async fn handle_review(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { +async fn handle_review(args: &acp::ExtRequest) -> ExtResult { match args.method.as_ref() { "x.ai/review/comment" => { let request: CommentRequest = parse_params(args)?; diff --git a/crates/codegen/kigi-shell/src/extensions/mod.rs b/crates/codegen/kigi-shell/src/extensions/mod.rs index 9dca80e..28d883f 100644 --- a/crates/codegen/kigi-shell/src/extensions/mod.rs +++ b/crates/codegen/kigi-shell/src/extensions/mod.rs @@ -1,7 +1,6 @@ pub mod auth; pub(crate) mod auth_gate; pub mod billing; -pub mod bundle; pub mod chat_conversation_history; pub mod code_nav; pub mod debug; @@ -28,7 +27,6 @@ pub mod search; pub mod session_admin; pub mod session_search; pub mod session_updates; -pub mod share; pub mod skills; pub mod suggest; pub mod task; diff --git a/crates/codegen/kigi-shell/src/extensions/session_admin.rs b/crates/codegen/kigi-shell/src/extensions/session_admin.rs index 3d23b2d..333748f 100644 --- a/crates/codegen/kigi-shell/src/extensions/session_admin.rs +++ b/crates/codegen/kigi-shell/src/extensions/session_admin.rs @@ -4,8 +4,8 @@ //! persistent or shared agent state but are not part of the per-turn prompt //! lifecycle: //! -//! - `x.ai/session/rename` rename a session locally + remote -//! - `x.ai/session/delete` delete a session locally + remote +//! - `x.ai/session/rename` rename a session locally +//! - `x.ai/session/delete` delete a session locally //! - `x.ai/session/update_mcp_servers` mid-session MCP server swap //! - `x.ai/session/fork` fork a session into a new one //! - `x.ai/internal/reload_all_mcp_servers` config hot-reload, all sessions @@ -77,7 +77,8 @@ async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR } if req.kind == SessionKind::Chat { - return rename_chat_conversation(agent, &req.session_id, &req.title).await; + return Err(acp::Error::invalid_request() + .data("chat conversations are not available in kigi (local sessions only)")); } let session_id = acp::SessionId::new(Arc::from(req.session_id.as_str())); @@ -111,22 +112,6 @@ async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR // Send a SessionSummaryGenerated notification so the TUI updates its title notify_session_title(agent, session_id, &req.title).await; - if agent.is_writeback_storage() && agent.current_auth().is_some() { - use crate::remote::client::BackendClient; - use crate::session::export::ExportedMetadata; - - let mut metadata = ExportedMetadata::from_summary(summary); - metadata.title = Some(req.title.clone()); - metadata.updated_at = Some(chrono::Utc::now().to_rfc3339()); - if let Err(e) = BackendClient::new() - .with_auth_manager(agent.auth_manager.clone()) - .save_session_data(&req.session_id, &[], Some(&metadata)) - .await - { - tracing::warn!(?e, session_id = %req.session_id, "failed to sync renamed title to backend"); - } - } - // Hook 2: update session replica with summary (fire-and-forget) if let Some(client) = agent.session_registry_client() { let sid = req.session_id.to_string(); @@ -169,49 +154,6 @@ async fn notify_session_title(agent: &MvpAgent, session_id: acp::SessionId, titl } } -async fn rename_chat_conversation( - agent: &MvpAgent, - conversation_id: &str, - title: &str, -) -> ExtResult { - use crate::remote::{ConvError, UpdateConversationBody}; - - let Some(client) = agent.conversations_client() else { - return Err(acp::Error::invalid_request() - .data("chat session rename requires the conversations lane (OIDC + chat feature)")); - }; - - let body = UpdateConversationBody { - title: Some(title.to_owned()), - starred: None, - }; - client - .update_conversation(conversation_id, &body) - .await - .map_err(|e| match e { - ConvError::NoOauth => acp::Error::invalid_request() - .data("chat session rename requires xAI OAuth credentials"), - ConvError::Http { status: 404 } => acp::Error::invalid_request() - .data(format!("conversation not found: {conversation_id}")), - other => acp::Error::internal_error() - .data(format!("chat conversation rename failed: {other}")), - })?; - - // If this conversation is open live, notify clients of the new title. - let session_id = acp::SessionId::new(Arc::from(conversation_id)); - if agent.sessions.borrow().contains_key(&session_id) { - notify_session_title(agent, session_id, title).await; - } - - tracing::info!( - session_id = %conversation_id, - title = %title, - "Chat conversation renamed" - ); - - to_raw_response(&serde_json::json!({ "success": true })) -} - // session/delete /// Delete a session from history. @@ -229,31 +171,17 @@ async fn handle_session_delete(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR let req: DeleteRequest = parse_params(args)?; if req.kind == SessionKind::Chat { - return soft_delete_chat_conversation(agent, &req.session_id).await; + return Err(acp::Error::invalid_request() + .data("chat conversations are not available in kigi (local sessions only)")); } let session_id = acp::SessionId::new(Arc::from(req.session_id.as_str())); - // For writeback storage (non-ZDR): remote delete is authoritative for - // the cloud history and runs first; on failure no local bits are - // touched so the pager does not remove the row or toast success. - let needs_remote = agent.is_writeback_storage() && agent.current_auth().is_some(); - - // Shared delete: remote-first, then local disk + FTS eviction. - // Mirrored by the `grok sessions delete ` CLI path. - crate::session::persistence::delete_session_history( - &req.session_id, - req.cwd.as_deref(), - needs_remote, - agent.auth_manager.clone(), - ) - .await - .map_err(|e| { - if let crate::session::persistence::DeleteSessionError::Remote(_) = &e { - tracing::warn!(?e, session_id = %req.session_id, "failed to delete remote session data"); - } - acp::Error::internal_error().data(e.to_string()) - })?; + // Local disk + FTS eviction. Mirrored by the `kigi sessions delete ` + // CLI path. + crate::session::persistence::delete_session_history(&req.session_id, req.cwd.as_deref()) + .await + .map_err(|e| acp::Error::internal_error().data(e.to_string()))?; // If an in-memory live session exists for this id (e.g. the user // deleted history for a session that is still open in another agent @@ -269,35 +197,6 @@ async fn handle_session_delete(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR to_raw_response(&serde_json::json!({ "success": true })) } -async fn soft_delete_chat_conversation(agent: &MvpAgent, conversation_id: &str) -> ExtResult { - use crate::remote::ConvError; - - let Some(client) = agent.conversations_client() else { - return Err(acp::Error::invalid_request() - .data("chat session delete requires the conversations lane (OIDC + chat feature)")); - }; - - client - .soft_delete_conversation(conversation_id) - .await - .map_err(|e| match e { - ConvError::NoOauth => acp::Error::invalid_request() - .data("chat session delete requires xAI OAuth credentials"), - other => acp::Error::internal_error() - .data(format!("chat conversation soft-delete failed: {other}")), - })?; - - let session_id = acp::SessionId::new(Arc::from(conversation_id)); - if agent.sessions.borrow().contains_key(&session_id) { - agent.request_session_shutdown(&session_id); - agent.remove_session(&session_id); - } - - tracing::info!(session_id = %conversation_id, "Chat conversation soft-deleted"); - - to_raw_response(&serde_json::json!({ "success": true })) -} - // session/update_mcp_servers async fn handle_update_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { @@ -708,8 +607,7 @@ async fn handle_session_fork(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtRes let request: ForkSessionRequest = parse_params(args)?; - let agent_id = crate::util::agent_id::agent_id(); - let response = fork_session(request, &agent_id, Some(agent.auth_manager.clone())) + let response = fork_session(request) .await .map_err(|e| acp::Error::internal_error().data(e.to_string()))?; diff --git a/crates/codegen/kigi-shell/src/extensions/share.rs b/crates/codegen/kigi-shell/src/extensions/share.rs deleted file mode 100644 index cc743a4..0000000 --- a/crates/codegen/kigi-shell/src/extensions/share.rs +++ /dev/null @@ -1,204 +0,0 @@ -//! `x.ai/share_session` extension handler. -//! -//! Loads a local session, exports it, and asks the backend for a public -//! share URL. - -use agent_client_protocol as acp; - -use super::{ExtResult, parse_params, to_raw_response}; -use crate::agent::MvpAgent; -use crate::remote::client::BackendClient; -use crate::session::export::ExportedSession; -use crate::session::info::Info as SessionInfo; -use crate::session::persistence::list_summaries; -use crate::session::share::{ShareSessionRequest, ShareSessionResponse}; - -#[tracing::instrument(skip_all, fields(method = %args.method))] -pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { - match args.method.as_ref() { - "x.ai/share_session" => { - tracing::info!("handling share session request"); - handle_share_session(agent, args).await - } - _ => Err(acp::Error::method_not_found()), - } -} - -async fn handle_share_session(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult { - let request: ShareSessionRequest = parse_params(args)?; - - // Get auth - required for sharing. - let auth = require_xai_auth_for_share(&agent.auth_manager)?; - - // Remote settings / feature-flag gate: sharing_enabled defaults to false - // and is only enabled for eligible accounts. - let sharing_enabled = agent - .cfg - .borrow() - .remote_settings - .as_ref() - .and_then(|rs| rs.sharing_enabled) - .unwrap_or(false); - if !sharing_enabled { - return Err( - acp::Error::invalid_params().data("Session sharing is not available for your account.") - ); - } - - // Find session info by searching through summaries - let summaries = list_summaries(None).await.map_err(|e| { - acp::Error::internal_error().data(format!("Failed to list sessions: {}", e)) - })?; - - let summary = summaries - .iter() - .find(|s| s.info.id.0.as_ref() == request.session_id.as_str()) - .ok_or_else(|| acp::Error::resource_not_found(Some("Session not found".into())))?; - - let info = SessionInfo { - id: acp::SessionId::new(request.session_id.clone()), - cwd: summary.info.cwd.clone(), - }; - - // Load and export session - let exported = ExportedSession::from_local_session(&info) - .await - .map_err(|e| acp::Error::internal_error().data(format!("Failed to load session: {}", e)))?; - - // Check for empty session - if exported.messages.is_empty() { - return Err(acp::Error::invalid_params().data("No messages to share yet")); - } - - // Upload to backend and get share URL. - let client = BackendClient::new().with_auth_manager(agent.auth_manager.clone()); - let agent_id = crate::util::agent_id::agent_id(); - let share_url = client - .share_session(&exported, &agent_id) - .await - .map_err(|e| { - tracing::error!(error = %e, "Failed to share session with backend"); - acp::Error::internal_error().data(format!("Failed to share session: {}", e)) - })?; - - let response = ShareSessionResponse { share_url }; - to_raw_response(&response) -} - -fn require_xai_auth_for_share( - auth_manager: &crate::auth::AuthManager, -) -> Result { - super::auth_gate::require_xai_auth( - auth_manager, - "Authentication required to share session", - "Share session is disabled. Run `grok login` to authenticate.", - ) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::auth::KimiCodeConfig; - use crate::auth::{AuthMode, KimiAuth}; - use chrono::{Duration, Utc}; - use std::sync::Arc; - use tempfile::tempdir; - - fn make_auth_manager_with_token_expiring_in( - ttl: Duration, - ) -> (Arc, tempfile::TempDir) { - let dir = tempdir().expect("tempdir for share auth test"); - let mgr = Arc::new(crate::auth::AuthManager::new( - dir.path(), - KimiCodeConfig::default(), - )); - - let expires_at = Utc::now() + ttl; - - // We must explicitly set oidc_issuer to a first-party xAI issuer. - // Only OIDC tokens against https://auth.x.ai (or the local-dev equivalent) - // return true from is_xai_auth(). This is required for the share tests to - // exercise the happy path through require_xai_auth_for_share. - let auth = KimiAuth { - auth_mode: AuthMode::OAuth, - key: "test-key".into(), - expires_at: Some(expires_at), - create_time: Utc::now() - Duration::hours(1), - ..Default::default() - }; - mgr.hot_swap(auth); - (mgr, dir) - } - - #[test] - fn share_works_outside_the_5m_early_invalidation_window() { - let (mgr, _dir) = make_auth_manager_with_token_expiring_in(Duration::minutes(10)); - assert!(mgr.current().is_some()); - assert!(require_xai_auth_for_share(&mgr).is_ok()); - } - - #[test] - fn share_succeeds_inside_the_5m_early_invalidation_window() { - let (mgr, _dir) = make_auth_manager_with_token_expiring_in(Duration::seconds(1)); - // This is exactly the state that triggered the user bug: - assert!( - mgr.current().is_none(), - "current() drops the token inside the buffer" - ); - assert!(mgr.expired_auth().is_some()); - - // Now that we use current_or_expired(), this passes. - let res = require_xai_auth_for_share(&mgr); - assert!( - res.is_ok(), - "require_xai_auth_for_share must succeed for a still-valid buffered xAI token" - ); - } - - #[test] - fn share_fails_with_no_auth_at_all() { - let dir = tempdir().expect("tempdir"); - let mgr = Arc::new(crate::auth::AuthManager::new( - dir.path(), - KimiCodeConfig::default(), - )); - assert!(require_xai_auth_for_share(&mgr).is_err()); - } - - #[test] - fn share_rejects_non_xai_auth_with_actionable_grok_login_message() { - let dir = tempdir().expect("tempdir"); - let mgr = Arc::new(crate::auth::AuthManager::new( - dir.path(), - KimiCodeConfig::default(), - )); - - // API key is the simplest non-xAI credential (External and enterprise OIDC - // are also rejected the same way). - let non_xai = KimiAuth { - auth_mode: AuthMode::ApiKey, - key: "xai-test-key".into(), - create_time: Utc::now(), - ..Default::default() - }; - mgr.hot_swap(non_xai); - - let err = require_xai_auth_for_share(&mgr) - .expect_err("non-xAI accounts (API key, External, enterprise IdP) must be rejected"); - - // This is the key assertion the review asked for: we must test the *exact* - // actionable data string for the non-xAI path (distinct from the generic - // "Authentication required to share session" path). - let serialized = - serde_json::to_value(&err).expect("acp::Error serializes to JSON-RPC shape"); - let data = serialized - .get("data") - .and_then(|v| v.as_str()) - .expect("auth_required error carries a data string"); - - assert_eq!( - data, - "Share session is disabled. Run `grok login` to authenticate." - ); - } -} diff --git a/crates/codegen/kigi-shell/src/lib.rs b/crates/codegen/kigi-shell/src/lib.rs index 50d1bdb..a0a6dd8 100644 --- a/crates/codegen/kigi-shell/src/lib.rs +++ b/crates/codegen/kigi-shell/src/lib.rs @@ -28,7 +28,6 @@ pub mod managed_config; pub mod mcp_doctor; pub use kigi_models as models; pub mod plugin; -pub mod remote; pub mod sampling; pub mod session; pub mod terminal; diff --git a/crates/codegen/kigi-shell/src/remote/agent.rs b/crates/codegen/kigi-shell/src/remote/agent.rs deleted file mode 100644 index 38bcb6a..0000000 --- a/crates/codegen/kigi-shell/src/remote/agent.rs +++ /dev/null @@ -1,326 +0,0 @@ -//! Remote sandbox client for cli-chat-proxy. -//! -//! This module provides an HTTP client to interact with cli-chat-proxy -//! for managing sandbox sessions and environments via REST API. - -use std::sync::Arc; - -use crate::auth::{AuthManager, KimiCodeConfig}; -use anyhow::{Context, Result, bail}; -use serde::de::DeserializeOwned; - -// Re-export sandbox API types from cli-chat-proxy-types for convenience. -// Sorted alphabetically; see sandbox_types.rs for logical grouping. -pub use prod_mc_cli_chat_proxy_types::{ - SandboxCreateEnvironmentRequest, SandboxEnvironment, SandboxEnvironmentResponse, - SandboxEnvironmentVariable, SandboxEnvironmentWithMetadata, SandboxForkRequest, - SandboxForkResponse, SandboxForkedSession, SandboxHibernateResponse, - SandboxListEnvironmentsRequest, SandboxListEnvironmentsResponse, - SandboxListPreinstalledPackagesResponse, SandboxLogsExitCodes, SandboxLogsResponse, - SandboxMode, SandboxPreinstalledPackage, SandboxRestoreRequest, SandboxRestoreResponse, - SandboxSecretInput, SandboxStartRequest, SandboxStartResponse, SandboxStatusResponse, - SandboxTerminateRequest, SandboxUpdateEnvironmentRequest, -}; - -// ============================================================================ -// Sandbox Client -// ============================================================================ - -/// HTTP client for interacting with the sandbox API via cli-chat-proxy. -/// -/// Path parameters (`session_id`, `environment_id`) are interpolated directly -/// into URLs without percent-encoding. This is safe because these IDs are -/// UUIDs in practice. If ID formats ever change to include URL-unsafe -/// characters, the `format!()` calls should be updated to use percent-encoding. -pub struct SandboxClient { - client: reqwest::Client, - base_url: String, - auth_manager: Arc, -} - -impl SandboxClient { - pub fn new(base_url: impl Into, auth_manager: Arc) -> Self { - Self { - client: crate::http::shared_client(), - base_url: base_url.into(), - auth_manager, - } - } - - /// Returns the base URL. - pub fn base_url(&self) -> &str { - &self.base_url - } - - // Do not set Content-Type — callers use .json() and reqwest .header() appends. - async fn auth_headers( - &self, - builder: reqwest::RequestBuilder, - ) -> Result { - let auth = self - .auth_manager - .auth() - .await - .context("failed to resolve sandbox auth")?; - let mut builder = builder - .header("Authorization", format!("Bearer {}", auth.key)) - .header("x-userid", &auth.user_id) - .header("x-grok-client-version", kigi_version::VERSION); - - if let Some(email) = &auth.email { - builder = builder.header("x-email", email); - } - - builder = builder - .header( - "x-grok-client-identifier", - crate::http::process_client_identifier(), - ) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ); - - Ok(kigi_file_utils::trace_context::inject_trace_context_into_request(builder)) - } - - /// Check an HTTP response for errors, then deserialize the JSON body. - async fn parse_response( - response: reqwest::Response, - operation: &str, - ) -> Result { - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); - bail!("{operation} failed: {status} - {body}"); - } - response - .json() - .await - .with_context(|| format!("failed to parse {operation} response")) - } - - /// Check an HTTP response for errors, discarding the body. - async fn check_response(response: reqwest::Response, operation: &str) -> Result<()> { - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); - bail!("{operation} failed: {status} - {body}"); - } - Ok(()) - } - - /// Fork an existing sandbox session. - pub async fn fork_session(&self, request: &SandboxForkRequest) -> Result { - let url = format!("{}/sandbox/sessions/fork", self.base_url); - let response = self - .auth_headers(self.client.post(&url)) - .await? - .json(request) - .send() - .await - .context("failed to send fork session request")?; - Self::parse_response(response, "fork session").await - } - - /// Terminate a sandbox session. - pub async fn terminate_session( - &self, - session_id: &str, - request: &SandboxTerminateRequest, - ) -> Result<()> { - let mut url = format!("{}/sandbox/sessions/{}", self.base_url, session_id); - if let Some(env_id) = &request.environment_id { - url = format!("{}?environmentId={}", url, env_id); - } - - let response = self - .auth_headers(self.client.delete(&url)) - .await? - .send() - .await - .context("failed to send terminate session request")?; - - if response.status().as_u16() == 404 { - bail!("session not found: {session_id}"); - } - Self::check_response(response, "terminate session").await - } - - // ======================================================================== - // Session Lifecycle - // ======================================================================== - - /// Start a sandbox session (non-TUI). - pub async fn start_session( - &self, - request: &SandboxStartRequest, - ) -> Result { - let url = format!("{}/sandbox/sessions/start", self.base_url); - let response = self - .auth_headers(self.client.post(&url)) - .await? - .json(request) - .send() - .await - .context("failed to send start session request")?; - Self::parse_response(response, "start session").await - } - - /// Get sandbox session status. - pub async fn get_session_status(&self, session_id: &str) -> Result { - let url = format!("{}/sandbox/sessions/{}/status", self.base_url, session_id); - let response = self - .auth_headers(self.client.get(&url)) - .await? - .send() - .await - .context("failed to send get session status request")?; - Self::parse_response(response, "get session status").await - } - - /// Get sandbox session logs. - pub async fn get_session_logs(&self, session_id: &str) -> Result { - let url = format!("{}/sandbox/sessions/{}/logs", self.base_url, session_id); - let response = self - .auth_headers(self.client.get(&url)) - .await? - .send() - .await - .context("failed to send get session logs request")?; - Self::parse_response(response, "get session logs").await - } - - /// Hibernate a sandbox session (snapshot rootfs to GCS and terminate). - pub async fn hibernate_session(&self, session_id: &str) -> Result { - let url = format!( - "{}/sandbox/sessions/{}/hibernate", - self.base_url, session_id - ); - let response = self - .auth_headers(self.client.post(&url)) - .await? - .send() - .await - .context("failed to send hibernate session request")?; - Self::parse_response(response, "hibernate session").await - } - - /// Restore a previously hibernated sandbox session from its snapshot. - pub async fn restore_session( - &self, - session_id: &str, - request: &SandboxRestoreRequest, - ) -> Result { - let url = format!("{}/sandbox/sessions/{}/restore", self.base_url, session_id); - let response = self - .auth_headers(self.client.post(&url)) - .await? - .json(request) - .send() - .await - .context("failed to send restore session request")?; - Self::parse_response(response, "restore session").await - } - - // ======================================================================== - // Environment CRUD - // ======================================================================== - - /// List sandbox environments. - pub async fn list_environments( - &self, - request: &SandboxListEnvironmentsRequest, - ) -> Result { - let url = format!("{}/sandbox/environments", self.base_url); - let mut builder = self.auth_headers(self.client.get(&url)).await?; - if let Some(page) = request.page { - builder = builder.query(&[("page", page)]); - } - if let Some(page_size) = request.page_size { - builder = builder.query(&[("pageSize", page_size)]); - } - let response = builder - .send() - .await - .context("failed to send list environments request")?; - Self::parse_response(response, "list environments").await - } - - /// Create a new sandbox environment. - pub async fn create_environment( - &self, - request: &SandboxCreateEnvironmentRequest, - ) -> Result { - let url = format!("{}/sandbox/environments", self.base_url); - let response = self - .auth_headers(self.client.post(&url)) - .await? - .json(request) - .send() - .await - .context("failed to send create environment request")?; - Self::parse_response(response, "create environment").await - } - - /// Get a sandbox environment by ID. - pub async fn get_environment( - &self, - environment_id: &str, - ) -> Result { - let url = format!("{}/sandbox/environments/{}", self.base_url, environment_id); - let response = self - .auth_headers(self.client.get(&url)) - .await? - .send() - .await - .context("failed to send get environment request")?; - Self::parse_response(response, "get environment").await - } - - /// Update a sandbox environment. - pub async fn update_environment( - &self, - environment_id: &str, - request: &SandboxUpdateEnvironmentRequest, - ) -> Result { - let url = format!("{}/sandbox/environments/{}", self.base_url, environment_id); - let response = self - .auth_headers(self.client.put(&url)) - .await? - .json(request) - .send() - .await - .context("failed to send update environment request")?; - Self::parse_response(response, "update environment").await - } - - /// Delete a sandbox environment. - pub async fn delete_environment(&self, environment_id: &str) -> Result<()> { - let url = format!("{}/sandbox/environments/{}", self.base_url, environment_id); - let response = self - .auth_headers(self.client.delete(&url)) - .await? - .send() - .await - .context("failed to send delete environment request")?; - Self::check_response(response, "delete environment").await - } - - /// List preinstalled packages available for sandbox environments. - pub async fn list_preinstalled_packages( - &self, - ) -> Result { - let url = format!( - "{}/sandbox/environments/preinstalled-packages", - self.base_url - ); - let response = self - .auth_headers(self.client.get(&url)) - .await? - .send() - .await - .context("failed to send list preinstalled packages request")?; - Self::parse_response(response, "list preinstalled packages").await - } -} diff --git a/crates/codegen/kigi-shell/src/remote/chat_models_client.rs b/crates/codegen/kigi-shell/src/remote/chat_models_client.rs deleted file mode 100644 index 26c9518..0000000 --- a/crates/codegen/kigi-shell/src/remote/chat_models_client.rs +++ /dev/null @@ -1,202 +0,0 @@ -//! grok.com chat-product model catalog (`POST /rest/modes`) — the models -//! grok-web's chat picker shows, distinct from the CLI `/v1/models` build -//! catalog. Transport only; cache + ACP mapping live in -//! [`crate::agent::chat_modes`]. - -use std::sync::Arc; - -use serde::Deserialize; - -use crate::auth::AuthManager; - -const KIGI_WEB_URL: &str = "https://grok.com"; - -#[derive(Debug, Clone, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct Mode { - #[serde(default)] - pub id: String, - #[serde(default)] - pub title: String, - #[serde(default)] - pub description: String, - #[serde(default)] - pub badge_text: Option, - #[serde(default)] - pub availability: ModeAvailability, - #[serde(default)] - pub icon_hint: String, - #[serde(default)] - pub tags: Vec, -} - -impl Mode { - pub fn is_available(&self) -> bool { - self.availability.available.is_some() - } -} - -/// proto3-JSON oneof: exactly one field is present. -#[derive(Debug, Clone, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ModeAvailability { - #[serde(default)] - pub available: Option, - #[serde(default)] - pub unavailable: Option, - #[serde(default)] - pub requires_upgrade: Option, - #[serde(default)] - pub coming_soon: Option, -} - -#[derive(Debug, Clone, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ListModesResponse { - #[serde(default)] - pub modes: Vec, - #[serde(default)] - pub default_mode_id: String, -} - -#[derive(Debug, thiserror::Error)] -pub enum ChatModelsError { - #[error("no grok.com credentials")] - NoAuth, - #[error("request timed out")] - Timeout, - #[error("network error: {0}")] - Network(#[from] reqwest::Error), - #[error("request failed: {status}")] - Http { status: u16 }, - #[error("parse error: {0}")] - Parse(#[from] serde_json::Error), -} - -/// Stateless transport for `POST /rest/modes`; caching lives in -/// [`crate::agent::chat_modes::ChatModesManager`]. -pub struct ChatModelsClient { - http: reqwest::Client, - base_url: String, - auth: Arc, -} - -impl ChatModelsClient { - pub fn new(auth: Arc) -> Self { - let base_url = std::env::var("KIGI_MODES_BASE_URL") - .ok() - .filter(|s| !s.is_empty()) - .or_else(|| { - std::env::var("KIGI_CONVERSATIONS_BASE_URL") - .ok() - .filter(|s| !s.is_empty()) - }) - .or_else(|| { - std::env::var("KIGI_CODE_WEB_URL") - .ok() - .filter(|s| !s.is_empty()) - }) - .unwrap_or_else(|| KIGI_WEB_URL.to_string()); - Self { - http: crate::http::shared_client(), - base_url, - auth, - } - } - - /// Gated only on a valid grok.com bearer — deliberately NOT `is_xai_auth()` - /// (unlike workspaces/conversations), since `/rest/modes` is the public chat - /// endpoint and that gate would exclude API-key / cached-token chat users. - pub async fn list_modes(&self, locale: &str) -> Result { - let auth = self - .auth - .auth() - .await - .map_err(|_| ChatModelsError::NoAuth)?; - - let url = format!("{}/rest/modes", self.base_url); - let body = serde_json::json!({ "locale": locale }); - let mut builder = self - .http - .post(&url) - .json(&body) - .header("Authorization", format!("Bearer {}", auth.key)) - .header("x-userid", &auth.user_id) - .header("x-grok-client-version", kigi_version::VERSION) - .header( - "x-grok-client-identifier", - crate::http::process_client_identifier(), - ) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ) - .header(reqwest::header::ACCEPT, "application/json"); - if let Some(email) = &auth.email { - builder = builder.header("x-email", email); - } - let builder = kigi_file_utils::trace_context::inject_trace_context_into_request(builder); - - let response = builder.send().await?; - let status = response.status(); - if !status.is_success() { - return Err(ChatModelsError::Http { - status: status.as_u16(), - }); - } - - let bytes = response.bytes().await?; - let resp: ListModesResponse = serde_json::from_slice(&bytes)?; - Ok(resp) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn modes_parse_camelcase_wire() { - let json = serde_json::json!({ - "modes": [{ - "id": "auto", - "title": "Auto", - "description": "Picks the best model", - "badgeText": "New", - "availability": { "available": {} }, - "iconHint": "rocket", - "tags": ["TAG_PRIMARY"] - }, { - "id": "heavy", - "title": "Heavy", - "availability": { "requiresUpgrade": { "message": "Upgrade" } } - }], - "defaultModeId": "auto" - }); - let resp: ListModesResponse = serde_json::from_value(json).unwrap(); - assert_eq!(resp.modes.len(), 2); - assert_eq!(resp.default_mode_id, "auto"); - let auto = &resp.modes[0]; - assert_eq!(auto.id, "auto"); - assert_eq!(auto.title, "Auto"); - assert_eq!(auto.badge_text.as_deref(), Some("New")); - assert_eq!(auto.icon_hint, "rocket"); - assert_eq!(auto.tags, vec!["TAG_PRIMARY".to_string()]); - assert!(auto.is_available()); - assert!(!resp.modes[1].is_available()); - } - - #[test] - fn missing_fields_default_gracefully() { - let json = serde_json::json!({ "modes": [{ "id": "m1" }] }); - let resp: ListModesResponse = serde_json::from_value(json).unwrap(); - let m = &resp.modes[0]; - assert_eq!(m.id, "m1"); - assert!(m.title.is_empty()); - assert!(m.description.is_empty()); - assert!(m.badge_text.is_none()); - // No availability field on the wire → not selectable. - assert!(!m.is_available()); - assert!(resp.default_mode_id.is_empty()); - } -} diff --git a/crates/codegen/kigi-shell/src/remote/client.rs b/crates/codegen/kigi-shell/src/remote/client.rs deleted file mode 100644 index 06514a4..0000000 --- a/crates/codegen/kigi-shell/src/remote/client.rs +++ /dev/null @@ -1,2268 +0,0 @@ -//! HTTP client for backend CRUD operations. -use crate::auth::{KimiAuth, KimiCodeConfig}; -use crate::session::export::{ExportedMessage, ExportedMetadata, ExportedSession}; -use indexmap::IndexMap; -use prod_mc_cli_chat_proxy_types::SubagentBundle; -use serde::{Deserialize, Serialize}; -use std::time::Duration; -const KIGI_CODE_BACKEND_URL: &str = "https://code.kigi.com"; -const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30); -const KIGI_CODE_WEB_URL: &str = "https://grok.com"; -/// Build a share URL from a permission ID -pub fn share_url(permission_id: &str) -> String { - let web_url = - std::env::var("KIGI_CODE_WEB_URL").unwrap_or_else(|_| KIGI_CODE_WEB_URL.to_string()); - format!("{}/build/share/{}", web_url, permission_id) -} -fn add_cli_chat_proxy_headers_blocking( - builder: reqwest::blocking::RequestBuilder, - auth: &KimiAuth, - alpha_test_key: Option<&str>, - url: &str, -) -> reqwest::blocking::RequestBuilder { - let mut builder = builder - .header("Authorization", format!("Bearer {}", auth.key)) - .header("x-userid", &auth.user_id) - .header("x-grok-client-version", kigi_version::VERSION); - if let Some(email) = &auth.email { - builder = builder.header("x-email", email); - } - let _ = (alpha_test_key, url); - builder - .header( - "x-grok-client-identifier", - crate::http::process_client_identifier(), - ) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ) -} -async fn parse_json_response( - response: reqwest::Response, -) -> Result { - let bytes = response.bytes().await?; - serde_json::from_slice(&bytes).map_err(BackendError::from) -} -async fn add_bundle_fetch_headers( - builder: reqwest::RequestBuilder, - auth_manager: Option<&std::sync::Arc>, - deployment_key: Option<&str>, - alpha_test_key: Option<&str>, - url: &str, -) -> reqwest::RequestBuilder { - let resolved_auth = match auth_manager { - Some(am) => am.auth().await.ok(), - None => None, - }; - let mut credentials = crate::util::kigi_auth_credentials::KigiAuthCredentials::new( - resolved_auth.as_ref().map(|auth| auth.key.clone()), - ); - credentials.deployment_key = deployment_key.map(str::to_owned); - credentials.alpha_test_key = alpha_test_key.map(str::to_owned); - let mut builder = credentials - .apply(builder, url) - .header("x-grok-client-version", kigi_version::VERSION); - if deployment_key.is_none() - && let Some(auth) = &resolved_auth - { - builder = builder.header("x-userid", &auth.user_id); - if let Some(email) = &auth.email { - builder = builder.header("x-email", email); - } - } - builder = builder - .header( - "x-grok-client-identifier", - crate::http::process_client_identifier(), - ) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ); - kigi_file_utils::trace_context::inject_trace_context_into_request(builder) -} -/// Fetch the bundled subagent cache payload from cli-chat-proxy `GET /v1/subagents/bundle`. -/// -/// Uses the shell's standard proxy-backed auth model: deployment key auth takes -/// precedence when configured; otherwise user-session token auth is used. -pub async fn fetch_subagent_bundle( - cli_chat_proxy_base_url: &str, - auth_manager: Option<&std::sync::Arc>, - deployment_key: Option<&str>, - alpha_test_key: Option<&str>, -) -> Result { - let url = format!("{}/subagents/bundle", cli_chat_proxy_base_url); - let response = add_bundle_fetch_headers( - crate::http::shared_client() - .get(&url) - .timeout(std::time::Duration::from_secs(10)), - auth_manager, - deployment_key, - alpha_test_key, - &url, - ) - .await - .send() - .await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); - return Err(BackendError::RequestFailed { status, body }); - } - let bundle: SubagentBundle = parse_json_response(response).await?; - tracing::debug!( - version = % bundle.version, personas = bundle.personas.len(), roles = bundle - .roles.len(), agents = bundle.agents.len(), - "Fetched subagent bundle from cli-chat-proxy" - ); - Ok(bundle) -} -/// The result of fetching a bundle: either raw tar.gz bytes from the new -/// archive endpoint, or a parsed JSON bundle from the legacy endpoint. -#[derive(Debug)] -pub enum FetchedBundle { - Archive(Vec), - Legacy(SubagentBundle), -} -/// Fetch a bundle, trying the archive endpoint first and falling back to -/// legacy JSON on any non-success HTTP status. -pub async fn fetch_bundle( - cli_chat_proxy_base_url: &str, - auth_manager: Option<&std::sync::Arc>, - deployment_key: Option<&str>, - alpha_test_key: Option<&str>, -) -> Result { - fetch_bundle_inner( - cli_chat_proxy_base_url, - auth_manager, - deployment_key, - alpha_test_key, - ) - .await -} -async fn fetch_bundle_inner( - cli_chat_proxy_base_url: &str, - auth_manager: Option<&std::sync::Arc>, - deployment_key: Option<&str>, - alpha_test_key: Option<&str>, -) -> Result { - let archive_url = format!("{}/bundle/archive", cli_chat_proxy_base_url); - let raw_client = crate::http::shared_client(); - let client: reqwest_middleware::ClientWithMiddleware = if let Some(am) = auth_manager { - let provider: std::sync::Arc = std::sync::Arc::new( - crate::auth::credential_provider::ShellAuthCredentialProvider::new( - am.clone(), - deployment_key.map(str::to_owned), - alpha_test_key.map(str::to_owned), - ), - ); - crate::http::with_auth_retry(raw_client, provider) - } else { - reqwest_middleware::ClientBuilder::new(raw_client).build() - }; - let mut request = client - .get(&archive_url) - .timeout(std::time::Duration::from_secs(30)) - .header("x-grok-client-version", kigi_version::VERSION) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ); - if deployment_key.is_none() - && let Some(am) = auth_manager - && let Some(auth) = am.current() - { - request = request.header("x-userid", &auth.user_id); - if let Some(ref email) = auth.email { - request = request.header("x-email", email); - } - } - let archive_response = request.send().await.map_err(|e| match e { - reqwest_middleware::Error::Reqwest(e) => BackendError::Network(e), - reqwest_middleware::Error::Middleware(e) => BackendError::Auth(e.to_string()), - })?; - if archive_response.status().is_success() { - let bytes = archive_response.bytes().await?; - return Ok(FetchedBundle::Archive(bytes.to_vec())); - } - if archive_response.status() == reqwest::StatusCode::UNAUTHORIZED { - let body = archive_response.text().await.unwrap_or_default(); - return Err(BackendError::RequestFailed { status: 401, body }); - } - tracing::debug!( - status = % archive_response.status(), - "archive endpoint unavailable, falling back to legacy JSON" - ); - let bundle = fetch_subagent_bundle( - cli_chat_proxy_base_url, - auth_manager, - deployment_key, - alpha_test_key, - ) - .await?; - Ok(FetchedBundle::Legacy(bundle)) -} -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ShareResponse { - pub permission_id: String, -} -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct LoadDataResponse { - pub messages: Option>, - pub session: Option, -} -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct LoadedMessage { - pub id: String, - pub content: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub timestamp: Option, -} -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SessionInfo { - pub session_id: String, - pub title: Option, - pub cwd: Option, - pub status: Option, - pub created_at: Option, - pub updated_at: Option, - #[serde(default)] - pub metadata: Option, -} -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SaveDataRequest { - pub messages: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct UpsertSessionRequest { - pub session: SessionUpdate, - pub agent_id: String, -} -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SessionUpdate { - #[serde(skip_serializing_if = "Option::is_none")] - pub title: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub cwd: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub status: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} -#[derive(Debug, thiserror::Error)] -pub enum BackendError { - #[error("Network error: {0}")] - Network(#[from] reqwest::Error), - #[error("Request failed: {status} - {body}")] - RequestFailed { status: u16, body: String }, - #[error("Serialization error: {0}")] - Serialization(#[from] serde_json::Error), - #[error("Session not found: {session_id}")] - SessionNotFound { session_id: String }, - #[error("Hydration I/O error at {path}: {source}")] - Hydration { - path: std::path::PathBuf, - source: std::io::Error, - }, - #[error("Auth error: {0}")] - Auth(String), -} -pub struct BackendClient { - reqwest_client: reqwest::Client, - client: reqwest_middleware::ClientWithMiddleware, - base_url: String, - pub(crate) auth_manager: Option>, -} -impl Default for BackendClient { - fn default() -> Self { - Self::new() - } -} -impl BackendClient { - fn build_default_client() -> reqwest::Client { - reqwest::Client::builder() - .connect_timeout(Duration::from_secs(10)) - .timeout(DEFAULT_TIMEOUT) - .build() - .expect("failed to build HTTP client") - } - pub fn new() -> Self { - let reqwest_client = Self::build_default_client(); - Self { - client: reqwest_middleware::ClientBuilder::new(reqwest_client.clone()).build(), - reqwest_client, - base_url: std::env::var("KIGI_CODE_BACKEND_URL") - .unwrap_or_else(|_| KIGI_CODE_BACKEND_URL.to_string()), - auth_manager: None, - } - } - pub fn with_base_url(base_url: impl Into) -> Self { - let reqwest_client = Self::build_default_client(); - Self { - client: reqwest_middleware::ClientBuilder::new(reqwest_client.clone()).build(), - reqwest_client, - base_url: base_url.into(), - auth_manager: None, - } - } - /// Attach a live `AuthManager` so every request resolves a fresh token - /// instead of requiring the caller to pass `&KimiAuth`. - pub fn with_auth_manager(mut self, manager: std::sync::Arc) -> Self { - let credentials: std::sync::Arc = - std::sync::Arc::new( - crate::auth::credential_provider::ShellAuthCredentialProvider::new( - manager.clone(), - None, - None, - ), - ); - self.client = crate::http::with_auth_retry(self.reqwest_client.clone(), credentials); - self.auth_manager = Some(manager); - self - } - /// Resolve auth from the attached `AuthManager`. - async fn resolve_auth(&self) -> Result { - let manager = self - .auth_manager - .as_ref() - .ok_or_else(|| BackendError::Auth("No AuthManager configured".into()))?; - manager - .auth() - .await - .map_err(|e| BackendError::Auth(format!("{e}"))) - } - pub fn base_url(&self) -> &str { - &self.base_url - } - /// Upload session and create share link. - /// - /// The session data (`save_session_data`) is sent inline to the backend. - /// If the backend responds with 413 (payload too large), the error is - /// logged as a warning and the share continues — the caller is expected - /// to have already uploaded the data to GCS via a signed URL as a - /// fallback. - pub async fn share_session( - &self, - session: &ExportedSession, - agent_id: &str, - ) -> Result { - self.upsert_session(&session.session_id, &session.metadata, agent_id) - .await?; - match self - .save_session_data( - &session.session_id, - &session.messages, - Some(&session.metadata), - ) - .await - { - Ok(()) => {} - Err(BackendError::RequestFailed { status: 413, .. }) => { - tracing::warn!( - session_id = % session.session_id, - "Backend returned 413 for save_session_data; \ - session data should already be in GCS via signed URL" - ); - } - Err(e) => return Err(e), - } - let share_response = self.create_share_link(&session.session_id).await?; - Ok(share_url(&share_response.permission_id)) - } - /// Sync session to backend without creating a share link. - pub async fn sync_session( - &self, - session: &ExportedSession, - agent_id: &str, - ) -> Result<(), BackendError> { - self.upsert_session(&session.session_id, &session.metadata, agent_id) - .await?; - self.save_session_data( - &session.session_id, - &session.messages, - Some(&session.metadata), - ) - .await?; - Ok(()) - } - /// Build auth + identity headers (plain bearer; no token-auth marker). - async fn auth_header_map(&self) -> Result { - use reqwest::header::{HeaderMap, HeaderValue}; - let auth = self.resolve_auth().await?; - let mut headers = HeaderMap::new(); - let required = |value: &str, name: &str| -> Result { - HeaderValue::from_str(value) - .map_err(|e| BackendError::Auth(format!("invalid {name} header: {e}"))) - }; - headers.insert("x-userid", required(&auth.user_id, "x-userid")?); - if let Some(email) = &auth.email - && let Ok(v) = HeaderValue::from_str(email) - { - headers.insert("x-email", v); - } - if let Ok(v) = HeaderValue::from_str(&crate::http::process_client_identifier()) { - headers.insert("x-grok-client-identifier", v); - } - headers.insert( - crate::http::CLIENT_MODE_HEADER, - HeaderValue::from_static(crate::http::process_client_mode()), - ); - headers.insert( - "x-grok-client-version", - HeaderValue::from_static(kigi_version::VERSION), - ); - Ok(headers) - } - async fn send_with_auth( - &self, - builder: reqwest::RequestBuilder, - ) -> Result { - let headers = self.auth_header_map().await?; - let builder = kigi_file_utils::trace_context::inject_trace_context_into_request( - builder.headers(headers), - ); - let request = builder.build()?; - self.client.execute(request).await.map_err(|e| match e { - reqwest_middleware::Error::Reqwest(e) => BackendError::Network(e), - reqwest_middleware::Error::Middleware(e) => BackendError::Auth(e.to_string()), - }) - } - pub async fn upsert_session( - &self, - session_id: &str, - metadata: &ExportedMetadata, - agent_id: &str, - ) -> Result<(), BackendError> { - let url = format!("{}/sessions/{}", self.base_url, session_id); - let request = UpsertSessionRequest { - session: SessionUpdate { - title: metadata.title.clone(), - cwd: Some(metadata.cwd.clone()), - status: Some("active".to_string()), - metadata: serde_json::to_value(metadata).ok(), - }, - agent_id: agent_id.to_string(), - }; - let response = self - .send_with_auth(self.reqwest_client.put(&url).json(&request)) - .await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); - return Err(BackendError::RequestFailed { status, body }); - } - Ok(()) - } - pub async fn save_session_data( - &self, - session_id: &str, - messages: &[ExportedMessage], - metadata: Option<&ExportedMetadata>, - ) -> Result<(), BackendError> { - let url = format!("{}/sessions/{}/data", self.base_url, session_id); - let request = SaveDataRequest { - messages: messages.to_vec(), - metadata: metadata.and_then(|m| serde_json::to_value(m).ok()), - }; - let response = self - .send_with_auth(self.reqwest_client.post(&url).json(&request)) - .await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); - return Err(BackendError::RequestFailed { status, body }); - } - Ok(()) - } - /// List all sessions for the authenticated user. `GET /sessions` - pub async fn list_sessions(&self) -> Result, BackendError> { - let url = format!("{}/sessions", self.base_url); - let response = self.send_with_auth(self.reqwest_client.get(&url)).await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); - return Err(BackendError::RequestFailed { status, body }); - } - #[derive(Deserialize)] - struct ListResponse { - sessions: Vec, - } - let data: ListResponse = response.json().await?; - Ok(data.sessions) - } - pub async fn load_session_data( - &self, - session_id: &str, - ) -> Result { - let url = format!("{}/sessions/{}/data", self.base_url, session_id); - let response = self.send_with_auth(self.reqwest_client.get(&url)).await?; - if response.status().as_u16() == 404 { - return Err(BackendError::SessionNotFound { - session_id: session_id.to_string(), - }); - } - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); - return Err(BackendError::RequestFailed { status, body }); - } - let data: LoadDataResponse = response.json().await?; - Ok(data) - } - pub async fn create_share_link(&self, session_id: &str) -> Result { - let url = format!("{}/sessions/{}/share", self.base_url, session_id); - let response = self.send_with_auth(self.reqwest_client.post(&url)).await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); - return Err(BackendError::RequestFailed { status, body }); - } - let share_response: ShareResponse = response.json().await?; - Ok(share_response) - } - pub async fn delete_session_data(&self, session_id: &str) -> Result<(), BackendError> { - let url = format!("{}/sessions/{}/data", self.base_url, session_id); - let response = self - .send_with_auth(self.reqwest_client.delete(&url)) - .await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().await.unwrap_or_default(); - return Err(BackendError::RequestFailed { status, body }); - } - Ok(()) - } -} -/// Fetch remote settings from cli-chat-proxy `GET /v1/settings`. -/// -/// This is a blocking call intended for use in the early prefetch thread -/// (`std::thread::spawn`, no tokio runtime). Returns `None` on any error -/// so startup is never blocked by a settings fetch failure. -/// -/// Retries up to 2 times (3 attempts total) on transient errors (5xx, -/// network). 4xx and parse errors are not retried. -pub fn fetch_settings_blocking( - cli_chat_proxy_base_url: &str, - auth: &KimiAuth, - alpha_test_key: Option<&str>, -) -> Option { - let client = crate::http::shared_blocking_client(); - let url = format!("{}/settings", cli_chat_proxy_base_url); - for attempt in 0u64..3 { - if attempt > 0 { - std::thread::sleep(std::time::Duration::from_millis(500 * attempt)); - } - let request = - add_cli_chat_proxy_headers_blocking(client.get(&url), auth, alpha_test_key, &url); - match request.send() { - Ok(resp) if resp.status().is_success() => match resp.json() { - Ok(settings) => { - tracing::debug!("Fetched remote settings from cli-chat-proxy"); - return Some(settings); - } - Err(e) => { - tracing::warn!(attempt, "Failed to parse settings response: {e}"); - return None; - } - }, - Ok(resp) if resp.status().is_server_error() => { - tracing::warn!( - attempt, - status = resp.status().as_u16(), - "Settings fetch server error, retrying" - ); - continue; - } - Ok(resp) => { - tracing::warn!(status = resp.status().as_u16(), "Failed to fetch settings"); - return None; - } - Err(e) => { - tracing::warn!(attempt, "Settings fetch network error: {e}"); - continue; - } - } - } - tracing::error!("Settings fetch failed after 3 attempts"); - None -} -#[derive(Deserialize)] -struct LoginConfigResponse { - /// Tri-state: `Some` forces a transport; `None`/absent → client default. - #[serde(default)] - device_flow: Option, -} -/// Fetch `grok_build_login_device_flow` from cli-chat-proxy `GET /v1/login-config`. -/// -/// Unauthenticated (pre-login); `x-grok-agent-id` is the per-install bucketing key. -/// Best-effort: any error or unset flag returns `None` so the caller keeps the -/// loopback default. Caps at 1.5s with no retries since it's on the login path; -/// `agent_id()` runs on the blocking pool so the fetch never stalls the executor. -pub async fn fetch_login_device_flow(cli_chat_proxy_base_url: &str) -> Option { - let agent_id = tokio::task::spawn_blocking(crate::util::agent_id::agent_id) - .await - .ok()?; - let client = crate::http::shared_client(); - let url = format!("{}/login-config", cli_chat_proxy_base_url); - let response = client - .get(&url) - .timeout(std::time::Duration::from_millis(1500)) - .header("x-grok-agent-id", agent_id) - .header("x-grok-client-version", kigi_version::VERSION) - .header( - "x-grok-client-identifier", - crate::http::process_client_identifier(), - ) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ) - .send() - .await; - let resp = match response { - Ok(resp) if resp.status().is_success() => resp, - Ok(resp) => { - tracing::debug!(status = resp.status().as_u16(), "login-config fetch failed"); - return None; - } - Err(e) => { - tracing::debug!("login-config fetch error: {e}"); - return None; - } - }; - match resp.json::().await { - Ok(cfg) => { - tracing::debug!( - device_flow = ? cfg.device_flow, "Fetched remote login-config" - ); - cfg.device_flow - } - Err(e) => { - tracing::debug!("Failed to parse login-config response: {e}"); - None - } - } -} -/// Default context window (256k) when the remote endpoint doesn't provide one. -pub(crate) const DEFAULT_CONTEXT_WINDOW: u64 = 256_000; -#[derive(Debug, Deserialize)] -struct ModelsResponse { - data: Vec, -} -/// The models-fetch origin key for this endpoints/auth shape. Used as the -/// models disk-cache origin: cached entries embed absolute `base_url`s from -/// the backend(s) that served them, so a catalog fetched against one fetch -/// plan (env override, different set of platform credentials, a test's mock -/// server) must be a cache miss for any other. Encodes URLs and enabled -/// platform NAMES only — never credential values. -pub(crate) fn models_fetch_origin( - endpoints: &crate::agent::config::EndpointsConfig, - fetch_auth: crate::agent::models::ModelFetchAuth, - has_oauth: bool, - platform_keys: &crate::agent::models::PlatformApiKeys, -) -> String { - match fetch_auth { - crate::agent::models::ModelFetchAuth::CustomEndpoint => endpoints.resolve_models_list_url(), - crate::agent::models::ModelFetchAuth::Platforms => { - let parts: Vec = enabled_platforms(has_oauth, platform_keys) - .into_iter() - .map(|p| format!("{}={}", p.as_str(), platform_models_url(p, endpoints))) - .collect(); - format!("platforms[{}]", parts.join(";")) - } - } -} -/// The platforms with usable credentials, in registry order (kimi-code first -/// so "default model = first list item" favors the subscription). -fn enabled_platforms( - has_oauth: bool, - platform_keys: &crate::agent::models::PlatformApiKeys, -) -> Vec { - kigi_models::PlatformId::ALL - .into_iter() - .filter(|p| { - if p.uses_oauth() { - has_oauth - } else { - platform_keys.key_for(*p).is_some() - } - }) - .collect() -} -/// `{base}/models` for one platform. The subscription platform resolves its -/// base through the endpoints config (`cli_chat_proxy_base_url` override, -/// else `KIGI_CODE_BASE_URL` / production default via kigi-env); the open -/// platforms use their fixed bases. -fn platform_models_url( - platform: kigi_models::PlatformId, - endpoints: &crate::agent::config::EndpointsConfig, -) -> String { - let base = if platform.uses_oauth() { - endpoints.proxy_url() - } else { - platform.base_url() - }; - format!("{}/models", base.trim_end_matches('/')) -} -/// Fetch result: model entries + optional etag from the subscription platform. -pub struct FetchModelsResult { - pub models: Vec, - pub etag: Option, - /// The OAuth platform answered 401. The async layer forces a token - /// refresh and retries once (port of kimi-cli `refresh_managed_models`). - pub oauth_unauthorized: bool, -} -/// Fetch the model catalog (PRD F4). -/// -/// - Custom endpoint mode (`KIGI_MODELS_BASE_URL` / `models_list_url`): a -/// single OpenAI-compatible listing fetched with the BYOK key or session -/// bearer, parsed leniently ([`parse_remote_model_value`]). -/// - Otherwise, the fixed platform registry: `GET {base}/models` with -/// `Authorization: Bearer ` per enabled platform, -/// parsed per the F4 wire contract with capability derivation and the -/// `kimi-k` prefix filter for the open platforms. -/// -/// Succeeds when at least one platform delivers; per-platform failures are -/// logged (status codes only, never credentials). -pub(crate) fn fetch_models_blocking( - endpoints: &crate::agent::config::EndpointsConfig, - auth: Option<&KimiAuth>, - fetch_auth: crate::agent::models::ModelFetchAuth, - platform_keys: &crate::agent::models::PlatformApiKeys, -) -> Result { - match fetch_auth { - crate::agent::models::ModelFetchAuth::CustomEndpoint => { - fetch_custom_endpoint_models_blocking(endpoints, auth) - } - crate::agent::models::ModelFetchAuth::Platforms => { - fetch_platform_models_blocking(endpoints, auth, platform_keys) - } - } -} -fn fetch_custom_endpoint_models_blocking( - endpoints: &crate::agent::config::EndpointsConfig, - auth: Option<&KimiAuth>, -) -> Result { - let client = crate::http::shared_blocking_client(); - let url = endpoints.resolve_models_list_url(); - let inference_base_url = endpoints.resolve_inference_base_url(); - tracing::info!("Fetching models from custom endpoint {}", url); - let api_key = crate::agent::auth_method::read_xai_api_key_env() - .or_else(|_| { - auth.map(|a| a.key.clone()) - .ok_or(std::env::VarError::NotPresent) - }) - .map_err(|_| { - BackendError::Auth("No API key for custom models endpoint. Set XAI_API_KEY.".into()) - })?; - let request = client - .get(&url) - .header("Authorization", format!("Bearer {}", api_key)); - let response = request.send()?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().unwrap_or_default(); - tracing::warn!("Failed to fetch models: {} - {}", status, body); - return Err(BackendError::RequestFailed { status, body }); - } - let etag = response - .headers() - .get("etag") - .and_then(|v| v.to_str().ok()) - .map(|s| s.to_string()); - let models_response: ModelsResponse = response.json()?; - tracing::info!("Fetched {} models from {}", models_response.data.len(), url); - let mut models = Vec::with_capacity(models_response.data.len()); - for (idx, value) in models_response.data.into_iter().enumerate() { - match parse_remote_model_value(&value, &inference_base_url) { - Some(model) => models.push(model), - None => { - tracing::warn!( - "Skipping model at index {}: missing required field ('model' or 'context_window') or invalid types", - idx - ) - } - } - } - Ok(FetchModelsResult { - models, - etag, - oauth_unauthorized: false, - }) -} -/// Registry fetch across all platforms with usable credentials. -fn fetch_platform_models_blocking( - endpoints: &crate::agent::config::EndpointsConfig, - auth: Option<&KimiAuth>, - platform_keys: &crate::agent::models::PlatformApiKeys, -) -> Result { - let enabled = enabled_platforms(auth.is_some(), platform_keys); - if enabled.is_empty() { - return Err(BackendError::Auth( - "No platform credentials: log in with `kigi login` or configure a moonshot API key \ - (KIGI_MOONSHOT_API_KEY or [platforms.*] in ~/.kigi/config.toml)." - .into(), - )); - } - - let mut models = Vec::new(); - let mut etag = None; - let mut oauth_unauthorized = false; - let mut successes = 0usize; - let mut last_error: Option = None; - for platform in &enabled { - let bearer = if platform.uses_oauth() { - auth.map(|a| a.key.clone()) - .expect("enabled_platforms gated on auth presence") - } else { - platform_keys - .key_for(*platform) - .expect("enabled_platforms gated on key presence") - .to_owned() - }; - match fetch_one_platform_models(*platform, endpoints, &bearer) { - Ok((platform_models, platform_etag)) => { - tracing::info!( - platform = platform.as_str(), - count = platform_models.len(), - "platform models fetch succeeded" - ); - successes += 1; - if platform.uses_oauth() { - etag = platform_etag; - } - models.extend(platform_models); - } - Err(e) => { - if platform.uses_oauth() - && matches!(&e, BackendError::RequestFailed { status: 401, .. }) - { - oauth_unauthorized = true; - } - tracing::warn!( - platform = platform.as_str(), - error = %e, - "platform models fetch failed" - ); - last_error = Some(e); - } - } - } - - if successes == 0 { - // All enabled platforms failed. When the failure includes an OAuth - // 401, return `Ok` with the flag set (and no models) so the async - // layer can force a token refresh and retry — an `Err` would drop - // the signal. Non-401 failures propagate as the last error. - if oauth_unauthorized { - return Ok(FetchModelsResult { - models: Vec::new(), - etag: None, - oauth_unauthorized: true, - }); - } - return Err(last_error.unwrap_or_else(|| { - BackendError::Auth("no platform models fetch was attempted".into()) - })); - } - Ok(FetchModelsResult { - models, - etag, - oauth_unauthorized, - }) -} -/// `GET {base}/models` for one platform (PRD F4 wire contract): -/// `Authorization: Bearer ` → `{data:[{id, context_length, -/// supports_reasoning, supports_image_in, supports_video_in, display_name?}]}`. -/// Applies the platform's `kimi-k` prefix filter and capability derivation, -/// and keys each entry `{platform_id}/{model_id}`. -fn fetch_one_platform_models( - platform: kigi_models::PlatformId, - endpoints: &crate::agent::config::EndpointsConfig, - bearer: &str, -) -> Result<(Vec, Option), BackendError> { - let client = crate::http::shared_blocking_client(); - let url = platform_models_url(platform, endpoints); - tracing::info!(platform = platform.as_str(), url = %url, "fetching platform models"); - let response = client - .get(&url) - .header("Authorization", format!("Bearer {}", bearer)) - .send()?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let body = response.text().unwrap_or_default(); - return Err(BackendError::RequestFailed { status, body }); - } - let etag = response - .headers() - .get("etag") - .and_then(|v| v.to_str().ok()) - .map(|s| s.to_string()); - let listing: kigi_models::WireModelsResponse = response.json()?; - let total = listing.data.len(); - let filtered = kigi_models::filter_allowed_models(platform, listing.data); - if filtered.len() != total { - tracing::info!( - platform = platform.as_str(), - total, - kept = filtered.len(), - "applied platform model-prefix filter" - ); - } - let base_url = if platform.uses_oauth() { - endpoints.proxy_url() - } else { - platform.base_url() - }; - let models = filtered - .into_iter() - .map(|wire| platform_wire_model_to_entry(platform, wire, &base_url)) - .collect(); - Ok((models, etag)) -} -/// Map one F4 wire model to a catalog entry config. -/// -/// SECURITY: the entry carries only env-var NAMES (`env_key`) for the open -/// platforms — never key values — because raw fetched entries are persisted -/// to the models disk cache. Config-file keys are stamped in-memory later by -/// `resolve_model_list`'s platform-credentials layer. -fn platform_wire_model_to_entry( - platform: kigi_models::PlatformId, - wire: kigi_models::WireModel, - base_url: &str, -) -> crate::agent::config::ModelEntryConfig { - let capabilities = wire.capabilities(); - let context_window = std::num::NonZeroU64::new(wire.context_length).unwrap_or_else(|| { - tracing::debug!( - model = %wire.id, - default = DEFAULT_CONTEXT_WINDOW, - "platform model missing context_length; using default" - ); - std::num::NonZeroU64::new(DEFAULT_CONTEXT_WINDOW).expect("non-zero") - }); - let env_key = (!platform.uses_oauth()) - .then(|| crate::agent::config::EnvKeys::new(platform.api_key_env_names().iter().copied())); - crate::agent::config::ModelEntryConfig { - id: Some(platform.managed_model_key(&wire.id)), - name: Some(wire.display_name.clone().unwrap_or_else(|| wire.id.clone())), - model: wire.id, - base_url: base_url.to_owned(), - description: None, - max_completion_tokens: None, - temperature: None, - top_p: None, - api_key: None, - env_key, - api_backend: Default::default(), - auth_scheme: None, - reasoning_effort: None, - supports_reasoning_effort: false, - reasoning_efforts: Vec::new(), - capabilities, - extra_headers: IndexMap::new(), - context_window, - auto_compact_threshold_percent: None, - system_prompt_label: None, - api_base_url: None, - use_concise: false, - agent_type: crate::agent::config::default_agent_type(), - inference_idle_timeout_secs: None, - max_retries: None, - hidden: false, - // Subscription models require the OAuth session; open-platform - // models are usable by API-key users. - supported_in_api: !platform.uses_oauth(), - supports_backend_search: false, - compactions_remaining: None, - compaction_at_tokens: None, - show_model_fingerprint: false, - stream_tool_calls: None, - laziness_detector: Default::default(), - } -} -/// Parse a single model entry from the /models-v2 response. -/// Used by both initial model fetch and session-resume metadata refresh. -pub fn parse_remote_model_value( - value: &serde_json::Value, - default_base_url: &str, -) -> Option { - let obj = value.as_object()?; - let meta = obj.get("_meta").and_then(|v| v.as_object()); - let id = get_string(obj, "id"); - let model = get_string(obj, "model") - .or_else(|| get_string(obj, "modelId")) - .or_else(|| id.clone()) - .or_else(|| meta.and_then(|m| get_string(m, "model"))) - .or_else(|| meta.and_then(|m| get_string(m, "modelId")))?; - let base_url = get_string(obj, "baseUrl") - .or_else(|| get_string(obj, "base_url")) - .unwrap_or_else(|| default_base_url.to_owned()); - let name = get_string(obj, "name").or_else(|| Some(model.clone())); - let context_window = get_u64(obj, "contextWindow") - .or_else(|| get_u64(obj, "context_window")) - .or_else(|| meta.and_then(|m| get_u64(m, "contextWindow"))) - .or_else(|| meta.and_then(|m| get_u64(m, "totalContextTokens"))) - .unwrap_or(DEFAULT_CONTEXT_WINDOW); - let context_window = std::num::NonZeroU64::new(context_window)?; - let agent_type = get_string(obj, "systemPromptType") - .or_else(|| get_string(obj, "system_prompt_type")) - .or_else(|| get_string(obj, "agent_type")) - .or_else(|| get_string(obj, "agentType")) - .or_else(|| meta.and_then(|m| get_string(m, "agentType"))) - .or_else(|| meta.and_then(|m| get_string(m, "agent_type"))) - .unwrap_or_else(crate::agent::config::default_agent_type); - let api_backend = get_string(obj, "apiBackend") - .or_else(|| get_string(obj, "api_backend")) - .and_then(|s| match s.as_str() { - "responses" => Some(crate::sampling::ApiBackend::Responses), - "chat_completions" => Some(crate::sampling::ApiBackend::ChatCompletions), - "messages" => Some(crate::sampling::ApiBackend::Messages), - _ => None, - }) - .unwrap_or_default(); - Some(crate::agent::config::ModelEntryConfig { - id, - model, - base_url, - name, - description: get_string(obj, "description"), - max_completion_tokens: get_u64(obj, "maxCompletionTokens") - .or_else(|| get_u64(obj, "max_completion_tokens")) - .and_then(|v| u32::try_from(v).ok()), - temperature: get_f64(obj, "temperature").map(|v| v as f32), - top_p: get_f64(obj, "topP").or_else(|| get_f64(obj, "top_p")).map(|v| v as f32), - api_key: get_string(obj, "apiKey").or_else(|| get_string(obj, "api_key")), - env_key: get_env_keys(obj, "envKey").or_else(|| get_env_keys(obj, "env_key")), - api_backend, - context_window, - auto_compact_threshold_percent: get_u64(obj, "autoCompactThresholdPercent") - .or_else(|| get_u64(obj, "auto_compact_threshold_percent")) - .and_then(|v| u8::try_from(v).ok()), - system_prompt_label: get_string(obj, "systemPromptLabel") - .or_else(|| get_string(obj, "system_prompt_label")) - .filter(|s| !s.trim().is_empty()), - extra_headers: get_string_map(obj, "extraHeaders"), - api_base_url: get_string(obj, "apiBaseUrl") - .or_else(|| get_string(obj, "api_base_url")), - use_concise: obj - .get("useConcise") - .or_else(|| obj.get("use_concise")) - .and_then(|v| v.as_bool()) - .unwrap_or(false), - agent_type, - inference_idle_timeout_secs: get_u64(obj, "inferenceIdleTimeoutSecs") - .or_else(|| get_u64(obj, "inference_idle_timeout_secs")), - max_retries: get_u64(obj, "maxRetries") - .or_else(|| get_u64(obj, "max_retries")) - .and_then(|v| u32::try_from(v).ok()), - hidden: obj - .get("hidden") - .or_else(|| meta.and_then(|m| m.get("hidden"))) - .and_then(|v| v.as_bool()) - .unwrap_or(false), - supported_in_api: obj - .get("supportedInApi") - .or_else(|| obj.get("supported_in_api")) - .or_else(|| meta.and_then(|m| m.get("supportedInApi"))) - .and_then(|v| v.as_bool()) - .unwrap_or(true), - auth_scheme: None, - reasoning_effort: get_string(obj, "reasoningEffort") - .or_else(|| get_string(obj, "reasoning_effort")) - .or_else(|| meta.and_then(|m| get_string(m, "reasoningEffort"))) - .and_then(|s| s.parse().ok()), - supports_reasoning_effort: obj - .get("supportsReasoningEffort") - .or_else(|| obj.get("supports_reasoning_effort")) - .or_else(|| meta.and_then(|m| m.get("supportsReasoningEffort"))) - .and_then(|v| v.as_bool()) - .unwrap_or(false), - reasoning_efforts: obj - .get("reasoningEfforts") - .or_else(|| obj.get("reasoning_efforts")) - .or_else(|| meta.and_then(|m| m.get("reasoningEfforts"))) - .and_then(|v| v.as_array()) - .map(|arr| kigi_sampling_types::parse_reasoning_effort_options(arr)) - .unwrap_or_default(), - capabilities: obj - .get("capabilities") - .and_then(|v| { - serde_json::from_value::>(v.clone()).ok() - }) - .unwrap_or_default(), - supports_backend_search: obj - .get("supportsBackendSearch") - .or_else(|| obj.get("supports_backend_search")) - .or_else(|| meta.and_then(|m| m.get("supportsBackendSearch"))) - .and_then(|v| v.as_bool()) - .unwrap_or(false), - compactions_remaining: obj - .get("compactionsRemaining") - .or_else(|| obj.get("compactions_remaining")) - .or_else(|| meta.and_then(|m| m.get("compactionsRemaining"))) - .and_then(parse_compactions_remaining) - .or_else(|| { - obj - .get("sendCompactionsRemaining") - .or_else(|| obj.get("send_compactions_remaining")) - .or_else(|| meta.and_then(|m| m.get("sendCompactionsRemaining"))) - .and_then(|v| v.as_bool()) - .map(kigi_sampling_types::CompactionsRemaining::Dynamic) - }), - compaction_at_tokens: obj - .get("compactionAtTokens") - .or_else(|| obj.get("compaction_at_tokens")) - .or_else(|| meta.and_then(|m| m.get("compactionAtTokens"))) - .and_then(parse_compaction_at_tokens), - show_model_fingerprint: obj - .get("showModelFingerprint") - .or_else(|| obj.get("show_model_fingerprint")) - .or_else(|| meta.and_then(|m| m.get("showModelFingerprint"))) - .and_then(|v| v.as_bool()) - .unwrap_or(false), - stream_tool_calls: obj - .get("streamToolCalls") - .or_else(|| obj.get("stream_tool_calls")) - .and_then(|v| v.as_bool()), - laziness_detector: get_object(obj, "lazinessDetector") - .or_else(|| get_object(obj, "laziness_detector")) - .or_else(|| meta.and_then(|m| get_object(m, "lazinessDetector"))) - .and_then(|v| match serde_json::from_value::< - crate::agent::config::LazinessDetectorPerModelConfig, - >(v.clone()) { - Ok(cfg) => Some(cfg), - Err(e) => { - tracing::warn!( - error = % e, - "Failed to deserialize laziness_detector block from remote model; falling back to default" - ); - None - } - }) - .unwrap_or_default(), - }) -} -fn get_string(obj: &serde_json::Map, key: &str) -> Option { - obj.get(key).and_then(|v| v.as_str()).map(|s| s.to_string()) -} -/// Parse `env_key` / `envKey` as a single string or a string array. -fn get_env_keys( - obj: &serde_json::Map, - key: &str, -) -> Option { - let v = obj.get(key)?; - if let Some(s) = v.as_str() { - return Some(crate::agent::config::EnvKeys::single(s)); - } - if let Some(arr) = v.as_array() { - let mut names = Vec::with_capacity(arr.len()); - for item in arr { - let Some(s) = item.as_str() else { - tracing::warn!( - key, - "env_key array has a non-string element; ignoring env_key" - ); - return None; - }; - if !s.is_empty() { - names.push(s.to_owned()); - } - } - if names.is_empty() { - return None; - } - return Some(crate::agent::config::EnvKeys::new(names)); - } - None -} -fn parse_compaction_at_tokens( - v: &serde_json::Value, -) -> Option { - use kigi_sampling_types::CompactionAtTokens; - v.as_bool() - .map(CompactionAtTokens::Enabled) - .or_else(|| v.as_u64().map(CompactionAtTokens::Fixed)) -} -fn parse_compactions_remaining( - v: &serde_json::Value, -) -> Option { - use kigi_sampling_types::CompactionsRemaining; - v.as_bool().map(CompactionsRemaining::Dynamic).or_else(|| { - v.as_u64() - .and_then(|n| u8::try_from(n).ok()) - .map(CompactionsRemaining::Fixed) - }) -} -fn get_u64(obj: &serde_json::Map, key: &str) -> Option { - obj.get(key).and_then(|v| v.as_u64()) -} -fn get_f64(obj: &serde_json::Map, key: &str) -> Option { - obj.get(key).and_then(|v| v.as_f64()) -} -fn get_object<'a>( - obj: &'a serde_json::Map, - key: &str, -) -> Option<&'a serde_json::Value> { - obj.get(key).filter(|v| v.is_object()) -} -fn get_string_map( - obj: &serde_json::Map, - key: &str, -) -> IndexMap { - obj.get(key) - .and_then(|v| v.as_object()) - .map(|map| { - map.iter() - .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string()))) - .collect() - }) - .unwrap_or_default() -} -#[cfg(test)] -mod tests { - use super::*; - use axum::{ - Router, - extract::{Path, State}, - http::{HeaderMap, StatusCode}, - routing::get, - }; - use std::sync::{Arc, Mutex}; - #[test] - fn login_config_response_parses_tristate() { - let parse = |s: &str| { - serde_json::from_str::(s) - .unwrap() - .device_flow - }; - assert_eq!(parse(r#"{"device_flow": true}"#), Some(true)); - assert_eq!(parse(r#"{"device_flow": false}"#), Some(false)); - assert_eq!(parse(r#"{"device_flow": null}"#), None); - assert_eq!(parse("{}"), None, "absent flag must parse as unset"); - } - #[test] - fn get_env_keys_parses_strings_and_rejects_non_strings() { - use crate::agent::config::EnvKeys; - let parse = |v: serde_json::Value| { - let obj = serde_json::json!({ "env_key" : v }); - get_env_keys(obj.as_object().unwrap(), "env_key") - }; - assert_eq!(parse(serde_json::json!("A")), Some(EnvKeys::single("A"))); - assert_eq!( - parse(serde_json::json!(["A", "B"])), - Some(EnvKeys::new(["A", "B"])) - ); - assert_eq!(parse(serde_json::json!(["A", 123])), None); - assert_eq!(parse(serde_json::json!([])), None); - } - fn header_str(headers: &HeaderMap, name: &str) -> Option { - headers - .get(name) - .and_then(|v| v.to_str().ok()) - .map(str::to_owned) - } - #[derive(Debug, Default, Clone)] - struct LoginConfigHeaders { - authorization: Option, - user_id: Option, - email: Option, - agent_id: Option, - client_identifier: Option, - client_version: Option, - } - #[derive(Clone)] - struct LoginConfigServerState { - status_code: StatusCode, - body: String, - seen: Arc>>, - } - /// Mock cli-chat-proxy serving `GET /v1/login-config` with a fixed status + - /// raw body, recording the request headers it saw. - async fn start_login_config_server( - status_code: StatusCode, - body: String, - ) -> ( - String, - Arc>>, - tokio::task::JoinHandle<()>, - ) { - let seen = Arc::new(Mutex::new(Vec::new())); - let state = LoginConfigServerState { - status_code, - body, - seen: seen.clone(), - }; - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let app = Router::new() - .route( - "/v1/login-config", - get( - |State(state): State, headers: HeaderMap| async move { - state.seen.lock().unwrap().push(LoginConfigHeaders { - authorization: header_str(&headers, "authorization"), - user_id: header_str(&headers, "x-userid"), - email: header_str(&headers, "x-email"), - agent_id: header_str(&headers, "x-grok-agent-id"), - client_identifier: header_str(&headers, "x-grok-client-identifier"), - client_version: header_str(&headers, "x-grok-client-version"), - }); - (state.status_code, state.body) - }, - ), - ) - .with_state(state); - let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (format!("{base}/v1"), seen, handle) - } - #[tokio::test] - async fn fetch_login_device_flow_parses_2xx_bodies() { - for (body, expected) in [ - (r#"{"device_flow": true}"#, Some(true)), - (r#"{"device_flow": false}"#, Some(false)), - (r#"{"device_flow": null}"#, None), - (r#"{}"#, None), - (r#"{"other": 1}"#, None), - ] { - let (base, _seen, server) = - start_login_config_server(StatusCode::OK, body.to_string()).await; - let got = fetch_login_device_flow(&base).await; - server.abort(); - assert_eq!(got, expected, "body {body:?}"); - } - } - #[tokio::test] - async fn fetch_login_device_flow_errors_return_none() { - for (status, body) in [ - (StatusCode::NOT_FOUND, r#"{"device_flow": true}"#), - ( - StatusCode::INTERNAL_SERVER_ERROR, - r#"{"device_flow": true}"#, - ), - (StatusCode::OK, "not json"), - ] { - let (base, _seen, server) = start_login_config_server(status, body.to_string()).await; - let got = fetch_login_device_flow(&base).await; - server.abort(); - assert_eq!(got, None, "status {status}, body {body:?}"); - } - } - #[tokio::test] - async fn fetch_login_device_flow_sends_only_unauthenticated_headers() { - let (base, seen, server) = - start_login_config_server(StatusCode::OK, r#"{"device_flow": true}"#.to_string()).await; - let got = fetch_login_device_flow(&base).await; - server.abort(); - assert_eq!(got, Some(true)); - let seen = seen.lock().unwrap(); - let h = seen - .last() - .expect("server should have received one request"); - assert!( - h.agent_id.as_deref().is_some_and(|v| !v.is_empty()), - "must send x-grok-agent-id (the bucketing key)" - ); - assert!( - h.client_identifier.is_some(), - "must send x-grok-client-identifier" - ); - assert!( - h.client_version.is_some(), - "must send x-grok-client-version" - ); - assert_eq!(h.authorization, None, "must not send Authorization"); - assert_eq!(h.user_id, None, "must not send x-userid"); - assert_eq!(h.email, None, "must not send x-email"); - } - #[derive(Debug, Default, Clone)] - struct SeenHeaders { - authorization: Option, - token_auth: Option, - user_id: Option, - email: Option, - alpha_test_key: Option, - client_version: Option, - } - #[derive(Clone)] - struct BundleServerState { - body: serde_json::Value, - status_code: StatusCode, - seen_headers: Arc>>, - } - async fn start_bundle_server( - status_code: StatusCode, - body: serde_json::Value, - ) -> ( - String, - Arc>>, - tokio::task::JoinHandle<()>, - ) { - let seen_headers = Arc::new(Mutex::new(Vec::new())); - let state = BundleServerState { - body, - status_code, - seen_headers: seen_headers.clone(), - }; - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let app = Router::new() - .route( - "/v1/subagents/bundle", - get( - |State(state): State, headers: HeaderMap| async move { - state.seen_headers.lock().unwrap().push(SeenHeaders { - authorization: headers - .get("authorization") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - token_auth: headers - .get("x-xai-token-auth") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - user_id: headers - .get("x-userid") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - email: headers - .get("x-email") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - alpha_test_key: { - let _ = &headers; - None - }, - client_version: headers - .get("x-grok-client-version") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - }); - (state.status_code, axum::Json(state.body)) - }, - ), - ) - .route( - "/forward/{tail}", - get( - |Path(_tail): Path, - State(state): State, - headers: HeaderMap| async move { - state.seen_headers.lock().unwrap().push(SeenHeaders { - authorization: headers - .get("authorization") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - token_auth: headers - .get("x-xai-token-auth") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - user_id: headers - .get("x-userid") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - email: headers - .get("x-email") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - alpha_test_key: { - let _ = &headers; - None - }, - client_version: headers - .get("x-grok-client-version") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned), - }); - (state.status_code, axum::Json(state.body)) - }, - ), - ) - .with_state(state); - let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (format!("{base}/v1"), seen_headers, handle) - } - fn test_auth() -> KimiAuth { - KimiAuth { - key: "token".to_string(), - auth_mode: crate::auth::AuthMode::OAuth, - create_time: chrono::Utc::now(), - user_id: "user-1".to_string(), - email: Some("test@example.com".to_string()), - refresh_token: None, - expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - expires_in: Some(3600), - scope: None, - token_type: None, - } - } - fn test_auth_manager() -> Arc { - let dir = tempfile::tempdir().unwrap(); - let mgr = crate::auth::AuthManager::new(dir.path(), crate::auth::KimiCodeConfig::default()); - mgr.hot_swap(test_auth()); - std::mem::forget(dir); - Arc::new(mgr) - } - #[tokio::test(flavor = "current_thread")] - async fn fetch_subagent_bundle_success() { - let body = serde_json::json!( - { "version" : "bundle-v1", "personas" : { "researcher" : "persona" }, "roles" - : { "reviewer" : "role" }, "agents" : { "default" : "agent" } } - ); - let (proxy_base_url, seen_headers, server) = - start_bundle_server(axum::http::StatusCode::OK, body).await; - let am = test_auth_manager(); - let bundle = fetch_subagent_bundle(&proxy_base_url, Some(&am), None, None) - .await - .unwrap(); - assert_eq!(bundle.version, "bundle-v1"); - assert_eq!( - bundle.personas.get("researcher"), - Some(&"persona".to_string()) - ); - assert_eq!(bundle.roles.get("reviewer"), Some(&"role".to_string())); - assert_eq!(bundle.agents.get("default"), Some(&"agent".to_string())); - let headers = seen_headers.lock().unwrap(); - let headers = headers.last().unwrap(); - assert_eq!(headers.authorization.as_deref(), Some("Bearer token")); - assert_eq!(headers.token_auth, None, "no token-auth marker header"); - assert_eq!(headers.user_id.as_deref(), Some("user-1")); - assert_eq!(headers.email.as_deref(), Some("test@example.com")); - assert_eq!(headers.alpha_test_key, None); - assert!(headers.client_version.is_some()); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - async fn fetch_subagent_bundle_uses_deployment_key_without_user_headers() { - let body = serde_json::json!( - { "version" : "bundle-v1", "personas" : {}, "roles" : {}, "agents" : {} } - ); - let (proxy_base_url, seen_headers, server) = - start_bundle_server(axum::http::StatusCode::OK, body).await; - let am = test_auth_manager(); - let bundle = fetch_subagent_bundle(&proxy_base_url, Some(&am), Some("deploy-key"), None) - .await - .unwrap(); - assert_eq!(bundle.version, "bundle-v1"); - let headers = seen_headers.lock().unwrap(); - let headers = headers.last().unwrap(); - assert_eq!(headers.authorization.as_deref(), Some("Bearer deploy-key")); - assert_eq!(headers.token_auth, None); - assert_eq!(headers.user_id, None); - assert_eq!(headers.email, None); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - async fn fetch_subagent_bundle_http_failure() { - let (proxy_base_url, _seen_headers, server) = start_bundle_server( - axum::http::StatusCode::UNAUTHORIZED, - serde_json::json!({ "error" : "unauthorized" }), - ) - .await; - let am = test_auth_manager(); - let error = fetch_subagent_bundle(&proxy_base_url, Some(&am), None, None) - .await - .unwrap_err(); - assert!(matches!( - error, - BackendError::RequestFailed { status: 401, .. } - )); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - async fn fetch_subagent_bundle_parse_failure() { - let (proxy_base_url, _seen_headers, server) = start_bundle_server( - axum::http::StatusCode::OK, - serde_json::json!({ "version" : 42 }), - ) - .await; - let am = test_auth_manager(); - let error = fetch_subagent_bundle(&proxy_base_url, Some(&am), None, None) - .await - .unwrap_err(); - assert!(matches!(error, BackendError::Serialization(_))); - server.abort(); - } - #[test] - fn parse_openai_format_uses_id_field() { - let value = serde_json::json!( - { "id" : "grok-3", "object" : "model", "owned_by" : "xai", "context_window" : - 131072 } - ); - let result = parse_remote_model_value(&value, "https://api.x.ai/v1").unwrap(); - assert_eq!(result.model, "grok-3"); - assert_eq!(result.base_url, "https://api.x.ai/v1"); - assert_eq!(result.name.as_deref(), Some("grok-3")); - } - #[test] - fn parse_model_field_takes_priority_over_id() { - let value = serde_json::json!( - { "id" : "display-key", "model" : "actual-model-id", "name" : "Display Name", - "context_window" : 131072 } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert_eq!(result.model, "actual-model-id"); - assert_eq!(result.name.as_deref(), Some("Display Name")); - } - #[test] - fn parse_reads_reasoning_effort_fields() { - use kigi_sampling_types::ReasoningEffort; - let value = serde_json::json!( - { "model" : "grok-4.5", "context_window" : 1_000_000, - "supports_reasoning_effort" : true, "reasoning_effort" : "high" } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert!(result.supports_reasoning_effort); - assert_eq!(result.reasoning_effort, Some(ReasoningEffort::High)); - let value = serde_json::json!( - { "model" : "grok-4.5", "contextWindow" : 1_000_000, - "supportsReasoningEffort" : true, "reasoningEffort" : "xhigh" } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert!(result.supports_reasoning_effort); - assert_eq!(result.reasoning_effort, Some(ReasoningEffort::Xhigh)); - let value = serde_json::json!({ "model" : "x", "context_window" : 256_000 }); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert!(!result.supports_reasoning_effort); - assert!(result.reasoning_effort.is_none()); - } - #[test] - fn parse_reads_reasoning_efforts_list() { - use kigi_sampling_types::ReasoningEffort; - let value = serde_json::json!( - { "model" : "grok-4.5", "context_window" : 1_000_000, "reasoning_efforts" : - [{ "id" : "deep", "value" : "xhigh", "label" : "Deep" }, { "value" : - "quantum" }, "low",] } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert_eq!(result.reasoning_efforts.len(), 2); - assert_eq!(result.reasoning_efforts[0].id, "deep"); - assert_eq!(result.reasoning_efforts[0].value, ReasoningEffort::Xhigh); - assert_eq!(result.reasoning_efforts[1].value, ReasoningEffort::Low); - for value in [ - serde_json::json!( - { "model" : "m", "context_window" : 256_000, "reasoningEfforts" : [{ - "value" : "high" }] } - ), - serde_json::json!( - { "model" : "m", "context_window" : 256_000, "_meta" : { - "reasoningEfforts" : [{ "value" : "high" }] } } - ), - ] { - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert_eq!(result.reasoning_efforts.len(), 1); - assert_eq!(result.reasoning_efforts[0].value, ReasoningEffort::High); - } - let value = serde_json::json!({ "model" : "x", "context_window" : 256_000 }); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert!(result.reasoning_efforts.is_empty()); - } - #[test] - fn parse_reads_meta_fallback_fields() { - let value = serde_json::json!( - { "_meta" : { "model" : "meta-model-id", "contextWindow" : 131072, - "agentType" : "concise" } } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert_eq!(result.model, "meta-model-id"); - assert_eq!( - result.context_window, - std::num::NonZeroU64::new(131072).unwrap() - ); - assert_eq!(result.agent_type, "concise"); - } - #[test] - fn parse_remote_model_value_no_laziness_detector_block_yields_default() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert_eq!( - result.laziness_detector, - crate::agent::config::LazinessDetectorPerModelConfig::default() - ); - } - #[test] - fn parse_remote_model_value_parses_camelcase_key() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { - "enabled" : true, "max_nudges_per_session" : 2, "idle_threshold_ms" : 12_000, - "min_confidence" : 0.75, }, } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - let expected = crate::agent::config::LazinessDetectorPerModelConfig { - enabled: true, - max_nudges_per_session: 2, - idle_threshold_ms: Some(12_000), - min_confidence: Some(0.75), - include_reasoning: None, - }; - assert_eq!(result.laziness_detector, expected); - } - #[test] - fn parse_remote_model_value_parses_snake_case_laziness_detector() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, "laziness_detector" : { - "enabled" : true, "max_nudges_per_session" : 3, "idle_threshold_ms" : 8_000, - "min_confidence" : 0.6, }, } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - let expected = crate::agent::config::LazinessDetectorPerModelConfig { - enabled: true, - max_nudges_per_session: 3, - idle_threshold_ms: Some(8_000), - min_confidence: Some(0.6), - include_reasoning: None, - }; - assert_eq!(result.laziness_detector, expected); - } - #[test] - fn parse_remote_model_value_parses_meta_laziness_detector() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, "_meta" : { - "lazinessDetector" : { "enabled" : true, "max_nudges_per_session" : 1, - "idle_threshold_ms" : 15_000, "min_confidence" : 0.9, }, }, } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - let expected = crate::agent::config::LazinessDetectorPerModelConfig { - enabled: true, - max_nudges_per_session: 1, - idle_threshold_ms: Some(15_000), - min_confidence: Some(0.9), - include_reasoning: None, - }; - assert_eq!(result.laziness_detector, expected); - } - #[test] - fn parse_remote_model_value_partial_block_uses_field_defaults() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { - "enabled" : true, }, } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - let expected = crate::agent::config::LazinessDetectorPerModelConfig { - enabled: true, - max_nudges_per_session: 0, - idle_threshold_ms: None, - min_confidence: None, - include_reasoning: None, - }; - assert_eq!(result.laziness_detector, expected); - } - #[test] - fn parse_remote_model_value_malformed_block_falls_back_to_default() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { - "enabled" : true, "max_nudges_per_session" : "abc", }, } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert_eq!( - result.laziness_detector, - crate::agent::config::LazinessDetectorPerModelConfig::default() - ); - } - #[test] - fn parse_remote_model_value_non_object_value_falls_back_to_default() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : - "not-an-object", } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert_eq!( - result.laziness_detector, - crate::agent::config::LazinessDetectorPerModelConfig::default() - ); - } - #[test] - fn parse_remote_model_value_top_level_camelcase_wins_over_snake_case() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { - "enabled" : true, "max_nudges_per_session" : 7, }, "laziness_detector" : { - "enabled" : false, "max_nudges_per_session" : 99, }, } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - let expected = crate::agent::config::LazinessDetectorPerModelConfig { - enabled: true, - max_nudges_per_session: 7, - idle_threshold_ms: None, - min_confidence: None, - include_reasoning: None, - }; - assert_eq!(result.laziness_detector, expected); - } - /// `include_reasoning: false` parses cleanly under the per-model - /// `lazinessDetector` block (camelCase wrapper, snake_case inner — - /// matching the existing field-naming convention used for the - /// sibling `min_confidence`, `idle_threshold_ms`, etc.). - #[test] - fn parse_remote_model_value_parses_include_reasoning_under_camelcase_wrapper() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { - "enabled" : true, "include_reasoning" : false, }, } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert_eq!(result.laziness_detector.include_reasoning, Some(false)); - } - #[test] - fn parse_remote_model_value_parses_include_reasoning_under_snake_case_wrapper() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, "laziness_detector" : { - "enabled" : true, "include_reasoning" : true, }, } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert_eq!(result.laziness_detector.include_reasoning, Some(true)); - } - #[test] - fn parse_remote_model_value_omitted_include_reasoning_defaults_to_none() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { - "enabled" : true, "max_nudges_per_session" : 2, }, } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert_eq!( - result.laziness_detector.include_reasoning, None, - "absent include_reasoning defers to harness default via None", - ); - } - #[test] - fn parse_remote_model_value_top_level_wins_over_meta() { - let value = serde_json::json!( - { "model" : "grok-4", "context_window" : 256_000, "lazinessDetector" : { - "enabled" : true, "max_nudges_per_session" : 5, }, "_meta" : { - "lazinessDetector" : { "enabled" : false, "max_nudges_per_session" : 99, }, - }, } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - let expected = crate::agent::config::LazinessDetectorPerModelConfig { - enabled: true, - max_nudges_per_session: 5, - idle_threshold_ms: None, - min_confidence: None, - include_reasoning: None, - }; - assert_eq!(result.laziness_detector, expected); - } - #[test] - fn parse_reads_show_model_fingerprint_field() { - let value = serde_json::json!( - { "model" : "grok-build", "context_window" : 256_000, - "show_model_fingerprint" : true } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert!(result.show_model_fingerprint); - let value = serde_json::json!( - { "model" : "grok-build", "contextWindow" : 256_000, "showModelFingerprint" : - true } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert!(result.show_model_fingerprint); - let value = serde_json::json!( - { "model" : "grok-build", "context_window" : 256_000, "_meta" : { - "showModelFingerprint" : true } } - ); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert!(result.show_model_fingerprint); - let value = serde_json::json!({ "model" : "x", "context_window" : 256_000 }); - let result = parse_remote_model_value(&value, "https://default.url").unwrap(); - assert!(!result.show_model_fingerprint); - } - #[test] - fn get_object_returns_none_for_non_object_values() { - let value = serde_json::json!( - { "string" : "hello", "number" : 42, "bool" : true, "array" : [1, 2, 3], - "null" : null, } - ); - let obj = value.as_object().unwrap(); - assert!(get_object(obj, "string").is_none()); - assert!(get_object(obj, "number").is_none()); - assert!(get_object(obj, "bool").is_none()); - assert!(get_object(obj, "array").is_none()); - assert!(get_object(obj, "null").is_none()); - assert!(get_object(obj, "missing").is_none()); - } - #[test] - fn get_object_returns_some_for_actual_object() { - let value = serde_json::json!({ "nested" : { "a" : 1, "b" : "two" }, }); - let obj = value.as_object().unwrap(); - let nested = get_object(obj, "nested").expect("nested key should resolve to object"); - assert!(nested.is_object()); - assert_eq!(nested["a"], serde_json::json!(1)); - assert_eq!(nested["b"], serde_json::json!("two")); - } - fn endpoints( - proxy: &str, - models_base_url: Option<&str>, - models_list_url: Option<&str>, - ) -> crate::agent::config::EndpointsConfig { - crate::agent::config::EndpointsConfig { - cli_chat_proxy_base_url: Some(proxy.to_owned()), - models_base_url: models_base_url.map(|s| s.to_owned()), - models_list_url: models_list_url.map(|s| s.to_owned()), - ..Default::default() - } - } - #[test] - fn inference_url_defaults_to_proxy() { - let ep = endpoints("https://proxy.kigi.com/v1", None, None); - assert_eq!(ep.resolve_inference_base_url(), "https://proxy.kigi.com/v1"); - } - #[test] - fn inference_url_uses_models_base_url() { - let ep = endpoints( - "https://proxy.kigi.com/v1", - Some("https://enterprise.acme.com/v1"), - None, - ); - assert_eq!( - ep.resolve_inference_base_url(), - "https://enterprise.acme.com/v1" - ); - } - #[test] - fn inference_url_base_url_wins_over_proxy() { - let ep = endpoints( - "https://proxy.kigi.com/v1", - Some("https://inference.acme.com/v1"), - Some("https://registry.acme.com/api/models"), - ); - assert_eq!( - ep.resolve_inference_base_url(), - "https://inference.acme.com/v1" - ); - } - #[test] - fn list_url_defaults_to_proxy_models() { - let ep = endpoints("https://proxy.kigi.com/v1", None, None); - assert_eq!( - ep.resolve_models_list_url(), - "https://proxy.kigi.com/v1/models" - ); - } - #[test] - fn list_url_derived_from_base_url() { - let ep = endpoints( - "https://proxy.kigi.com/v1", - Some("https://api.x.ai/v1"), - None, - ); - assert_eq!(ep.resolve_models_list_url(), "https://api.x.ai/v1/models"); - } - #[test] - fn list_url_explicit_overrides_derivation() { - let ep = endpoints( - "https://proxy.kigi.com/v1", - Some("https://inference.acme.com/v1"), - Some("https://registry.acme.com/api/list-models"), - ); - assert_eq!( - ep.resolve_models_list_url(), - "https://registry.acme.com/api/list-models" - ); - } - /// INVARIANT: each platform's `/models` URL matches its registry base — - /// kimi-code → the subscription proxy (config override respected, else the - /// kigi-env default), moonshot platforms → their fixed bases — and the - /// cache-origin key encodes the enabled fetch plan without any secrets. - #[test] - #[serial_test::serial] - fn platform_models_urls_and_fetch_origin() { - use crate::agent::config::EndpointsConfig; - use crate::agent::models::{ModelFetchAuth, PlatformApiKeys}; - for k in [ - "KIGI_CLI_CHAT_PROXY_BASE_URL", - "KIGI_CODE_BASE_URL", - "KIGI_MODELS_LIST_URL", - ] { - unsafe { std::env::remove_var(k) }; - } - let cfg = EndpointsConfig::from_config_value(&toml::Value::Table(Default::default())); - assert_eq!( - platform_models_url(kigi_models::PlatformId::KimiCode, &cfg), - "https://api.kimi.com/coding/v1/models" - ); - assert_eq!( - platform_models_url(kigi_models::PlatformId::MoonshotCn, &cfg), - "https://api.moonshot.cn/v1/models" - ); - assert_eq!( - platform_models_url(kigi_models::PlatformId::MoonshotAi, &cfg), - "https://api.moonshot.ai/v1/models" - ); - // Proxy override re-points the subscription platform only. - let proxied = EndpointsConfig::from_config_value( - &toml::from_str( - r#"[endpoints] - cli_chat_proxy_base_url = "https://proxy.acme.example/v1""#, - ) - .unwrap(), - ); - assert_eq!( - platform_models_url(kigi_models::PlatformId::KimiCode, &proxied), - "https://proxy.acme.example/v1/models" - ); - assert_eq!( - platform_models_url(kigi_models::PlatformId::MoonshotCn, &proxied), - "https://api.moonshot.cn/v1/models" - ); - - // Origin key: OAuth-only plan lists kimi-code only; adding a moonshot - // key changes the plan (→ cache miss); the key VALUE never appears. - let oauth_only = models_fetch_origin( - &cfg, - ModelFetchAuth::Platforms, - true, - &PlatformApiKeys::default(), - ); - assert_eq!( - oauth_only, - "platforms[kimi-code=https://api.kimi.com/coding/v1/models]" - ); - let with_cn = models_fetch_origin( - &cfg, - ModelFetchAuth::Platforms, - true, - &crate::agent::models::PlatformApiKeys::test_keys(Some("sk-secret-cn"), None), - ); - assert_ne!( - oauth_only, with_cn, - "enabling a platform must change the origin" - ); - assert!(with_cn.contains("moonshot-cn=https://api.moonshot.cn/v1/models")); - assert!( - !with_cn.contains("sk-secret-cn"), - "origin key must never embed credential values" - ); - - // Custom endpoint mode → the explicit list URL verbatim. - let custom = EndpointsConfig::from_config_value( - &toml::from_str( - r#"[endpoints] - models_base_url = "https://models.acme.com/v1""#, - ) - .unwrap(), - ); - assert_eq!( - models_fetch_origin( - &custom, - ModelFetchAuth::CustomEndpoint, - false, - &PlatformApiKeys::default(), - ), - "https://models.acme.com/v1/models" - ); - } - /// REGRESSION: `grok setup` must send the deployment key to - /// the proxy, never the inference endpoint. - #[test] - #[serial_test::serial] - fn deployment_config_url_uses_cli_chat_proxy_when_not_overridden() { - use crate::agent::config::EndpointsConfig; - for k in [ - "KIGI_CLI_CHAT_PROXY_BASE_URL", - "KIGI_MANAGED_CONFIG_URL", - "KIGI_XAI_API_BASE_URL", - ] { - unsafe { std::env::remove_var(k) }; - } - unsafe { std::env::set_var("KIGI_DEPLOYMENT_KEY", "xai-token-ENTERPRISE") }; - let managed: toml::Value = toml::from_str( - r#"[endpoints] - deployment_key = "xai-token-ENTERPRISE" - xai_api_base_url = "https://inference.acme-corp.example/xai/v1""#, - ) - .unwrap(); - let url = EndpointsConfig::from_config_value(&managed).resolve_managed_config_url(); - assert_eq!(url, "https://api.kimi.com/coding/v1/deployment/config"); - assert!( - !url.contains("acme-corp"), - "deployment key would be sent to the inference host: {url}" - ); - let pinned: toml::Value = toml::from_str( - r#"[endpoints] - xai_api_base_url = "https://inference.acme-corp.example/xai/v1" - cli_chat_proxy_base_url = "https://proxy.acme-corp.example/v1""#, - ) - .unwrap(); - assert_eq!( - EndpointsConfig::from_config_value(&pinned).resolve_managed_config_url(), - "https://proxy.acme-corp.example/v1/deployment/config" - ); - unsafe { std::env::remove_var("KIGI_DEPLOYMENT_KEY") }; - } - #[derive(Clone)] - struct DualBundleServerState { - archive_status: StatusCode, - archive_bytes: Vec, - legacy_status: StatusCode, - legacy_body: serde_json::Value, - } - async fn start_dual_bundle_server( - state: DualBundleServerState, - ) -> (String, tokio::task::JoinHandle<()>) { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port()); - let app = Router::new() - .route( - "/v1/bundle/archive", - get(|State(state): State| async move { - (state.archive_status, state.archive_bytes) - }), - ) - .route( - "/v1/subagents/bundle", - get(|State(state): State| async move { - (state.legacy_status, axum::Json(state.legacy_body)) - }), - ) - .with_state(state); - let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (format!("{base}/v1"), handle) - } - #[tokio::test(flavor = "current_thread")] - async fn fetch_bundle_returns_archive_on_success() { - let archive_bytes = b"fake-tar-gz-bytes".to_vec(); - let (proxy_base_url, server) = start_dual_bundle_server(DualBundleServerState { - archive_status: StatusCode::OK, - archive_bytes: archive_bytes.clone(), - legacy_status: StatusCode::OK, - legacy_body: serde_json::json!( - { "version" : "v1", "personas" : {}, "roles" : {}, "agents" : {} } - ), - }) - .await; - let am = test_auth_manager(); - let result = fetch_bundle(&proxy_base_url, Some(&am), None, None) - .await - .unwrap(); - match result { - FetchedBundle::Archive(bytes) => assert_eq!(bytes, archive_bytes), - FetchedBundle::Legacy(_) => panic!("expected Archive variant"), - } - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - async fn fetch_bundle_falls_back_on_archive_404() { - let (proxy_base_url, server) = start_dual_bundle_server(DualBundleServerState { - archive_status: StatusCode::NOT_FOUND, - archive_bytes: Vec::new(), - legacy_status: StatusCode::OK, - legacy_body: serde_json::json!( - { "version" : "v1", "personas" : { "r" : "p" }, "roles" : {}, - "agents" : {} } - ), - }) - .await; - let am = test_auth_manager(); - let result = fetch_bundle(&proxy_base_url, Some(&am), None, None) - .await - .unwrap(); - match result { - FetchedBundle::Legacy(bundle) => { - assert_eq!(bundle.version, "v1"); - assert_eq!(bundle.personas.get("r"), Some(&"p".to_string())); - } - FetchedBundle::Archive(_) => panic!("expected Legacy variant"), - } - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - async fn fetch_bundle_falls_back_on_archive_503() { - let (proxy_base_url, server) = start_dual_bundle_server(DualBundleServerState { - archive_status: StatusCode::SERVICE_UNAVAILABLE, - archive_bytes: Vec::new(), - legacy_status: StatusCode::OK, - legacy_body: serde_json::json!( - { "version" : "v1", "personas" : {}, "roles" : {}, "agents" : {} } - ), - }) - .await; - let am = test_auth_manager(); - let result = fetch_bundle(&proxy_base_url, Some(&am), None, None) - .await - .unwrap(); - match &result { - FetchedBundle::Legacy(bundle) => assert_eq!(bundle.version, "v1"), - FetchedBundle::Archive(_) => panic!("expected Legacy variant"), - } - server.abort(); - } - /// `BackendClient::save_session_data` resolves auth from the attached - /// `AuthManager` and sends the token as `Bearer ` on the wire. - /// This is the writeback path used on every session flush. - #[tokio::test(flavor = "current_thread")] - async fn backend_client_resolves_auth_from_auth_manager() { - let captured_auth = Arc::new(Mutex::new(None::)); - let captured = captured_auth.clone(); - let app = Router::new().route( - "/sessions/{id}/data", - axum::routing::post(move |headers: HeaderMap| async move { - *captured.lock().unwrap() = headers - .get("authorization") - .and_then(|v| v.to_str().ok()) - .map(str::to_owned); - StatusCode::OK - }), - ); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - let am = test_auth_manager(); - let client = BackendClient::with_base_url(format!("http://{addr}")).with_auth_manager(am); - client - .save_session_data("test-session", &[], None) - .await - .unwrap(); - let sent = captured_auth - .lock() - .unwrap() - .clone() - .expect("server must receive Authorization header"); - assert_eq!(sent, "Bearer token", "must use token from AuthManager"); - server.abort(); - } - #[tokio::test(flavor = "current_thread")] - async fn fetch_bundle_propagates_legacy_error_after_fallback() { - let (proxy_base_url, server) = start_dual_bundle_server(DualBundleServerState { - archive_status: StatusCode::NOT_FOUND, - archive_bytes: Vec::new(), - legacy_status: StatusCode::UNAUTHORIZED, - legacy_body: serde_json::json!({ "error" : "unauthorized" }), - }) - .await; - let am = test_auth_manager(); - let error = fetch_bundle(&proxy_base_url, Some(&am), None, None) - .await - .unwrap_err(); - assert!(matches!( - error, - BackendError::RequestFailed { status: 401, .. } - )); - server.abort(); - } - /// Regression: reqwest .header() appends — duplicate - /// or overlapping headers cause Cloudflare to reject the request. - #[tokio::test(flavor = "current_thread")] - async fn auth_headers_do_not_collide_with_json() { - let client = - BackendClient::with_base_url("http://localhost").with_auth_manager(test_auth_manager()); - let auth_headers = client.auth_header_map().await.unwrap(); - assert!( - !auth_headers.contains_key("content-type"), - "content-type in auth map would overwrite .json()" - ); - let request = reqwest::Client::new() - .put("http://localhost/sessions/test") - .json(&serde_json::json!({ "test" : true })) - .headers(auth_headers) - .build() - .unwrap(); - for name in request.headers().keys() { - let count = request.headers().get_all(name).iter().count(); - assert_eq!(count, 1, "duplicate header {name}"); - } - } -} diff --git a/crates/codegen/kigi-shell/src/remote/conversations_client.rs b/crates/codegen/kigi-shell/src/remote/conversations_client.rs deleted file mode 100644 index a84212b..0000000 --- a/crates/codegen/kigi-shell/src/remote/conversations_client.rs +++ /dev/null @@ -1,305 +0,0 @@ -use std::sync::Arc; - -use serde::{Deserialize, Serialize}; - -use crate::auth::{AuthManager, KimiAuth}; - -const KIGI_WEB_URL: &str = "https://grok.com"; - -#[derive(Debug, Clone, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct Conversation { - #[serde(default)] - pub conversation_id: String, - #[serde(default)] - pub title: String, - #[serde(default)] - pub starred: bool, - #[serde(default)] - pub create_time: Option, - #[serde(default)] - pub modify_time: Option, - #[serde(default)] - pub workspaces: Vec, -} - -#[derive(Debug, Clone, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct Workspace { - #[serde(default)] - pub workspace_id: String, -} - -#[derive(Debug, Clone, Default)] -pub struct ConvQuery { - pub page_size: i64, - pub page_token: Option, - pub search_query: Option, - pub workspace_id: Option, -} - -#[derive(Debug, Clone, Default)] -pub struct ListConversationsPage { - pub conversations: Vec, - pub next_page_token: Option, -} - -/// Body for `PUT /rest/app-chat/conversations/{id}` (grok-web `chatUpdateConversation`). -#[derive(Debug, Clone, Default, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdateConversationBody { - #[serde(skip_serializing_if = "Option::is_none")] - pub title: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub starred: Option, -} - -#[derive(Debug, thiserror::Error)] -pub enum ConvError { - #[error("no OAuth credentials for conversations:read")] - NoOauth, - #[error("network error: {0}")] - Network(#[from] reqwest::Error), - #[error("request failed: {status}")] - Http { status: u16 }, - #[error("parse error: {0}")] - Parse(#[from] serde_json::Error), -} - -#[derive(Debug, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -struct ListConversationsResponseWire { - #[serde(default)] - conversations: Vec, - #[serde(default)] - next_page_token: Option, - #[serde(default)] - text_search_matches: Vec, -} - -#[derive(Debug, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -struct ListConversationsMatchWire { - #[serde(default)] - conversation: Option, -} - -pub struct ConversationsClient { - http: reqwest::Client, - base_url: String, - auth: Arc, -} - -impl ConversationsClient { - pub fn new(auth: Arc) -> Self { - let base_url = std::env::var("KIGI_CONVERSATIONS_BASE_URL") - .ok() - .filter(|s| !s.is_empty()) - .or_else(|| { - std::env::var("KIGI_CODE_WEB_URL") - .ok() - .filter(|s| !s.is_empty()) - }) - .unwrap_or_else(|| KIGI_WEB_URL.to_string()); - Self { - http: crate::http::shared_client(), - base_url, - auth, - } - } - - async fn require_xai_auth(&self) -> Result { - let auth = self.auth.auth().await.map_err(|_| ConvError::NoOauth)?; - if !auth.is_session_auth() { - return Err(ConvError::NoOauth); - } - Ok(auth) - } - - fn apply_auth_headers( - &self, - builder: reqwest::RequestBuilder, - auth: &KimiAuth, - ) -> reqwest::RequestBuilder { - let mut builder = builder - .header("Authorization", format!("Bearer {}", auth.key)) - .header("x-userid", &auth.user_id) - .header("x-grok-client-version", kigi_version::VERSION) - .header( - "x-grok-client-identifier", - crate::http::process_client_identifier(), - ) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ) - .header(reqwest::header::ACCEPT, "application/json"); - if let Some(email) = &auth.email { - builder = builder.header("x-email", email); - } - kigi_file_utils::trace_context::inject_trace_context_into_request(builder) - } - - pub async fn list_conversations( - &self, - q: &ConvQuery, - ) -> Result { - let auth = self.require_xai_auth().await?; - - let url = format!("{}/rest/app-chat/conversations", self.base_url); - let mut query: Vec<(&str, String)> = vec![("pageSize", q.page_size.to_string())]; - if let Some(token) = q.page_token.as_deref().filter(|s| !s.is_empty()) { - query.push(("pageToken", token.to_owned())); - } - if let Some(search) = q.search_query.as_deref().filter(|s| !s.is_empty()) { - query.push(("searchQuery", search.to_owned())); - } - if let Some(workspace) = q.workspace_id.as_deref().filter(|s| !s.is_empty()) { - query.push(("workspaceId", workspace.to_owned())); - } - - let builder = self.apply_auth_headers(self.http.get(&url).query(&query), &auth); - - let response = builder.send().await?; - let status = response.status(); - if !status.is_success() { - return Err(ConvError::Http { - status: status.as_u16(), - }); - } - - let bytes = response.bytes().await?; - let wire: ListConversationsResponseWire = serde_json::from_slice(&bytes)?; - - let searching = q.search_query.as_deref().is_some_and(|s| !s.is_empty()); - // During an active search, results come exclusively from - // `text_search_matches`. Never fall back to `wire.conversations` here: - // an empty match set means "no hits", and the server may return - // recent/unfiltered conversations in `conversations` that are NOT search - // matches — surfacing those would be wrong. - let conversations = if searching { - wire.text_search_matches - .into_iter() - .filter_map(|m| m.conversation) - .collect() - } else { - wire.conversations - }; - - Ok(ListConversationsPage { - conversations, - next_page_token: wire.next_page_token.filter(|t| !t.is_empty()), - }) - } - - /// `PUT /rest/app-chat/conversations/{conversation_id}` — rename and/or star. - pub async fn update_conversation( - &self, - conversation_id: &str, - body: &UpdateConversationBody, - ) -> Result<(), ConvError> { - let auth = self.require_xai_auth().await?; - let url = format!( - "{}/rest/app-chat/conversations/{}", - self.base_url, - urlencoding::encode(conversation_id) - ); - let builder = self - .apply_auth_headers(self.http.put(&url), &auth) - .json(body); - - let response = builder.send().await?; - let status = response.status(); - if !status.is_success() { - return Err(ConvError::Http { - status: status.as_u16(), - }); - } - Ok(()) - } - - /// `DELETE /rest/app-chat/conversations/soft/{conversation_id}` — soft-delete. - pub async fn soft_delete_conversation(&self, conversation_id: &str) -> Result<(), ConvError> { - let auth = self.require_xai_auth().await?; - let url = format!( - "{}/rest/app-chat/conversations/soft/{}", - self.base_url, - urlencoding::encode(conversation_id) - ); - let builder = self.apply_auth_headers(self.http.delete(&url), &auth); - - let response = builder.send().await?; - let status = response.status(); - // 404 = already soft-deleted; keep deletion idempotent like the - // build path's `classify_remote_delete`. - if !status.is_success() && status.as_u16() != 404 { - return Err(ConvError::Http { - status: status.as_u16(), - }); - } - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn conversation_parses_camelcase_wire() { - let json = serde_json::json!({ - "conversations": [{ - "conversationId": "conv_abc", - "title": "Compare GPU vendors", - "starred": true, - "createTime": "2026-06-18T17:30:00Z", - "modifyTime": "2026-06-18T18:02:00Z", - "workspaces": [{ "workspaceId": "ws_9f3a" }] - }], - "nextPageToken": "tok2" - }); - let wire: ListConversationsResponseWire = serde_json::from_value(json).unwrap(); - assert_eq!(wire.conversations.len(), 1); - let c = &wire.conversations[0]; - assert_eq!(c.conversation_id, "conv_abc"); - assert_eq!(c.title, "Compare GPU vendors"); - assert!(c.starred); - assert_eq!(c.modify_time.as_deref(), Some("2026-06-18T18:02:00Z")); - assert_eq!(c.workspaces[0].workspace_id, "ws_9f3a"); - assert_eq!(wire.next_page_token.as_deref(), Some("tok2")); - } - - #[test] - fn missing_fields_default_gracefully() { - let json = serde_json::json!({ "conversations": [{ "conversationId": "c1" }] }); - let wire: ListConversationsResponseWire = serde_json::from_value(json).unwrap(); - let c = &wire.conversations[0]; - assert_eq!(c.conversation_id, "c1"); - assert!(c.title.is_empty()); - assert!(c.modify_time.is_none()); - assert!(c.create_time.is_none()); - assert!(c.workspaces.is_empty()); - assert!(wire.next_page_token.is_none()); - } - - #[test] - fn update_body_serializes_only_set_fields() { - let title_only = UpdateConversationBody { - title: Some("New title".into()), - starred: None, - }; - assert_eq!( - serde_json::to_value(&title_only).unwrap(), - serde_json::json!({ "title": "New title" }) - ); - - let both = UpdateConversationBody { - title: Some("T".into()), - starred: Some(true), - }; - assert_eq!( - serde_json::to_value(&both).unwrap(), - serde_json::json!({ "title": "T", "starred": true }) - ); - } -} diff --git a/crates/codegen/kigi-shell/src/remote/mod.rs b/crates/codegen/kigi-shell/src/remote/mod.rs deleted file mode 100644 index b46663d..0000000 --- a/crates/codegen/kigi-shell/src/remote/mod.rs +++ /dev/null @@ -1,37 +0,0 @@ -//! Remote storage client for the backend. - -pub mod agent; -pub mod chat_models_client; -pub mod client; -pub mod conversations_client; -pub mod pull; -#[cfg(test)] -mod pull_smoke_test; -pub mod sync; -pub mod workspaces_client; - -pub use agent::{ - SandboxClient, SandboxCreateEnvironmentRequest, SandboxEnvironment, SandboxEnvironmentResponse, - SandboxEnvironmentVariable, SandboxEnvironmentWithMetadata, SandboxForkRequest, - SandboxForkResponse, SandboxForkedSession, SandboxHibernateResponse, - SandboxListEnvironmentsRequest, SandboxListEnvironmentsResponse, - SandboxListPreinstalledPackagesResponse, SandboxLogsExitCodes, SandboxLogsResponse, - SandboxMode, SandboxPreinstalledPackage, SandboxRestoreRequest, SandboxRestoreResponse, - SandboxSecretInput, SandboxStartRequest, SandboxStartResponse, SandboxStatusResponse, - SandboxTerminateRequest, SandboxUpdateEnvironmentRequest, -}; -pub use chat_models_client::{ - ChatModelsClient, ChatModelsError, ListModesResponse, Mode, ModeAvailability, -}; -pub use client::{ - BackendClient, BackendError, FetchModelsResult, FetchedBundle, fetch_bundle, - fetch_login_device_flow, fetch_settings_blocking, fetch_subagent_bundle, share_url, -}; -pub(crate) use client::{DEFAULT_CONTEXT_WINDOW, fetch_models_blocking, models_fetch_origin}; -pub use conversations_client::{ - ConvError, ConvQuery, Conversation, ConversationsClient, ListConversationsPage, - UpdateConversationBody, -}; -pub use pull::{PullResult, pull_session_to_local}; -pub use sync::RemoteSync; -pub use workspaces_client::{ListWorkspacesPage, Workspace, WorkspacesClient, WsError, WsQuery}; diff --git a/crates/codegen/kigi-shell/src/remote/pull.rs b/crates/codegen/kigi-shell/src/remote/pull.rs deleted file mode 100644 index b9b723b..0000000 --- a/crates/codegen/kigi-shell/src/remote/pull.rs +++ /dev/null @@ -1,772 +0,0 @@ -//! Pull-on-miss: fetch a session from the backend and hydrate local JSONL storage. - -use crate::remote::client::{BackendClient, BackendError}; - -#[derive(Debug)] -pub enum PullResult { - /// Written to local storage. The [`Info`] cwd comes from the backend (may differ from caller's). - Hydrated(crate::session::info::Info), - /// Not found on the backend. - NotFound, -} - -/// Fetch a session from the backend and hydrate local JSONL storage. -pub async fn pull_session_to_local( - session_id: &str, - client: &BackendClient, -) -> Result { - let loaded = match client.load_session_data(session_id).await { - Ok(resp) => resp, - Err(BackendError::SessionNotFound { .. }) => return Ok(PullResult::NotFound), - Err(e) => return Err(e), - }; - - let remote = match loaded.session.as_ref() { - Some(s) => s, - None => return Ok(PullResult::NotFound), - }; - - // cwd required for local dir placement; null means pre-writeback session. - let cwd = match remote.cwd.as_ref() { - Some(cwd) => cwd, - None => { - tracing::warn!(session_id, "Cannot pull session: backend has cwd=null"); - return Ok(PullResult::NotFound); - } - }; - - let info = crate::session::info::Info { - id: agent_client_protocol::SessionId::new(std::sync::Arc::from(session_id)), - cwd: cwd.clone(), - }; - let dir = crate::session::persistence::session_dir(&info); - - let num_messages = hydrate::write_to_dir(&dir, &loaded)?; - - tracing::info!(session_id, %cwd, num_messages, "Pulled session from backend"); - - Ok(PullResult::Hydrated(info)) -} - -pub(crate) mod hydrate { - use std::path::Path; - use std::sync::Arc; - - use crate::remote::client::{BackendError, LoadDataResponse, LoadedMessage, SessionInfo}; - use crate::session::info::Info; - use crate::session::persistence::{CHAT_FORMAT_VERSION, Summary, default_model_id}; - - fn io_err(path: &Path, source: std::io::Error) -> BackendError { - BackendError::Hydration { - path: path.to_path_buf(), - source, - } - } - - /// Write all session files to `dir`. - pub(super) fn write_to_dir( - dir: &Path, - loaded: &LoadDataResponse, - ) -> Result { - let remote = loaded - .session - .as_ref() - .expect("caller checked session.is_some()"); - - let info = Info { - id: agent_client_protocol::SessionId::new(Arc::from(remote.session_id.as_str())), - cwd: remote.cwd.clone().expect("caller verified cwd is Some"), - }; - - std::fs::create_dir_all(dir).map_err(|e| io_err(dir, e))?; - - let num_messages = loaded.messages.as_ref().map_or(0, |m| m.len()); - let mut num_chat_messages = 0; - - if let Some(ref messages) = loaded.messages { - write_updates(dir, messages)?; - num_chat_messages = rebuild_chat_history(dir)?; - } - - write_summary(dir, &info, remote, num_messages, num_chat_messages)?; - write_remote_origin_marker(dir); - - Ok(num_messages) - } - - fn write_summary( - dir: &Path, - info: &Info, - remote: &SessionInfo, - num_messages: usize, - num_chat_messages: usize, - ) -> Result<(), BackendError> { - let meta = remote.metadata.as_ref(); - - let model_id = meta - .and_then(|m| m.get("modelId")) - .and_then(|v| v.as_str()) - .map(agent_client_protocol::ModelId::new) - .unwrap_or_else(default_model_id); - - let parent_session_id = meta - .and_then(|m| m.get("parentSessionId")) - .and_then(|v| v.as_str()) - .map(String::from); - - let summary = Summary { - info: info.clone(), - session_summary: remote.title.clone().unwrap_or_default(), - created_at: parse_rfc3339_or_now(remote.created_at.as_deref()), - updated_at: parse_rfc3339_or_now(remote.updated_at.as_deref()), - num_messages, - num_chat_messages, - current_model_id: model_id, - parent_session_id, - forked_at: None, - collection_id: None, - next_trace_turn: 0, - chat_format_version: CHAT_FORMAT_VERSION, - prompt_display_cwd: None, - session_kind: None, - fork_context_source: None, - fork_parent_prompt_id: None, - inherited_prefix_len: None, - hidden: None, - source_workspace_dir: None, - git_root_dir: None, - git_remotes: Vec::new(), - head_commit: None, - head_branch: None, - request_id: None, - // Record the *local* kigi_home (where this hydrated copy lives), - // not the original remote session's, since reconstruction runs locally. - kigi_home: crate::session::persistence::kigi_home_string(), - last_active_at: None, - generated_title: None, - title_is_manual: false, - worktree_label: None, - agent_name: None, - // Hydrated locally — record the profile this process runs under. - sandbox_profile: kigi_sandbox::configured_profile_name().map(String::from), - reasoning_effort: None, - }; - - let json = serde_json::to_string_pretty(&summary)?; - write_file(&dir.join("summary.json"), json.as_bytes()) - } - - /// Convert backend JSON-RPC messages to local updates.jsonl (replayable methods only). - pub(super) fn write_updates( - dir: &Path, - messages: &[LoadedMessage], - ) -> Result<(), BackendError> { - use std::io::Write; - - let path = dir.join("updates.jsonl"); - let file = std::fs::File::create(&path).map_err(|e| io_err(&path, e))?; - let mut w = std::io::BufWriter::new(file); - - for msg in messages { - let parsed = match serde_json::from_str::(&msg.content) { - Ok(v) => v, - Err(_) => continue, - }; - if !is_session_update(&parsed) { - continue; - } - if let Some(line) = to_envelope_line(&parsed) { - let _ = w.write_all(line.as_bytes()); - let _ = w.write_all(b"\n"); - } - } - - w.flush().map_err(|e| io_err(&path, e)) - } - - /// Rebuild `chat_history.jsonl` from `updates.jsonl` so pulled sessions are continuable. - fn rebuild_chat_history(dir: &Path) -> Result { - use crate::session::storage::UpdatesIterator; - use std::io::{Seek, Write}; - - let updates_path = dir.join("updates.jsonl"); - let Some(iter) = - UpdatesIterator::open(&updates_path).map_err(|e| io_err(&updates_path, e))? - else { - return Ok(0); - }; - - let chat_path = dir.join("chat_history.jsonl"); - let file = std::fs::File::create(&chat_path).map_err(|e| io_err(&chat_path, e))?; - let mut writer = std::io::BufWriter::new(file); - let mut reducer = ChatReducer::new(); - - for result in iter { - let update = match result { - Ok(u) => u, - Err(_) => continue, - }; - - for item in reducer.process(&update) { - if let Ok(line) = serde_json::to_string(&item) { - let _ = writer.write_all(line.as_bytes()); - let _ = writer.write_all(b"\n"); - } - } - - // CompactionCheckpoint: truncate file and reset - if reducer.should_truncate() { - reducer.clear_truncate_flag(); - let _ = writer.seek(std::io::SeekFrom::Start(0)); - let _ = writer.get_mut().set_len(0); - } - } - - // Flush trailing state - for item in reducer.flush() { - if let Ok(line) = serde_json::to_string(&item) { - let _ = writer.write_all(line.as_bytes()); - let _ = writer.write_all(b"\n"); - } - } - - writer.flush().map_err(|e| io_err(&chat_path, e))?; - Ok(reducer.count()) - } - - use crate::sampling::{AssistantItem, ContentPart, ConversationItem, ToolCall}; - use agent_client_protocol as acp; - use std::collections::{HashMap, HashSet}; - - /// Reduces ACP session updates into conversation items. - /// - /// Turn boundaries: User→Agent flushes user, Agent→User flushes agent, - /// tool completion flushes agent before emitting result. - struct ChatReducer { - user_parts: Vec, - agent_text: String, - agent_tool_calls: Vec, - - in_user_turn: bool, - has_agent_content: bool, - needs_truncate: bool, - - tool_args: HashMap, - emitted_tool_results: HashSet, - item_count: usize, - } - - impl ChatReducer { - fn new() -> Self { - Self { - user_parts: Vec::new(), - agent_text: String::new(), - agent_tool_calls: Vec::new(), - in_user_turn: false, - has_agent_content: false, - needs_truncate: false, - tool_args: HashMap::new(), - emitted_tool_results: HashSet::new(), - item_count: 0, - } - } - - fn process( - &mut self, - update: &crate::session::storage::SessionUpdate, - ) -> Vec { - use crate::session::storage::SessionUpdate; - - match update { - SessionUpdate::Acp(n) => self.handle_acp(&n.update), - SessionUpdate::Xai(n) => self.handle_xai(&n.update), - } - } - - fn handle_acp(&mut self, update: &acp::SessionUpdate) -> Vec { - match update { - acp::SessionUpdate::UserMessageChunk(chunk) => self.on_user_chunk(chunk), - acp::SessionUpdate::AgentMessageChunk(chunk) => self.on_agent_chunk(chunk), - acp::SessionUpdate::ToolCall(tc) => self.on_tool_call(tc), - acp::SessionUpdate::ToolCallUpdate(tc) => self.on_tool_call_update(tc), - _ => Vec::new(), // AgentThoughtChunk, Retry, Plan not needed - } - } - - fn handle_xai( - &mut self, - update: &crate::extensions::notification::SessionUpdate, - ) -> Vec { - use crate::extensions::notification::SessionUpdate as XaiUpdate; - - match update { - XaiUpdate::CompactionCheckpoint(_) => { - self.reset(); - self.needs_truncate = true; - Vec::new() - } - _ => Vec::new(), // DiffReview, MemoryFlush, etc. not needed - } - } - - fn on_user_chunk(&mut self, chunk: &acp::ContentChunk) -> Vec { - let mut out = Vec::new(); - - if !self.in_user_turn { - out.extend(self.flush_agent()); - self.in_user_turn = true; - } - - match &chunk.content { - acp::ContentBlock::Text(t) => { - self.user_parts.push(ContentPart::Text { - text: std::sync::Arc::::from(t.text.clone()), - }); - } - acp::ContentBlock::Image(img) => { - if let Some(uri) = &img.uri { - self.user_parts.push(ContentPart::Image { - url: std::sync::Arc::::from(uri.clone()), - }); - } - } - _ => {} // Audio, Resource, etc. not needed for chat replay - } - - out - } - - fn on_agent_chunk(&mut self, chunk: &acp::ContentChunk) -> Vec { - let mut out = Vec::new(); - - if self.in_user_turn { - out.extend(self.flush_user()); - self.in_user_turn = false; - } - - if let acp::ContentBlock::Text(t) = &chunk.content { - self.agent_text.push_str(&t.text); - self.has_agent_content = true; - } - - out - } - - fn on_tool_call(&mut self, tc: &acp::ToolCall) -> Vec { - let id = tc.tool_call_id.0.to_string(); - let args = tc - .raw_input - .as_ref() - .map(|v| v.to_string()) - .unwrap_or_default(); - - self.tool_args.insert(id.clone(), args.clone()); - self.agent_tool_calls.push(ToolCall { - id: std::sync::Arc::::from(id), - name: tc.title.clone(), - arguments: std::sync::Arc::::from(args), - }); - - Vec::new() - } - - fn on_tool_call_update(&mut self, tc: &acp::ToolCallUpdate) -> Vec { - let id = tc.tool_call_id.0.to_string(); - self.maybe_backfill_args(&id, &tc.fields); - - if Self::is_completed(&tc.fields) && self.emitted_tool_results.insert(id.clone()) { - return self.emit_tool_result(&id, &tc.fields); - } - Vec::new() - } - - /// Backfill tool arguments from ToolCallUpdate if ToolCall didn't have them. - fn maybe_backfill_args(&mut self, id: &str, fields: &acp::ToolCallUpdateFields) { - let Some(raw) = &fields.raw_input else { return }; - let needs_backfill = self.tool_args.get(id).is_none_or(String::is_empty); - if !needs_backfill { - return; - } - - let args = raw.to_string(); - self.tool_args.insert(id.to_string(), args.clone()); - - if let Some(call) = self - .agent_tool_calls - .iter_mut() - .find(|c| c.id.as_ref() == id) - { - call.arguments = std::sync::Arc::::from(args); - } - } - - fn is_completed(fields: &acp::ToolCallUpdateFields) -> bool { - matches!( - fields.status, - Some(acp::ToolCallStatus::Completed | acp::ToolCallStatus::Failed) - ) - } - - fn emit_tool_result( - &mut self, - id: &str, - fields: &acp::ToolCallUpdateFields, - ) -> Vec { - let mut out = Vec::new(); - out.extend(self.flush_agent()); - - let content = extract_tool_result_text(fields); - let item = ConversationItem::tool_result(id.to_string(), content); - self.item_count += 1; - out.push(item); - out - } - - fn flush_user(&mut self) -> Option { - if self.user_parts.is_empty() { - return None; - } - let item = ConversationItem::user_with_parts(std::mem::take(&mut self.user_parts)); - self.item_count += 1; - Some(item) - } - - fn flush_agent(&mut self) -> Option { - if !self.has_agent_content && self.agent_tool_calls.is_empty() { - return None; - } - let item = ConversationItem::Assistant(AssistantItem { - content: std::sync::Arc::::from(std::mem::take(&mut self.agent_text)), - tool_calls: std::mem::take(&mut self.agent_tool_calls), - model_id: None, - model_fingerprint: None, - reasoning_effort: None, - }); - self.has_agent_content = false; - self.item_count += 1; - Some(item) - } - - fn flush(&mut self) -> Vec { - let mut out = Vec::new(); - out.extend(self.flush_user()); - out.extend(self.flush_agent()); - out - } - - fn reset(&mut self) { - self.user_parts.clear(); - self.agent_text.clear(); - self.agent_tool_calls.clear(); - self.tool_args.clear(); - self.emitted_tool_results.clear(); - self.in_user_turn = false; - self.has_agent_content = false; - self.item_count = 0; - } - - fn should_truncate(&self) -> bool { - self.needs_truncate - } - - fn clear_truncate_flag(&mut self) { - self.needs_truncate = false; - } - - fn count(&self) -> usize { - self.item_count - } - } - - /// Extract displayable text from a completed ToolCallUpdate. - fn extract_tool_result_text(fields: &agent_client_protocol::ToolCallUpdateFields) -> String { - if let Some(content) = &fields.content { - let text: String = content - .iter() - .filter_map(|c| match c { - agent_client_protocol::ToolCallContent::Content( - agent_client_protocol::Content { - content: agent_client_protocol::ContentBlock::Text(t), - .. - }, - ) => Some(t.text.as_str()), - _ => None, - }) - .collect::>() - .join(""); - if !text.is_empty() { - return text; - } - } - if let Some(raw) = &fields.raw_output { - return raw.to_string(); - } - String::new() - } - - fn write_remote_origin_marker(dir: &Path) { - let _ = std::fs::write( - dir.join(".remote_origin"), - format!("pulled_at={}\n", chrono::Utc::now().to_rfc3339()), - ); - } - - /// Replayable JSON-RPC methods (excludes metadata like `prompt_complete`). - const REPLAYABLE_METHODS: &[&str] = &["session/update", "_x.ai/session/update"]; - - fn is_session_update(json_rpc: &serde_json::Value) -> bool { - json_rpc - .get("method") - .and_then(|v| v.as_str()) - .is_some_and(|m| REPLAYABLE_METHODS.contains(&m)) - } - - fn to_envelope_line(json_rpc: &serde_json::Value) -> Option { - let method = json_rpc.get("method").and_then(|v| v.as_str())?; - let params = json_rpc.get("params").cloned().unwrap_or_default(); - - serde_json::to_string(&serde_json::json!({ - "timestamp": 0u64, - "method": method, - "params": params, - })) - .ok() - } - - fn parse_rfc3339_or_now(s: Option<&str>) -> chrono::DateTime { - s.and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok()) - .map(|dt| dt.with_timezone(&chrono::Utc)) - .unwrap_or_else(chrono::Utc::now) - } - - fn write_file(path: &Path, data: &[u8]) -> Result<(), BackendError> { - std::fs::write(path, data).map_err(|e| io_err(path, e)) - } -} - -#[cfg(test)] -mod tests { - use crate::remote::client::LoadedMessage; - - #[test] - fn hydrate_writes_valid_updates_jsonl() { - let tmp = tempfile::TempDir::new().unwrap(); - let messages = vec![ - LoadedMessage { - id: "1".into(), - content: r#"{"method":"session/update","params":{"update":"hello"}}"#.into(), - timestamp: None, - }, - LoadedMessage { - id: "2".into(), - content: r#"{"method":"session/update","params":{"update":"world"}}"#.into(), - timestamp: None, - }, - ]; - - super::hydrate::write_updates(tmp.path(), &messages).unwrap(); - - let content = std::fs::read_to_string(tmp.path().join("updates.jsonl")).unwrap(); - let lines: Vec<&str> = content.lines().collect(); - assert_eq!(lines.len(), 2); - - for line in &lines { - let v: serde_json::Value = serde_json::from_str(line).unwrap(); - assert_eq!(v["timestamp"], 0); - assert_eq!(v["method"], "session/update"); - assert!(v["params"].is_object()); - } - } - - #[test] - fn rebuild_chat_history_merges_chunks() { - use crate::session::export::ExportedMessage; - use agent_client_protocol::{ContentBlock, ContentChunk, SessionUpdate, TextContent}; - use std::sync::Arc; - - // Build ACP notifications matching the RemoteSync path - let sid = agent_client_protocol::SessionId::new(Arc::from("test")); - let notifications = [ - agent_client_protocol::SessionNotification::new( - sid.clone(), - SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text( - TextContent::new("hello "), - ))), - ), - agent_client_protocol::SessionNotification::new( - sid.clone(), - SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text( - TextContent::new("world"), - ))), - ), - agent_client_protocol::SessionNotification::new( - sid.clone(), - SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text( - TextContent::new("hi back"), - ))), - ), - ]; - - // Serialize through ExportedMessage (writeback path) - let messages: Vec = notifications - .iter() - .map(|n| { - let exported = ExportedMessage::from_notification(n); - LoadedMessage { - id: "x".into(), - content: exported.content, - timestamp: None, - } - }) - .collect(); - - let data = crate::remote::client::LoadDataResponse { - messages: Some(messages), - session: Some(crate::remote::client::SessionInfo { - session_id: "test".into(), - title: None, - cwd: Some("/tmp".into()), - status: None, - created_at: None, - updated_at: None, - metadata: None, - }), - }; - let tmp = tempfile::TempDir::new().unwrap(); - super::hydrate::write_to_dir(tmp.path(), &data).unwrap(); - - let chat = std::fs::read_to_string(tmp.path().join("chat_history.jsonl")).unwrap(); - let items: Vec = chat - .lines() - .filter(|l| !l.is_empty()) - .filter_map(|l| serde_json::from_str(l).ok()) - .collect(); - - assert_eq!(items.len(), 2, "should have 1 user + 1 agent item"); - assert!(matches!( - &items[0], - crate::sampling::ConversationItem::User(_) - )); - assert!(matches!( - &items[1], - crate::sampling::ConversationItem::Assistant(_) - )); - if let crate::sampling::ConversationItem::User(u) = &items[0] { - let text: String = u - .content - .iter() - .filter_map(|p| match p { - crate::sampling::ContentPart::Text { text } => Some(text.as_ref()), - _ => None, - }) - .collect(); - assert_eq!(text, "hello world"); - } - } - - #[test] - fn rebuild_chat_history_preserves_user_images() { - use crate::session::export::ExportedMessage; - use agent_client_protocol::{ - ContentBlock, ContentChunk, ImageContent, SessionUpdate, TextContent, - }; - use std::sync::Arc; - - let sid = agent_client_protocol::SessionId::new(Arc::from("test")); - let notifications = [ - agent_client_protocol::SessionNotification::new( - sid.clone(), - SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text( - TextContent::new("look at this"), - ))), - ), - agent_client_protocol::SessionNotification::new( - sid.clone(), - SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Image( - ImageContent::new(String::new(), String::new()) - .uri(Some("data:image/png;base64,abc".into())), - ))), - ), - agent_client_protocol::SessionNotification::new( - sid.clone(), - SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text( - TextContent::new("I see an image"), - ))), - ), - ]; - - let messages: Vec = notifications - .iter() - .map(|n| LoadedMessage { - id: "x".into(), - content: ExportedMessage::from_notification(n).content, - timestamp: None, - }) - .collect(); - - let data = crate::remote::client::LoadDataResponse { - messages: Some(messages), - session: Some(crate::remote::client::SessionInfo { - session_id: "test".into(), - title: None, - cwd: Some("/tmp".into()), - status: None, - created_at: None, - updated_at: None, - metadata: None, - }), - }; - let tmp = tempfile::TempDir::new().unwrap(); - super::hydrate::write_to_dir(tmp.path(), &data).unwrap(); - - let chat = std::fs::read_to_string(tmp.path().join("chat_history.jsonl")).unwrap(); - let items: Vec = chat - .lines() - .filter(|l| !l.is_empty()) - .filter_map(|l| serde_json::from_str(l).ok()) - .collect(); - - assert_eq!(items.len(), 2); - if let crate::sampling::ConversationItem::User(u) = &items[0] { - assert_eq!(u.content.len(), 2, "should have text + image parts"); - assert!(matches!( - &u.content[0], - crate::sampling::ContentPart::Text { .. } - )); - assert!(matches!( - &u.content[1], - crate::sampling::ContentPart::Image { .. } - )); - } else { - panic!("expected User item"); - } - } - - #[test] - fn hydrate_skips_invalid_messages() { - let tmp = tempfile::TempDir::new().unwrap(); - let messages = vec![ - LoadedMessage { - id: "1".into(), - content: r#"{"method":"session/update","params":{}}"#.into(), - timestamp: None, - }, - LoadedMessage { - id: "bad".into(), - content: "not valid json".into(), - timestamp: None, - }, - LoadedMessage { - id: "3".into(), - content: r#"{"method":"session/update","params":{"x":1}}"#.into(), - timestamp: None, - }, - ]; - - super::hydrate::write_updates(tmp.path(), &messages).unwrap(); - - let content = std::fs::read_to_string(tmp.path().join("updates.jsonl")).unwrap(); - let lines: Vec<&str> = content.lines().collect(); - assert_eq!(lines.len(), 2, "invalid message should be skipped"); - } -} diff --git a/crates/codegen/kigi-shell/src/remote/pull_smoke_test.rs b/crates/codegen/kigi-shell/src/remote/pull_smoke_test.rs deleted file mode 100644 index 9cbbc10..0000000 --- a/crates/codegen/kigi-shell/src/remote/pull_smoke_test.rs +++ /dev/null @@ -1,127 +0,0 @@ -//! Push → pull round-trip smoke test against the live backend. -//! -//! Run with: `cargo test -p kigi-shell -- pull_smoke --ignored --nocapture` - -#[cfg(test)] -mod tests { - use crate::auth::KimiAuth; - use crate::remote::client::BackendClient; - use crate::session::storage::{JsonlStorageAdapter, StorageAdapter}; - use std::collections::BTreeMap; - use std::sync::Arc; - - fn load_prod_auth() -> Option { - let path = crate::util::kigi_home::kigi_home().join("auth.json"); - let contents = std::fs::read_to_string(&path).ok()?; - let store: BTreeMap = serde_json::from_str(&contents).ok()?; - let scope = crate::auth::KimiCodeConfig::default().auth_scope(); - crate::auth::lookup_auth(&store, &scope) - } - - /// Full round-trip using the real RemoteSync production code path: - /// create RemoteSync → queue ACP notifications → flush → verify on - /// backend → pull back → verify local hydration + storage adapter load. - #[tokio::test] - #[ignore] - async fn smoke_push_pull_round_trip() { - use crate::remote::sync::RemoteSync; - use crate::session::export::ExportedMetadata; - use agent_client_protocol::{ - ContentBlock, ContentChunk, SessionNotification, SessionUpdate, TextContent, - }; - - let auth = load_prod_auth().expect("No auth.json — run `grok login`"); - let am = Arc::new(crate::auth::AuthManager::new( - &crate::util::kigi_home::kigi_home(), - crate::auth::KimiCodeConfig::default(), - )); - am.hot_swap(auth); - let client = BackendClient::new().with_auth_manager(am.clone()); - - let session_id = format!("test-rt-{}", uuid::Uuid::new_v4()); - let test_cwd = "/tmp/smoke-test".to_string(); - let test_title = "Push-Pull Round Trip Test"; - - // PUSH via RemoteSync (real production path) - let metadata = ExportedMetadata { - title: Some(test_title.into()), - cwd: test_cwd.clone(), - model_id: Some("grok-3".into()), - created_at: Some(chrono::Utc::now().to_rfc3339()), - updated_at: Some(chrono::Utc::now().to_rfc3339()), - total_messages: None, - parent_session_id: None, - session_kind: None, - subagent_type: None, - subagent_persona: None, - subagent_role: None, - fork_context_source: None, - subagent_depth: None, - }; - - let sync = RemoteSync::new( - session_id.clone(), - metadata, - BackendClient::new().with_auth_manager(am.clone()), - ); - - let sid = agent_client_protocol::SessionId::new(Arc::from(session_id.as_str())); - sync.queue(SessionNotification::new( - sid.clone(), - SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text( - TextContent::new("Hello from smoke test — user".to_string()), - ))), - )); - sync.queue(SessionNotification::new( - sid.clone(), - SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text( - TextContent::new("Hello from smoke test — agent".to_string()), - ))), - )); - sync.flush(); - tokio::time::sleep(tokio::time::Duration::from_secs(3)).await; - - // Verify backend has cwd, title, messages - let loaded = client - .load_session_data(&session_id) - .await - .expect("load after push failed"); - let remote = loaded.session.as_ref().expect("no session row"); - assert_eq!(remote.cwd.as_deref(), Some(test_cwd.as_str())); - assert_eq!(remote.title.as_deref(), Some(test_title)); - assert!(loaded.messages.as_ref().map_or(0, |m| m.len()) >= 2); - - // PULL back to local - let result = crate::remote::pull_session_to_local(&session_id, &client) - .await - .expect("pull failed"); - let pulled = match result { - crate::remote::PullResult::Hydrated(info) => info, - crate::remote::PullResult::NotFound => panic!("pull returned NotFound"), - }; - assert_eq!(pulled.cwd, test_cwd); - - // Verify local storage loads - let local_dir = crate::session::persistence::session_dir(&pulled); - assert!(local_dir.join("summary.json").exists()); - assert!(local_dir.join("updates.jsonl").exists()); - - let storage = JsonlStorageAdapter::default(); - let data = storage - .load_session_without_updates(&pulled) - .await - .expect("storage load failed"); - assert_eq!(data.summary.session_summary, test_title); - - // Verify chat_history has both turns - let chat = - std::fs::read_to_string(local_dir.join("chat_history.jsonl")).unwrap_or_default(); - assert!(chat.contains("user"), "chat_history missing user turn"); - assert!(chat.contains("agent"), "chat_history missing agent turn"); - - // Cleanup - drop(sync); - let _ = client.delete_session_data(&session_id).await; - let _ = std::fs::remove_dir_all(&local_dir); - } -} diff --git a/crates/codegen/kigi-shell/src/remote/sync.rs b/crates/codegen/kigi-shell/src/remote/sync.rs deleted file mode 100644 index 7abb690..0000000 --- a/crates/codegen/kigi-shell/src/remote/sync.rs +++ /dev/null @@ -1,167 +0,0 @@ -//! Writeback push: async queue that flushes session updates to the backend. -//! -//! `RemoteSync` runs a background tokio task that buffers ACP notifications -//! and flushes them to the backend via [`BackendClient::save_session_data()`]. -//! -//! ## Backpressure -//! -//! When the buffer exceeds [`MAX_PENDING`], the task attempts an emergency -//! flush. If that also fails (network down), the oldest messages are dropped -//! to prevent unbounded memory growth. -//! -//! ## Drop behavior -//! -//! When `RemoteSync` is dropped, the sender half of the channel closes and -//! the background task exits. **Pending buffered messages are lost.** This -//! is acceptable because the local JSONL files are the source of truth — -//! writeback is best-effort. - -use crate::remote::BackendClient; -use crate::session::export::{ExportedMessage, ExportedMetadata}; -use agent_client_protocol as acp; -use tokio::sync::mpsc; - -/// Max buffered notifications before triggering an emergency flush. -/// Sized to keep memory under ~50MB even with large notifications. -const MAX_PENDING: usize = 512; - -/// How many oldest messages to drop when an emergency flush fails. -/// Dropping a batch (not one-by-one) avoids repeated failed flushes. -const DROP_BATCH_SIZE: usize = 64; - -enum SyncMsg { - Queue(Box), - Flush, - SetTitle(String), - SetModelId(String), -} - -#[derive(Clone)] -pub struct RemoteSync { - tx: mpsc::UnboundedSender, -} - -impl RemoteSync { - /// Metadata is included on every flush to keep the backend session row current. - pub(crate) fn new( - session_id: String, - metadata: ExportedMetadata, - client: BackendClient, - ) -> Self { - let (tx, rx) = mpsc::unbounded_channel(); - tokio::spawn(sync_task(session_id, metadata, client, rx)); - Self { tx } - } - - pub fn queue(&self, notification: acp::SessionNotification) { - let _ = self.tx.send(SyncMsg::Queue(Box::new(notification))); - } - - pub fn flush(&self) { - let _ = self.tx.send(SyncMsg::Flush); - } - - pub fn set_title(&self, title: String) { - let _ = self.tx.send(SyncMsg::SetTitle(title)); - } - - pub fn set_model_id(&self, model_id: String) { - let _ = self.tx.send(SyncMsg::SetModelId(model_id)); - } -} - -async fn do_flush( - client: &BackendClient, - session_id: &str, - metadata: &ExportedMetadata, - pending: &mut Vec, -) -> bool { - if pending.is_empty() { - return true; - } - - let messages: Vec = pending - .iter() - .map(ExportedMessage::from_notification) - .collect(); - - match client - .save_session_data(session_id, &messages, Some(metadata)) - .await - { - Ok(()) => { - tracing::debug!(count = pending.len(), "Writeback: synced"); - pending.clear(); - - // Link session to agent so the relay can route requests to it. - if let Err(e) = client - .upsert_session(session_id, metadata, &crate::util::agent_id::agent_id()) - .await - { - tracing::warn!(error = %e, "Writeback: failed to upsert session"); - } - - true - } - Err(e) => { - tracing::warn!(error = %e, pending = pending.len(), "Writeback: flush failed"); - false - } - } -} - -async fn sync_task( - session_id: String, - mut metadata: ExportedMetadata, - client: BackendClient, - mut rx: mpsc::UnboundedReceiver, -) { - let mut pending: Vec = Vec::new(); - - while let Some(msg) = rx.recv().await { - match msg { - SyncMsg::Queue(n) => { - if pending.len() >= MAX_PENDING { - tracing::warn!( - pending = pending.len(), - "Writeback: buffer full, attempting emergency flush" - ); - - metadata.updated_at = Some(chrono::Utc::now().to_rfc3339()); - if !do_flush(&client, &session_id, &metadata, &mut pending).await { - let dropped = pending.drain(0..DROP_BATCH_SIZE.min(pending.len())).count(); - tracing::error!( - dropped = dropped, - "Writeback: emergency flush failed, dropping oldest messages" - ); - } - } - pending.push(*n); - } - SyncMsg::Flush => { - metadata.updated_at = Some(chrono::Utc::now().to_rfc3339()); - do_flush(&client, &session_id, &metadata, &mut pending).await; - } - SyncMsg::SetTitle(title) => { - metadata.title = Some(title); - metadata.updated_at = Some(chrono::Utc::now().to_rfc3339()); - if let Err(e) = client - .save_session_data(&session_id, &[], Some(&metadata)) - .await - { - tracing::warn!(?e, "Writeback: failed to sync title to backend"); - } - } - SyncMsg::SetModelId(id) => { - metadata.model_id = Some(id); - metadata.updated_at = Some(chrono::Utc::now().to_rfc3339()); - if let Err(e) = client - .save_session_data(&session_id, &[], Some(&metadata)) - .await - { - tracing::warn!(?e, "Writeback: failed to sync model_id to backend"); - } - } - } - } -} diff --git a/crates/codegen/kigi-shell/src/remote/workspaces_client.rs b/crates/codegen/kigi-shell/src/remote/workspaces_client.rs deleted file mode 100644 index a3ebce3..0000000 --- a/crates/codegen/kigi-shell/src/remote/workspaces_client.rs +++ /dev/null @@ -1,176 +0,0 @@ -use std::sync::Arc; - -use serde::Deserialize; - -use crate::auth::AuthManager; - -const KIGI_WEB_URL: &str = "https://grok.com"; - -#[derive(Debug, Clone, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct Workspace { - #[serde(default)] - pub workspace_id: String, - #[serde(default)] - pub name: String, - #[serde(default)] - pub create_time: Option, - #[serde(default)] - pub kind: Option, -} - -#[derive(Debug, Clone, Default)] -pub struct WsQuery { - pub page_size: i64, - pub page_token: Option, - pub query: Option, - pub kind: Option, -} - -#[derive(Debug, Clone, Default)] -pub struct ListWorkspacesPage { - pub workspaces: Vec, - pub next_page_token: Option, -} - -#[derive(Debug, thiserror::Error)] -pub enum WsError { - #[error("no OAuth credentials for workspaces:read")] - NoOauth, - #[error("network error: {0}")] - Network(#[from] reqwest::Error), - #[error("request failed: {status}")] - Http { status: u16 }, - #[error("parse error: {0}")] - Parse(#[from] serde_json::Error), -} - -#[derive(Debug, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -struct ListWorkspacesResponseWire { - #[serde(default)] - workspaces: Vec, - #[serde(default)] - next_page_token: Option, -} - -pub struct WorkspacesClient { - http: reqwest::Client, - base_url: String, - auth: Arc, -} - -impl WorkspacesClient { - pub fn new(auth: Arc) -> Self { - let base_url = first_nonempty_env(&[ - "KIGI_WORKSPACES_BASE_URL", - "KIGI_CONVERSATIONS_BASE_URL", - "KIGI_CODE_WEB_URL", - ]) - .unwrap_or_else(|| KIGI_WEB_URL.to_string()); - Self { - http: crate::http::shared_client(), - base_url, - auth, - } - } - - pub async fn list_workspaces(&self, q: &WsQuery) -> Result { - let auth = self.auth.auth().await.map_err(|_| WsError::NoOauth)?; - if !auth.is_session_auth() { - return Err(WsError::NoOauth); - } - - let url = format!("{}/rest/workspaces", self.base_url); - let mut query: Vec<(&str, String)> = vec![("pageSize", q.page_size.to_string())]; - if let Some(token) = q.page_token.as_deref().filter(|s| !s.is_empty()) { - query.push(("pageToken", token.to_owned())); - } - if let Some(search) = q.query.as_deref().filter(|s| !s.is_empty()) { - query.push(("query", search.to_owned())); - } - if let Some(kind) = q.kind.as_deref().filter(|s| !s.is_empty()) { - query.push(("kind", kind.to_owned())); - } - - let mut builder = self - .http - .get(&url) - .query(&query) - .header("Authorization", format!("Bearer {}", auth.key)) - .header("x-userid", &auth.user_id) - .header("x-grok-client-version", kigi_version::VERSION) - .header( - "x-grok-client-identifier", - crate::http::process_client_identifier(), - ) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ) - .header(reqwest::header::ACCEPT, "application/json"); - if let Some(email) = &auth.email { - builder = builder.header("x-email", email); - } - let builder = kigi_file_utils::trace_context::inject_trace_context_into_request(builder); - - let response = builder.send().await?; - let status = response.status(); - if !status.is_success() { - return Err(WsError::Http { - status: status.as_u16(), - }); - } - - let bytes = response.bytes().await?; - let wire: ListWorkspacesResponseWire = serde_json::from_slice(&bytes)?; - - Ok(ListWorkspacesPage { - workspaces: wire.workspaces, - next_page_token: wire.next_page_token.filter(|t| !t.is_empty()), - }) - } -} - -fn first_nonempty_env(keys: &[&str]) -> Option { - keys.iter() - .find_map(|k| std::env::var(k).ok().filter(|s| !s.is_empty())) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn workspace_parses_camelcase_wire() { - let json = serde_json::json!({ - "workspaces": [{ - "workspaceId": "ws_9f3a", - "name": "GPU vendor research", - "createTime": "2026-06-18T17:30:00Z", - "kind": "WORKSPACE_KIND_IMAGINE" - }], - "nextPageToken": "tok2" - }); - let wire: ListWorkspacesResponseWire = serde_json::from_value(json).unwrap(); - assert_eq!(wire.workspaces.len(), 1); - let w = &wire.workspaces[0]; - assert_eq!(w.workspace_id, "ws_9f3a"); - assert_eq!(w.name, "GPU vendor research"); - assert_eq!(w.create_time.as_deref(), Some("2026-06-18T17:30:00Z")); - assert_eq!(w.kind.as_deref(), Some("WORKSPACE_KIND_IMAGINE")); - assert_eq!(wire.next_page_token.as_deref(), Some("tok2")); - } - - #[test] - fn missing_fields_default_gracefully() { - let json = serde_json::json!({ "workspaces": [{ "workspaceId": "w1" }] }); - let wire: ListWorkspacesResponseWire = serde_json::from_value(json).unwrap(); - let w = &wire.workspaces[0]; - assert_eq!(w.workspace_id, "w1"); - assert!(w.name.is_empty()); - assert!(w.create_time.is_none()); - assert!(w.kind.is_none()); - assert!(wire.next_page_token.is_none()); - } -} diff --git a/crates/codegen/kigi-shell/src/sampling/error.rs b/crates/codegen/kigi-shell/src/sampling/error.rs index c01d9f9..ef5be69 100644 --- a/crates/codegen/kigi-shell/src/sampling/error.rs +++ b/crates/codegen/kigi-shell/src/sampling/error.rs @@ -17,26 +17,36 @@ use agent_client_protocol as acp; /// see this code and show a user-friendly upgrade message instead. pub const RATE_LIMITED_ERROR_CODE: i32 = -32003; -/// OAuth / session rate-limit copy (personal plan upgrade path). -pub const RATE_LIMITED_USER_MESSAGE_OAUTH: &str = - "You\u{2019}ve hit the rate limit for your plan. Upgrade your account or try again later."; +/// Subscription (OAuth) rate-limit copy. PRD Q3: the official Kimi CLI and +/// Kigi draw on the SAME subscription quota, so the message says so — a user +/// who also runs `kimi` should understand why the limit arrived early. +/// Deliberately promises no reset duration; the quota window is server-side. +pub static RATE_LIMITED_USER_MESSAGE_OAUTH: std::sync::LazyLock = + std::sync::LazyLock::new(|| { + format!( + "You\u{2019}ve hit the usage limit of your Kimi subscription. Note that Kigi and the \ + official Kimi CLI share the same subscription quota. Upgrade your plan at {} or try \ + again later.", + kigi_env::upgrade_page_url() + ) + }); -/// API key / team rate-limit copy. Personal grok.com upgrades do not raise API -/// team limits; admins purchase credits or a higher spend-based tier. -/// See https://docs.x.ai/developers/rate-limits#rate-limit-tiers -pub const RATE_LIMITED_USER_MESSAGE_API_KEY: &str = "You\u{2019}ve hit your team\u{2019}s API rate limit. Ask a team admin to purchase more credits for higher limits, or try again later. See https://docs.x.ai/developers/rate-limits#rate-limit-tiers"; +/// Moonshot API-key rate-limit copy. Platform keys are tier-limited (RPM/TPM); +/// raising the tier happens in the Moonshot Open Platform console, not via a +/// Kimi subscription. +pub const RATE_LIMITED_USER_MESSAGE_API_KEY: &str = "You\u{2019}ve hit the rate limit for your Moonshot API key. Check your tier\u{2019}s limits in the Moonshot Open Platform console (platform.moonshot.ai or platform.moonshot.cn), or try again later."; /// Pick rate-limit copy from the *active* auth method. /// /// Pass the real `is_api_key_auth` flag (pager `AppView`, `AuthMethodKind::is_api_key` -/// for the selected method). Do **not** decide from `has_xai_api_key_env()` alone: +/// for the selected method). Do **not** decide from the key env var alone: /// when both an env key and a cached OAuth session exist, auth prefers the /// cached session over the API key. pub fn rate_limited_user_message(is_api_key_auth: bool) -> &'static str { if is_api_key_auth { RATE_LIMITED_USER_MESSAGE_API_KEY } else { - RATE_LIMITED_USER_MESSAGE_OAUTH + RATE_LIMITED_USER_MESSAGE_OAUTH.as_str() } } @@ -318,20 +328,20 @@ mod tests { fn rate_limited_user_message_oauth_vs_api_key() { assert_eq!( rate_limited_user_message(false), - RATE_LIMITED_USER_MESSAGE_OAUTH + RATE_LIMITED_USER_MESSAGE_OAUTH.as_str() ); assert_eq!( rate_limited_user_message(true), RATE_LIMITED_USER_MESSAGE_API_KEY ); - assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains("Upgrade your account")); - assert!(RATE_LIMITED_USER_MESSAGE_API_KEY.contains("team")); - assert!(RATE_LIMITED_USER_MESSAGE_API_KEY.contains("credits")); - assert!( - RATE_LIMITED_USER_MESSAGE_API_KEY - .contains("https://docs.x.ai/developers/rate-limits#rate-limit-tiers") - ); - assert!(!RATE_LIMITED_USER_MESSAGE_API_KEY.contains("Upgrade your account")); + // PRD Q3: the subscription copy must state the shared quota with the + // official Kimi CLI and point at the upgrade page. + assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains("official Kimi CLI")); + assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains("same subscription quota")); + assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains(kigi_env::upgrade_page_url())); + // API-key copy points at the Moonshot platform, not the subscription. + assert!(RATE_LIMITED_USER_MESSAGE_API_KEY.contains("Moonshot")); + assert!(!RATE_LIMITED_USER_MESSAGE_API_KEY.contains("subscription quota")); } #[test] @@ -341,7 +351,6 @@ mod tests { message: "Rate limit exceeded".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; let acp_err = map_sampling_err_to_acp(err); assert_eq!(acp_err.code, acp::ErrorCode::from(RATE_LIMITED_ERROR_CODE)); @@ -359,7 +368,6 @@ mod tests { message: "Rate limit exceeded".into(), model_metadata: None, retry_after_secs: Some(60), - should_retry: None, }; assert_eq!(err.retry_after(), Some(60)); let acp_err = map_sampling_err_to_acp(err); @@ -374,14 +382,12 @@ mod tests { message: "limited".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; let server_err = SamplingError::Api { status: StatusCode::INTERNAL_SERVER_ERROR, message: "oops".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; let rate_acp = map_sampling_err_to_acp(rate_err); let server_acp = map_sampling_err_to_acp(server_err); @@ -398,7 +404,6 @@ mod tests { message: "bad token".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; let acp_err = map_sampling_err_to_acp(err); assert_eq!(acp_err.code, acp::Error::auth_required().code); @@ -420,7 +425,6 @@ mod tests { .into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; let acp_err = map_sampling_err_to_acp(err); assert_ne!( @@ -476,7 +480,6 @@ mod tests { message: "The model 'grok-build' requires a Grok subscription.".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; let acp_err = map_sampling_err_to_acp(err); let data = acp_err.data.unwrap(); @@ -501,7 +504,6 @@ mod tests { message: "The model 'grok-build' requires a Grok subscription.".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; let acp_err = map_sampling_err_to_acp(err); let data = acp_err.data.unwrap(); @@ -522,7 +524,6 @@ mod tests { message: "Content violates usage guidelines.".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }; let acp_err = map_sampling_err_to_acp(err); let data = acp_err.data.unwrap(); diff --git a/crates/codegen/kigi-shell/src/session/acp_session.rs b/crates/codegen/kigi-shell/src/session/acp_session.rs index 1de09a7..e883834 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session.rs @@ -679,8 +679,6 @@ pub(crate) struct SessionActor { pub(crate) origin_client: Option, /// Feedback manager for signal tracking and feedback request heuristics pub(crate) feedback_manager: Arc, - /// Cancellation token for the feedback sync loop (None if no feedback client) - pub(crate) sync_loop_cancel: Option, /// The fully-built Agent: owns the ToolBridge, system prompt, policies, /// and the AgentDefinition. Replaces the old `tool_bridge` + `agent_definition` fields. /// Wrapped in `RefCell` for mid-session mutation (skill refresh, prompt regen). diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs index ffe943d..199aa52 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs @@ -1,5 +1,5 @@ use super::*; -use crate::remote::DEFAULT_CONTEXT_WINDOW; +use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW; use kigi_chat_state::conversation_util::replace_or_insert_system_head; impl SessionActor { pub(super) async fn handle_set_session_model( @@ -72,7 +72,6 @@ impl SessionActor { existing.auth_type, ), alpha_test_key: existing.alpha_test_key, - client_version: sampling_config.client_version.clone(), }); self.model_auth_facts.replace(None); self.signals_handle() diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs index e6b16ab..3f2d587 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs @@ -691,7 +691,6 @@ impl SessionActor { crate::agent::config::finalize_image_describe_sampler_config( resolved_describe, &active_session_config, - self.client_identifier.clone(), Some(self.max_retries), ); let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| { diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/recap.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/recap.rs index 4c62c55..d5ad677 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/recap.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/recap.rs @@ -3,7 +3,7 @@ use super::*; -use crate::remote::DEFAULT_CONTEXT_WINDOW; +use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW; impl SessionActor { /// Handle a /btw side question — single-turn model call using the diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/reminders.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/reminders.rs index 0d78ade..a96b0cb 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/reminders.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/reminders.rs @@ -134,7 +134,7 @@ pub(super) fn build_todo_gate_reminder(pending: &[&str], unbacked_in_progress: & /// (which is disabled). Extracted from `spawn_session_actor` so the /// precedence rules are unit-testable. Named `resolve_*` to match the /// sibling precedence helpers in `crate::util::config` -/// (`resolve_zdr_access_enabled`, `resolve_restore_code`, …). +/// (`resolve_restore_code`, …). pub(crate) fn resolve_reminder_policy( remote: Option<&crate::util::config::RemoteSettings>, todo_gate: bool, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/run_loop.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/run_loop.rs index cf3851a..e14c114 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/run_loop.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/run_loop.rs @@ -205,8 +205,7 @@ pub(super) async fn run_session( .emit_buffered(notification). await; } if let Some(tx) = respond_to { let _ = tx.send(()); } } } } } maybe_completion = completion_rx.recv() => { let - Some((prompt_id, result)) = maybe_completion else { if let Some(cancel) = & - session.sync_loop_cancel { cancel.cancel(); } cleanup_session_scratch(& + Some((prompt_id, result)) = maybe_completion else { cleanup_session_scratch(& session); return; }; if let Some(notification) = replay_buffer.flush() { session.emit_buffered(notification). await; } let (turn_succeeded, infra_pause_message) = SessionActor::post_turn_goal_degradation_plan(& @@ -265,8 +264,7 @@ pub(super) async fn run_session( { let model_id = session.current_model_id(). await; if let Some(signals) = session .signals_handle().snapshot(). await { - } } if let - Some(cancel) = & session.sync_loop_cancel { cancel.cancel(); } session + } } session .feedback_manager.shutdown(). await; if ! session .startup_hints.is_subagent { session.persist_background_task_manifest(). await; } cleanup_session_scratch(& session); return; }; match cmd { @@ -327,8 +325,7 @@ pub(super) async fn run_session( ::agent::config::try_resolve_model_credentials(model_name.as_str(), existing .api_key.as_deref()) { session.chat_state_handle .update_credentials(kigi_chat_state::Credentials { api_key : r.api_key, - auth_type : r.auth_type, alpha_test_key : existing.alpha_test_key, - client_version : existing.client_version, }); } session.model_auth_facts + auth_type : r.auth_type, alpha_test_key : existing.alpha_test_key, }); } session.model_auth_facts .replace(None); } } SessionCommand::GetCurrentModel { responds_to } => { let model = session.chat_state_handle.get_sampling_config(). await .map(| c | c .model).unwrap_or_default(); let _ = responds_to.send(model); } @@ -697,7 +694,7 @@ pub(super) async fn run_session( await; session.send_hook_execution("session_start", None, None, & results). await; } } SessionCommand::GetFeedbackContext { turn_number, responds_to } => { let s = session.clone(); tokio::task::spawn_local(async move { use - prod_mc_cli_chat_proxy_types::feedback_types::FeedbackToolOutcome; let + crate::session::feedback_types::FeedbackToolOutcome; let turn_idx = turn_number.and_then(| n | usize::try_from(n).ok()); let (last_user_message, last_assistant_message) = match turn_idx { Some(n) => { let conv = s.chat_state_handle.get_conversation(). await; @@ -809,8 +806,7 @@ pub(super) async fn run_session( "MEMORY_SUBAGENT_SKIP: skipping on_session_end for subagent session"); } session.maybe_run_dream(). await; let telem = session.memory .telemetry_snapshot(); session.emit_memory_session_summary(& telem, - total_chunks_at_end, session_end_result); if let Some(cancel) = & session - .sync_loop_cancel { cancel.cancel(); } session.feedback_manager + total_chunks_at_end, session_end_result); session.feedback_manager .shutdown(). await; if ! session.startup_hints .is_subagent { session.persist_background_task_manifest(). await; } cleanup_session_scratch(& session); return; } } } diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs index f26b26d..75bb2f9 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs @@ -49,7 +49,7 @@ impl SessionTokenAuthGate { is_session_based: auth_method_id .is_some_and(crate::agent::auth_method::is_session_based_method), model_byok, - endpoint_is_first_party: crate::util::is_first_party_xai_url(base_url), + endpoint_is_first_party: crate::util::is_first_party_url(base_url), } } fn active(self) -> bool { @@ -314,22 +314,11 @@ impl SessionActor { auth_scheme, extra_headers, context_window: cfg.context_window.get(), - client_version: creds.client_version, reasoning_effort: cfg.reasoning_effort, force_http1: false, max_retries: Some(self.max_retries), stream_tool_calls: cfg.stream_tool_calls.unwrap_or(false), idle_timeout_secs: None, - client_identifier: self.client_identifier.clone(), - deployment_id: crate::managed_config::resolve_deployment_id( - crate::managed_config::resolve_deployment_key().as_deref(), - ), - user_id: self - .auth_manager - .as_ref() - .and_then(|am| am.current_or_expired()) - .filter(|a| a.is_session_auth()) - .map(|a| a.user_id), origin_client: self.origin_client.clone(), attribution_callback: self.attribution_callback.clone(), bearer_resolver: if use_bearer_resolver { @@ -482,7 +471,6 @@ impl SessionActor { &endpoints, session_key.as_deref(), creds.alpha_test_key.clone(), - creds.client_version.clone(), ) } /// Resolve a dedicated sampler for the Auto-mode classifier model `slug`, @@ -499,7 +487,6 @@ impl SessionActor { crate::agent::config::stamp_session_local_sampler_fields( &mut cfg, &active_session_config, - self.client_identifier.clone(), Some(self.max_retries), ); let model = cfg.model.clone(); diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/session_setup.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/session_setup.rs index e764984..947ae7f 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/session_setup.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/session_setup.rs @@ -323,10 +323,10 @@ impl SessionActor { /// Check if the session has been idle and proactively refresh model metadata. /// /// Called at the start of each turn. If idle > `IDLE_REFRESH_THRESHOLD_SECS`, - /// fetches `/models-v2` from cli-chat-proxy and updates the cached + /// fetches `/models` from cli-chat-proxy and updates the cached /// context_window / max_completion_tokens if remote settings changed them. /// - /// Skipped for BYOK users (no remote settings, no `/models-v2`). + /// Skipped for BYOK users (no remote settings, no `/models`). pub(super) async fn maybe_refresh_model_metadata_on_resume(&self) { if !self.is_session_based_auth() { return; @@ -353,7 +353,7 @@ impl SessionActor { tracing::info!( idle_secs, threshold_secs = Self::IDLE_REFRESH_THRESHOLD_SECS, - "Session resumed after idle — refreshing model metadata from cli-chat-proxy" + "Session resumed after idle — refreshing model metadata" ); let creds = self.chat_state_handle.get_credentials().await; let Some(ref am) = self.auth_manager else { @@ -370,27 +370,21 @@ impl SessionActor { ); let middleware_client = crate::http::with_auth_retry(crate::http::shared_client(), provider); - let url = format!("{}/models-v2", base_url); + let url = format!("{}/models", base_url); let parse_models_response = |json: serde_json::Value| -> Option<(std::num::NonZeroU64, Option)> { let data = json.get("data")?.as_array()?; for entry in data { - let parsed = crate::remote::client::parse_remote_model_value(entry, base_url)?; + let parsed = + crate::agent::models_fetch::parse_remote_model_value(entry, base_url)?; if parsed.model == *current_model { return Some((parsed.context_window, parsed.max_completion_tokens)); } } None }; - #[allow(unused_mut)] - let mut request = middleware_client + let request = middleware_client .get(&url) - .header("X-XAI-Token-Auth", "xai-grok-cli") - .header("x-grok-client-version", kigi_version::VERSION) - .header( - crate::http::CLIENT_MODE_HEADER, - crate::http::process_client_mode(), - ) .timeout(std::time::Duration::from_secs(5)); let response = match request.send().await { Ok(r) => r, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/slash_exec.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/slash_exec.rs index f7568e9..0366642 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/slash_exec.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/slash_exec.rs @@ -828,7 +828,7 @@ impl SessionActor { ); let model_id = sampling_config.map(|c| c.model); let resolved_model_id = model_metadata.resolved_model_id; - let client_version = credentials.client_version; + let client_version = Some(kigi_version::VERSION.to_string()); use crate::session::feedback_manager::{SessionFeedbackData, SubmitOutcome}; let outcome = self diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs index 1b6b7e3..28c5f45 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs @@ -3,7 +3,7 @@ //! the MCP auto-restart wiring (`SessionRestartActions`). #![allow(clippy::items_after_test_module)] use super::*; -use crate::remote::DEFAULT_CONTEXT_WINDOW; +use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW; /// Partition CLI `--allow` rules under the pin: blanket catch-all allows /// (`Allow(Any)` `*` / `**`, plus bare/match-all Bash/MCP/WebFetch grants — see /// `resolution::is_catchall_allow`) substitute for the blocked `--yolo`, so drop them when @@ -123,10 +123,7 @@ pub(crate) async fn spawn_session_actor( codebase_indexes: std::sync::Arc>, code_nav_enabled: bool, fs_watch_caps: fs_watch::FsWatchCapabilities, - feedback_proxy_url: Option, - feedback_user_token: Option, - feedback_alpha_test_key: Option, - deployment_key: Option, + feedback_base_url: Option, client_terminal_capable: bool, client_fs_capable: bool, gateway_enabled: std::sync::Arc, @@ -141,7 +138,6 @@ pub(crate) async fn spawn_session_actor( persisted_goal_mode: Option, persisted_announcement_state: Option, memory_config: Option, - loc_tracking_enabled: bool, feedback_flags: crate::session::feedback_manager::FeedbackFlags, managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle, managed_mcp_expires_at: Option>, @@ -879,36 +875,30 @@ pub(crate) async fn spawn_session_actor( } persist_chat_history_jsonl_sync(&session_info, &conversation); chat_state_handle.replace_conversation(conversation); - let feedback_client = feedback_proxy_url.map(|base_url| { - let mut client = - crate::agent::feedback_client::FeedbackClient::new(base_url, feedback_user_token) - .with_alpha_test_key(feedback_alpha_test_key) - .with_deployment_key(deployment_key); - if let Some(am) = auth_manager.as_ref() { - client = client.with_auth_manager(am.clone()); - } - client - }); + let feedback_client = match (feedback_base_url, auth_manager.as_ref()) { + (Some(base_url), Some(am)) => Some( + crate::agent::feedback_client::FeedbackClient::new(base_url, am.clone()) + .with_session_id(session_info.id.0.to_string()), + ), + _ => None, + }; let has_feedback_client = feedback_client.is_some(); tracing::info!( session_id = % session_info.id.0, has_feedback_client = has_feedback_client, "Creating feedback manager" ); let feedback_client_type = match client_type { - ClientType::GrokTUI => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Tui, - ClientType::GrokWeb => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Web, - ClientType::Nebula => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Nebula, - ClientType::Extension => { - prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Extension - } - ClientType::Generic => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Agent, - ClientType::Desktop => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Desktop, - ClientType::GrokPager => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Tui, + ClientType::GrokTUI => crate::session::feedback_types::ClientType::Tui, + ClientType::GrokWeb => crate::session::feedback_types::ClientType::Web, + ClientType::Nebula => crate::session::feedback_types::ClientType::Nebula, + ClientType::Extension => crate::session::feedback_types::ClientType::Extension, + ClientType::Generic => crate::session::feedback_types::ClientType::Agent, + ClientType::Desktop => crate::session::feedback_types::ClientType::Desktop, + ClientType::GrokPager => crate::session::feedback_types::ClientType::Tui, }; let feedback_config = FeedbackManagerConfig { feedback_enabled: feedback_flags.enabled, client_type: feedback_client_type, - loc_tracking_enabled, ..Default::default() }; let feedback_manager = Arc::new(FeedbackManager::new( @@ -930,11 +920,6 @@ pub(crate) async fn spawn_session_actor( } signals_handle.set_primary_model(&primary_model_id); signals_handle.set_tracing_config(inference_idle_timeout_secs); - let sync_loop_cancel = if has_feedback_client { - Some(tokio_util::sync::CancellationToken::new()) - } else { - None - }; let force_compact = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); let resolved_workspace_root = kigi_workspace::session::git::find_git_root_from_path( std::path::Path::new(&session_info.cwd), @@ -1141,7 +1126,6 @@ pub(crate) async fn spawn_session_actor( client_identifier: session_client_identifier.clone(), origin_client: origin_client.clone(), feedback_manager: feedback_manager.clone(), - sync_loop_cancel: sync_loop_cancel.clone(), agent: std::cell::RefCell::new(agent), last_reported_branch: Arc::new(Mutex::new(None)), git_head_enabled: fs_watch_caps.git_head, @@ -1375,18 +1359,6 @@ pub(crate) async fn spawn_session_actor( } }); } - if let Some(cancel) = sync_loop_cancel { - tracing::info!(session_id = % session_info.id.0, "Spawning feedback sync loop"); - let fm = feedback_manager.clone(); - tokio::spawn(async move { - fm.run_sync_loop(cancel).await; - }); - } else { - tracing::debug!( - session_id = % session_info.id.0, - "No feedback client available, skipping sync loop" - ); - } { use agent_client_protocol::Client as _; use kigi_tools::implementations::grok_build::ask_user_question::{ @@ -1600,10 +1572,7 @@ pub(crate) async fn spawn_session_on_thread( codebase_indexes: std::sync::Arc>, code_nav_enabled: bool, fs_watch_caps: fs_watch::FsWatchCapabilities, - feedback_proxy_url: Option, - feedback_user_token: Option, - feedback_alpha_test_key: Option, - deployment_key: Option, + feedback_base_url: Option, client_terminal_capable: bool, client_fs_capable: bool, gateway_enabled: std::sync::Arc, @@ -1618,7 +1587,6 @@ pub(crate) async fn spawn_session_on_thread( persisted_goal_mode: Option, persisted_announcement_state: Option, memory_config: Option, - loc_tracking_enabled: bool, feedback_flags: crate::session::feedback_manager::FeedbackFlags, managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle, managed_mcp_expires_at: Option>, @@ -1751,10 +1719,7 @@ pub(crate) async fn spawn_session_on_thread( codebase_indexes, code_nav_enabled, fs_watch_caps, - feedback_proxy_url, - feedback_user_token, - feedback_alpha_test_key, - deployment_key, + feedback_base_url, client_terminal_capable, client_fs_capable, gateway_enabled, @@ -1769,7 +1734,6 @@ pub(crate) async fn spawn_session_on_thread( persisted_goal_mode, persisted_announcement_state, memory_config, - loc_tracking_enabled, feedback_flags, managed_mcp_handle, managed_mcp_expires_at, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs index 1fa3055..00b2d8b 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs @@ -850,15 +850,11 @@ async fn set_session_model_invalidates_byok_memo_for_same_model_id() { auth_scheme: Default::default(), extra_headers: Default::default(), context_window: 256_000, - client_version: None, force_http1: false, max_retries: None, stream_tool_calls: false, idle_timeout_secs: None, - client_identifier: None, reasoning_effort: None, - deployment_id: None, - user_id: None, origin_client: None, attribution_callback: None, bearer_resolver: None, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs index 15f9cc1..70c0af1 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs @@ -47,15 +47,11 @@ async fn persist_ack_waits_for_disk_flush_before_success() { auth_scheme: Default::default(), extra_headers: Default::default(), context_window: 100_000, - client_version: None, force_http1: false, max_retries: None, stream_tool_calls: false, idle_timeout_secs: None, - client_identifier: None, reasoning_effort: None, - deployment_id: None, - user_id: None, origin_client: None, attribution_callback: None, bearer_resolver: None, @@ -193,7 +189,6 @@ async fn persist_ack_waits_for_disk_flush_before_success() { client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-session")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(test_agent_default().await), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, @@ -338,15 +333,11 @@ async fn first_turn_memory_injection_persists_to_chat_history() { api_backend: Default::default(), auth_scheme: Default::default(), context_window: 100_000, - client_version: None, force_http1: false, max_retries: None, stream_tool_calls: false, idle_timeout_secs: None, - client_identifier: None, reasoning_effort: None, - deployment_id: None, - user_id: None, origin_client: None, attribution_callback: None, bearer_resolver: None, @@ -470,15 +461,11 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history() api_backend: Default::default(), auth_scheme: Default::default(), context_window: 100_000, - client_version: None, force_http1: false, max_retries: None, stream_tool_calls: false, idle_timeout_secs: None, - client_identifier: None, reasoning_effort: None, - deployment_id: None, - user_id: None, origin_client: None, attribution_callback: None, bearer_resolver: None, @@ -641,7 +628,6 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history() client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-session")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(test_agent_default().await), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, @@ -890,7 +876,6 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() { client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-session")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(agent), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, @@ -1729,15 +1714,11 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() { auth_scheme: Default::default(), extra_headers: Default::default(), context_window: 100_000, - client_version: None, force_http1: false, max_retries: Some(0), stream_tool_calls: false, idle_timeout_secs: Some(60), - client_identifier: None, reasoning_effort: None, - deployment_id: None, - user_id: None, origin_client: None, attribution_callback: None, bearer_resolver: None, @@ -1876,7 +1857,6 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() { client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-session")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(agent), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs index f894422..ca2144e 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs @@ -43,7 +43,7 @@ async fn test_last_api_request_at_idle_detection() { /// End-to-end test for `maybe_refresh_model_metadata_on_resume`. /// /// Simulates a session idle for >10 minutes, then verifies the function -/// fetches `/models-v2`, parses the response, and updates `context_window` +/// fetches `/models`, parses the response, and updates `context_window` /// and `max_completion_tokens` in the sampling config. #[tokio::test(flavor = "current_thread")] async fn test_e2e_idle_resume_refreshes_model_metadata() { @@ -52,7 +52,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { local .run_until(async { let app = axum::Router::new().route( - "/v1/models-v2", + "/v1/models", get(|| async { axum::Json(serde_json::json!( { "data" : [{ "model" : "test-model", "name" : "Test Model", @@ -117,7 +117,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { api_key: Some("test-key".to_string()), auth_type: Default::default(), alpha_test_key: None, - client_version: None, }); tokio::time::sleep(std::time::Duration::from_millis(50)).await; let actor = SessionActor { @@ -219,7 +218,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-session")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(test_agent_default().await), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, @@ -321,12 +319,12 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { assert_eq!( cfg_after.context_window, std::num::NonZeroU64::new(300_000).unwrap(), - "context_window should be updated to 300K from /models-v2" + "context_window should be updated to 300K from /models" ); assert_eq!( cfg_after.max_completion_tokens, Some(16384), - "max_completion_tokens should be updated to 16384 from /models-v2" + "max_completion_tokens should be updated to 16384 from /models" ); }) .await; diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs index a6952ab..5a42214 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs @@ -152,7 +152,6 @@ async fn create_test_actor( client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-session")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(test_agent_default().await), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, @@ -591,7 +590,6 @@ async fn create_test_actor_with_memory( client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-memory")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(test_agent_default().await), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, @@ -1168,7 +1166,7 @@ async fn test_compact_on_error_no_trigger_when_tokens_within_new_window() { /// End-to-end test for `maybe_refresh_model_metadata_on_resume`. /// /// Simulates a session idle for >10 minutes, then verifies the function -/// fetches `/models-v2`, parses the response, and updates `context_window` +/// fetches `/models`, parses the response, and updates `context_window` /// and `max_completion_tokens` in the sampling config. #[tokio::test(flavor = "current_thread")] async fn test_e2e_idle_resume_refreshes_model_metadata() { @@ -1177,7 +1175,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { local .run_until(async { let app = axum::Router::new().route( - "/v1/models-v2", + "/v1/models", get(|| async { axum::Json(serde_json::json!( { "data" : [{ "model" : "test-model", "name" : "Test Model", @@ -1241,7 +1239,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { api_key: Some("test-key".to_string()), auth_type: Default::default(), alpha_test_key: None, - client_version: None, }); tokio::time::sleep(std::time::Duration::from_millis(50)).await; let actor = SessionActor { @@ -1346,7 +1343,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-session")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(test_agent_default().await), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, @@ -1448,12 +1444,12 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() { assert_eq!( cfg_after.context_window, std::num::NonZeroU64::new(300_000).unwrap(), - "context_window should be updated to 300K from /models-v2" + "context_window should be updated to 300K from /models" ); assert_eq!( cfg_after.max_completion_tokens, Some(16384), - "max_completion_tokens should be updated to 16384 from /models-v2" + "max_completion_tokens should be updated to 16384 from /models" ); }) .await; diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs index d7fda78..7689481 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs @@ -211,7 +211,6 @@ async fn create_test_actor_with_memory( client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-memory")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(test_agent_default().await), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs index 5848b5b..6a1e885 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs @@ -166,7 +166,7 @@ async fn actor_with_proxy( let cfg = crate::agent::config::Config { endpoints: crate::agent::config::EndpointsConfig { - cli_chat_proxy_base_url: Some(proxy_base.to_string()), + coding_api_base_url: Some(proxy_base.to_string()), ..Default::default() }, ..Default::default() diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs index 3fce531..ed64fde 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs @@ -157,7 +157,6 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-session")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(test_agent_default().await), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs index 95ff153..6b63542 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs @@ -271,7 +271,6 @@ pub(crate) async fn create_test_actor_ex( client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-session")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(test_agent_default().await), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, diff --git a/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs b/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs index 0562fbf..0dffe18 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs @@ -64,9 +64,6 @@ async fn web_search_uses_model_override_from_config_end_to_end() { entry, crate::agent::config::resolve_credentials(entry, None), None, - None, - None, - None, ); let web_search_sampling = crate::tools::config::web_search_sampling_config(resolved); diff --git a/crates/codegen/kigi-shell/src/session/acp_types.rs b/crates/codegen/kigi-shell/src/session/acp_types.rs index 5e5254b..62f3f0b 100644 --- a/crates/codegen/kigi-shell/src/session/acp_types.rs +++ b/crates/codegen/kigi-shell/src/session/acp_types.rs @@ -88,11 +88,11 @@ pub struct ClientFeedbackInput { pub session_id: String, /// Type of client submitting feedback - pub client_type: prod_mc_cli_chat_proxy_types::feedback_types::ClientType, + pub client_type: crate::session::feedback_types::ClientType, /// Rating type (thumbs, stars, nps) #[serde(default)] - pub rating_type: Option, + pub rating_type: Option, /// Rating value (interpretation depends on rating_type): /// - thumbs: -1 (down), 0 (neutral), 1 (up) @@ -113,7 +113,7 @@ pub struct ClientFeedbackInput { /// Context type for the feedback #[serde(default)] - pub context_type: Option, + pub context_type: Option, /// 0-based turn number this feedback is about. #[serde(default, alias = "turnNumber")] @@ -134,7 +134,7 @@ pub struct ClientFeedbackInput { /// Terminal environment snapshot from the client. #[serde(default)] - pub terminal_info: Option, + pub terminal_info: Option, } impl ClientFeedbackInput { @@ -144,10 +144,10 @@ impl ClientFeedbackInput { /// - stars: 1 to 5 /// - nps: 0 to 10 fn clamp_rating_value( - rating_type: Option, + rating_type: Option, rating_value: Option, ) -> Option { - use prod_mc_cli_chat_proxy_types::feedback_types::RatingType; + use crate::session::feedback_types::RatingType; match (rating_type, rating_value) { (Some(RatingType::Thumbs), Some(v)) => Some(v.clamp(-1, 1)), @@ -175,8 +175,8 @@ impl ClientFeedbackInput { resolved_model_id: Option, model_fingerprint: Option, turn_number: Option, - ) -> prod_mc_cli_chat_proxy_types::feedback_types::FeedbackSubmission { - use prod_mc_cli_chat_proxy_types::feedback_types::FeedbackContent; + ) -> crate::session::feedback_types::FeedbackSubmission { + use crate::session::feedback_types::FeedbackContent; let clamped_rating_value = Self::clamp_rating_value(self.rating_type, self.rating_value); let content = match ( @@ -577,7 +577,7 @@ pub struct SessionInfoResponse { pub struct FeedbackContext { pub last_user_message: Option, pub last_assistant_message: Option, - pub tool_outcomes: Vec, + pub tool_outcomes: Vec, pub compaction_count: i64, pub context_window_usage: u8, pub context_tokens_used: u64, @@ -655,14 +655,14 @@ mod tests { let input: ClientFeedbackInput = serde_json::from_str(json).unwrap(); assert_eq!( input.client_type, - prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Desktop + crate::session::feedback_types::ClientType::Desktop ); assert_eq!(input.session_id, "sess-1"); let submission = input.to_submission(Some("grok-3".into()), None, None, Some(5)); assert_eq!( submission.client_type, - prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Desktop + crate::session::feedback_types::ClientType::Desktop ); assert_eq!(submission.client_type.to_string(), "desktop"); } diff --git a/crates/codegen/kigi-shell/src/session/compaction.rs b/crates/codegen/kigi-shell/src/session/compaction.rs index d867ffe..31f7795 100644 --- a/crates/codegen/kigi-shell/src/session/compaction.rs +++ b/crates/codegen/kigi-shell/src/session/compaction.rs @@ -7,7 +7,7 @@ //! lives alongside the primary one in `acp_session.rs`. use super::SessionActor; use super::is_project_instructions; -use crate::remote::DEFAULT_CONTEXT_WINDOW; +use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW; use crate::session::compaction_config::{ AsyncCompactionCache, SUPPRESS_NONE, SUPPRESS_STICKY, SUPPRESS_TURN, SUPPRESS_UNTIL_SUCCESS, }; @@ -2242,7 +2242,6 @@ mod inline_auto_compact_flow_tests { client_identifier: None, origin_client: None, feedback_manager: Arc::new(FeedbackManager::local_only("test-session")), - sync_loop_cancel: None, agent: std::cell::RefCell::new(test_agent_default().await), last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)), git_head_enabled: false, diff --git a/crates/codegen/kigi-shell/src/session/feedback.rs b/crates/codegen/kigi-shell/src/session/feedback.rs index 4c5385e..e668fc5 100644 --- a/crates/codegen/kigi-shell/src/session/feedback.rs +++ b/crates/codegen/kigi-shell/src/session/feedback.rs @@ -10,9 +10,7 @@ use super::signals::SessionSignals; use crate::util::probabilistic_sample; // Re-export shared feedback API wire types to avoid duplication -pub use prod_mc_cli_chat_proxy_types::feedback_types::{ - FeedbackHeuristicsConfig, FeedbackMode, TierConfig, -}; +pub use crate::session::feedback_types::{FeedbackHeuristicsConfig, FeedbackMode, TierConfig}; /// Feedback request tier with associated probability and criteria. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -288,7 +286,7 @@ impl FeedbackHeuristics { /// Create a heuristics evaluator from a remote feedback-heuristics config. pub fn from_config(config: &FeedbackHeuristicsConfig) -> Self { - use prod_mc_cli_chat_proxy_types::feedback_types::parse_feedback_mode_str; + use crate::session::feedback_types::parse_feedback_mode_str; Self { enabled: config.enabled, @@ -343,7 +341,7 @@ impl FeedbackHeuristics { /// Update the heuristics configuration from a loaded config. /// Preserves the triggered_tiers state and request tracking. pub fn update_config(&mut self, config: &FeedbackHeuristicsConfig) { - use prod_mc_cli_chat_proxy_types::feedback_types::parse_feedback_mode_str; + use crate::session::feedback_types::parse_feedback_mode_str; self.enabled = config.enabled; diff --git a/crates/codegen/kigi-shell/src/session/feedback_manager.rs b/crates/codegen/kigi-shell/src/session/feedback_manager.rs index 26d349b..f1d3e9a 100644 --- a/crates/codegen/kigi-shell/src/session/feedback_manager.rs +++ b/crates/codegen/kigi-shell/src/session/feedback_manager.rs @@ -3,57 +3,48 @@ //! This manager coordinates: //! - Signal tracking via SessionSignalsHandle //! - Heuristics evaluation to determine when to request feedback -//! - Periodic sync of signals to the feedback/analytics backend -//! - Background loading of feedback configuration from the backend -//! - Creating feedback request records when triggered -//! - Sending feedback request notifications to clients +//! - Local persistence of every feedback record +//! - Forwarding text feedback to the Kimi Code feedback endpoint for +//! subscription (OAuth) sessions //! //! ## Usage //! ```ignore //! // Create the manager when a session starts -//! let manager = FeedbackManager::new(session_id, feedback_api_url, user_token); +//! let manager = FeedbackManager::new(session_id, feedback_client, config); //! //! // Get the signals handle to pass around for event tracking //! let signals = manager.signals_handle(); //! -//! // Spawn the background sync task (also loads config) -//! tokio::spawn(manager.run_sync_loop()); -//! //! // Track events //! signals.increment_turn(); //! signals.record_tool_call("read_file"); //! //! // Check for feedback after each turn -//! // This also records the request with the feedback API if triggered //! if let Some(request) = manager.maybe_request_feedback(None).await { //! // Send FeedbackRequest notification to client //! } //! ``` -use std::ops::ControlFlow; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use tokio::sync::RwLock; -use crate::agent::feedback_client::{ - FeedbackApiError, FeedbackClient, signals_to_update, snapshot_to_turn_delta, -}; +use crate::agent::feedback_client::FeedbackClient; use crate::session::feedback::{ FeedbackEvaluation, FeedbackHeuristics, FeedbackRequest, FeedbackTier, TriggerCondition, }; -use crate::session::signals::{SessionSignalsActor, SessionSignalsHandle, TurnDeltaSnapshot}; +use crate::session::signals::{SessionSignalsActor, SessionSignalsHandle}; -use prod_mc_cli_chat_proxy_types::feedback_types::{ - ClientType, ContextType, CreateFeedbackRequestInput, FeedbackContent, FeedbackMode, - FeedbackSubmission, FeedbackToolOutcome, +use crate::session::feedback_types::{ + ClientType, FeedbackContent, FeedbackMode, FeedbackSubmission, FeedbackToolOutcome, }; use crate::session::persistence::{LocalFeedbackEntry, PersistenceMsg, UserFeedbackEntry}; pub(crate) enum SubmitOutcome { Submitted, - /// No server configured for this session. + /// Persisted locally only: no subscription session, or a rating-only + /// record with no text content for the Kimi feedback endpoint. LocalOnly, /// Server request failed. Failed(anyhow::Error), @@ -70,8 +61,9 @@ pub(crate) fn new_submission( s } -/// Pipeline: persist → strip → submit. Callers merge `KIGI_USER_METADATA` and -/// set `submission.request_id`. +/// Pipeline: persist locally → forward text content to the Kimi feedback +/// endpoint (subscription sessions only). Callers merge `KIGI_USER_METADATA` +/// and set `submission.request_id`. pub(crate) async fn submit_feedback_workflow( submission: &mut FeedbackSubmission, feedback_client: Option<&FeedbackClient>, @@ -96,42 +88,33 @@ pub(crate) async fn submit_feedback_workflow( } } - let telemetry_model_id = submission.model_id.clone(); let telemetry_rating_value = submission.rating_value; - let telemetry_session_id = submission.session_id.clone(); let has_feedback_text = submission .feedback_text .as_ref() .is_some_and(|t| !t.is_empty()); - let request_id = submission.request_id.clone(); - let appearance_id = request_id.clone(); + let appearance_id = submission.request_id.clone(); - // Keep client-enriched triage fields; do not strip_metadata (Slack shows Option fields when set). - - let outcome = if let Some(client) = feedback_client { - let result = if let Some(req_id) = request_id { - with_one_shot_auth_retry(client, || async { - client - .complete_request(&req_id, submission) - .await - .map(|_| ()) - }) - .await - } else { - with_one_shot_auth_retry(client, || async { - client.submit_feedback(submission).await.map(|_| ()) - }) - .await - }; - match result { - Ok(()) => SubmitOutcome::Submitted, - Err(e) => { - tracing::warn!(error = %e, "feedback submission failed"); - SubmitOutcome::Failed(e) + // Only text-bearing feedback goes over the wire: the Kimi endpoint takes + // a `content` string (kimi-cli slash.py parity); ratings stay local. + let outcome = match (feedback_client, &submission.feedback_text) { + (Some(client), Some(text)) if !text.is_empty() => { + let model = submission + .model_id + .as_deref() + .or(submission.resolved_model_id.as_deref()); + match client + .submit_feedback(&submission.session_id, text, model) + .await + { + Ok(()) => SubmitOutcome::Submitted, + Err(e) => { + tracing::warn!(error = %e, "feedback submission failed"); + SubmitOutcome::Failed(e) + } } } - } else { - SubmitOutcome::LocalOnly + _ => SubmitOutcome::LocalOnly, }; { @@ -172,18 +155,13 @@ pub struct FeedbackFlags { /// Configuration for the feedback manager. #[derive(Debug, Clone)] pub struct FeedbackManagerConfig { - /// Interval for syncing signals to the analytics backend (default: 30s) + /// Interval for the signals actor's periodic bookkeeping tick. pub sync_interval: Duration, /// Whether user-facing feedback features are enabled (popups, `/feedback`, /// ratings). Gated by `KIGI_FEEDBACK_ENABLED`. pub feedback_enabled: bool, /// Client type (Agent, Tui, Web, Extension) pub client_type: ClientType, - /// Whether LOC attribution tracking is enabled for this session. - /// Propagated into every `SessionTurnDelta` so the server can - /// distinguish "tracking off" (zeros are noise) from "tracking on, - /// no code changed" (zeros are real data). - pub loc_tracking_enabled: bool, } impl Default for FeedbackManagerConfig { @@ -192,7 +170,6 @@ impl Default for FeedbackManagerConfig { sync_interval: Duration::from_secs(60), feedback_enabled: false, client_type: ClientType::Agent, - loc_tracking_enabled: false, } } } @@ -205,19 +182,17 @@ pub struct FeedbackManager { signals_handle: SessionSignalsHandle, /// Feedback heuristics evaluator heuristics: Arc>, - /// REST client for the feedback/analytics backend + /// Client for the Kimi Code feedback endpoint (subscription sessions). feedback_client: Option, /// Configuration config: FeedbackManagerConfig, - /// Whether config has been loaded from server - config_loaded: Arc, } impl FeedbackManager { /// Create a new feedback manager for a session. /// - /// If `feedback_client` is None, signal syncing is disabled but local - /// tracking and heuristics evaluation still work. + /// If `feedback_client` is None, submissions stay local but tracking and + /// heuristics evaluation still work. pub fn new( session_id: impl Into, feedback_client: Option, @@ -243,7 +218,6 @@ impl FeedbackManager { heuristics: Arc::new(RwLock::new(FeedbackHeuristics::new())), feedback_client, config, - config_loaded: Arc::new(AtomicBool::new(false)), } } @@ -267,13 +241,14 @@ impl FeedbackManager { self.config.feedback_enabled } - /// REST client for the feedback/analytics backend, if configured. + /// Client for the Kimi Code feedback endpoint, if this is a subscription + /// session. pub fn feedback_client(&self) -> Option<&FeedbackClient> { self.feedback_client.as_ref() } /// Client type for this session (Agent, Tui, Web, etc.). - pub fn client_type(&self) -> prod_mc_cli_chat_proxy_types::feedback_types::ClientType { + pub fn client_type(&self) -> ClientType { self.config.client_type } @@ -330,47 +305,6 @@ impl FeedbackManager { .await } - /// Check if config has been loaded from the server. - pub fn is_config_loaded(&self) -> bool { - self.config_loaded.load(Ordering::Relaxed) - } - - /// Load feedback heuristics config from the backend. - /// This is called automatically in run_sync_loop but can be called manually. - /// Does not block - errors are logged and defaults are used. - #[tracing::instrument(name = "feedback.load_config", skip_all, fields( - session_id = %self.session_id, - ))] - pub async fn load_config(&self) { - let Some(client) = &self.feedback_client else { - return; // No client, use defaults - }; - - if self.config.feedback_enabled { - match client.get_feedback_config().await { - Ok(config) => { - let mut heuristics = self.heuristics.write().await; - heuristics.update_config(&config); - self.config_loaded.store(true, Ordering::Relaxed); - tracing::info!( - session_id = %self.session_id, - config_id = %config.config_id, - config_version = config.config_version, - enabled = config.enabled, - "Loaded feedback heuristics config from server" - ); - } - Err(e) => { - tracing::warn!( - session_id = %self.session_id, - error = %e, - "Failed to load feedback heuristics config, using defaults" - ); - } - } - } - } - /// Evaluate heuristics and return a FeedbackRequest if one should be sent. /// /// Call this after each turn to check if feedback should be requested. @@ -378,9 +312,6 @@ impl FeedbackManager { /// - No tier criteria are met /// - The tier was already triggered this session /// - Probabilistic sampling says no - /// - /// When a request is triggered, this method also creates a record via the - /// feedback API for tracking and analytics. #[tracing::instrument(name = "feedback.maybe_request_feedback", skip_all, fields( session_id = %self.session_id, ))] @@ -395,7 +326,6 @@ impl FeedbackManager { let signals = self.signals_handle.snapshot().await?; let mut heuristics = self.heuristics.write().await; - // Check if heuristics are globally enabled (from server config) if !heuristics.is_enabled() { return None; } @@ -422,12 +352,10 @@ impl FeedbackManager { tier = ?request.tier, trigger_type = %request.trigger_type, feedback_mode = ?request.feedback_mode, + prompt_id = ?prompt_id, "Feedback request triggered" ); - self.record_feedback_request(&request, trigger_condition, feedback_mode, prompt_id) - .await; - return Some(request); } @@ -449,11 +377,6 @@ impl FeedbackManager { /// `x.ai/debug/trigger_feedback` ACP extension method to exercise /// the full feedback notification ↔ response flow without needing a /// real session that meets tier criteria. - /// - /// When a `feedback_client` is configured, the request is also recorded - /// via the feedback API — exactly like a real trigger — so that the - /// subsequent `complete_request` / `dismiss_request` round-trip from the - /// client works end-to-end. #[tracing::instrument(name = "feedback.force_feedback_request", skip_all, fields( session_id = %self.session_id, ))] @@ -481,96 +404,7 @@ impl FeedbackManager { // Manual/debug triggers are always dismissible regardless of tier config, // since they exist for developer testing, not real user feedback collection. - let request = FeedbackRequest::with_mode( - self.session_id.clone(), - condition.clone(), - mode, - true, - None, - ); - - self.record_feedback_request(&request, &condition, mode, None) - .await; - - request - } - - /// Record a feedback request via the feedback API. - /// - /// This is a best-effort operation — errors are logged but do not - /// prevent the request from being sent to the client. - #[tracing::instrument(name = "feedback.record_feedback_request", skip_all, fields( - session_id = %self.session_id, - ))] - async fn record_feedback_request( - &self, - request: &FeedbackRequest, - trigger_condition: &TriggerCondition, - feedback_mode: FeedbackMode, - prompt_id: Option, - ) { - let Some(client) = &self.feedback_client else { - return; - }; - - let input = CreateFeedbackRequestInput { - request_id: request.request_id.clone(), - session_id: self.session_id.clone(), - client_type: self.config.client_type, - feedback_mode, - feedback_prompt: Some(request.prompt.clone()), - priority: tier_to_priority(trigger_condition.tier), - trigger_type: request.trigger_type.clone(), - trigger_reason: Some(trigger_condition.trigger_reason()), - context_type: Some(ContextType::Session), - context_message_ids: vec![], - expires_at: None, - experiment_id: None, - trigger_condition: serde_json::to_value(trigger_condition).ok(), - prompt_id, - }; - - match with_one_shot_auth_retry(client, || client.create_feedback_request(&input)).await { - Ok(response) => { - tracing::debug!( - request_id = %response.request_id, - "Feedback request recorded with feedback API" - ); - } - Err(e) => { - tracing::warn!( - request_id = %request.request_id, - error = %e, - "Failed to record feedback request (continuing anyway)" - ); - } - } - } - - /// Capture a turn-end snapshot and send the delta to the analytics backend. - /// - /// Call this once per user turn, after the agent has finished all tool-call - /// rounds and produced a final response (i.e. alongside `record_turn_complete`). - /// Intermediate tool-call steps within the same turn do NOT need their own - /// call — the signals actor accumulates tool calls, errors, and latency - /// continuously, so the single snapshot at turn end captures the full diff. - /// - /// The caller provides a pre-captured `TurnDeltaSnapshot` (taken exactly - /// once inside the session actor). This avoids double-advancing the delta - /// baseline. If the snapshot is `None` (e.g. the signals actor was shut - /// down), this is a no-op. - /// - /// The delta is converted and sent asynchronously to the backend. Errors - /// are logged but never block the turn flow. - /// - /// Load feedback heuristics config on startup. - /// This should be spawned as a background task. - #[tracing::instrument(skip_all, fields(session_id = %self.session_id))] - pub async fn run_sync_loop(self: Arc, cancel: tokio_util::sync::CancellationToken) { - // Load config in background (non-blocking, errors logged) - self.load_config().await; - cancel.cancelled().await; - tracing::debug!("Feedback sync loop cancelled"); + FeedbackRequest::with_mode(self.session_id.clone(), condition, mode, true, None) } /// Shutdown the manager: shuts down the signals actor. @@ -579,173 +413,6 @@ impl FeedbackManager { } } -// Auth outcome handler used by run_sync_loop on 401. - -/// Max consecutive failed sync ticks tolerated before stopping the loop. -/// ~10 minutes at the default 60s interval. -const MAX_CONSECUTIVE_AUTH_FAILURES: u8 = 10; - -/// telemetry `reason` discriminators on the `signals sync loop stopped permanently` -/// event. Pinned because alerts filter on these strings. -const REASON_AUTH_PERMANENT_FAILURE: &str = "auth_permanent_failure"; -const REASON_NO_CLIENT_OR_REFRESHER: &str = "no_client_or_refresher"; - -const LOG_TITLE_TRANSIENT: &str = "signals sync transient auth failure"; -const LOG_TITLE_STOPPED_PERMANENTLY: &str = "signals sync loop stopped permanently"; - -/// Classification of one 401-recovery attempt. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum SyncAuthOutcome { - /// Refresh + retry succeeded. - Recovered, - /// Refresh or retry failed transiently (lock timeout, network, sibling - /// race, post-refresh 5xx). Increment the counter and retry next tick. - Transient, - /// IdP confirmed a terminal failure (`invalid_grant` / `invalid_client`). - /// Only re-login will recover. - Permanent, - /// No client or no refresher configured — nothing to retry. - Unrecoverable, -} - -fn handle_auth_outcome( - outcome: SyncAuthOutcome, - consecutive_auth_failures: &mut u8, - session_id: &str, -) -> ControlFlow<()> { - match outcome { - SyncAuthOutcome::Recovered => { - *consecutive_auth_failures = 0; - tracing::info!( - session_id = %session_id, - "Signal sync recovered after token refresh" - ); - ControlFlow::Continue(()) - } - SyncAuthOutcome::Transient => { - *consecutive_auth_failures = consecutive_auth_failures.saturating_add(1); - tracing::warn!( - session_id = %session_id, - consecutive_failures = *consecutive_auth_failures, - max = MAX_CONSECUTIVE_AUTH_FAILURES, - "Signals sync transient auth failure" - ); - kigi_log::unified_log::warn( - LOG_TITLE_TRANSIENT, - Some(session_id), - Some(serde_json::json!({ - "consecutive_failures": *consecutive_auth_failures, - "max": MAX_CONSECUTIVE_AUTH_FAILURES, - })), - ); - if *consecutive_auth_failures >= MAX_CONSECUTIVE_AUTH_FAILURES { - tracing::warn!( - session_id = %session_id, - consecutive_failures = *consecutive_auth_failures, - "Signals sync loop stopped: consecutive transient auth failures" - ); - kigi_log::unified_log::warn( - "signals sync loop stopped: consecutive transient auth failures", - Some(session_id), - Some(serde_json::json!({ - "consecutive_failures": *consecutive_auth_failures, - "max": MAX_CONSECUTIVE_AUTH_FAILURES, - })), - ); - ControlFlow::Break(()) - } else { - ControlFlow::Continue(()) - } - } - SyncAuthOutcome::Permanent => { - tracing::warn!( - session_id = %session_id, - reason = REASON_AUTH_PERMANENT_FAILURE, - "Signals sync loop stopped: IdP confirmed permanent auth failure" - ); - kigi_log::unified_log::warn( - LOG_TITLE_STOPPED_PERMANENTLY, - Some(session_id), - Some(serde_json::json!({ "reason": REASON_AUTH_PERMANENT_FAILURE })), - ); - ControlFlow::Break(()) - } - SyncAuthOutcome::Unrecoverable => { - tracing::warn!( - session_id = %session_id, - reason = REASON_NO_CLIENT_OR_REFRESHER, - "Signals sync loop stopped: no client or no refresher configured" - ); - kigi_log::unified_log::warn( - LOG_TITLE_STOPPED_PERMANENTLY, - Some(session_id), - Some(serde_json::json!({ "reason": REASON_NO_CLIENT_OR_REFRESHER })), - ); - ControlFlow::Break(()) - } - } -} - -/// Check if an error is an HTTP 401 Unauthorized response. -/// -/// Uses typed downcast on [`FeedbackApiError`] instead of string matching, -/// so it stays correct even if error messages change. -fn is_auth_error(error: &anyhow::Error) -> bool { - error - .downcast_ref::() - .is_some_and(|e| e.is_unauthorized()) -} - -/// Check if an error is an HTTP 403 Forbidden response. -/// -/// 403 from the signals endpoint means the session does not belong to the -/// current user — a permanent condition that will never self-resolve. -fn is_forbidden_error(error: &anyhow::Error) -> bool { - error - .downcast_ref::() - .is_some_and(|e| e.is_forbidden()) -} - -/// Run `op` once; on 401, wait for an in-flight refresh to land, then -/// retry once. Prefers waiting for the proactive-refresh task or -/// main-request-path recovery over driving a `ServerRejected` refresh -/// itself, avoiding the 401-amplification pattern during token-expiry -/// windows. -async fn with_one_shot_auth_retry( - client: &FeedbackClient, - mut op: F, -) -> anyhow::Result -where - F: FnMut() -> Fut, - Fut: std::future::Future>, -{ - match op().await { - Ok(v) => Ok(v), - Err(e) if is_auth_error(&e) => { - // 1. Wait briefly for the proactive refresh or main-path - // recovery to land a fresh token. - let refreshed = client.wait_for_token_refresh(Duration::from_secs(3)).await; - // 2. If nobody refreshed, drive our own recovery as fallback. - if refreshed || client.try_refresh_credentials().await { - op().await - } else { - Err(e) - } - } - Err(e) => Err(e), - } -} - -/// Convert a FeedbackTier to a priority value (1-10, higher = more important). -fn tier_to_priority(tier: crate::session::feedback::FeedbackTier) -> i32 { - use crate::session::feedback::FeedbackTier; - match tier { - FeedbackTier::Tier1 => 5, // Standard engagement - FeedbackTier::Tier2 => 6, // Complex session with recovery - FeedbackTier::Tier3 => 7, // Recovery from friction - } -} - #[cfg(test)] mod tests { use super::*; @@ -786,83 +453,6 @@ mod tests { manager.shutdown().await; } - #[test] - fn test_is_auth_error_detects_401() { - use crate::agent::feedback_client::FeedbackApiError; - let err: anyhow::Error = FeedbackApiError { - status: reqwest::StatusCode::UNAUTHORIZED, - context: "Signals update", - body: "Invalid or expired credentials".to_string(), - } - .into(); - assert!(is_auth_error(&err)); - } - - #[test] - fn test_is_auth_error_ignores_other_statuses() { - use crate::agent::feedback_client::FeedbackApiError; - let err_500: anyhow::Error = FeedbackApiError { - status: reqwest::StatusCode::INTERNAL_SERVER_ERROR, - context: "Signals update", - body: "oops".to_string(), - } - .into(); - assert!(!is_auth_error(&err_500)); - - let err_403: anyhow::Error = FeedbackApiError { - status: reqwest::StatusCode::FORBIDDEN, - context: "Signals update", - body: "ZDR team".to_string(), - } - .into(); - assert!(!is_auth_error(&err_403)); - } - - #[test] - fn test_is_auth_error_ignores_non_api_errors() { - assert!(!is_auth_error(&anyhow::anyhow!("network timeout"))); - assert!(!is_auth_error(&anyhow::anyhow!("connection refused"))); - } - - #[test] - fn test_is_forbidden_error_detects_403() { - use crate::agent::feedback_client::FeedbackApiError; - let err: anyhow::Error = FeedbackApiError { - status: reqwest::StatusCode::FORBIDDEN, - context: "Signals update", - body: "Access denied: session does not belong to this user".to_string(), - } - .into(); - assert!(is_forbidden_error(&err)); - } - - #[test] - fn test_is_forbidden_error_ignores_other_statuses() { - use crate::agent::feedback_client::FeedbackApiError; - let err_401: anyhow::Error = FeedbackApiError { - status: reqwest::StatusCode::UNAUTHORIZED, - context: "Signals update", - body: "Invalid credentials".to_string(), - } - .into(); - assert!(!is_forbidden_error(&err_401)); - assert!(!is_forbidden_error(&anyhow::anyhow!("network error"))); - } - - #[test] - fn test_is_auth_error_works_through_anyhow_conversion() { - use crate::agent::feedback_client::FeedbackApiError; - // Verify the FeedbackApiError survives anyhow::Error round-trip - // (this is the actual path: send_json returns FeedbackApiError.into()) - let api_err = FeedbackApiError { - status: reqwest::StatusCode::UNAUTHORIZED, - context: "Signals update", - body: "token expired".to_string(), - }; - let anyhow_err: anyhow::Error = api_err.into(); - assert!(is_auth_error(&anyhow_err)); - } - #[tokio::test] async fn test_feedback_manager_disabled() { let config = FeedbackManagerConfig { @@ -895,158 +485,21 @@ mod tests { assert!(snapshot.is_none(), "Signals actor should be shut down"); } - // ── handle_auth_outcome tests ────────────────────────────────────────── - - /// 9 transient failures must NOT break, and a subsequent `Recovered` - /// must reset the counter to 0. - #[test] - fn test_sync_loop_continues_through_transient_auth_failures() { - let mut counter: u8 = 0; - for _ in 0..(MAX_CONSECUTIVE_AUTH_FAILURES - 1) { - let flow = handle_auth_outcome(SyncAuthOutcome::Transient, &mut counter, "s"); - assert_eq!(flow, ControlFlow::Continue(())); - } - assert_eq!(counter, MAX_CONSECUTIVE_AUTH_FAILURES - 1); - let flow = handle_auth_outcome(SyncAuthOutcome::Recovered, &mut counter, "s"); - assert_eq!(flow, ControlFlow::Continue(())); - assert_eq!(counter, 0, "Recovered must reset the counter"); - } - - /// Exactly `MAX_CONSECUTIVE_AUTH_FAILURES` consecutive `Transient` - /// outcomes break the loop. - #[test] - fn test_sync_loop_breaks_after_max_transient_auth_failures() { - let mut counter: u8 = 0; - for i in 0..(MAX_CONSECUTIVE_AUTH_FAILURES - 1) { - let flow = handle_auth_outcome(SyncAuthOutcome::Transient, &mut counter, "s"); - assert_eq!( - flow, - ControlFlow::Continue(()), - "iteration {i} should still continue" - ); - } - // The 10th (== MAX_CONSECUTIVE_AUTH_FAILURES) transient breaks. - let flow = handle_auth_outcome(SyncAuthOutcome::Transient, &mut counter, "s"); - assert_eq!(flow, ControlFlow::Break(())); - assert_eq!(counter, MAX_CONSECUTIVE_AUTH_FAILURES); - } - - /// `Permanent` breaks immediately and does not bump the counter. - #[test] - fn test_sync_loop_breaks_immediately_on_permanent_failure() { - let mut counter: u8 = 0; - let flow = handle_auth_outcome(SyncAuthOutcome::Permanent, &mut counter, "s"); - assert_eq!(flow, ControlFlow::Break(())); - assert_eq!(counter, 0); - } - - /// 5 transient → 1 recovered → 5 transient must not break. - #[test] - fn test_sync_loop_counter_resets_on_successful_sync() { - let mut counter: u8 = 0; - for _ in 0..5 { - assert_eq!( - handle_auth_outcome(SyncAuthOutcome::Transient, &mut counter, "s"), - ControlFlow::Continue(()) - ); - } - assert_eq!(counter, 5); - assert_eq!( - handle_auth_outcome(SyncAuthOutcome::Recovered, &mut counter, "s"), - ControlFlow::Continue(()) - ); - assert_eq!(counter, 0, "Recovered must reset the counter"); - for _ in 0..5 { - assert_eq!( - handle_auth_outcome(SyncAuthOutcome::Transient, &mut counter, "s"), - ControlFlow::Continue(()) - ); - } - assert_eq!(counter, 5, "second burst should be re-counted from zero"); - } - - /// `Unrecoverable` breaks the loop and does not bump the counter. - #[test] - fn test_sync_loop_breaks_on_unrecoverable() { - let mut counter: u8 = 0; - let flow = handle_auth_outcome(SyncAuthOutcome::Unrecoverable, &mut counter, "s"); - assert_eq!(flow, ControlFlow::Break(())); - assert_eq!(counter, 0); - } - - /// `FeedbackClient::is_auth_permanently_failed` reflects the attached - /// `AuthManager`'s `permanent_failure()` cache (record → true, - /// age-out → false). + /// A rating-only submission (no text) must not hit the network even when + /// no client is configured — the workflow reports LocalOnly. #[tokio::test] - async fn test_is_auth_permanently_failed_reads_auth_manager() { - use crate::agent::feedback_client::FeedbackClient; - use crate::auth::error::RefreshTokenFailedReason; - use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig}; - use std::sync::Arc; + async fn test_rating_only_submission_stays_local() { + use crate::session::feedback_types::RatingType; - let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); - let client = FeedbackClient::new("http://example/v1", None).with_auth_manager(am.clone()); - - assert!(!client.is_auth_permanently_failed()); - - // The tombstone is scoped to the live credential's refresh token. - am.hot_swap(KimiAuth { - key: "tok".into(), - refresh_token: Some("rt".into()), - expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)), - ..KimiAuth::test_default() - }); - am.record_permanent_failure("rt".to_string(), RefreshTokenFailedReason::Other.into()); - assert!(client.is_auth_permanently_failed()); - - am.force_permanent_failure_aged_out(); - assert!(!client.is_auth_permanently_failed()); - } - - /// With no `AuthManager` attached, `is_auth_permanently_failed` is false. - #[test] - fn test_is_auth_permanently_failed_without_auth_manager() { - use crate::agent::feedback_client::FeedbackClient; - let client = FeedbackClient::new("http://example/v1", None); - assert!(!client.is_auth_permanently_failed()); - } - - /// `has_token_refresher` requires BOTH an `AuthManager` AND a refresher - /// wired in. Without this, a static-deployment-key session would be - /// mis-classified as recoverable. - #[tokio::test] - async fn test_has_token_refresher_requires_refresher_attached() { - use crate::agent::feedback_client::FeedbackClient; - use crate::auth::{AuthManager, KimiCodeConfig}; - use std::sync::Arc; - - struct NoOpRefresher; - #[async_trait::async_trait] - impl crate::auth::refresh::TokenRefresher for NoOpRefresher { - async fn refresh( - &self, - _reason: crate::auth::refresh::RefreshReason, - ) -> crate::auth::refresh::RefreshOutcome { - crate::auth::refresh::RefreshOutcome::TransientFailure { - message: "noop".into(), - } - } - } - - let dir = tempfile::tempdir().unwrap(); - let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default())); - - let bare = FeedbackClient::new("http://example/v1", None); - assert!(!bare.has_token_refresher()); - - let with_am = FeedbackClient::new("http://example/v1", None).with_auth_manager(am.clone()); - assert!( - !with_am.has_token_refresher(), - "AuthManager without a refresher must NOT be reported as recoverable" + let mut submission = new_submission( + "sess-local".into(), + ClientType::Tui, + FeedbackContent::Rating { + rating_type: RatingType::Thumbs, + rating_value: 1, + }, ); - - am.set_refresher(std::sync::Arc::new(NoOpRefresher)); - assert!(with_am.has_token_refresher()); + let outcome = submit_feedback_workflow(&mut submission, None, None, false).await; + assert!(matches!(outcome, SubmitOutcome::LocalOnly)); } } diff --git a/crates/codegen/kigi-shell/src/session/feedback_types.rs b/crates/codegen/kigi-shell/src/session/feedback_types.rs new file mode 100644 index 0000000..9e68932 --- /dev/null +++ b/crates/codegen/kigi-shell/src/session/feedback_types.rs @@ -0,0 +1,763 @@ +//! Local feedback data types. +//! +//! Formerly the wire contract with the deleted xAI cli-chat-proxy feedback +//! backend; now these types only back the LOCAL feedback records persisted in +//! the session store and the heuristics that decide when to solicit feedback. +//! The only remaining network surface is the Kimi Code `POST {base}/feedback` +//! call in [`crate::agent::feedback_client`], which sends a small flat JSON +//! body — none of these types go over the wire anymore, so the proxy-only +//! null-column fields (experiment/comparison/preference plumbing) are gone. + +use serde::{Deserialize, Deserializer, Serialize}; + +pub use kigi_shared::session::FeedbackTerminalInfo; + +/// Type of client submitting feedback. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ClientType { + /// Terminal/CLI agent + #[default] + Agent, + /// Terminal UI + Tui, + /// Web interface + Web, + /// IDE extension (VS Code, JetBrains, etc.) + Extension, + /// Remote workspace / hosted agent client (wire value `nebula`). + Nebula, + /// Desktop (Electron app) + Desktop, +} + +impl std::fmt::Display for ClientType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ClientType::Agent => write!(f, "agent"), + ClientType::Tui => write!(f, "tui"), + ClientType::Web => write!(f, "web"), + ClientType::Extension => write!(f, "extension"), + ClientType::Nebula => write!(f, "nebula"), + ClientType::Desktop => write!(f, "desktop"), + } + } +} + +/// Type of feedback being submitted. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FeedbackType { + /// Numeric rating only + #[default] + Rating, + /// Free-form text only + Text, + /// Both rating and text + RatingWithText, + /// Model preference comparison + ModelPreference, + /// Bug report + BugReport, + /// Feature request + FeatureRequest, +} + +impl std::fmt::Display for FeedbackType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + FeedbackType::Rating => write!(f, "rating"), + FeedbackType::Text => write!(f, "text"), + FeedbackType::RatingWithText => write!(f, "rating_with_text"), + FeedbackType::ModelPreference => write!(f, "model_preference"), + FeedbackType::BugReport => write!(f, "bug_report"), + FeedbackType::FeatureRequest => write!(f, "feature_request"), + } + } +} + +/// Type of rating scale used. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RatingType { + /// Thumbs up/down (-1, 0, 1) + Thumbs, + /// Star rating (1-5) + Stars, + /// Net Promoter Score (0-10) + Nps, +} + +impl std::fmt::Display for RatingType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + RatingType::Thumbs => write!(f, "thumbs"), + RatingType::Stars => write!(f, "stars"), + RatingType::Nps => write!(f, "nps"), + } + } +} + +/// Context type for what the feedback is about. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ContextType { + /// Feedback about a specific message + Message, + /// Feedback about the overall session/conversation + Session, + /// Feedback about a specific feature + Feature, + /// Feedback about tool usage + ToolUse, + /// General feedback + General, +} + +impl std::fmt::Display for ContextType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ContextType::Message => write!(f, "message"), + ContextType::Session => write!(f, "session"), + ContextType::Feature => write!(f, "feature"), + ContextType::ToolUse => write!(f, "tool_use"), + ContextType::General => write!(f, "general"), + } + } +} + +/// Type of feedback mode requested. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FeedbackMode { + /// Thumbs up/down + Thumbs, + /// Star rating (1-5) + Stars, + /// Free-form text + Text, + /// Thumbs up/down with optional text comment + ThumbsText, + /// Star rating with optional text comment + StarsText, + /// Model comparison + Comparison, + /// Multi-question survey + Survey, + /// Net Promoter Score (0-10) + Nps, + /// NPS with optional text comment + NpsText, +} + +impl std::fmt::Display for FeedbackMode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + FeedbackMode::Thumbs => write!(f, "thumbs"), + FeedbackMode::Stars => write!(f, "stars"), + FeedbackMode::Text => write!(f, "text"), + FeedbackMode::ThumbsText => write!(f, "thumbs_text"), + FeedbackMode::StarsText => write!(f, "stars_text"), + FeedbackMode::Comparison => write!(f, "comparison"), + FeedbackMode::Survey => write!(f, "survey"), + FeedbackMode::Nps => write!(f, "nps"), + FeedbackMode::NpsText => write!(f, "nps_text"), + } + } +} + +/// Parse a feedback mode string to FeedbackMode enum. +pub fn parse_feedback_mode_str(s: &str) -> FeedbackMode { + match s { + "thumbs" => FeedbackMode::Thumbs, + "stars" => FeedbackMode::Stars, + "text" => FeedbackMode::Text, + "thumbs_text" => FeedbackMode::ThumbsText, + "stars_text" => FeedbackMode::StarsText, + "comparison" => FeedbackMode::Comparison, + "survey" => FeedbackMode::Survey, + "nps" => FeedbackMode::Nps, + "nps_text" => FeedbackMode::NpsText, + _ => FeedbackMode::Thumbs, + } +} + +/// Allowed `feedback_type` + value-field combinations. Construct submissions +/// via [`FeedbackSubmission::with_content`]. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub enum FeedbackContent { + Rating { + rating_type: RatingType, + rating_value: i32, + }, + Text(String), + RatingWithText { + rating_type: RatingType, + rating_value: i32, + text: String, + }, +} + +impl FeedbackContent { + fn apply_to(self, s: &mut FeedbackSubmission) { + s.rating_type = None; + s.rating_value = None; + s.feedback_text = None; + match self { + Self::Rating { + rating_type, + rating_value, + } => { + s.feedback_type = FeedbackType::Rating; + s.rating_type = Some(rating_type); + s.rating_value = Some(rating_value); + } + Self::Text(text) => { + s.feedback_type = FeedbackType::Text; + s.feedback_text = Some(text); + } + Self::RatingWithText { + rating_type, + rating_value, + text, + } => { + s.feedback_type = FeedbackType::RatingWithText; + s.rating_type = Some(rating_type); + s.rating_value = Some(rating_value); + s.feedback_text = Some(text); + } + } + } +} + +fn empty_string_as_none<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + let opt = Option::::deserialize(deserializer)?; + Ok(opt.filter(|s| !s.is_empty())) +} + +/// A user feedback record. Persisted locally in the session store; the text +/// content is forwarded to the Kimi Code feedback endpoint for subscription +/// sessions. Construct via [`FeedbackSubmission::with_content`]; the `Default` +/// impl exists for builder-style construction and test fixtures and does not +/// produce a valid submission on its own (empty `session_id`). +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FeedbackSubmission { + /// Session ID this feedback is for + pub session_id: String, + + /// Type of client submitting feedback + pub client_type: ClientType, + + /// Type of feedback being submitted + pub feedback_type: FeedbackType, + + /// Turn number within the session (optional) + #[serde(skip_serializing_if = "Option::is_none")] + pub turn_number: Option, + + /// Rating type (if applicable) + #[serde(skip_serializing_if = "Option::is_none")] + pub rating_type: Option, + + /// Rating value (interpretation depends on rating_type) + /// - thumbs: -1 (down), 0 (neutral), 1 (up) + /// - stars: 1-5 + /// - nps: 0-10 + #[serde(skip_serializing_if = "Option::is_none")] + pub rating_value: Option, + + /// Free-form feedback text + #[serde(skip_serializing_if = "Option::is_none")] + pub feedback_text: Option, + + /// Feedback categories (e.g., ["accuracy", "speed", "helpfulness"]) + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub feedback_categories: Vec, + + /// Model ID used for the response being rated + #[serde(skip_serializing_if = "Option::is_none")] + pub model_id: Option, + + /// Server-resolved model ID from the actual chat completion response. + #[serde(skip_serializing_if = "Option::is_none")] + pub resolved_model_id: Option, + + /// Checkpoint fingerprint from the inference provider (`system_fingerprint`). + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "empty_string_as_none" + )] + pub model_fingerprint: Option, + + /// Context type for the feedback + #[serde(skip_serializing_if = "Option::is_none")] + pub context_type: Option, + + /// Feedback request ID (set when responding to a solicited request) + #[serde(skip_serializing_if = "Option::is_none")] + pub request_id: Option, + + /// Client version + #[serde(skip_serializing_if = "Option::is_none")] + pub client_version: Option, + + /// Shell (kigi-shell) version + #[serde(skip_serializing_if = "Option::is_none")] + pub shell_version: Option, + + /// Additional metadata as JSON + #[serde(skip_serializing_if = "Option::is_none")] + pub metadata: Option, + + /// Last user message at feedback time. + #[serde( + default, + skip_serializing_if = "Option::is_none", + alias = "last_user_turn" + )] + pub last_user_message: Option, + + /// Last assistant response at feedback time. + #[serde( + default, + skip_serializing_if = "Option::is_none", + alias = "last_assistant_turn" + )] + pub last_assistant_message: Option, + + /// Per-tool call counts for the rated turn. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub tool_outcomes: Vec, + + /// Session working directory. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub session_cwd: Option, + + /// Number of compactions in the session. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub compaction_count: Option, + + /// Context window usage percentage (0–100). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub context_window_usage: Option, + + /// Raw context tokens used at feedback time. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub context_tokens_used: Option, + + /// Raw model context window token limit at feedback time. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub context_window_tokens: Option, + + /// Terminal environment snapshot at feedback time (brand, multiplexer, SSH, etc.). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub terminal_info: Option, +} + +impl FeedbackSubmission { + /// Construct from typed content; set optional fields after. + pub fn with_content( + session_id: String, + client_type: ClientType, + content: FeedbackContent, + ) -> Self { + let mut s = Self { + session_id, + client_type, + ..Default::default() + }; + content.apply_to(&mut s); + s + } + + /// Merge a JSON object into `metadata`, inserting if absent. + pub fn merge_metadata(&mut self, extra: serde_json::Value) { + match &mut self.metadata { + Some(existing) if existing.is_object() => { + if let (Some(dst), Some(src)) = (existing.as_object_mut(), extra.as_object()) { + for (k, v) in src { + dst.insert(k.clone(), v.clone()); + } + } + } + _ => { + self.metadata = Some(extra); + } + } + } +} + +/// Per-tool call/failure counts for a single tool in a turn. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FeedbackToolOutcome { + pub tool_name: String, + pub calls: u32, + pub failures: u32, +} + +/// Configuration for a single feedback tier. +/// +/// Each tier has specific thresholds and conditions that must be met +/// for feedback to be requested at that tier. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub struct TierConfig { + /// Whether this tier is enabled + pub enabled: bool, + /// Sample rate (0.0 to 1.0, e.g., 0.0005 = 0.05%) + pub sample_rate: f64, + /// Minimum turns required to trigger + pub min_turns: i64, + /// Minimum tool calls required (Tier 1 & 2) + #[serde(default)] + pub min_tool_calls: i64, + /// Minimum compactions required (Tier 1 & 2) + #[serde(default)] + pub min_compactions: i64, + /// Minimum errors required (Tier 2 only) + #[serde(default)] + pub min_errors: i64, + /// Whether cancellations disqualify this tier (Tier 1) + #[serde(default)] + pub no_cancellations: bool, + /// Whether cancellation is required (Tier 3) + #[serde(default)] + pub requires_cancellation: bool, + /// Whether revert is required (Tier 3) + #[serde(default)] + pub requires_revert: bool, + /// Whether at least one of cancellation/revert is required (Tier 3) + #[serde(default)] + pub requires_recovery: bool, + /// Feedback mode to use when this tier triggers + pub feedback_mode: FeedbackMode, + /// Whether feedback requests from this tier are dismissible (non-intrusive) + #[serde(default = "default_true")] + pub dismissible: bool, + /// Prompt text shown to users when this tier's feedback is requested + #[serde(default)] + pub prompt: String, + /// Max times this tier can trigger per session (0 = unlimited) + #[serde(default = "default_one")] + pub max_triggers: i32, +} + +impl Default for TierConfig { + fn default() -> Self { + Self { + enabled: true, + sample_rate: 0.0005, + min_turns: 10, + min_tool_calls: 5, + min_compactions: 2, + min_errors: 0, + no_cancellations: false, + requires_cancellation: false, + requires_revert: false, + requires_recovery: false, + feedback_mode: FeedbackMode::Thumbs, + dismissible: true, + prompt: String::new(), + max_triggers: 1, + } + } +} + +/// Configuration for feedback heuristics. +/// +/// Formerly fetched from the proxy backend; now purely local — the built-in +/// [`Default`] is the only production source, kept as a struct so tests and +/// future config surfaces can tune the tiers. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub struct FeedbackHeuristicsConfig { + /// Unique configuration identifier + pub config_id: String, + /// Configuration version (monotonically increasing) + pub config_version: i64, + + // === Global Settings === + /// Master enable/disable switch for all feedback collection + pub enabled: bool, + /// Minimum seconds between feedback requests (cooldown period) + #[serde(default = "default_cooldown_seconds")] + pub cooldown_seconds: i64, + /// Maximum feedback requests per session + #[serde(default = "default_max_requests")] + pub max_requests_per_session: i64, + + // === Tier 1: Standard Engagement === + /// Whether Tier 1 is enabled + #[serde(default = "default_true")] + pub tier1_enabled: bool, + /// Sample rate for Tier 1 (0.0-1.0) + #[serde(default = "default_tier1_sample_rate")] + pub tier1_sample_rate: f64, + /// Minimum turns for Tier 1 + #[serde(default = "default_tier1_min_turns")] + pub tier1_min_turns: i64, + /// Minimum tool calls for Tier 1 + #[serde(default = "default_tier1_min_tool_calls")] + pub tier1_min_tool_calls: i64, + /// Minimum compactions for Tier 1 + #[serde(default = "default_tier1_min_compactions")] + pub tier1_min_compactions: i64, + /// Whether Tier 1 requires no cancellations + #[serde(default = "default_true")] + pub tier1_no_cancellations: bool, + /// Feedback mode for Tier 1 + #[serde(default = "default_feedback_mode_thumbs")] + pub tier1_feedback_mode: String, + /// Whether Tier 1 feedback requests are dismissible + #[serde(default = "default_true")] + pub tier1_dismissible: bool, + /// Prompt text shown to users when Tier 1 feedback is requested + #[serde(default = "default_tier1_prompt")] + pub tier1_prompt: String, + /// Max times Tier 1 can trigger per session (0 = unlimited) + #[serde(default = "default_one")] + pub tier1_max_triggers: i32, + + // === Tier 2: Complex Session with Recovery === + /// Whether Tier 2 is enabled + #[serde(default = "default_true")] + pub tier2_enabled: bool, + /// Sample rate for Tier 2 (0.0-1.0) + #[serde(default = "default_tier2_sample_rate")] + pub tier2_sample_rate: f64, + /// Minimum turns for Tier 2 + #[serde(default = "default_tier2_min_turns")] + pub tier2_min_turns: i64, + /// Minimum tool calls for Tier 2 + #[serde(default = "default_tier2_min_tool_calls")] + pub tier2_min_tool_calls: i64, + /// Minimum compactions for Tier 2 + #[serde(default = "default_tier2_min_compactions")] + pub tier2_min_compactions: i64, + /// Minimum errors for Tier 2 + #[serde(default = "default_tier2_min_errors")] + pub tier2_min_errors: i64, + /// Feedback mode for Tier 2 + #[serde(default = "default_feedback_mode_thumbs_text")] + pub tier2_feedback_mode: String, + /// Whether Tier 2 feedback requests are dismissible + #[serde(default = "default_true")] + pub tier2_dismissible: bool, + /// Prompt text shown to users when Tier 2 feedback is requested + #[serde(default = "default_tier2_prompt")] + pub tier2_prompt: String, + /// Max times Tier 2 can trigger per session (0 = unlimited) + #[serde(default = "default_one")] + pub tier2_max_triggers: i32, + + // === Tier 3: Recovery from Friction === + /// Whether Tier 3 is enabled + #[serde(default = "default_true")] + pub tier3_enabled: bool, + /// Sample rate for Tier 3 (0.0-1.0) + #[serde(default = "default_tier3_sample_rate")] + pub tier3_sample_rate: f64, + /// Minimum turns for Tier 3 + #[serde(default = "default_tier3_min_turns")] + pub tier3_min_turns: i64, + /// Whether Tier 3 requires at least one cancellation + #[serde(default)] + pub tier3_requires_cancellation: bool, + /// Whether Tier 3 requires at least one revert + #[serde(default)] + pub tier3_requires_revert: bool, + /// Whether Tier 3 requires recovery (cancellation OR revert) + #[serde(default = "default_true")] + pub tier3_requires_recovery: bool, + /// Feedback mode for Tier 3 + #[serde(default = "default_feedback_mode_stars_text")] + pub tier3_feedback_mode: String, + /// Whether Tier 3 feedback requests are dismissible + #[serde(default = "default_true")] + pub tier3_dismissible: bool, + /// Prompt text shown to users when Tier 3 feedback is requested + #[serde(default = "default_tier3_prompt")] + pub tier3_prompt: String, + /// Max times Tier 3 can trigger per session (0 = unlimited) + #[serde(default = "default_one")] + pub tier3_max_triggers: i32, +} + +impl Default for FeedbackHeuristicsConfig { + fn default() -> Self { + Self { + config_id: "default".to_string(), + config_version: 1, + enabled: true, + cooldown_seconds: 300, + max_requests_per_session: 3, + // Tier 1 + tier1_enabled: true, + tier1_sample_rate: 0.0005, + tier1_min_turns: 10, + tier1_min_tool_calls: 5, + tier1_min_compactions: 2, + tier1_no_cancellations: true, + tier1_feedback_mode: "thumbs".to_string(), + tier1_dismissible: true, + tier1_prompt: default_tier1_prompt(), + tier1_max_triggers: 1, + // Tier 2 + tier2_enabled: true, + tier2_sample_rate: 0.0002, + tier2_min_turns: 15, + tier2_min_tool_calls: 10, + tier2_min_compactions: 3, + tier2_min_errors: 1, + tier2_feedback_mode: "thumbs_text".to_string(), + tier2_dismissible: true, + tier2_prompt: default_tier2_prompt(), + tier2_max_triggers: 1, + // Tier 3 + tier3_enabled: true, + tier3_sample_rate: 0.0001, + tier3_min_turns: 20, + tier3_requires_cancellation: false, + tier3_requires_revert: false, + tier3_requires_recovery: true, + tier3_feedback_mode: "stars_text".to_string(), + tier3_dismissible: true, + tier3_prompt: default_tier3_prompt(), + tier3_max_triggers: 1, + } + } +} + +impl FeedbackHeuristicsConfig { + /// Get the Tier 1 configuration as a TierConfig. + pub fn tier1_config(&self) -> TierConfig { + TierConfig { + enabled: self.tier1_enabled, + sample_rate: self.tier1_sample_rate, + min_turns: self.tier1_min_turns, + min_tool_calls: self.tier1_min_tool_calls, + min_compactions: self.tier1_min_compactions, + min_errors: 0, + no_cancellations: self.tier1_no_cancellations, + requires_cancellation: false, + requires_revert: false, + requires_recovery: false, + feedback_mode: parse_feedback_mode_str(&self.tier1_feedback_mode), + dismissible: self.tier1_dismissible, + prompt: self.tier1_prompt.clone(), + max_triggers: self.tier1_max_triggers, + } + } + + /// Get the Tier 2 configuration as a TierConfig. + pub fn tier2_config(&self) -> TierConfig { + TierConfig { + enabled: self.tier2_enabled, + sample_rate: self.tier2_sample_rate, + min_turns: self.tier2_min_turns, + min_tool_calls: self.tier2_min_tool_calls, + min_compactions: self.tier2_min_compactions, + min_errors: self.tier2_min_errors, + no_cancellations: false, + requires_cancellation: false, + requires_revert: false, + requires_recovery: false, + feedback_mode: parse_feedback_mode_str(&self.tier2_feedback_mode), + dismissible: self.tier2_dismissible, + prompt: self.tier2_prompt.clone(), + max_triggers: self.tier2_max_triggers, + } + } + + /// Get the Tier 3 configuration as a TierConfig. + pub fn tier3_config(&self) -> TierConfig { + TierConfig { + enabled: self.tier3_enabled, + sample_rate: self.tier3_sample_rate, + min_turns: self.tier3_min_turns, + min_tool_calls: 0, + min_compactions: 0, + min_errors: 0, + no_cancellations: false, + requires_cancellation: self.tier3_requires_cancellation, + requires_revert: self.tier3_requires_revert, + requires_recovery: self.tier3_requires_recovery, + feedback_mode: parse_feedback_mode_str(&self.tier3_feedback_mode), + dismissible: self.tier3_dismissible, + prompt: self.tier3_prompt.clone(), + max_triggers: self.tier3_max_triggers, + } + } +} + +// Default value functions for serde +fn default_true() -> bool { + true +} +fn default_cooldown_seconds() -> i64 { + 300 +} +fn default_max_requests() -> i64 { + 3 +} +fn default_tier1_sample_rate() -> f64 { + 0.0005 +} +fn default_tier1_min_turns() -> i64 { + 10 +} +fn default_tier1_min_tool_calls() -> i64 { + 5 +} +fn default_tier1_min_compactions() -> i64 { + 2 +} +fn default_tier2_sample_rate() -> f64 { + 0.0002 +} +fn default_tier2_min_turns() -> i64 { + 15 +} +fn default_tier2_min_tool_calls() -> i64 { + 10 +} +fn default_tier2_min_compactions() -> i64 { + 3 +} +fn default_tier2_min_errors() -> i64 { + 1 +} +fn default_tier3_sample_rate() -> f64 { + 0.0001 +} +fn default_tier3_min_turns() -> i64 { + 20 +} +fn default_feedback_mode_thumbs() -> String { + "thumbs".to_string() +} +fn default_feedback_mode_thumbs_text() -> String { + "thumbs_text".to_string() +} +fn default_feedback_mode_stars_text() -> String { + "stars_text".to_string() +} +fn default_tier1_prompt() -> String { + "You've been having a productive session! Would you mind sharing quick feedback?".to_string() +} +fn default_tier2_prompt() -> String { + "You've worked through a complex session. Your feedback would help us improve.".to_string() +} +fn default_tier3_prompt() -> String { + "Thanks for sticking with us through that session. Got a moment to share feedback?".to_string() +} +fn default_one() -> i32 { + 1 +} diff --git a/crates/codegen/kigi-shell/src/session/fork.rs b/crates/codegen/kigi-shell/src/session/fork.rs index af50e8f..bb6d0fd 100644 --- a/crates/codegen/kigi-shell/src/session/fork.rs +++ b/crates/codegen/kigi-shell/src/session/fork.rs @@ -3,9 +3,7 @@ //! Forks a saved session to a new working directory with a new session ID. //! This creates new session files but does not start the session. -use crate::remote::BackendClient; -const FORK_LOG: &str = "xai_fork"; -use crate::session::export::ExportedMetadata; +const FORK_LOG: &str = "kigi_fork"; use crate::session::info::Info; use crate::session::storage::{CopySessionOptions, JsonlStorageAdapter}; use crate::util::kigi_home::kigi_home; @@ -62,11 +60,7 @@ fn generate_fork_session_id(_source_id: &str) -> String { } /// Fork a saved session to a new working directory. -pub async fn fork_session( - request: ForkSessionRequest, - agent_id: &str, - auth_manager: Option>, -) -> io::Result { +pub async fn fork_session(request: ForkSessionRequest) -> io::Result { let t0 = std::time::Instant::now(); let root_dir = kigi_home(); @@ -114,32 +108,6 @@ pub async fn fork_session( let copy_ms = t0.elapsed().as_millis() as u64; - // Writeback session to backend (fire-and-forget). - // This is telemetry-grade: the local fork works without it. All fork - // state lives locally (session files on disk), and the caller does not - // depend on synchronous backend registration. The backend eventually - // learns about the session when the background task completes. - // Spawning removes the network round-trip (~200-400ms) from the - // critical path. - if let Some(am) = auth_manager { - let sid = new_session_id.clone(); - let cwd = request.new_cwd.clone(); - let parent = request.source_session_id.clone(); - let model = request.new_model_id.clone(); - let aid = agent_id.to_string(); - tokio::spawn(async move { - if let Err(e) = - sync_forked_session_to_backend(&sid, &cwd, parent, model, &aid, am).await - { - tracing::warn!( - session_id = %sid, - error = %e, - "Failed to register forked session with backend (background)" - ); - } - }); - } - let total_ms = t0.elapsed().as_millis() as u64; tracing::info!( target: FORK_LOG, @@ -149,7 +117,7 @@ pub async fn fork_session( total_ms, chat_copied = result.chat_messages_copied, updates_copied = result.updates_copied, - "FORK_COPY: session data copied (backend sync spawned in background)" + "FORK_COPY: session data copied" ); Ok(ForkSessionResponse { @@ -163,43 +131,6 @@ pub async fn fork_session( }) } -/// Sync a forked session to the backend (for writeback mode). -async fn sync_forked_session_to_backend( - session_id: &str, - cwd: &str, - parent_session_id: String, - model_id: Option, - agent_id: &str, - auth_manager: std::sync::Arc, -) -> Result<(), Box> { - let client = BackendClient::new().with_auth_manager(auth_manager); - let metadata = ExportedMetadata { - title: None, // Will be generated later when session runs - cwd: cwd.to_string(), - model_id, - created_at: Some(chrono::Utc::now().to_rfc3339()), - updated_at: Some(chrono::Utc::now().to_rfc3339()), - total_messages: Some(0), - parent_session_id: Some(parent_session_id), - session_kind: None, - subagent_type: None, - subagent_persona: None, - subagent_role: None, - fork_context_source: None, - subagent_depth: None, - }; - - client - .upsert_session(session_id, &metadata, agent_id) - .await?; - tracing::info!( - session_id = %session_id, - "Forked session registered with backend" - ); - - Ok(()) -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/codegen/kigi-shell/src/session/helpers/session_compact.rs b/crates/codegen/kigi-shell/src/session/helpers/session_compact.rs index d323b60..56db105 100644 --- a/crates/codegen/kigi-shell/src/session/helpers/session_compact.rs +++ b/crates/codegen/kigi-shell/src/session/helpers/session_compact.rs @@ -743,7 +743,6 @@ mod classify_tests { message: "test".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, })) }; assert!(det(StatusCode::BAD_REQUEST)); @@ -836,7 +835,6 @@ mod classify_tests { .into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }))); } #[test] @@ -866,7 +864,6 @@ mod classify_tests { message: "bad payload".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }) else { panic!("expected Deterministic for 400"); }; @@ -878,7 +875,6 @@ mod classify_tests { message: "upstream blip".into(), model_metadata: None, retry_after_secs: None, - should_retry: None, }) else { panic!("expected Transient for 500"); }; @@ -1595,15 +1591,11 @@ mod reasoning_compaction_regression_tests { auth_scheme: Default::default(), extra_headers: Default::default(), context_window: 256_000, - client_version: None, force_http1: false, max_retries: None, stream_tool_calls: false, idle_timeout_secs: None, - client_identifier: None, reasoning_effort: None, - deployment_id: None, - user_id: None, origin_client: None, attribution_callback: None, bearer_resolver: None, diff --git a/crates/codegen/kigi-shell/src/session/mod.rs b/crates/codegen/kigi-shell/src/session/mod.rs index 8e58fea..398434e 100644 --- a/crates/codegen/kigi-shell/src/session/mod.rs +++ b/crates/codegen/kigi-shell/src/session/mod.rs @@ -12,6 +12,7 @@ pub mod two_pass; pub use self::acp_session::*; pub use self::acp_types::*; pub use self::commands::*; +pub use self::feedback_types::{ClientType, FeedbackTerminalInfo, RatingType}; pub use self::fork::{ForkSessionRequest, ForkSessionResponse, fork_session}; pub use self::handle::*; pub use self::persistence::{ @@ -19,13 +20,9 @@ pub use self::persistence::{ resolve_local_session_any_cwd, session_exists_by_id, session_exists_for_cwd, }; pub use self::result::{Empty, ExtMethodResult}; -pub use self::share::{ShareSessionRequest, ShareSessionResponse}; pub use kigi_fsnotify::{ FsConfig, FsEvent, FsEventKind, FsEventSource, FsNotifyError, GitMetaKind, }; -pub use prod_mc_cli_chat_proxy_types::feedback_types::{ - ClientType, FeedbackTerminalInfo, RatingType, -}; /// `false` twin: this template is not compiled into this build, so no /// template matches. Keeps ungated call sites compiling in both /// configurations. @@ -273,18 +270,6 @@ pub struct ClientFsConfig { pub mode: ClientFsMode, } /// Share session request/response types -pub mod share { - /// Request to share a session via URL - #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] - pub struct ShareSessionRequest { - pub session_id: String, - } - /// Response containing the shareable URL - #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] - pub struct ShareSessionResponse { - pub share_url: String, - } -} /// Proxy config for the session registry client. /// Shared between `acp_session` (slash commands) and `persistence` (title generation). #[derive(Clone)] @@ -303,6 +288,7 @@ pub(crate) mod events; pub mod export; pub mod feedback; pub mod feedback_manager; +pub mod feedback_types; pub mod file_system; pub mod fork; pub(crate) mod fs_watch; diff --git a/crates/codegen/kigi-shell/src/session/persistence.rs b/crates/codegen/kigi-shell/src/session/persistence.rs index 62896cc..b2209e4 100644 --- a/crates/codegen/kigi-shell/src/session/persistence.rs +++ b/crates/codegen/kigi-shell/src/session/persistence.rs @@ -5,8 +5,6 @@ use std::sync::Arc; use crate::config::StorageMode; -use crate::remote::RemoteSync; - use crate::sampling::Client as OaiCompatClient; use crate::sampling::ConversationItem; use crate::session::export::ExportedMetadata; @@ -105,7 +103,7 @@ pub struct UserFeedbackEntry { pub dismissed: bool, /// The full submission payload (omitted when dismissed) #[serde(skip_serializing_if = "Option::is_none")] - pub submission: Option, + pub submission: Option, } /// Helper for `#[serde(skip_serializing_if)]` on bool fields. @@ -116,14 +114,13 @@ pub(crate) fn is_false(v: &bool) -> bool { #[cfg(test)] mod feedback_tests { use super::*; - use prod_mc_cli_chat_proxy_types::feedback_types::{ + use crate::session::feedback_types::{ ClientType, FeedbackSubmission, FeedbackType, RatingType, }; fn make_submission(thumbs_up: bool) -> FeedbackSubmission { FeedbackSubmission { session_id: "session-abc".into(), - user_id: None, client_type: ClientType::Tui, feedback_type: if thumbs_up { FeedbackType::Rating @@ -139,22 +136,13 @@ mod feedback_tests { Some("could be better".into()) }, feedback_categories: vec![], - message_id: None, model_id: Some("grok-3-fast".into()), resolved_model_id: Some("grok-4.5".into()), model_fingerprint: None, context_type: None, - feature_name: None, - tool_name: None, - experiment_id: None, - comparison_id: None, - preferred_model_id: None, - preference_strength: None, - preference_reasons: vec![], request_id: None, client_version: None, shell_version: None, - extension_host: None, metadata: None, last_user_message: None, last_assistant_message: None, @@ -165,7 +153,6 @@ mod feedback_tests { context_tokens_used: None, context_window_tokens: None, terminal_info: None, - unified_log_url: None, } } @@ -1392,7 +1379,6 @@ struct SessionPersistence { /// Pending ACP notification for merging consecutive text chunks pending_notification: Option, rx: mpsc::UnboundedReceiver, - remote_sync: Option, /// Session title generation lifecycle. summary: crate::session::summary::SummaryGenerator, registry_title_sync: Option, @@ -1497,20 +1483,12 @@ impl SessionPersistence { } } - /// Flush any pending merged ACP notification to disk and remote sync. + /// Flush any pending merged ACP notification to disk. async fn flush_pending(&mut self) { // Write any pending merged ACP notification if let Some(notification) = self.pending_notification.take() { - self.write_update(&SessionUpdate::Acp(Box::new(notification.clone()))) + self.write_update(&SessionUpdate::Acp(Box::new(notification))) .await; - // HTTP-based remote sync (Writeback mode) - if let Some(sync) = &self.remote_sync { - sync.queue(notification); - } - } - // Flush HTTP sync - if let Some(sync) = &self.remote_sync { - sync.flush(); } } @@ -1549,12 +1527,8 @@ impl SessionPersistence { SessionUpdate::Acp(notification) => { // ACP notifications use merging to coalesce consecutive text chunks if let Some(to_write) = self.maybe_merge_notification(¬ification) { - self.write_update(&SessionUpdate::Acp(Box::new(to_write.clone()))) + self.write_update(&SessionUpdate::Acp(Box::new(to_write))) .await; - // HTTP-based remote sync (Writeback mode) - if let Some(sync) = &self.remote_sync { - sync.queue(to_write); - } } } SessionUpdate::Xai(_) => { @@ -1602,9 +1576,6 @@ impl SessionPersistence { { tracing::warn!(?e, "failed to update current model"); } - if let Some(sync) = &self.remote_sync { - sync.set_model_id(model_id.0.to_string()); - } } PersistenceMsg::PlanState(state) => { if let Err(e) = self.storage.write_plan_state(&self.info, &state).await { @@ -1660,9 +1631,6 @@ impl SessionPersistence { &self.info, &title, ); - if let Some(sync) = &self.remote_sync { - sync.set_title(title.clone()); - } if let Some(reg) = self.registry_title_sync.as_ref() && !reg.suppress_for_zdr { @@ -1901,69 +1869,6 @@ fn collect_session_files_recursive(base: &Path, dir: &Path, files: &mut Vec>, -) -> io::Result> { - match storage_mode { - StorageMode::Local => Ok(None), - StorageMode::Writeback => { - let auth_manager = auth_manager.ok_or_else(|| { - io::Error::new( - io::ErrorKind::PermissionDenied, - "Writeback storage mode requires authentication. Run 'grok login' first.", - ) - })?; - if auth_manager.current_or_expired().is_some() { - // ZDR was an xAI team concept; nothing gates remote sync here. - } else { - tracing::warn!( - "writeback: no auth loaded yet, ZDR check skipped (backend enforces server-side)" - ); - } - tracing::info!("Writeback mode enabled, syncing to backend"); - let client = - crate::remote::BackendClient::new().with_auth_manager(auth_manager.clone()); - let metadata = ExportedMetadata::from_summary(summary); - Ok(Some(RemoteSync::new( - summary.info.id.to_string(), - metadata, - client, - ))) - } - } -} - -/// Pull a session from the backend if not found locally. Returns the pulled -/// session's [`Info`] (cwd may differ from caller's on different machines), -/// or `None` if not found or on error. -async fn try_pull_from_remote(info: &Info, client: &crate::remote::BackendClient) -> Option { - // BackendClient resolves auth internally via its auth_manager. - client.auth_manager.as_ref()?; - - tracing::info!(session_id = %info.id, "Session not found locally, trying backend"); - - match crate::remote::pull_session_to_local(&info.id.0, client).await { - Ok(crate::remote::PullResult::Hydrated(pulled_info)) => { - tracing::info!( - session_id = %info.id, - pulled_cwd = %pulled_info.cwd, - "Pulled session from backend" - ); - Some(pulled_info) - } - Ok(crate::remote::PullResult::NotFound) => { - tracing::debug!(session_id = %info.id, "Session not found on backend either"); - None - } - Err(e) => { - tracing::warn!(session_id = %info.id, error = %e, "Backend pull failed"); - None - } - } -} - /// Map a persistence `io::Error` into an `acp::Error` with a human-friendly /// `message` and a stable `data.code` for log aggregation. pub(crate) fn io_error_to_acp(e: &io::Error) -> acp::Error { @@ -2066,7 +1971,6 @@ pub(crate) async fn new( let info_clone = info.clone(); let storage: Arc = Arc::from(storage); - let remote_sync = init_remote_sync(&summary, storage_mode, auth_manager)?; let handle = PersistenceHandle { tx: tx.clone(), noop: false, @@ -2078,7 +1982,6 @@ pub(crate) async fn new( storage: storage.clone(), pending_notification: None, rx, - remote_sync: remote_sync.clone(), summary: crate::session::summary::SummaryGenerator::new( crate::session::summary::SummaryConfig { sampling_client, @@ -2147,7 +2050,6 @@ pub async fn new_with_explicit_dir( storage: storage.clone(), pending_notification: None, rx, - remote_sync: None, summary: crate::session::summary::SummaryGenerator::new( crate::session::summary::SummaryConfig { sampling_client, @@ -2195,101 +2097,12 @@ pub struct PersistedInfoLight { pub goal_mode_state: Option, } -/// On NotFound, try pulling from backend. Returns pulled info or the original error. -async fn pull_on_miss( - info: &Info, - client: &crate::remote::BackendClient, - err: io::Error, -) -> io::Result { - if err.kind() != io::ErrorKind::NotFound { - return Err(err); - } - try_pull_from_remote(info, client).await.ok_or(err) -} - -#[expect(dead_code, reason = "wired when session restore flow calls load")] -pub(crate) async fn load( - info: &Info, - sampling_client: OaiCompatClient, - storage_mode: StorageMode, - auth_manager: Option>, - backend: Option<&crate::remote::BackendClient>, - gateway: Option, - session_summary_model: String, - registry_title_sync: Option, -) -> io::Result<(PersistedInfo, PersistenceHandle)> { - let root_dir = kigi_home(); - let storage: Box = Box::new(JsonlStorageAdapter::with_root(root_dir)); - - let (persisted, loaded_info) = match storage.load_session(info).await { - Ok(p) => (p, info.clone()), - Err(e) => match backend { - Some(client) => { - let pulled = pull_on_miss(info, client, e).await?; - let p = storage.load_session(&pulled).await?; - (p, pulled) - } - None => return Err(e), - }, - }; - // Touch on load too: resuming must reset the worktree's gc expiry clock. - touch_worktree_for_session(&loaded_info).await; - - let persisted_info = PersistedInfo { - summary: persisted.summary, - chat_history: persisted.chat_history, - updates: persisted.updates, - plan_state: persisted.plan_state, - rewind_points: persisted.rewind_points, - signals: persisted.signals, - }; - - let (tx, rx) = mpsc::unbounded_channel::(); - - let storage: Arc = Arc::from(storage); - let remote_sync = init_remote_sync(&persisted_info.summary, storage_mode, auth_manager)?; - - let has_title = !persisted_info.summary.display_title().is_empty(); - let handle = PersistenceHandle { - tx: tx.clone(), - noop: false, - }; - tokio::task::spawn(async move { - let mut summary_gen = crate::session::summary::SummaryGenerator::new( - crate::session::summary::SummaryConfig { - sampling_client, - model: session_summary_model, - persistence_tx: tx, - }, - ); - if has_title { - summary_gen.mark_done(); - } - let persistence = SessionPersistence { - info: loaded_info, - storage: storage.clone(), - pending_notification: None, - rx, - remote_sync: remote_sync.clone(), - summary: summary_gen, - registry_title_sync, - gateway, - }; - persistence.run().await; - }); - - Ok((persisted_info, handle)) -} - -/// Like `load`, but doesn't load updates into memory. +/// Loads a session for streaming updates without reading them into memory. /// Instead, provides the path to the updates file for streaming reads. /// Use this for memory-efficient session loading when replaying updates. pub(crate) async fn load_light( info: &Info, sampling_client: OaiCompatClient, - storage_mode: StorageMode, - auth_manager: Option>, - backend: Option<&crate::remote::BackendClient>, gateway: Option, session_summary_model: String, registry_title_sync: Option, @@ -2298,16 +2111,9 @@ pub(crate) async fn load_light( let storage: Box = Box::new(JsonlStorageAdapter::with_root(root_dir.clone())); - let (persisted, loaded_info) = match storage.load_session_without_updates(info).await { - Ok(p) => (p, info.clone()), - Err(e) => match backend { - Some(client) => { - let pulled = pull_on_miss(info, client, e).await?; - let p = storage.load_session_without_updates(&pulled).await?; - (p, pulled) - } - None => return Err(e), - }, + let (persisted, loaded_info) = { + let p = storage.load_session_without_updates(info).await?; + (p, info.clone()) }; // Touch on load too: resuming must reset the worktree's gc expiry clock. touch_worktree_for_session(&loaded_info).await; @@ -2330,7 +2136,6 @@ pub(crate) async fn load_light( let (tx, rx) = mpsc::unbounded_channel::(); let storage: Arc = Arc::from(storage); - let remote_sync = init_remote_sync(&persisted_info.summary, storage_mode, auth_manager)?; let has_title = !persisted_info.summary.display_title().is_empty(); let handle = PersistenceHandle { @@ -2353,7 +2158,6 @@ pub(crate) async fn load_light( storage: storage.clone(), pending_notification: None, rx, - remote_sync: remote_sync.clone(), summary: summary_gen, registry_title_sync, gateway, @@ -2373,100 +2177,58 @@ pub async fn list_summaries(cwd: Option<&str>) -> io::Result> { } /// Failure modes of [`delete_session_history`]. -/// -/// Kept distinct so callers can surface a precise message: a remote -/// failure is reported separately from a local-disk failure because the -/// remote delete runs first and aborts the whole operation (see the doc -/// on [`delete_session_history`]). #[derive(Debug, thiserror::Error)] pub enum DeleteSessionError { /// Listing local summaries (to resolve the on-disk session dir) failed. #[error("failed to list sessions: {0}")] List(#[source] io::Error), - /// The remote (writeback) copy could not be deleted; local bits were - /// left untouched so the operation can be retried. - #[error("failed to delete remote session data: {0}")] - Remote(#[source] crate::remote::client::BackendError), /// The local on-disk session directory could not be removed. #[error("failed to delete session: {0}")] Local(#[source] io::Error), } -/// Where a session copy was actually removed by [`delete_session_history`]. +/// Whether a session copy was removed by [`delete_session_history`]. /// -/// Both fields are `false` when nothing existed to delete (still a +/// `local_removed` is `false` when nothing existed to delete (still a /// success). Callers use [`Self::any_removed`] to decide between a -/// "deleted" and a "not found" message without conflating a remote-only -/// delete with a no-op. +/// "deleted" and a "not found" message. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct SessionDeletion { /// A local on-disk session directory was found and removed. pub local_removed: bool, - /// A remote (writeback) copy was found and removed. `false` when - /// `needs_remote` was not set, or the remote copy was already absent - /// (the backend returned `404`). - pub remote_removed: bool, } impl SessionDeletion { - /// `true` when a copy was removed from at least one location. + /// `true` when the local session directory was removed. pub fn any_removed(self) -> bool { - self.local_removed || self.remote_removed + self.local_removed } } -/// Permanently delete a session's history: the remote (writeback) copy -/// when `needs_remote`, the local on-disk session directory, and the -/// FTS search-index entry. +/// Permanently delete a session's history: the local on-disk session +/// directory and the FTS search-index entry. /// -/// Idempotent: a session that is missing locally (e.g. remote-only) -/// still succeeds, and a remote `404` (copy already gone) is treated as -/// success rather than an error. When `needs_remote` is set the remote -/// delete runs *first* and is authoritative — only on its success (or a -/// `404`) are the local bits removed. This ordering prevents a partial -/// delete where the local copy is nuked but the remote copy lingers and -/// re-appears on the next session list. +/// Idempotent: a session that is missing locally still succeeds. /// -/// Returns a [`SessionDeletion`] recording which copies (local / remote) -/// were actually removed; both fields `false` means nothing existed -/// (still `Ok`). +/// Returns a [`SessionDeletion`] recording whether a local copy was +/// removed; `false` means nothing existed (still `Ok`). pub async fn delete_session_history( session_id: &str, cwd: Option<&str>, - needs_remote: bool, - auth_manager: Arc, ) -> Result { let sid = acp::SessionId::new(Arc::from(session_id)); - // Resolve the local session info, scoping to cwd if provided. A - // remote-only session won't be found here — that's fine, the remote - // delete (if applicable) still runs. + // Resolve the local session info, scoping to cwd if provided. let summaries = list_summaries(cwd) .await .map_err(DeleteSessionError::List)?; - let local_info = summaries + let Some(info) = summaries .iter() .find(|s| s.info.id == sid) - .map(|s| s.info.clone()); - - // Remote delete first (authoritative for cloud history). A genuine - // failure aborts before any local mutation so the row does not - // reappear; a `404` means the copy is already gone, so deletion stays - // idempotent and falls through to local cleanup. - let remote_removed = if needs_remote { - let result = crate::remote::client::BackendClient::new() - .with_auth_manager(auth_manager) - .delete_session_data(session_id) - .await; - classify_remote_delete(result)? - } else { - false - }; - - let Some(info) = local_info else { + .map(|s| s.info.clone()) + else { return Ok(SessionDeletion { local_removed: false, - remote_removed, }); }; @@ -2481,85 +2243,22 @@ pub async fn delete_session_history( Ok(SessionDeletion { local_removed: true, - remote_removed, }) } -/// Classify a remote `delete_session_data` result, reporting whether a -/// remote copy was actually removed: a `2xx` means a copy was deleted -/// (`Ok(true)`), a `404` means it was already gone so deletion stays -/// idempotent (`Ok(false)`), and any other backend error aborts the -/// delete (`Err`) so local bits are left untouched and it can be retried. -fn classify_remote_delete( - result: Result<(), crate::remote::client::BackendError>, -) -> Result { - use crate::remote::client::BackendError; - match result { - Ok(()) => Ok(true), - Err(BackendError::RequestFailed { status: 404, .. }) => Ok(false), - Err(e) => Err(DeleteSessionError::Remote(e)), - } -} - #[cfg(test)] mod delete_session_history_tests { - use super::{DeleteSessionError, SessionDeletion, classify_remote_delete}; - use crate::remote::client::BackendError; + use super::SessionDeletion; #[test] - fn remote_ok_reports_removed() { - assert!( - classify_remote_delete(Ok(())).unwrap(), - "a 2xx delete must report that a remote copy was removed" - ); - } - - #[test] - fn remote_404_is_treated_as_already_deleted() { - let removed = classify_remote_delete(Err(BackendError::RequestFailed { - status: 404, - body: "not found".into(), - })) - .expect("a 404 means the remote copy is gone — deletion must stay idempotent"); - assert!( - !removed, - "a 404 must report that nothing was removed remotely" - ); - } - - #[test] - fn remote_non_404_request_failure_aborts() { - let res = classify_remote_delete(Err(BackendError::RequestFailed { - status: 500, - body: "boom".into(), - })); - assert!(matches!(res, Err(DeleteSessionError::Remote(_)))); - } - - #[test] - fn remote_auth_failure_aborts() { - let res = classify_remote_delete(Err(BackendError::Auth("denied".into()))); - assert!(matches!(res, Err(DeleteSessionError::Remote(_)))); - } - - #[test] - fn any_removed_reflects_either_location() { + fn any_removed_reflects_local_removal() { assert!(!SessionDeletion::default().any_removed()); assert!( SessionDeletion { local_removed: true, - remote_removed: false, } .any_removed() ); - assert!( - SessionDeletion { - local_removed: false, - remote_removed: true, - } - .any_removed(), - "a remote-only delete must count as removed" - ); } } diff --git a/crates/codegen/kigi-shell/src/session/storage/jsonl/tests.rs b/crates/codegen/kigi-shell/src/session/storage/jsonl/tests.rs index 722a700..a69723f 100644 --- a/crates/codegen/kigi-shell/src/session/storage/jsonl/tests.rs +++ b/crates/codegen/kigi-shell/src/session/storage/jsonl/tests.rs @@ -1085,7 +1085,7 @@ async fn test_load_prompts_only_large_session() { #[tokio::test] async fn test_append_feedback_creates_file_and_persists() { use crate::session::persistence::{LocalFeedbackEntry, UserFeedbackEntry}; - use prod_mc_cli_chat_proxy_types::feedback_types::{ + use crate::session::feedback_types::{ ClientType, FeedbackSubmission, FeedbackType, RatingType, }; let temp_dir = TempDir::new().unwrap(); @@ -1101,7 +1101,6 @@ async fn test_append_feedback_creates_file_and_persists() { dismissed: false, submission: Some(FeedbackSubmission { session_id: "test-session-123".into(), - user_id: None, client_type: ClientType::Tui, feedback_type: FeedbackType::Rating, turn_number: Some(3), @@ -1109,22 +1108,13 @@ async fn test_append_feedback_creates_file_and_persists() { rating_value: Some(1), feedback_text: None, feedback_categories: vec![], - message_id: None, model_id: Some("grok-3-fast".into()), resolved_model_id: Some("grok-4.5".into()), model_fingerprint: None, context_type: None, - feature_name: None, - tool_name: None, - experiment_id: None, - comparison_id: None, - preferred_model_id: None, - preference_strength: None, - preference_reasons: vec![], request_id: None, client_version: None, shell_version: None, - extension_host: None, metadata: None, last_user_message: None, last_assistant_message: None, @@ -1135,7 +1125,6 @@ async fn test_append_feedback_creates_file_and_persists() { context_tokens_used: None, context_window_tokens: None, terminal_info: None, - unified_log_url: None, }), }); adapter.append_feedback(&info, &user_entry).await.unwrap(); diff --git a/crates/codegen/kigi-shell/src/session/unified_list/cursor.rs b/crates/codegen/kigi-shell/src/session/unified_list/cursor.rs index 41b745a..2722f65 100644 --- a/crates/codegen/kigi-shell/src/session/unified_list/cursor.rs +++ b/crates/codegen/kigi-shell/src/session/unified_list/cursor.rs @@ -3,7 +3,6 @@ use std::cmp::{Ordering, Reverse}; use base64::Engine as _; use serde::{Deserialize, Serialize}; -use super::PartialReason; use super::envelope::SessionKind; use super::row::UnifiedRow; @@ -11,10 +10,6 @@ use super::row::UnifiedRow; pub(super) struct CompositeCursor { #[serde(default, skip_serializing_if = "Option::is_none")] pub boundary: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub conv_page_token: Option, - #[serde(default)] - pub conv_page_drained: bool, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -42,52 +37,21 @@ impl CompositeCursor { } } -pub(super) enum ConvLane { - Skipped, - Degraded(PartialReason), - Page { - rows: Vec, - next_token: Option, - frontier: Option, - }, -} - -pub(super) fn conv_frontier(raw_rows: &[UnifiedRow], has_more: bool) -> Option { - if !has_more { - return None; - } - raw_rows - .iter() - .max_by(|a, b| cmp_total_order(a, b)) - .map(boundary_of) -} - pub(super) struct Paginated { pub candidates: Vec, pub emit_count: usize, pub next_cursor: Option, - pub partial: Option, } -pub(super) fn merge_and_paginate( +/// Sort local rows newest-first, resume after the cursor boundary, and cut +/// one page. `next_cursor` is set only when rows remain past the page. +pub(super) fn paginate( local: Vec, - conv: ConvLane, cursor: &CompositeCursor, limit: usize, ) -> Paginated { - let (conv_rows, conv_next_token, conv_fetched, conv_frontier, partial) = match conv { - ConvLane::Skipped => (Vec::new(), None, false, None, None), - ConvLane::Degraded(reason) => (Vec::new(), None, false, None, Some(reason)), - ConvLane::Page { - rows, - next_token, - frontier, - } => (rows, next_token, true, frontier, None), - }; - let mut keyed: Vec<(SortKey, UnifiedRow)> = local .into_iter() - .chain(conv_rows) .map(|row| (row_sort_key(&row), row)) .collect(); @@ -98,46 +62,12 @@ pub(super) fn merge_and_paginate( keyed.sort_by(|(a, _), (b, _)| a.cmp(b)); - let mut emit_count = keyed.len().min(limit); - if let Some(frontier) = &conv_frontier { - let fkey = boundary_sort_key(frontier); - let frontier_count = keyed - .iter() - .take_while(|(k, _)| k.cmp(&fkey) != Ordering::Greater) - .count(); - emit_count = emit_count.min(frontier_count); - } + let emit_count = keyed.len().min(limit); let new_boundary = (emit_count > 0).then(|| boundary_of(&keyed[emit_count - 1].1)); + let has_more = keyed.len() > emit_count; - let tail = &keyed[emit_count..]; - let local_has_more = tail.iter().any(|(_, r)| r.kind == SessionKind::Build); - let conv_in_tail = tail.iter().any(|(_, r)| r.kind == SessionKind::Chat); - - let (next_conv_token, next_conv_drained, conv_has_more) = if conv_fetched { - if conv_in_tail { - (cursor.conv_page_token.clone(), false, true) - } else { - let has_more = conv_next_token.is_some(); - (conv_next_token, true, has_more) - } - } else if partial.is_some() && cursor.conv_page_token.is_some() && new_boundary.is_some() { - ( - cursor.conv_page_token.clone(), - cursor.conv_page_drained, - true, - ) - } else { - ( - cursor.conv_page_token.clone(), - cursor.conv_page_drained, - false, - ) - }; - - let next_cursor = (local_has_more || conv_has_more).then(|| CompositeCursor { + let next_cursor = has_more.then(|| CompositeCursor { boundary: new_boundary.or_else(|| cursor.boundary.clone()), - conv_page_token: next_conv_token, - conv_page_drained: next_conv_drained, }); let candidates: Vec = keyed.into_iter().map(|(_, row)| row).collect(); @@ -146,7 +76,6 @@ pub(super) fn merge_and_paginate( candidates, emit_count, next_cursor, - partial, } } @@ -205,12 +134,8 @@ pub(super) fn cmp_total_order(a: &UnifiedRow, b: &UnifiedRow) -> Ordering { #[cfg(test)] mod tests { use super::*; - use crate::remote::Conversation; use crate::session::merge::MergedSession; - use crate::session::unified_list::{ - conversation_to_row, facet_registry, merged_session_to_row, - }; - use std::collections::BTreeSet; + use crate::session::unified_list::{facet_registry, merged_session_to_row}; fn local(id: &str, ts: &str) -> UnifiedRow { let m = MergedSession { @@ -236,336 +161,85 @@ mod tests { merged_session_to_row(m, facet_registry()) } - fn conv(id: &str, ts: &str) -> UnifiedRow { - let c = Conversation { - conversation_id: id.into(), - title: "t".into(), - modify_time: Some(ts.into()), - ..Conversation::default() - }; - conversation_to_row(c, facet_registry()) - } - - struct ConvSource { - rows: Vec, - page_size: usize, - } - - impl ConvSource { - fn new(mut rows: Vec, page_size: usize) -> Self { - rows.sort_by(cmp_total_order); - Self { rows, page_size } - } - - fn page(&self, token: Option<&str>) -> ConvLane { - if self.rows.is_empty() { - return ConvLane::Page { - rows: Vec::new(), - next_token: None, - frontier: None, - }; - } - let idx = token - .and_then(|t| t.strip_prefix('p')) - .and_then(|n| n.parse::().ok()) - .unwrap_or(0); - let start = idx * self.page_size; - let end = (start + self.page_size).min(self.rows.len()); - let rows = self.rows.get(start..end).unwrap_or(&[]).to_vec(); - let next_token = (end < self.rows.len()).then(|| format!("p{}", idx + 1)); - let frontier = conv_frontier(&rows, next_token.is_some()); - ConvLane::Page { - rows, - next_token, - frontier, - } - } - } - - fn walk_all(local_window: &[UnifiedRow], conv: &ConvSource, limit: usize) -> Vec { - let mut cursor = CompositeCursor::default(); - let mut emitted: Vec = Vec::new(); - for _ in 0..1000 { - let lane = conv.page(cursor.conv_page_token.as_deref()); - let result = merge_and_paginate(local_window.to_vec(), lane, &cursor, limit); - emitted.extend( - result.candidates[..result.emit_count] - .iter() - .map(|r| r.legacy.session_id.clone()), - ); - match result.next_cursor { - Some(c) => cursor = c, - None => return emitted, - } - } - panic!("pagination did not terminate"); - } - - fn ids(rows: &[UnifiedRow]) -> Vec { - rows.iter().map(|r| r.legacy.session_id.clone()).collect() + fn ids(p: &Paginated) -> Vec { + p.candidates[..p.emit_count] + .iter() + .map(|r| r.legacy.session_id.clone()) + .collect() } #[test] - fn cursor_round_trips() { - let cur = CompositeCursor { + fn cursor_roundtrip_boundary_only() { + let c = CompositeCursor { boundary: Some(BoundaryKey { - updated_at: "2026-06-01T00:00:00Z".into(), - kind: SessionKind::Chat, - session_id: "conv_1".into(), - }), - conv_page_token: Some("p3".into()), - conv_page_drained: true, - }; - let decoded = CompositeCursor::decode(Some(&cur.encode())); - assert_eq!(decoded.conv_page_token.as_deref(), Some("p3")); - assert!(decoded.conv_page_drained); - let b = decoded.boundary.unwrap(); - assert_eq!(b.session_id, "conv_1"); - assert_eq!(b.kind, SessionKind::Chat); - } - - #[test] - fn malformed_cursor_decodes_to_fresh_first_page() { - for bad in [Some("not base64 !!!"), Some(""), None] { - let c = CompositeCursor::decode(bad); - assert!(c.boundary.is_none()); - assert!(c.conv_page_token.is_none()); - assert!(!c.conv_page_drained); - } - } - - #[test] - fn multi_page_walk_equals_single_fetch_window() { - let local_window = vec![ - local("l1", "2026-06-10T00:00:00Z"), - local("l2", "2026-06-08T00:00:00Z"), - local("l3", "2026-06-04T00:00:00Z"), - local("l4", "2026-05-30T00:00:00Z"), - ]; - let conv_rows = vec![ - conv("c1", "2026-06-09T00:00:00Z"), - conv("c2", "2026-06-07T00:00:00Z"), - conv("c3", "2026-06-06T00:00:00Z"), - conv("c4", "2026-06-03T00:00:00Z"), - conv("c5", "2026-05-29T00:00:00Z"), - ]; - - let mut expected_all = local_window.clone(); - expected_all.extend(conv_rows.clone()); - expected_all.sort_by(cmp_total_order); - let expected_ids = ids(&expected_all); - - for &limit in &[1usize, 2, 3, 5, 7, 100] { - for &page_size in &[1usize, 2, 3] { - let source = ConvSource::new(conv_rows.clone(), page_size); - let got = walk_all(&local_window, &source, limit); - let unique: BTreeSet<&String> = got.iter().collect(); - assert_eq!( - unique.len(), - got.len(), - "duplicate emitted (limit={limit}, page_size={page_size}): {got:?}" - ); - assert_eq!( - got, expected_ids, - "walk != single fetch (limit={limit}, page_size={page_size})" - ); - } - } - } - - #[test] - fn equal_updated_at_tie_break_no_drop_or_dup() { - let ts = "2026-06-01T00:00:00Z"; - let local_window = vec![local("l_same", ts), local("l_old", "2026-05-01T00:00:00Z")]; - let conv_rows = vec![conv("c_same", ts), conv("c_old", "2026-05-15T00:00:00Z")]; - - let mut expected_all = local_window.clone(); - expected_all.extend(conv_rows.clone()); - expected_all.sort_by(cmp_total_order); - let expected_ids = ids(&expected_all); - assert_eq!(expected_ids[0], "l_same"); - assert_eq!(expected_ids[1], "c_same"); - - for &limit in &[1usize, 2, 3] { - let source = ConvSource::new(conv_rows.clone(), 1); - let got = walk_all(&local_window, &source, limit); - let unique: BTreeSet<&String> = got.iter().collect(); - assert_eq!(unique.len(), got.len(), "dup at limit={limit}: {got:?}"); - assert_eq!( - got, expected_ids, - "tie-break walk mismatch at limit={limit}" - ); - } - } - - #[test] - fn partial_conv_page_is_not_advanced_until_drained() { - let local_window = vec![ - local("l1", "2026-06-10T00:00:00Z"), - local("l2", "2026-06-08T00:00:00Z"), - ]; - let conv_rows = vec![ - conv("c1", "2026-06-09T00:00:00Z"), - conv("c2", "2026-06-07T00:00:00Z"), - ]; - let source = ConvSource::new(conv_rows.clone(), 2); - let got = walk_all(&local_window, &source, 1); - - let mut expected_all = local_window.clone(); - expected_all.extend(conv_rows.clone()); - expected_all.sort_by(cmp_total_order); - assert_eq!(got, ids(&expected_all)); - } - - #[test] - fn whole_page_filtered_out_does_not_drop_later_match() { - let local_window = vec![ - local("l1", "2026-06-10T00:00:00Z"), - local("l2", "2026-06-01T00:00:00Z"), - ]; - let raw = vec![ - conv("c1_drop", "2026-06-09T00:00:00Z"), - conv("c2_drop", "2026-06-05T00:00:00Z"), - conv("c3_ok", "2026-06-03T00:00:00Z"), - ]; - let source = ConvSource::new(raw, 1); - - let mut cursor = CompositeCursor::default(); - let mut emitted: Vec = Vec::new(); - for _ in 0..1000 { - let lane = match source.page(cursor.conv_page_token.as_deref()) { - ConvLane::Page { - rows, - next_token, - frontier, - } => ConvLane::Page { - rows: rows - .into_iter() - .filter(|r| r.legacy.session_id.contains("ok")) - .collect(), - next_token, - frontier, - }, - other => other, - }; - let result = merge_and_paginate(local_window.clone(), lane, &cursor, 2); - emitted.extend( - result.candidates[..result.emit_count] - .iter() - .map(|r| r.legacy.session_id.clone()), - ); - match result.next_cursor { - Some(c) => cursor = c, - None => break, - } - } - - let mut expected = local_window.clone(); - expected.push(conv("c3_ok", "2026-06-03T00:00:00Z")); - expected.sort_by(cmp_total_order); - assert_eq!(emitted, ids(&expected)); - assert!( - emitted.iter().any(|id| id == "c3_ok"), - "the later matching conversation must not be dropped" - ); - } - - #[test] - fn local_only_when_conversations_skipped() { - let local_window = vec![ - local("l1", "2026-06-10T00:00:00Z"), - local("l2", "2026-06-08T00:00:00Z"), - ]; - let result = merge_and_paginate( - local_window.clone(), - ConvLane::Skipped, - &CompositeCursor::default(), - 10, - ); - assert_eq!(result.emit_count, 2); - assert!(result.partial.is_none()); - assert!(result.next_cursor.is_none()); - } - - #[test] - fn degraded_lane_sets_partial_and_returns_local() { - let local_window = vec![local("l1", "2026-06-10T00:00:00Z")]; - let result = merge_and_paginate( - local_window, - ConvLane::Degraded(PartialReason::Timeout), - &CompositeCursor::default(), - 10, - ); - assert_eq!(result.partial, Some(PartialReason::Timeout)); - assert_eq!(result.emit_count, 1); - } - - #[test] - fn degraded_mid_walk_with_progress_keeps_live_conv_token() { - let cursor = CompositeCursor { - boundary: Some(BoundaryKey { - updated_at: "2026-06-15T00:00:00Z".into(), + updated_at: "2026-02-01T00:00:00Z".into(), kind: SessionKind::Build, - session_id: "z_newer".into(), + session_id: "a".into(), }), - conv_page_token: Some("p2".into()), - conv_page_drained: true, }; - let result = merge_and_paginate( - vec![local("l1", "2026-06-10T00:00:00Z")], - ConvLane::Degraded(PartialReason::Timeout), - &cursor, - 10, - ); - assert_eq!(result.emit_count, 1, "the local row is emitted (progress)"); - assert_eq!(result.partial, Some(PartialReason::Timeout)); - let next = result - .next_cursor - .expect("progress + live conv token must keep the continuation"); - assert_eq!(next.conv_page_token.as_deref(), Some("p2")); - assert_eq!( - next.boundary.as_ref().map(|b| b.session_id.as_str()), - Some("l1") - ); + let decoded = CompositeCursor::decode(Some(&c.encode())); + let b = decoded.boundary.expect("boundary survives roundtrip"); + assert_eq!(b.session_id, "a"); + assert_eq!(b.updated_at, "2026-02-01T00:00:00Z"); } #[test] - fn degraded_mid_walk_with_no_progress_terminates() { - let cursor = CompositeCursor { - boundary: Some(BoundaryKey { - updated_at: "2026-06-10T00:00:00Z".into(), - kind: SessionKind::Build, - session_id: "l1".into(), - }), - conv_page_token: Some("p2".into()), - conv_page_drained: true, - }; - let result = merge_and_paginate( - vec![local("l1", "2026-06-10T00:00:00Z")], - ConvLane::Degraded(PartialReason::Timeout), - &cursor, - 10, - ); - assert_eq!( - result.emit_count, 0, - "local lane is exhausted (no progress)" - ); - assert_eq!(result.partial, Some(PartialReason::Timeout)); + fn decode_garbage_yields_default() { assert!( - result.next_cursor.is_none(), - "a zero-progress degraded page must terminate, not re-emit an identical cursor" + CompositeCursor::decode(Some("!!!not-base64!!!")) + .boundary + .is_none() ); + assert!(CompositeCursor::decode(None).boundary.is_none()); + assert!(CompositeCursor::decode(Some("")).boundary.is_none()); } #[test] - fn degraded_first_page_with_no_token_does_not_fabricate_a_cursor() { - let result = merge_and_paginate( - Vec::new(), - ConvLane::Degraded(PartialReason::Error), - &CompositeCursor::default(), - 10, - ); - assert_eq!(result.partial, Some(PartialReason::Error)); - assert!(result.next_cursor.is_none()); + fn paginate_sorts_newest_first_and_cuts_page() { + let rows = vec![ + local("old", "2026-01-01T00:00:00Z"), + local("new", "2026-03-01T00:00:00Z"), + local("mid", "2026-02-01T00:00:00Z"), + ]; + let page = paginate(rows, &CompositeCursor::default(), 2); + assert_eq!(ids(&page), vec!["new", "mid"]); + assert!(page.next_cursor.is_some(), "a third row remains"); + } + + #[test] + fn paginate_resumes_after_boundary_without_duplicates() { + let rows: Vec = vec![ + local("a", "2026-03-01T00:00:00Z"), + local("b", "2026-02-01T00:00:00Z"), + local("c", "2026-01-01T00:00:00Z"), + ]; + let first = paginate(rows.clone(), &CompositeCursor::default(), 2); + assert_eq!(ids(&first), vec!["a", "b"]); + let cursor = first.next_cursor.expect("more rows remain"); + let second = paginate(rows, &cursor, 2); + assert_eq!(ids(&second), vec!["c"]); + assert!(second.next_cursor.is_none(), "list is exhausted"); + } + + #[test] + fn paginate_exact_page_has_no_next_cursor() { + let rows = vec![ + local("a", "2026-03-01T00:00:00Z"), + local("b", "2026-02-01T00:00:00Z"), + ]; + let page = paginate(rows, &CompositeCursor::default(), 2); + assert_eq!(ids(&page).len(), 2); + assert!(page.next_cursor.is_none()); + } + + #[test] + fn paginate_ties_break_stably_by_session_id() { + let ts = "2026-02-01T00:00:00Z"; + let rows = vec![local("b", ts), local("a", ts), local("c", ts)]; + let first = paginate(rows.clone(), &CompositeCursor::default(), 2); + assert_eq!(ids(&first), vec!["a", "b"]); + let cursor = first.next_cursor.expect("one row remains"); + let second = paginate(rows, &cursor, 2); + assert_eq!(ids(&second), vec!["c"]); } } diff --git a/crates/codegen/kigi-shell/src/session/unified_list/facets.rs b/crates/codegen/kigi-shell/src/session/unified_list/facets.rs index 129d709..6d029da 100644 --- a/crates/codegen/kigi-shell/src/session/unified_list/facets.rs +++ b/crates/codegen/kigi-shell/src/session/unified_list/facets.rs @@ -4,7 +4,6 @@ use serde::Serialize; use super::envelope::{FacetMap, FacetValue, SessionKind}; use super::row::UnifiedRow; -use crate::remote::Conversation; use crate::session::merge::MergedSession; pub const KIND_FACET_KEY: &str = "kind"; @@ -44,25 +43,6 @@ impl NormalizedItem { starred: false, } } - - pub fn from_conversation(c: &Conversation) -> Self { - Self { - kind: SessionKind::Chat, - cwd: String::new(), - repo_name: None, - branch: None, - worktree_label: None, - git_root_dir: None, - source_workspace_dir: None, - workspace_ids: c - .workspaces - .iter() - .map(|w| w.workspace_id.clone()) - .filter(|id| !id.is_empty()) - .collect(), - starred: c.starred, - } - } } #[derive(Debug, Default)] @@ -393,7 +373,7 @@ pub struct FacetSummaryValue { #[cfg(test)] mod tests { use super::*; - use crate::session::unified_list::{conversation_to_row, merged_session_to_row}; + use crate::session::unified_list::merged_session_to_row; fn local_row(session_id: &str, repo: Option<&str>, branch: Option<&str>) -> UnifiedRow { let m = MergedSession { @@ -419,74 +399,6 @@ mod tests { merged_session_to_row(m, &build_facet_registry()) } - fn conv_row(conversation_id: &str, workspaces: &[&str]) -> UnifiedRow { - let c = Conversation { - conversation_id: conversation_id.into(), - title: "t".into(), - modify_time: Some("2026-06-01T00:00:00Z".into()), - workspaces: workspaces - .iter() - .map(|w| crate::remote::conversations_client::Workspace { - workspace_id: (*w).into(), - }) - .collect(), - ..Conversation::default() - }; - conversation_to_row(c, &build_facet_registry()) - } - - fn conv_row_starred(conversation_id: &str, starred: bool) -> UnifiedRow { - let c = Conversation { - conversation_id: conversation_id.into(), - title: "t".into(), - modify_time: Some("2026-06-01T00:00:00Z".into()), - starred, - ..Conversation::default() - }; - conversation_to_row(c, &build_facet_registry()) - } - - #[test] - fn project_facet_only_on_conversations() { - let reg = build_facet_registry(); - let conv = NormalizedItem::from_conversation(&Conversation { - conversation_id: "c1".into(), - workspaces: vec![crate::remote::conversations_client::Workspace { - workspace_id: "ws_9f3a".into(), - }], - ..Conversation::default() - }); - let facets = reg.extract_all(&conv); - assert!(matches!( - facets.get(WORKSPACE_FACET_KEY), - Some(FacetValue::Many(v)) if v == &[serde_json::json!("ws_9f3a")] - )); - let local = NormalizedItem::from_merged(&MergedSession { - session_id: "s".into(), - summary: String::new(), - first_prompt: None, - updated_at: String::new(), - created_at: String::new(), - cwd: "/x".into(), - hostname: None, - source: "local".into(), - model_id: None, - num_messages: 0, - last_active_at: None, - branch: Some("main".into()), - repo_name: Some("xai".into()), - worktree_label: None, - git_root_dir: None, - git_remotes: Vec::new(), - source_workspace_dir: None, - session_kind: None, - }); - let lf = reg.extract_all(&local); - assert!(!lf.contains_key(WORKSPACE_FACET_KEY)); - assert!(matches!(lf.get(REPO_FACET_KEY), Some(FacetValue::One(_)))); - assert!(matches!(lf.get(BRANCH_FACET_KEY), Some(FacetValue::One(_)))); - } - #[test] fn project_pushdown_single_value_sets_workspace_id() { let reg = build_facet_registry(); @@ -513,121 +425,6 @@ mod tests { assert!(q.workspace_id.is_none()); } - #[test] - fn project_filter_is_partition_aware_keeps_local_rows() { - let reg = build_facet_registry(); - let rows = vec![ - local_row("local-1", Some("xai"), Some("main")), - conv_row("conv-match", &["ws_9f3a"]), - conv_row("conv-other", &["ws_zzz"]), - ]; - let mut filters = BTreeMap::new(); - filters.insert( - WORKSPACE_FACET_KEY.to_owned(), - vec![serde_json::json!("ws_9f3a")], - ); - let kept = reg.apply_in_memory_filters(&filters, rows); - let ids: Vec<&str> = kept.iter().map(|r| r.legacy.session_id.as_str()).collect(); - assert!(ids.contains(&"local-1")); - assert!(ids.contains(&"conv-match")); - assert!(!ids.contains(&"conv-other")); - } - - #[test] - fn repo_filter_is_partition_aware_keeps_conversation_rows() { - let reg = build_facet_registry(); - let rows = vec![ - local_row("local-xai", Some("xai"), Some("main")), - local_row("local-other", Some("other"), Some("main")), - conv_row("conv-1", &["ws_9f3a"]), - ]; - let mut filters = BTreeMap::new(); - filters.insert(REPO_FACET_KEY.to_owned(), vec![serde_json::json!("xai")]); - let kept = reg.apply_in_memory_filters(&filters, rows); - let ids: Vec<&str> = kept.iter().map(|r| r.legacy.session_id.as_str()).collect(); - assert!(ids.contains(&"local-xai")); - assert!(!ids.contains(&"local-other")); - assert!(ids.contains(&"conv-1")); - } - - #[test] - fn pushdown_and_in_memory_project_filter_agree() { - let reg = build_facet_registry(); - let convs = vec![conv_row("a", &["ws_1"]), conv_row("b", &["ws_2"])]; - let mut filters = BTreeMap::new(); - filters.insert( - WORKSPACE_FACET_KEY.to_owned(), - vec![serde_json::json!("ws_1")], - ); - let in_memory = reg.apply_in_memory_filters(&filters, convs); - let ids: Vec<&str> = in_memory - .iter() - .map(|r| r.legacy.session_id.as_str()) - .collect(); - assert_eq!(ids, ["a"]); - let mut q = SourceQuery::default(); - reg.apply_pushdown(&filters, &mut q); - assert_eq!(q.workspace_id.as_deref(), Some("ws_1")); - } - - #[test] - fn starred_facet_present_only_for_starred_conversations() { - let reg = build_facet_registry(); - let starred = NormalizedItem::from_conversation(&Conversation { - conversation_id: "c1".into(), - starred: true, - ..Conversation::default() - }); - assert!(matches!( - reg.extract_all(&starred).get(STARRED_FACET_KEY), - Some(FacetValue::One(serde_json::Value::Bool(true))) - )); - let plain = NormalizedItem::from_conversation(&Conversation { - conversation_id: "c2".into(), - starred: false, - ..Conversation::default() - }); - assert!(!reg.extract_all(&plain).contains_key(STARRED_FACET_KEY)); - let local = NormalizedItem::from_merged(&MergedSession { - session_id: "s".into(), - summary: String::new(), - first_prompt: None, - updated_at: String::new(), - created_at: String::new(), - cwd: "/x".into(), - hostname: None, - source: "local".into(), - model_id: None, - num_messages: 0, - last_active_at: None, - branch: None, - repo_name: None, - worktree_label: None, - git_root_dir: None, - git_remotes: Vec::new(), - source_workspace_dir: None, - session_kind: None, - }); - assert!(!reg.extract_all(&local).contains_key(STARRED_FACET_KEY)); - } - - #[test] - fn starred_filter_is_partition_aware_keeps_local_rows() { - let reg = build_facet_registry(); - let rows = vec![ - local_row("local-1", Some("xai"), Some("main")), - conv_row_starred("conv-starred", true), - conv_row_starred("conv-plain", false), - ]; - let mut filters = BTreeMap::new(); - filters.insert(STARRED_FACET_KEY.to_owned(), vec![serde_json::json!(true)]); - let kept = reg.apply_in_memory_filters(&filters, rows); - let ids: Vec<&str> = kept.iter().map(|r| r.legacy.session_id.as_str()).collect(); - assert!(ids.contains(&"local-1")); - assert!(ids.contains(&"conv-starred")); - assert!(!ids.contains(&"conv-plain")); - } - fn local_row_with_git( session_id: &str, git_root: Option<&str>, @@ -656,49 +453,6 @@ mod tests { merged_session_to_row(m, &build_facet_registry()) } - #[test] - fn git_path_facets_present_only_for_local_rows() { - let reg = build_facet_registry(); - let local = NormalizedItem::from_merged(&MergedSession { - session_id: "s".into(), - summary: String::new(), - first_prompt: None, - updated_at: String::new(), - created_at: String::new(), - cwd: "/x".into(), - hostname: None, - source: "local".into(), - model_id: None, - num_messages: 0, - last_active_at: None, - branch: None, - repo_name: None, - worktree_label: None, - git_root_dir: Some("/Users/me/xai".into()), - git_remotes: Vec::new(), - source_workspace_dir: Some("/Users/me/xai-main".into()), - session_kind: Some("worktree".into()), - }); - let f = reg.extract_all(&local); - assert!(matches!( - f.get(GIT_ROOT_FACET_KEY), - Some(FacetValue::One(serde_json::Value::String(s))) if s == "/Users/me/xai" - )); - assert!(matches!( - f.get(SOURCE_WORKSPACE_FACET_KEY), - Some(FacetValue::One(serde_json::Value::String(s))) if s == "/Users/me/xai-main" - )); - - // Conversations carry no local git enrichment. - let conv = NormalizedItem::from_conversation(&Conversation { - conversation_id: "c1".into(), - ..Conversation::default() - }); - let cf = reg.extract_all(&conv); - assert!(!cf.contains_key(GIT_ROOT_FACET_KEY)); - assert!(!cf.contains_key(SOURCE_WORKSPACE_FACET_KEY)); - } - #[test] fn git_root_filter_keeps_matching_local_rows() { let reg = build_facet_registry(); diff --git a/crates/codegen/kigi-shell/src/session/unified_list/mod.rs b/crates/codegen/kigi-shell/src/session/unified_list/mod.rs index 0bede93..114b462 100644 --- a/crates/codegen/kigi-shell/src/session/unified_list/mod.rs +++ b/crates/codegen/kigi-shell/src/session/unified_list/mod.rs @@ -3,8 +3,7 @@ mod envelope; mod facets; mod row; use crate::agent::session_registry_client::SessionRegistryClient; -use crate::remote::{ConvError, ConvQuery, ConversationsClient}; -use cursor::{CompositeCursor, ConvLane, Paginated, merge_and_paginate}; +use cursor::{CompositeCursor, Paginated, paginate}; pub use envelope::{FacetMap, FacetValue, SessionKind, SessionMetaEnvelope}; pub use facets::{ BRANCH_FACET_KEY, BranchFacet, CWD_FACET_KEY, CwdFacet, FacetProvider, FacetRegistry, @@ -13,62 +12,18 @@ pub use facets::{ SOURCE_WORKSPACE_FACET_KEY, STARRED_FACET_KEY, SourceQuery, SourceWorkspaceFacet, StarredFacet, WORKSPACE_FACET_KEY, WORKTREE_FACET_KEY, WorkspaceFacet, WorktreeFacet, build_facet_registry, }; -pub use row::{ - ExtSupersetRow, RowMeta, SessionInfo, UnifiedRow, conversation_to_row, merged_session_to_row, -}; +pub use row::{ExtSupersetRow, RowMeta, SessionInfo, UnifiedRow, merged_session_to_row}; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; use std::sync::LazyLock; pub const DEFAULT_LIMIT: usize = 30; const CONV_PAGE_HEADROOM: usize = 5; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum PartialReason { - Timeout, - Error, - NoOauth, -} -impl PartialReason { - fn as_str(self) -> &'static str { - match self { - PartialReason::Timeout => "timeout", - PartialReason::Error => "error", - PartialReason::NoOauth => "no_oauth", - } - } -} static FACET_REGISTRY: LazyLock = LazyLock::new(build_facet_registry); pub fn facet_registry() -> &'static FacetRegistry { &FACET_REGISTRY } -/// Hard-off in release builds so they can't enable the -/// conversations lane via env. -pub fn conversations_lane_enabled() -> bool { - if true { - return false; - } - std::env::var("KIGI_SESSION_LIST_CONVERSATIONS") - .ok() - .is_some_and(|v| { - !matches!( - v.trim().to_ascii_lowercase().as_str(), - "" | "0" | "false" | "off" | "no" - ) - }) -} -/// Env lane (desktop `KIGI_SESSION_LIST_CONVERSATIONS`) OR process-wide -/// `--chat` (`KIGI_CHAT_MODE`); hard-off in release builds. -/// The single predicate `MvpAgent::conversations_client()` keys on. -pub fn conversations_lane_active() -> bool { - conversations_lane_enabled() || crate::agent::chat_modes::process_chat_mode_enabled() -} -/// Parse `x.ai/session/list` params and, under process-wide chat mode, force -/// the conversations-only `kind` facet (see [`force_kind_chat`]). pub fn parse_list_req(raw: &str) -> Result { - let mut req: ListReq = serde_json::from_str(raw)?; - if crate::agent::chat_modes::process_chat_mode_enabled() { - force_kind_chat(&mut req); - } - Ok(req) + serde_json::from_str(raw) } #[derive(Debug, Default, Deserialize)] #[serde(rename_all = "camelCase")] @@ -88,7 +43,6 @@ pub struct UnifiedListResult { pub rows: Vec, pub next_cursor: Option, pub facets: FacetSummary, - pub conversations_partial: Option, } #[derive(Debug, Default)] struct ParsedMeta { @@ -131,33 +85,8 @@ fn value_list(v: &serde_json::Value) -> Vec { other => vec![other.clone()], } } -/// Rewrite `req` so the `kind` facet filter is exactly `["chat"]`. -/// -/// REPLACES any client-sent `kind` allow-list (a union with `"build"` would -/// re-enable the local lane); every other facet filter and `_meta` key is -/// left untouched. -pub fn force_kind_chat(req: &mut ListReq) { - let mut meta = match req.meta.take() { - Some(serde_json::Value::Object(map)) => map, - _ => serde_json::Map::new(), - }; - let mut filters = match meta.remove("x.ai/facetFilters") { - Some(serde_json::Value::Object(map)) => map, - _ => serde_json::Map::new(), - }; - filters.insert( - KIND_FACET_KEY.to_owned(), - serde_json::json!([SessionKind::Chat.as_str()]), - ); - meta.insert( - "x.ai/facetFilters".to_owned(), - serde_json::Value::Object(filters), - ); - req.meta = Some(serde_json::Value::Object(meta)); -} pub async fn build_unified_list( registry_client: Option<&SessionRegistryClient>, - conversations_client: Option<&ConversationsClient>, req: ListReq, ) -> UnifiedListResult { let reg = facet_registry(); @@ -171,13 +100,11 @@ pub async fn build_unified_list( let cursor = CompositeCursor::decode(req.cursor.as_deref()); let mut source_query = SourceQuery::default(); reg.apply_pushdown(&facet_filters, &mut source_query); - let exclude_conversations = excludes_conversations(&facet_filters); let exclude_build = excludes_build(&facet_filters); let over = (limit * 3).max(100); - let local_fut = async { - if exclude_build { - return Vec::new(); - } + let local_rows = if exclude_build { + Vec::new() + } else { crate::session::merge::fetch_merged( registry_client, req.cwd.as_deref(), @@ -189,84 +116,17 @@ pub async fn build_unified_list( .map(|m| merged_session_to_row(m, reg)) .collect::>() }; - let conv_fut = async { - if exclude_conversations { - return ConvLane::Skipped; - } - let Some(client) = conversations_client else { - return ConvLane::Skipped; - }; - let q = ConvQuery { - page_size: (limit + CONV_PAGE_HEADROOM) as i64, - page_token: cursor.conv_page_token.clone(), - search_query: query.clone(), - workspace_id: source_query.workspace_id.clone(), - }; - match tokio::time::timeout( - crate::session::merge::REMOTE_TIMEOUT, - client.list_conversations(&q), - ) - .await - { - Ok(Ok(page)) => { - let next_token = page.next_page_token; - let rows: Vec = page - .conversations - .into_iter() - .map(|c| conversation_to_row(c, reg)) - .collect(); - let frontier = cursor::conv_frontier(&rows, next_token.is_some()); - ConvLane::Page { - rows, - next_token, - frontier, - } - } - Ok(Err(ConvError::NoOauth)) => ConvLane::Degraded(PartialReason::NoOauth), - Ok(Err(e)) => { - tracing::warn!("conversation list failed: {e}"); - ConvLane::Degraded(PartialReason::Error) - } - Err(_) => { - tracing::warn!("conversation list timed out"); - ConvLane::Degraded(PartialReason::Timeout) - } - } - }; - let (local_rows, conv_lane) = tokio::join!(local_fut, conv_fut); - { - let (conv_lane_status, conv_rows) = match &conv_lane { - ConvLane::Skipped => ("skipped", 0), - ConvLane::Degraded(reason) => (reason.as_str(), 0), - ConvLane::Page { rows, .. } => ("ok", rows.len()), - }; - tracing::debug!( - local_lane_skipped = exclude_build, - local_rows = local_rows.len(), - conv_lane = conv_lane_status, - conv_rows, - "session list lanes" - ); - } + tracing::debug!( + local_lane_skipped = exclude_build, + local_rows = local_rows.len(), + "session list" + ); let local_rows = reg.apply_in_memory_filters(&facet_filters, local_rows); - let conv_lane = match conv_lane { - ConvLane::Page { - rows, - next_token, - frontier, - } => ConvLane::Page { - rows: reg.apply_in_memory_filters(&facet_filters, rows), - next_token, - frontier, - }, - other => other, - }; let Paginated { candidates, emit_count, next_cursor, - partial, - } = merge_and_paginate(local_rows, conv_lane, &cursor, limit); + } = paginate(local_rows, &cursor, limit); let mut rows = candidates; rows.truncate(emit_count); let facets = reg.summarize_window(&rows); @@ -274,15 +134,6 @@ pub async fn build_unified_list( rows, next_cursor: next_cursor.map(|c| c.encode()), facets, - conversations_partial: partial, - } -} -fn excludes_conversations(filters: &BTreeMap>) -> bool { - match filters.get(KIND_FACET_KEY) { - Some(allowed) if !allowed.is_empty() => !allowed - .iter() - .any(|v| v.as_str() == Some(SessionKind::Chat.as_str())), - _ => false, } } /// Mirror of [`excludes_conversations`]: `true` when a non-empty `kind` @@ -307,21 +158,12 @@ pub struct ExtListResponse { pub struct ExtListResponseMeta { #[serde(rename = "x.ai/facets")] pub facets: FacetSummary, - #[serde(rename = "x.ai/partial")] - pub partial: PartialInfo, -} -#[derive(Debug, Clone, Serialize)] -pub struct PartialInfo { - pub conversations: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub reason: Option<&'static str>, } pub fn ext_list_response(result: UnifiedListResult) -> ExtListResponse { let UnifiedListResult { rows, next_cursor, facets, - conversations_partial, } = result; ExtListResponse { sessions: rows @@ -329,13 +171,7 @@ pub fn ext_list_response(result: UnifiedListResult) -> ExtListResponse { .map(UnifiedRow::into_ext_superset) .collect(), next_cursor, - meta: ExtListResponseMeta { - facets, - partial: PartialInfo { - conversations: conversations_partial.is_some(), - reason: conversations_partial.map(PartialReason::as_str), - }, - }, + meta: ExtListResponseMeta { facets }, } } #[cfg(test)] @@ -494,74 +330,8 @@ mod tests { ); filters } - #[test] - fn excludes_build_mirrors_excludes_conversations() { - assert!(excludes_build(&kind_filter(&["chat"]))); - assert!(!excludes_conversations(&kind_filter(&["chat"]))); - assert!(!excludes_build(&kind_filter(&["build"]))); - assert!(excludes_conversations(&kind_filter(&["build"]))); - assert!(!excludes_build(&kind_filter(&["build", "chat"]))); - assert!(!excludes_conversations(&kind_filter(&["build", "chat"]))); - assert!(!excludes_build(&kind_filter(&[]))); - assert!(!excludes_conversations(&kind_filter(&[]))); - assert!(!excludes_build(&BTreeMap::new())); - assert!(!excludes_conversations(&BTreeMap::new())); - } /// The forced `kind` REPLACES a client-sent `kind: ["build"]` (never /// unions), so the local lane stays excluded. - #[test] - fn forced_kind_replaces_client_build_filter() { - let mut req = ListReq { - meta: Some(serde_json::json!({ "x.ai/facetFilters" : { "kind" : ["build"] }, })), - ..ListReq::default() - }; - force_kind_chat(&mut req); - let parsed = ParsedMeta::parse(req.meta.as_ref()); - assert_eq!( - parsed.facet_filters.get(KIND_FACET_KEY), - Some(&vec![serde_json::json!("chat")]), - "forced kind must replace the client filter, not union with it" - ); - assert!(excludes_build(&parsed.facet_filters)); - assert!(!excludes_conversations(&parsed.facet_filters)); - } - #[test] - fn forced_kind_preserves_other_facets() { - let mut req = ListReq { - meta: Some(serde_json::json!( - { "x.ai/facetFilters" : { "kind" : ["build"], "starred" : [true], - "workspace" : ["w1"] }, "x.ai/query" : "antelope", "x.ai/limit" : 5, - } - )), - ..ListReq::default() - }; - force_kind_chat(&mut req); - let parsed = ParsedMeta::parse(req.meta.as_ref()); - assert_eq!( - parsed.facet_filters.get(KIND_FACET_KEY), - Some(&vec![serde_json::json!("chat")]) - ); - assert_eq!( - parsed.facet_filters.get("starred"), - Some(&vec![serde_json::json!(true)]) - ); - assert_eq!( - parsed.facet_filters.get("workspace"), - Some(&vec![serde_json::json!("w1")]) - ); - assert_eq!(parsed.query.as_deref(), Some("antelope")); - assert_eq!(parsed.limit, Some(5)); - } - #[test] - fn forced_kind_creates_facet_filters_when_meta_absent() { - let mut req = ListReq::default(); - force_kind_chat(&mut req); - let parsed = ParsedMeta::parse(req.meta.as_ref()); - assert_eq!( - parsed.facet_filters.get(KIND_FACET_KEY), - Some(&vec![serde_json::json!("chat")]) - ); - } fn xai_auth_manager(dir: &std::path::Path) -> std::sync::Arc { let am = std::sync::Arc::new(crate::auth::AuthManager::new( dir, @@ -598,191 +368,4 @@ mod tests { }); addr } - /// A client-sent `kind: ["build"]` rewritten by [`force_kind_chat`] - /// yields conversations only. - #[tokio::test] - #[serial_test::serial] - async fn forced_kind_serves_conversations_only() { - let addr = spawn_conversations_stub( - serde_json::json!( - { "conversations" : [{ "conversationId" : "c1", "title" : "Hello", - "modifyTime" : "2026-07-01T00:00:00Z" }, { "conversationId" : "c2", - "title" : "", "modifyTime" : "2026-07-02T00:00:00Z" },], } - ) - .to_string(), - ) - .await; - let _env = kigi_test_support::EnvGuard::set( - "KIGI_CONVERSATIONS_BASE_URL", - format!("http://{addr}"), - ); - let home = tempfile::tempdir().expect("tempdir"); - let client = ConversationsClient::new(xai_auth_manager(home.path())); - let mut req = ListReq { - meta: Some(serde_json::json!({ "x.ai/facetFilters" : { "kind" : ["build"] }, })), - ..ListReq::default() - }; - force_kind_chat(&mut req); - let result = build_unified_list(None, Some(&client), req).await; - let ids: Vec<&str> = result - .rows - .iter() - .map(|r| r.legacy.session_id.as_str()) - .collect(); - assert_eq!(ids, ["c2", "c1"], "conversations only, newest first"); - assert!( - result - .rows - .iter() - .all(|r| r.legacy.source == "conversation"), - "no build row may survive the forced kind filter" - ); - assert_eq!(result.conversations_partial, None); - } - /// A degraded conversations lane (no OAuth) surfaces through - /// `conversations_partial` instead of failing the list. - #[tokio::test] - #[serial_test::serial] - async fn degraded_conversations_lane_reports_no_oauth() { - let home = tempfile::tempdir().expect("tempdir"); - let auth = std::sync::Arc::new(crate::auth::AuthManager::new( - home.path(), - crate::auth::KimiCodeConfig::default(), - )); - let client = ConversationsClient::new(auth); - let mut req = ListReq::default(); - force_kind_chat(&mut req); - let result = build_unified_list(None, Some(&client), req).await; - assert!(result.rows.is_empty()); - assert_eq!(result.conversations_partial, Some(PartialReason::NoOauth)); - } - /// Build-mode canary: with no conversations client the lane is skipped — - /// not degraded. - #[tokio::test] - async fn non_chat_list_without_client_skips_conversations_lane() { - let req = ListReq { - cwd: Some("/nonexistent/unified-list-canary".into()), - ..ListReq::default() - }; - let result = build_unified_list(None, None, req).await; - assert_eq!( - result.conversations_partial, None, - "no client ⇒ lane skipped, never reported as degraded" - ); - assert!(result.rows.is_empty()); - } - /// Desktop env lane stays env-gated; process chat mode is feature-gated. - #[test] - #[serial_test::serial] - fn conversations_lane_env_gating_matrix() { - { - let _off = kigi_test_support::EnvGuard::unset("KIGI_SESSION_LIST_CONVERSATIONS"); - assert!(!conversations_lane_enabled()); - } - { - let _on = kigi_test_support::EnvGuard::set("KIGI_SESSION_LIST_CONVERSATIONS", "1"); - assert!(!conversations_lane_enabled()); - } - { - let _off = kigi_test_support::EnvGuard::set("KIGI_SESSION_LIST_CONVERSATIONS", "0"); - assert!(!conversations_lane_enabled()); - } - } - /// Truth table for `conversations_lane_active`: desktop env lane OR - /// process chat mode, hard-off in release builds. - #[test] - #[serial_test::serial] - fn conversations_lane_active_truth_table() { - use crate::agent::chat_modes::KIGI_CHAT_MODE_ENV; - let _chat_off = kigi_test_support::EnvGuard::unset(KIGI_CHAT_MODE_ENV); - let _desktop_off = kigi_test_support::EnvGuard::unset("KIGI_SESSION_LIST_CONVERSATIONS"); - assert!( - !conversations_lane_active(), - "no env ⇒ lane off (Build-mode default)" - ); - { - let _desktop = kigi_test_support::EnvGuard::set("KIGI_SESSION_LIST_CONVERSATIONS", "1"); - assert!(!conversations_lane_active()); - } - { - let _chat = kigi_test_support::EnvGuard::set(KIGI_CHAT_MODE_ENV, "1"); - assert!( - !conversations_lane_active(), - "process chat mode must enable the lane (chat feature only)" - ); - } - } - /// `parse_list_req` forces the conversations-only `kind` exactly when - /// process chat mode is on; otherwise the client request is untouched. - #[test] - #[serial_test::serial] - fn parse_list_req_forces_kind_under_process_chat_mode_only() { - use crate::agent::chat_modes::KIGI_CHAT_MODE_ENV; - let raw = serde_json::json!( - { "_meta" : { "x.ai/facetFilters" : { "kind" : ["build"], "starred" : [true] - } }, } - ) - .to_string(); - { - let _off = kigi_test_support::EnvGuard::unset(KIGI_CHAT_MODE_ENV); - let req = parse_list_req(&raw).expect("parse"); - let parsed = ParsedMeta::parse(req.meta.as_ref()); - assert_eq!( - parsed.facet_filters.get(KIND_FACET_KEY), - Some(&vec![serde_json::json!("build")]), - "non-chat: client kind filter untouched" - ); - } - { - let _on = kigi_test_support::EnvGuard::set(KIGI_CHAT_MODE_ENV, "1"); - let req = parse_list_req(&raw).expect("parse"); - let parsed = ParsedMeta::parse(req.meta.as_ref()); - let expected = if false { "chat" } else { "build" }; - assert_eq!( - parsed.facet_filters.get(KIND_FACET_KEY), - Some(&vec![serde_json::json!(expected)]) - ); - assert_eq!( - parsed.facet_filters.get("starred"), - Some(&vec![serde_json::json!(true)]), - "other facets pass through" - ); - } - } - /// Wire pin for the cross-crate `x.ai/partial` envelope the pager parses: - /// the serialized reason strings must not drift (the pager maps unknown - /// reasons to a generic retry notice, masking a rename). - #[test] - fn ext_list_response_serializes_partial_reasons() { - for (reason, wire) in [ - (PartialReason::NoOauth, "no_oauth"), - (PartialReason::Timeout, "timeout"), - (PartialReason::Error, "error"), - ] { - let value = serde_json::to_value(ext_list_response(UnifiedListResult { - rows: Vec::new(), - next_cursor: None, - facets: facet_registry().summarize_window(&[]), - conversations_partial: Some(reason), - })) - .expect("serialize"); - assert_eq!( - value["_meta"]["x.ai/partial"], - serde_json::json!({ "conversations" : - true, "reason" : wire }) - ); - } - let healthy = serde_json::to_value(ext_list_response(UnifiedListResult { - rows: Vec::new(), - next_cursor: None, - facets: facet_registry().summarize_window(&[]), - conversations_partial: None, - })) - .expect("serialize"); - assert_eq!( - healthy["_meta"]["x.ai/partial"], - serde_json::json!({ "conversations" : - false }) - ); - } } diff --git a/crates/codegen/kigi-shell/src/session/unified_list/row.rs b/crates/codegen/kigi-shell/src/session/unified_list/row.rs index 00458d1..47fc549 100644 --- a/crates/codegen/kigi-shell/src/session/unified_list/row.rs +++ b/crates/codegen/kigi-shell/src/session/unified_list/row.rs @@ -2,7 +2,6 @@ use serde::Serialize; use super::envelope::{FacetMap, SessionKind, SessionMetaEnvelope}; use super::facets::{FacetRegistry, NormalizedItem}; -use crate::remote::Conversation; use crate::session::merge::MergedSession; #[derive(Debug, Clone)] @@ -73,44 +72,6 @@ pub fn merged_session_to_row(m: MergedSession, reg: &FacetRegistry) -> UnifiedRo } } -pub fn conversation_to_row(c: Conversation, reg: &FacetRegistry) -> UnifiedRow { - let facets = reg.extract_all(&NormalizedItem::from_conversation(&c)); - let Conversation { - conversation_id, - title, - modify_time, - create_time, - .. - } = c; - let legacy = MergedSession { - session_id: conversation_id, - summary: title.clone(), - first_prompt: None, - updated_at: modify_time.as_deref().unwrap_or_default().to_owned(), - created_at: create_time.unwrap_or_default(), - cwd: String::new(), - hostname: None, - source: "conversation".to_string(), - model_id: None, - num_messages: 0, - last_active_at: modify_time.clone(), - branch: None, - repo_name: None, - worktree_label: None, - git_root_dir: None, - git_remotes: Vec::new(), - source_workspace_dir: None, - session_kind: None, - }; - UnifiedRow { - kind: SessionKind::Chat, - legacy, - title, - updated_at: modify_time, - facets, - } -} - fn effective_local_ts(m: &MergedSession) -> Option { m.last_active_at .as_deref() @@ -151,47 +112,4 @@ pub struct SessionInfo { mod tests { use super::*; use crate::session::unified_list::facet_registry; - - #[test] - fn conversation_row_uses_conversation_id_as_session_id() { - let c = Conversation { - conversation_id: "conv_abc123".into(), - title: "Compare GPU vendors".into(), - modify_time: Some("2026-06-18T18:02:00Z".into()), - create_time: Some("2026-06-18T17:30:00Z".into()), - ..Conversation::default() - }; - let row = conversation_to_row(c, facet_registry()); - assert_eq!(row.legacy.session_id, "conv_abc123"); - assert_eq!(row.kind, SessionKind::Chat); - assert_eq!(row.legacy.source, "conversation"); - assert_eq!(row.legacy.cwd, ""); - - let ext = serde_json::to_value(row.clone().into_ext_superset()).unwrap(); - assert_eq!(ext["sessionId"], "conv_abc123"); - assert_eq!(ext["cwd"], ""); - assert_eq!(ext["source"], "conversation"); - assert_eq!(ext["_meta"]["x.ai/session"]["kind"], "chat"); - // Chat rows have no local git enrichment (fields omitted). - assert!(ext.get("gitRootDir").is_none()); - assert!(ext.get("gitRemotes").is_none()); - assert!(ext.get("sourceWorkspaceDir").is_none()); - assert!(ext.get("sessionKind").is_none()); - - let bare = serde_json::to_value(row.into_session_info()).unwrap(); - assert_eq!(bare["sessionId"], "conv_abc123"); - } - - #[test] - fn conversation_missing_modify_time_still_resumable() { - let c = Conversation { - conversation_id: "conv_no_time".into(), - title: "Untitled".into(), - ..Conversation::default() - }; - let row = conversation_to_row(c, facet_registry()); - assert_eq!(row.legacy.session_id, "conv_no_time"); - assert!(row.updated_at.is_none()); - assert_eq!(row.legacy.updated_at, ""); - } } diff --git a/crates/codegen/kigi-shell/src/session/worktree.rs b/crates/codegen/kigi-shell/src/session/worktree.rs index 3731f3f..238f31a 100644 --- a/crates/codegen/kigi-shell/src/session/worktree.rs +++ b/crates/codegen/kigi-shell/src/session/worktree.rs @@ -387,7 +387,7 @@ async fn resume_local_session_in_worktree( source_workspace_dir: Some(resolved_source_cwd.to_owned()), ..Default::default() }; - let fork_resp = match fork_session(fork_req, agent_id, auth_manager).await { + let fork_resp = match fork_session(fork_req).await { Ok(r) => r, Err(e) => { cleanup_worktree_on_failure(resolved_source_cwd, &wt_resp.worktree_path).await; diff --git a/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs b/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs index 2b8e067..5595b4b 100644 --- a/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs +++ b/crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs @@ -47,15 +47,11 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap) -> SubagentSpawnCon auth_scheme: Default::default(), extra_headers: Default::default(), context_window: 256_000, - client_version: None, force_http1: false, max_retries: None, stream_tool_calls: false, idle_timeout_secs: None, - client_identifier: None, reasoning_effort: None, - deployment_id: None, - user_id: None, origin_client: None, attribution_callback: None, bearer_resolver: None, diff --git a/crates/codegen/kigi-shell/src/tools/config.rs b/crates/codegen/kigi-shell/src/tools/config.rs index 0b6d466..55eb3db 100644 --- a/crates/codegen/kigi-shell/src/tools/config.rs +++ b/crates/codegen/kigi-shell/src/tools/config.rs @@ -205,7 +205,7 @@ impl ShellToolsetConfig { pub fn new(base: Option, sampling_config: Option) -> Self { let default_base = SamplerConfig { api_key: None, - base_url: "https://api.x.ai/v1".to_string(), + base_url: kigi_env::coding_api_base_url(), model: String::new(), max_completion_tokens: None, temperature: None, @@ -214,15 +214,11 @@ impl ShellToolsetConfig { auth_scheme: Default::default(), extra_headers: indexmap::IndexMap::new(), context_window: 256_000, - client_version: None, reasoning_effort: None, force_http1: false, max_retries: None, stream_tool_calls: false, idle_timeout_secs: None, - client_identifier: None, - deployment_id: None, - user_id: None, origin_client: None, // Default base for the in-process web-search tool config. // Real `SamplerConfig`s (e.g. from `sampling_config_for_model`) diff --git a/crates/codegen/kigi-shell/src/util/config/resolve/auto_mode.rs b/crates/codegen/kigi-shell/src/util/config/resolve/auto_mode.rs index a474f6c..6f45ae4 100644 --- a/crates/codegen/kigi-shell/src/util/config/resolve/auto_mode.rs +++ b/crates/codegen/kigi-shell/src/util/config/resolve/auto_mode.rs @@ -67,8 +67,7 @@ fn resolve_auto_permission_mode_layers( } /// Resolve whether the **auto** permission mode feature (`PermissionMode::Auto`, -/// the LLM/heuristic classifier) is enabled. Full chain mirroring -/// [`resolve_zdr_access_enabled`](super::resolve_zdr_access_enabled): +/// the LLM/heuristic classifier) is enabled. Full precedence chain: /// /// requirements > env (`KIGI_AUTO_PERMISSION_MODE`) > `[auto_mode] enabled` in /// `config.toml` > managed > remote settings (`auto_mode.enabled`, coerced diff --git a/crates/codegen/kigi-shell/src/util/config/resolve/features.rs b/crates/codegen/kigi-shell/src/util/config/resolve/features.rs index 5b4fc77..748c63f 100644 --- a/crates/codegen/kigi-shell/src/util/config/resolve/features.rs +++ b/crates/codegen/kigi-shell/src/util/config/resolve/features.rs @@ -1,28 +1,6 @@ use crate::util::config::RemoteSettings; use toml::Value as TomlValue; -/// Resolve whether ZDR users are allowed to use the product. -/// -/// Precedence: requirements > env > config.toml > managed > remote settings > default (false). -pub fn resolve_zdr_access_enabled( - requirements: Option<&TomlValue>, - user: Option<&TomlValue>, - managed: Option<&TomlValue>, - remote: Option<&RemoteSettings>, -) -> bool { - use crate::agent::config::BoolFlag; - fn from_toml(v: Option<&TomlValue>) -> Option { - v?.get("features")?.get("zdr_access_enabled")?.as_bool() - } - BoolFlag::env("KIGI_ZDR_ACCESS_ENABLED") - .requirement(from_toml(requirements)) - .config(from_toml(user)) - .managed(from_toml(managed)) - .feature_flag(remote.and_then(|r| r.zdr_access_enabled)) - .resolve() - .value -} - /// Whether model-catalog (`/v1/models`) and remote-settings (`/v1/settings`) /// fetches from xAI backends are allowed, including the deployment-config sync /// bundled into the startup prefetch (the background managed-config sync has diff --git a/crates/codegen/kigi-shell/tests/common/mod.rs b/crates/codegen/kigi-shell/tests/common/mod.rs index 9cace66..a705090 100644 --- a/crates/codegen/kigi-shell/tests/common/mod.rs +++ b/crates/codegen/kigi-shell/tests/common/mod.rs @@ -43,15 +43,11 @@ pub fn test_sampler_config( .map(|(k, v)| (k.to_string(), v.to_string())) .collect(), context_window: 256_000, - client_version: None, force_http1: false, max_retries: None, stream_tool_calls: false, idle_timeout_secs: None, - client_identifier: None, reasoning_effort: None, - deployment_id: None, - user_id: None, origin_client: None, attribution_callback: None, bearer_resolver: None, diff --git a/crates/codegen/kigi-shell/tests/git_contention_e2e.rs b/crates/codegen/kigi-shell/tests/git_contention_e2e.rs index 1174be6..010a30e 100644 --- a/crates/codegen/kigi-shell/tests/git_contention_e2e.rs +++ b/crates/codegen/kigi-shell/tests/git_contention_e2e.rs @@ -431,7 +431,7 @@ fn git_rebase_refresh_storm_e2e() { // serve HTTP and never read the process environment. unsafe { std::env::set_var("KIGI_SHARE_DIR", kigi_home.path()); - std::env::set_var("KIGI_CLI_CHAT_PROXY_BASE_URL", server.url()); + std::env::set_var("KIGI_CODE_BASE_URL", server.url()); std::env::set_var("KIGI_XAI_API_BASE_URL", server.url()); std::env::set_var("XAI_API_KEY", "test-key-for-ci"); std::env::set_var("KIGI_TELEMETRY_ENABLED", "false"); diff --git a/crates/codegen/kigi-shell/tests/session_load_perf.rs b/crates/codegen/kigi-shell/tests/session_load_perf.rs index a9e0bad..afba299 100644 --- a/crates/codegen/kigi-shell/tests/session_load_perf.rs +++ b/crates/codegen/kigi-shell/tests/session_load_perf.rs @@ -625,7 +625,7 @@ async fn full_session_load_e2e() { std::env::set_var("KIGI_SHARE_DIR", kigi_home.path()); std::env::set_var("KIGI_INSTRUMENTATION", "log"); std::env::set_var("KIGI_INSTRUMENTATION_LOG", &instr_log); - std::env::set_var("KIGI_CLI_CHAT_PROXY_BASE_URL", server.url()); + std::env::set_var("KIGI_CODE_BASE_URL", server.url()); std::env::set_var("KIGI_XAI_API_BASE_URL", server.url()); std::env::set_var("XAI_API_KEY", "test-key-for-ci"); std::env::set_var("KIGI_TELEMETRY_ENABLED", "false"); diff --git a/crates/codegen/kigi-shell/tests/signed_managed_config/common.rs b/crates/codegen/kigi-shell/tests/signed_managed_config/common.rs index 07b59af..4c90e25 100644 --- a/crates/codegen/kigi-shell/tests/signed_managed_config/common.rs +++ b/crates/codegen/kigi-shell/tests/signed_managed_config/common.rs @@ -147,7 +147,7 @@ pub fn signed_dk_body( requirements: Option<&str>, ) -> String { let payload = SignedPayload { - version: prod_mc_cli_chat_proxy_types::SIGNED_PAYLOAD_VERSION, + version: kigi_config::signed_policy::SIGNED_PAYLOAD_VERSION, deployment_id: Some(deployment_id.to_owned()), team_id: None, managed_config: managed.map(str::to_owned), diff --git a/crates/codegen/kigi-shell/tests/signed_managed_config_extended.rs b/crates/codegen/kigi-shell/tests/signed_managed_config_extended.rs index e4928e5..ae8b8aa 100644 --- a/crates/codegen/kigi-shell/tests/signed_managed_config_extended.rs +++ b/crates/codegen/kigi-shell/tests/signed_managed_config_extended.rs @@ -38,7 +38,7 @@ async fn sync_fail_closed_policy(home: &std::path::Path, kp: &ring::signature::E /// provisioned key with no config row. fn signed_dk_empty_body(kp: &ring::signature::Ed25519KeyPair, deployment_id: &str) -> String { let payload = SignedPayload { - version: prod_mc_cli_chat_proxy_types::SIGNED_PAYLOAD_VERSION, + version: kigi_config::signed_policy::SIGNED_PAYLOAD_VERSION, deployment_id: Some(deployment_id.to_owned()), team_id: None, managed_config: None, diff --git a/crates/codegen/kigi-shell/tests/test_built_binary_e2e.rs b/crates/codegen/kigi-shell/tests/test_built_binary_e2e.rs index 16f0d75..91ed5a6 100644 --- a/crates/codegen/kigi-shell/tests/test_built_binary_e2e.rs +++ b/crates/codegen/kigi-shell/tests/test_built_binary_e2e.rs @@ -1267,7 +1267,7 @@ impl ConfigTestHarness { home, workdir: git_workdir(), env: vec![ - ("KIGI_CLI_CHAT_PROXY_BASE_URL".into(), server.url()), + ("KIGI_CODE_BASE_URL".into(), server.url()), ("KIGI_TELEMETRY_ENABLED".into(), "false".into()), ("KIGI_FEEDBACK_ENABLED".into(), "false".into()), ("KIGI_TRACE_UPLOAD".into(), "false".into()), diff --git a/crates/codegen/kigi-shell/tests/test_leader_soak.rs b/crates/codegen/kigi-shell/tests/test_leader_soak.rs index 9ee02a2..c733071 100644 --- a/crates/codegen/kigi-shell/tests/test_leader_soak.rs +++ b/crates/codegen/kigi-shell/tests/test_leader_soak.rs @@ -136,7 +136,7 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() { // code reads these process-globals (same pattern as session_load_perf). unsafe { std::env::set_var("KIGI_SHARE_DIR", kigi_home.path()); - std::env::set_var("KIGI_CLI_CHAT_PROXY_BASE_URL", server.url()); + std::env::set_var("KIGI_CODE_BASE_URL", server.url()); std::env::set_var("KIGI_XAI_API_BASE_URL", server.url()); std::env::set_var("XAI_API_KEY", "test-key-for-ci"); std::env::set_var("KIGI_TELEMETRY_ENABLED", "false"); diff --git a/crates/codegen/kigi-shell/tests/test_settings_refresh.rs b/crates/codegen/kigi-shell/tests/test_settings_refresh.rs deleted file mode 100644 index 498bbaa..0000000 --- a/crates/codegen/kigi-shell/tests/test_settings_refresh.rs +++ /dev/null @@ -1,258 +0,0 @@ -//! Integration test: MockInferenceServer `/v1/settings` endpoint and -//! remote settings settings refresh infrastructure. -//! -//! Tests the mock endpoint directly (no binary needed) and verifies -//! the `fetch_settings_blocking` client round-trips correctly with -//! runtime-mutated mock settings. -//! -//! Run locally: -//! ```bash -//! cargo test -p kigi-shell --test test_settings_refresh -//! ``` - -use std::future::Future; - -use kigi_shell::util::config::RemoteSettings; -use kigi_test_support::*; - -async fn with_local_set(f: F) -where - F: FnOnce() -> Fut, - Fut: Future, -{ - tokio::task::LocalSet::new().run_until(f()).await; -} - -/// Verify the mock `/v1/settings` endpoint returns 404 when no settings -/// are configured (the default). This preserves backward compatibility: -/// existing tests that never call `set_settings` see a 404, and -/// `fetch_settings_blocking` returns `None`. -#[tokio::test] -async fn test_settings_endpoint_returns_404_when_unconfigured() { - with_local_set(|| async { - let server = MockInferenceServer::start() - .await - .expect("start mock server"); - - let resp = reqwest::get(format!("{}/settings", server.url())) - .await - .expect("request failed"); - assert_eq!(resp.status(), 404); - }) - .await; -} - -/// Verify the mock `/v1/settings` endpoint returns configured settings -/// and that `set_settings` runtime mutation is reflected immediately. -#[tokio::test] -async fn test_settings_endpoint_returns_configured_settings() { - with_local_set(|| async { - let server = MockInferenceServer::start() - .await - .expect("start mock server"); - - // Configure initial settings - server.set_settings(RemoteSettings { - tips: Some(vec!["tip_v1".into()]), - leader_mode: Some(false), - ..Default::default() - }); - - // Fetch and verify - let resp = reqwest::get(format!("{}/settings", server.url())) - .await - .expect("request failed"); - assert_eq!(resp.status(), 200); - let settings: RemoteSettings = resp.json().await.expect("parse failed"); - assert_eq!(settings.tips, Some(vec!["tip_v1".into()])); - assert_eq!(settings.leader_mode, Some(false)); - }) - .await; -} - -/// Verify that `set_settings` updates are visible to subsequent requests -/// (runtime mutation for multi-session test scenarios). -#[tokio::test] -async fn test_settings_endpoint_reflects_runtime_mutations() { - with_local_set(|| async { - let server = MockInferenceServer::start() - .await - .expect("start mock server"); - - // Initial settings - server.set_settings(RemoteSettings { - tips: Some(vec!["tip_v1".into()]), - ..Default::default() - }); - let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url())) - .await - .unwrap() - .json() - .await - .unwrap(); - assert_eq!(settings.tips, Some(vec!["tip_v1".into()])); - - // Mutate settings (simulating a remote feature flag change) - server.set_settings(RemoteSettings { - tips: Some(vec!["tip_v2".into()]), - leader_mode: Some(true), - ..Default::default() - }); - - // Subsequent request sees the updated values - let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url())) - .await - .unwrap() - .json() - .await - .unwrap(); - assert_eq!(settings.tips, Some(vec!["tip_v2".into()])); - assert_eq!(settings.leader_mode, Some(true)); - }) - .await; -} - -/// Verify `fetch_settings_blocking` round-trips through the mock server. -/// This is the actual client function used by `refresh_remote_settings`. -#[tokio::test] -async fn test_fetch_settings_blocking_round_trip() { - with_local_set(|| async { - let server = MockInferenceServer::start() - .await - .expect("start mock server"); - - // Without settings configured: returns None (404 from mock) - let auth = kigi_shell::auth::KimiAuth { - key: "test-key".into(), - ..Default::default() - }; - let result = tokio::task::spawn_blocking({ - let url = server.url().to_string(); - let auth = auth.clone(); - move || kigi_shell::remote::fetch_settings_blocking(&url, &auth, None) - }) - .await - .unwrap(); - assert!( - result.is_none(), - "Expected None when settings not configured" - ); - - // With settings configured: returns Some(settings) - server.set_settings(RemoteSettings { - tips: Some(vec!["fetched_tip".into()]), - ..Default::default() - }); - let result = tokio::task::spawn_blocking({ - let url = server.url().to_string(); - let auth = auth.clone(); - move || kigi_shell::remote::fetch_settings_blocking(&url, &auth, None) - }) - .await - .unwrap(); - let settings = result.expect("Expected Some when settings are configured"); - assert_eq!(settings.tips, Some(vec!["fetched_tip".into()])); - }) - .await; -} - -/// Verify the `doom_loop_recovery` settings object survives the -/// `/v1/settings` round-trip, that its absence deserializes to `None` (old -/// servers), and that a partial object keeps its unset fields `None`. -#[tokio::test] -async fn test_doom_loop_recovery_settings_round_trip() { - use kigi_shell::util::config::DoomLoopRecoverySettings; - - with_local_set(|| async { - let server = MockInferenceServer::start() - .await - .expect("start mock server"); - - // Absent from the payload ⇒ None on the client. - server.set_settings(RemoteSettings::default()); - let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url())) - .await - .unwrap() - .json() - .await - .unwrap(); - assert_eq!(settings.doom_loop_recovery, None); - - server.set_settings(RemoteSettings { - doom_loop_recovery: Some(DoomLoopRecoverySettings { - enabled: Some(true), - max_threshold: Some(16), - max_retries: Some(1), - }), - ..Default::default() - }); - let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url())) - .await - .unwrap() - .json() - .await - .unwrap(); - let recovery = settings.doom_loop_recovery.expect("object round-trips"); - assert_eq!(recovery.enabled, Some(true)); - assert_eq!(recovery.max_threshold, Some(16)); - assert_eq!(recovery.max_retries, Some(1)); - - // Partial object: only the set field comes through; the rest stay - // None so the resolver falls through per-field. - server.set_settings(RemoteSettings { - doom_loop_recovery: Some(DoomLoopRecoverySettings { - max_threshold: Some(32), - ..Default::default() - }), - ..Default::default() - }); - let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url())) - .await - .unwrap() - .json() - .await - .unwrap(); - let recovery = settings.doom_loop_recovery.expect("object round-trips"); - assert_eq!(recovery.enabled, None); - assert_eq!(recovery.max_threshold, Some(32)); - assert_eq!(recovery.max_retries, None); - }) - .await; -} - -/// Verify that the mock server's request log correctly tracks -/// GET /v1/settings requests for assertion in multi-session tests. -#[tokio::test] -async fn test_settings_requests_appear_in_request_log() { - with_local_set(|| async { - let server = MockInferenceServer::start() - .await - .expect("start mock server"); - server.set_settings(RemoteSettings::default()); - - assert_eq!(server.request_count(), 0); - - // First request - let _ = reqwest::get(format!("{}/settings", server.url())) - .await - .unwrap(); - let settings_reqs: Vec<_> = server - .requests() - .into_iter() - .filter(|r| r.method == "GET" && r.path.contains("/settings")) - .collect(); - assert_eq!(settings_reqs.len(), 1, "Expected 1 settings request"); - - // Second request (simulating /new refresh) - let _ = reqwest::get(format!("{}/settings", server.url())) - .await - .unwrap(); - let settings_reqs: Vec<_> = server - .requests() - .into_iter() - .filter(|r| r.method == "GET" && r.path.contains("/settings")) - .collect(); - assert_eq!(settings_reqs.len(), 2, "Expected 2 settings requests"); - }) - .await; -} diff --git a/crates/codegen/kigi-test-support/src/env.rs b/crates/codegen/kigi-test-support/src/env.rs index a41f16b..dba42e9 100644 --- a/crates/codegen/kigi-test-support/src/env.rs +++ b/crates/codegen/kigi-test-support/src/env.rs @@ -164,7 +164,7 @@ pub fn test_env_cmd_tokio( // prompt (the windows-x86_64 lifecycle "prompt timed out" failure). // Mirrors `leader.rs` and the pty-harness `env_for_pager`. .env("KIGI_SHARE_DIR", home.join(".kigi")) - .env("KIGI_CLI_CHAT_PROXY_BASE_URL", mock_url) + .env("KIGI_CODE_BASE_URL", mock_url) .env("KIGI_XAI_API_BASE_URL", mock_url) .env("XAI_API_KEY", "test-key-for-ci") .env("KIGI_TELEMETRY_ENABLED", "false") diff --git a/crates/codegen/kigi-test-support/src/leader.rs b/crates/codegen/kigi-test-support/src/leader.rs index 684c6d6..ba0193e 100644 --- a/crates/codegen/kigi-test-support/src/leader.rs +++ b/crates/codegen/kigi-test-support/src/leader.rs @@ -141,7 +141,7 @@ impl LeaderStdioClient { // leader subprocess inherits/forwards this env var, so every // (re-)elected leader binds the same sandboxed path. .env("KIGI_LEADER_SOCKET", home.join(".kigi").join("leader.sock")) - .env("KIGI_CLI_CHAT_PROXY_BASE_URL", server.url()) + .env("KIGI_CODE_BASE_URL", server.url()) .env("KIGI_XAI_API_BASE_URL", server.url()) .env("XAI_API_KEY", "test-key-for-ci") .env("KIGI_TELEMETRY_ENABLED", "false") diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/interactions.rs b/crates/codegen/kigi-tui/src/app/acp_handler/interactions.rs index eaaa09d..71767f7 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/interactions.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/interactions.rs @@ -80,8 +80,6 @@ pub(crate) fn handle_ask_user_question( let cmd = match kind { LocalQuestionKind::Fork { .. } => "/fork", LocalQuestionKind::NewSession => "/new", - LocalQuestionKind::CreditLimitUpsell => "credit-limit upsell", - LocalQuestionKind::FreeUsageUpsell => "SuperGrok upsell", LocalQuestionKind::AgentTypeMismatch { .. } => "model switch", LocalQuestionKind::ProjectSelect { .. } => "project select", }; diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/mod.rs b/crates/codegen/kigi-tui/src/app/acp_handler/mod.rs index d0401a4..fddb44e 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/mod.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/mod.rs @@ -67,8 +67,6 @@ use prompt_origin::{push_wake_end_marker, viewer_turn_anchor, wake_turn_elapsed} pub(crate) use subagent_activity::finalize_killed_subagent; use subagent_activity::{subagent_activity_label, sync_subagent_activity}; -#[cfg(test)] -pub(crate) use session_notification::apply_session_event_for_test; use session_notification::{ advance_reconnect_cursor, confirm_context_used, detect_plan_mode_change, drop_unexpected_replay, handle_session_notification, diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/session_notification.rs b/crates/codegen/kigi-tui/src/app/acp_handler/session_notification.rs index b77a93c..ccab2a8 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/session_notification.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/session_notification.rs @@ -342,8 +342,6 @@ pub(super) fn handle_session_notification(notif: &acp::ExtNotification, app: &mu restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, bg_tasks: std::collections::BTreeMap::new(), bg_tool_call_to_task: std::collections::HashMap::new(), scheduled_tasks: std::collections::HashMap::new(), @@ -363,7 +361,6 @@ pub(super) fn handle_session_notification(notif: &acp::ExtNotification, app: &mu child_view.set_input_mode(InputMode::Vim); child_view.is_subagent_view = true; child_view.active_pane = crate::views::agent::ActivePane::Scrollback; - child_view.set_sharing_enabled(agent.sharing_enabled); let usage_visible = agent .prompt .slash_controller @@ -1087,18 +1084,6 @@ pub(super) fn handle_child_session_notification( /// Apply a compaction or retry event to a session's activity state and scrollback. /// /// Shared between the root agent and child (subagent) notification paths. -/// Test-only shim so dispatch-level tests can replay real notification -/// sequences (e.g. `RetryState::Retrying` → `Exhausted`) through the -/// production handler — the Retrying arm clears the `in_flight_prompt` -/// rewind stash, which a fixture setting fields directly would miss. -#[cfg(test)] -pub(crate) fn apply_session_event_for_test( - update: &XaiSessionUpdate, - session: &mut AgentSession, - scrollback: &mut crate::scrollback::state::ScrollbackState, -) -> bool { - apply_session_event(update, session, scrollback, false) -} pub(super) fn apply_session_event( update: &XaiSessionUpdate, session: &mut AgentSession, @@ -1215,7 +1200,6 @@ pub(super) fn apply_retry_state( scrollback: &mut crate::scrollback::state::ScrollbackState, is_api_key_auth: bool, ) { - let mut is_credit_limit = false; let mut is_reauth = false; use kigi_shell::extensions::notification::RetryState; match retry { @@ -1238,14 +1222,7 @@ pub(super) fn apply_retry_state( session.set_retry_activity(None); session.rate_limited = *rate_limited; - is_credit_limit = super::super::dispatch::is_credit_limit_error(None, reason); - let is_free_usage = - *rate_limited && super::super::dispatch::is_free_usage_exhausted_error(reason); - if is_credit_limit { - session.credit_limit_blocked = true; - } else if is_free_usage { - session.free_usage_blocked = true; - } else if !*rate_limited && is_reauthable_failure(None, reason) { + if !*rate_limited && is_reauthable_failure(None, reason) { is_reauth = true; scrollback.push_block(RenderBlock::session_event(SessionEvent::ReAuthRequired)); } else { @@ -1268,10 +1245,7 @@ pub(super) fn apply_retry_state( if error_type == "encrypted_content_mismatch" { session.model_incompatible = true; } - is_credit_limit = super::super::dispatch::is_credit_limit_error(None, message); - if is_credit_limit { - session.credit_limit_blocked = true; - } else if is_reauthable_failure(Some(error_type.as_str()), message) { + if is_reauthable_failure(Some(error_type.as_str()), message) { is_reauth = true; scrollback.push_block(RenderBlock::session_event(SessionEvent::ReAuthRequired)); } else if error_type == "context_length" { @@ -1287,8 +1261,7 @@ pub(super) fn apply_retry_state( } } } - if is_credit_limit { - } else if !is_reauth { + if !is_reauth { session.in_flight_prompt = None; } } diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/settings.rs b/crates/codegen/kigi-tui/src/app/acp_handler/settings.rs index e845c60..e9073b0 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/settings.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/settings.rs @@ -107,23 +107,6 @@ pub(super) fn handle_settings_update(notif: &acp::ExtNotification, app: &mut App if let Some(v) = update.show_resolved_model { app.show_resolved_model = v; } - if let Some(v) = update.sharing_enabled { - app.sharing_enabled = v; - // Propagate to existing agents so slash-command registries stay - // in sync (same fan-out pattern used when creating new agents). - for agent in app.agents.values_mut() { - agent.set_sharing_enabled(v); - } - } - // Always recompute is_api_key_auth from the tier so a later Free/SuperGrok - // stamp does not leave API-key bypass / hidden `/usage` stuck. - if let Some(v) = update.subscription_tier_display { - let is_key = super::super::app_view::is_api_key_label(&v); - app.is_api_key_auth = is_key; - app.usage_visible = !is_key && app.team_name.is_none(); - app.subscription_tier = Some(v); - app.apply_tier_restrictions(); - } // TODO: extract resolve_session_picker_grouped helper (duplicates event_loop.rs:143-160) // Respect env var > config > remote precedence (mirrors event_loop.rs startup). if let Some(remote_val) = update.session_picker_grouped { @@ -142,33 +125,6 @@ pub(super) fn handle_settings_update(notif: &acp::ExtNotification, app: &mut App .unwrap_or(remote_val); app.session_picker_grouped = resolved; } - if let Some(v) = update.subscription_watch_interval_secs { - app.subscription_watch_interval_secs = Some(v); - } - - // Gate update logic: - // - allow_access == Some(true): explicitly granted → lift the gate - // - gate_message.is_some(): server sent a new message → impose/update - // - Neither condition met: don't touch the gate. In particular, - // allow_access=Some(false) without a gate_message must NOT clear the - // gate (gate_from_settings returns None when gate_message is absent, - // which would incorrectly lift an existing gate). - if update.allow_access == Some(true) { - let effs = app.lift_gate(); - app.pending_effects.extend(effs); - } else if let Some(msg) = update.gate_message.as_ref() - && !msg.is_empty() - { - // (An empty gate_message would only clear the gate message text, NOT - // access, so it intentionally does not touch the gate here.) - let effs = app.impose_gate(kigi_shell::auth::GateInfo { - message: msg.clone(), - url: update.gate_url.clone(), - label: update.gate_label.clone(), - }); - app.pending_effects.extend(effs); - } - // Load config layers once for tips + group_tool_verbs + // collapsed_edit_blocks resolution. Loaded unconditionally: the UI flags // re-resolve on every update (see below), and updates are rare (post-auth @@ -361,22 +317,10 @@ pub(super) struct PagerSettingsUpdate { #[serde(default)] show_resolved_model: Option, #[serde(default)] - sharing_enabled: Option, - #[serde(default)] session_picker_grouped: Option, #[serde(default)] tips: Option>, #[serde(default)] - gate_message: Option, - #[serde(default)] - gate_url: Option, - #[serde(default)] - gate_label: Option, - #[serde(default)] - allow_access: Option, - #[serde(default)] - subscription_tier_display: Option, - #[serde(default)] auto_permission_mode_enabled: Option, /// Soft-default permission mode. Presence-aware: omit = no update, /// `null` = recompute with remote=None, string = that soft-default. @@ -389,8 +333,6 @@ pub(super) struct PagerSettingsUpdate { group_tool_verbs: Option, #[serde(default)] collapsed_edit_blocks: Option, - #[serde(default)] - subscription_watch_interval_secs: Option, } /// Presence-aware string: omit → `None` (`#[serde(default)]`), null → diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/tests/mod.rs b/crates/codegen/kigi-tui/src/app/acp_handler/tests/mod.rs index 93ede86..1e4c130 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/tests/mod.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/tests/mod.rs @@ -33,8 +33,6 @@ pub(super) fn make_session(session_id: Option<&str>) -> AgentSession { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, @@ -257,17 +255,6 @@ pub(super) fn follow_ups_ext_with_prompt( std::sync::Arc::from(serde_json::value::to_raw_value(¶ms).unwrap()), ) } -pub(super) fn tier_settings_update(tier: &str) -> acp::ExtNotification { - acp::ExtNotification::new( - "x.ai/settings/update", - std::sync::Arc::from( - serde_json::value::to_raw_value( - &serde_json::json!({ "subscription_tier_display" : tier }), - ) - .unwrap(), - ), - ) -} pub(super) fn group_tool_verbs_settings_update( value: Option, ) -> acp::ExtNotification { diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/tests/plan_mode.rs b/crates/codegen/kigi-tui/src/app/acp_handler/tests/plan_mode.rs index 3406b97..5717257 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/tests/plan_mode.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/tests/plan_mode.rs @@ -119,7 +119,7 @@ let agent = app.agents.get_mut(&AgentId(0)).unwrap(); seed_pending_tool(agent, "create-plan-call", "CreatePlan"); agent.active_modal = Some(crate::views::modal::ActiveModal::CommandPalette { - entries: crate::views::modal::default_palette_entries(agent.sharing_enabled), + entries: crate::views::modal::default_palette_entries(), state: crate::views::picker::PickerState::input_active(), window: crate::views::modal_window::ModalWindowState::new(), }); diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/tests/queue_and_adoption.rs b/crates/codegen/kigi-tui/src/app/acp_handler/tests/queue_and_adoption.rs index 0d7c962..6c0caaa 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/tests/queue_and_adoption.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/tests/queue_and_adoption.rs @@ -1751,24 +1751,7 @@ } /// The credit-limit early return discards the popped adoption's buffer. - #[test] - fn credit_limit_response_discards_adoption_buffer() { - let mut app = app_with_running_p1_and_stashed_b1(); - let id = AgentId(0); - send_tool_call_update(&mut app, "b1", "bash-mode-1", None); - app.agents - .get_mut(&id) - .unwrap() - .session - .credit_limit_blocked = true; - - prompt_response(&mut app, "p1"); - let agent = app.agents.get(&id).unwrap(); - assert!(!app.pending_running_adoptions.contains_key(&id)); - assert!(agent.pending_adoption_updates.is_empty()); - } - - /// A stash whose pid replayed a durable terminal is discarded, never adopted. + /// A stash whose pid replayed a durable terminal is discarded, never adopted. #[test] fn terminal_in_replay_stash_is_discarded_not_adopted() { let mut app = app_with_running_p1_and_stashed_b1(); diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/tests/session_events.rs b/crates/codegen/kigi-tui/src/app/acp_handler/tests/session_events.rs index 45ee0ba..e382c61 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/tests/session_events.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/tests/session_events.rs @@ -157,7 +157,7 @@ apply_retry_state(&exhausted, &mut session, &mut scrollback, false); match last_session_event(&scrollback) { Some(SessionEvent::RetryFailed { error, .. }) => { - assert_eq!(error, RATE_LIMITED_USER_MESSAGE_OAUTH); + assert_eq!(error, RATE_LIMITED_USER_MESSAGE_OAUTH.as_str()); } other => panic!("expected OAuth rate-limit RetryFailed, got {other:?}"), } @@ -194,140 +194,8 @@ ); } - /// A rate-limit exhaustion whose flattened reason carries the - /// free-usage code sets both flags and pushes NO generic block (the - /// driver shows the paywall modal on PromptResponse; viewers keep no - /// marker). #[test] - fn retry_exhausted_free_usage_sets_paywall_flag_without_marker() { - let mut session = make_session(Some("s1")); - let mut scrollback = ScrollbackState::new(); - session.in_flight_prompt = Some(InFlightPrompt { - text: "try me again".into(), - images: Vec::new(), - scrollback_entry: EntryId::new(2), - chip_elements: Vec::new(), - }); - - apply_retry_state( - &RetryState::Exhausted { - attempts: 0, - reason: "API error (status 429 Too Many Requests): \ - subscription:free-usage-exhausted: You have used all your free usage." - .into(), - is_rate_limited: true, - }, - &mut session, - &mut scrollback, - false, - ); - assert!( - session.rate_limited, - "free-usage keeps rate_limited (TurnFailed/toast suppression)" - ); - assert!(session.free_usage_blocked); - assert_eq!( - scrollback.len(), - 0, - "no RetryFailed marker — the paywall modal shows instead" - ); - assert!( - session.in_flight_prompt.is_none(), - "free-usage exhaustion clears in_flight_prompt like other failures" - ); - } - - #[test] - fn apply_retry_state_credit_limit_exhausted_preserves_in_flight_prompt() { - let mut session = make_session(Some("s1")); - let mut scrollback = ScrollbackState::new(); - session.in_flight_prompt = Some(InFlightPrompt { - text: "stash me".into(), - images: Vec::new(), - scrollback_entry: EntryId::new(2), - chip_elements: Vec::new(), - }); - apply_retry_state( - &RetryState::Exhausted { - attempts: 3, - reason: "status 403: run out of credits".into(), - is_rate_limited: false, - }, - &mut session, - &mut scrollback, - false, - ); - assert!( - session.credit_limit_blocked, - "credit_limit_blocked must be set for credit-limit 403" - ); - assert!( - session.in_flight_prompt.is_some(), - "in_flight_prompt must be preserved so PromptResponse handler can stash it" - ); - assert_eq!(session.in_flight_prompt.unwrap().text, "stash me"); - } - - #[test] - fn apply_retry_state_credit_limit_failed_preserves_in_flight_prompt() { - let mut session = make_session(Some("s1")); - let mut scrollback = ScrollbackState::new(); - session.in_flight_prompt = Some(InFlightPrompt { - text: "stash me too".into(), - images: Vec::new(), - scrollback_entry: EntryId::new(3), - chip_elements: Vec::new(), - }); - apply_retry_state( - &RetryState::Failed { - error_type: "proxy_error".into(), - message: "status 403: run out of credits".into(), - }, - &mut session, - &mut scrollback, - false, - ); - assert!( - session.credit_limit_blocked, - "credit_limit_blocked must be set for credit-limit 403" - ); - assert!( - session.in_flight_prompt.is_some(), - "in_flight_prompt must be preserved so PromptResponse handler can stash it" - ); - assert_eq!(session.in_flight_prompt.unwrap().text, "stash me too"); - } - - #[test] - fn apply_retry_state_pool_402_sets_credit_limit_blocked() { - let mut session = make_session(Some("s1")); - let mut scrollback = ScrollbackState::new(); - session.in_flight_prompt = Some(InFlightPrompt { - text: "pool blocked".into(), - images: Vec::new(), - scrollback_entry: EntryId::new(5), - chip_elements: Vec::new(), - }); - apply_retry_state( - &RetryState::Failed { - error_type: "proxy_error".into(), - message: - "API error (status 402 Payment Required): Grok Build usage balance exhausted" - .into(), - }, - &mut session, - &mut scrollback, - false, - ); - assert!( - session.credit_limit_blocked, - "credit_limit_blocked must be set for pool 402 balance exhausted" - ); - assert!(session.in_flight_prompt.is_some()); - } - - #[test] - fn apply_retry_state_non_credit_limit_failed_clears_in_flight_prompt() { + fn apply_retry_state_generic_failed_clears_in_flight_prompt() { let mut session = make_session(Some("s1")); let mut scrollback = ScrollbackState::new(); session.in_flight_prompt = Some(InFlightPrompt { @@ -345,13 +213,9 @@ &mut scrollback, false, ); - assert!( - !session.credit_limit_blocked, - "credit_limit_blocked must NOT be set for non-credit-limit errors" - ); assert!( session.in_flight_prompt.is_none(), - "in_flight_prompt must be cleared for non-credit-limit errors" + "in_flight_prompt must be cleared for generic errors" ); } @@ -400,7 +264,6 @@ ), "auth 401 must surface the actionable re-auth prompt" ); - assert!(!session.credit_limit_blocked); } /// A recoverable auth failure preserves `in_flight_prompt` so the diff --git a/crates/codegen/kigi-tui/src/app/acp_handler/tests/settings.rs b/crates/codegen/kigi-tui/src/app/acp_handler/tests/settings.rs index 89e6c03..c49b171 100644 --- a/crates/codegen/kigi-tui/src/app/acp_handler/tests/settings.rs +++ b/crates/codegen/kigi-tui/src/app/acp_handler/tests/settings.rs @@ -1,41 +1,6 @@ #![cfg_attr(rustfmt, rustfmt::skip)] use super::*; - #[test] - fn settings_non_api_key_tier_clears_stale_api_key_flag() { - let mut app = make_app_with_agent("sess-stale-key"); - assert!(handle_ext_notification( - &tier_settings_update("API Key"), - &mut app - )); - assert!(app.is_api_key_auth); - assert!(!app.usage_visible); - assert!(app.tier_restricted_commands.is_empty()); - - // Later personal Free stamp must not keep the API-key bypass. - assert!(handle_ext_notification( - &tier_settings_update("Free"), - &mut app - )); - assert!(!app.is_api_key_auth); - assert!(app.usage_visible); - // Tier gating no longer exists; nothing gets re-restricted. - assert!(app.tier_restricted_commands.is_empty()); - - // A paid tier after API Key clears the api-key flag and tier limits. - let mut app = make_app_with_agent("sess-paid-tier"); - assert!(handle_ext_notification( - &tier_settings_update("API Key"), - &mut app - )); - assert!(handle_ext_notification( - &tier_settings_update("SuperGrok"), - &mut app - )); - assert!(!app.is_api_key_auth); - assert!(app.tier_restricted_commands.is_empty()); - } - #[test] fn settings_update_clearing_group_tool_verbs_reverts_to_default() { // Expected values come from the same chain the handler resolves, so the diff --git a/crates/codegen/kigi-tui/src/app/actions.rs b/crates/codegen/kigi-tui/src/app/actions.rs index e872a7d..fd6212e 100644 --- a/crates/codegen/kigi-tui/src/app/actions.rs +++ b/crates/codegen/kigi-tui/src/app/actions.rs @@ -56,10 +56,6 @@ pub enum Action { ExitSession, /// Exit session without double-press confirmation (e.g., from command palette). ExitSessionConfirmed, - /// Open grok.com in the browser for SuperGrok subscription upsell. - OpenSupergrokUrl, - /// Re-check subscription status via the shell's `x.ai/auth/check_subscription`. - CheckSubscription, /// Open an arbitrary URL in the system browser (with scheme validation). OpenUrl(String), /// Open grok.com managed connectors, appending session teamId when set. @@ -587,8 +583,6 @@ pub enum Action { TrustFolder, /// A spawned task completed. TaskComplete(TaskResult), - /// Share the current session via URL. - ShareSession, /// Show session info (ID, cwd, model, context usage) instantly. ShowSessionInfo, /// Show release notes in a modal. @@ -602,7 +596,7 @@ pub enum Action { }, /// Show detailed context usage (progress bar, token breakdown, stats). ShowContextInfo, - /// Show credit usage via /usage command. + /// Show Kimi usage/quota via the /usage command. ShowUsage, /// Commit a read-only list of the queued prompts as a system block /// (`/queue`). The surface minimal mode uses in place of the `QueuePane`. @@ -663,11 +657,6 @@ pub enum Action { TriggerDeepSearch, /// Force an immediate deep content search, skipping the debounce. ForceDeepSearch, - /// Show privacy and data retention status. - ShowPrivacyInfo, - SetCodingDataSharing { - opted_in: bool, - }, /// `/fork` slash command: parsed args produced by /// [`crate::slash::commands::fork::parse_fork_args`]. The dispatcher /// resolves the worktree question (via flag or the local @@ -1698,11 +1687,6 @@ pub enum Effect { tool_name: String, enabled: bool, }, - /// Share the current session via URL. - ShareSession { - agent_id: AgentId, - session_id: acp::SessionId, - }, /// Fetch and display session info via x.ai/session/info. ShowSessionInfo { agent_id: AgentId, @@ -1777,19 +1761,6 @@ pub enum Effect { }, /// Log out via `x.ai/auth/logout` (shell clears auth.json + in-memory state). Logout, - /// Re-check subscription status via `x.ai/auth/check_subscription`. - /// `verify` scopes the result to a deferred-gate verification (see - /// [`crate::app::subscription`]); `None` for generic checks. - CheckSubscription { verify: Option }, - /// One-shot subscription re-check triggered by a credit-limit 403. - /// If the tier changed, the stashed prompt is retried instead of - /// showing the upsell modal. - CreditLimitRecheck { agent_id: AgentId }, - /// Schedule a 5s timer that fires `TaskResult::PaywallCheckTick`. - SchedulePaywallCheck, - /// Schedule `TaskResult::GateVerifyTimeout { generation }` after - /// [`crate::app::subscription::GATE_VERIFY_TIMEOUT`]. - ScheduleGateVerifyTimeout { generation: u64 }, /// Log out then authenticate sequentially in one task. SwitchAccount { request_seq: u64, @@ -1808,13 +1779,6 @@ pub enum Effect { UnregisterActiveSession { session_id: acp::SessionId }, /// Quit the application. Quit, - /// Toggle coding data sharing via ACP. - SetCodingDataSharing { - agent_id: AgentId, - opted_in: bool, - /// Pre-toggle value to revert to on failure. - rollback_to_opted_in: bool, - }, /// Rename the current session. RenameSession { agent_id: AgentId, @@ -1871,16 +1835,9 @@ pub enum Effect { target_prompt_index: usize, mode: crate::views::rewind::RewindMode, }, - /// Fetch billing/credit usage from the agent's `x.ai/billing` extension. - /// When `silent` is true the result updates `credit_balance` without - /// pushing a system message into scrollback (used for automatic refreshes - /// on session init and after each turn). - FetchBilling { agent_id: AgentId, silent: bool }, - /// Fetch billing data at the app level (no agent required). - /// Used on startup to populate the welcome-screen credit warning. - FetchAppBilling, - /// Re-fetch remote settings to check subscription gate. - RefreshGate, + /// Fetch Kimi usage/quota rows from the agent's `x.ai/billing` + /// extension (`GET {base}/usages` shell-side) for the `/usage` view. + FetchUsage { agent_id: AgentId }, /// Spawn a debounce sleep task for shell suggestions. `agent_id` rides /// to the expiry so the fetch is built from the arming agent, not /// whatever view is active when the timer fires. @@ -2019,9 +1976,6 @@ pub enum TaskResult { /// Session list fetched for the welcome screen picker. SessionListLoaded { sessions: Vec, - /// Degraded conversations lane (`_meta["x.ai/partial"]`), surfaced - /// as an actionable picker notice instead of a silent empty list. - partial: Option, /// Echo of [`Effect::FetchSessionList::seq`]; stale results are dropped. seq: u64, /// Echo of [`Effect::FetchSessionList::query`]. `Some` marks the @@ -2256,16 +2210,6 @@ pub enum TaskResult { agent_id: AgentId, result: Result<(), String>, }, - /// Share session completed successfully. - ShareSessionComplete { - agent_id: AgentId, - share_url: String, - }, - /// Share session failed. - ShareSessionFailed { - agent_id: AgentId, - error: String, - }, /// Session info fetched successfully. SessionInfoComplete { agent_id: AgentId, @@ -2277,17 +2221,6 @@ pub enum TaskResult { agent_id: AgentId, error: String, }, - /// Coding data sharing preference updated. - CodingDataSharingUpdated { - agent_id: AgentId, - opted_in: bool, - }, - /// Coding data sharing update failed. - CodingDataSharingFailed { - agent_id: AgentId, - error: String, - rollback_to_opted_in: bool, - }, /// Session rename completed successfully. RenameSessionComplete { agent_id: AgentId, @@ -2405,25 +2338,6 @@ pub enum TaskResult { }, /// Shell acknowledged logout (auth cleared). LogoutComplete, - /// Shell responded to `x.ai/auth/check_subscription`. `verify` echoes - /// the generation from `Effect::CheckSubscription` for deferred-gate - /// verifications. - CheckSubscriptionComplete { - verify: Option, - meta: Option, - }, - /// Result of the credit-limit subscription re-check. If the tier - /// changed the stashed prompt is retried; otherwise the upsell is shown. - CreditLimitRecheckComplete { - agent_id: AgentId, - meta: Option, - }, - /// 5s paywall check timer fired -- time to send another check. - PaywallCheckTick, - /// The deferred-gate verification window expired. - GateVerifyTimeout { - generation: u64, - }, /// The 2-second "copied!" display timer expired. AuthCopiedTimeout, DeepSearchResults { @@ -2470,30 +2384,10 @@ pub enum TaskResult { agent_id: AgentId, error: String, }, - /// Billing data fetched from the agent. - BillingFetched { + /// Kimi usage/quota rows fetched for the `/usage` view. + UsageFetched { agent_id: AgentId, - balance: Option, - /// When true, update `credit_balance` silently (no scrollback message). - silent: bool, - /// Subscription tier piggybacked from remote settings. - subscription_tier: Option, - /// Auto top-up rule fetch result; `Unchanged` keeps any cached rule. - autotopup: crate::views::credit_bar::AutoTopupFetch, - }, - /// App-level billing data (welcome screen). - AppBillingFetched { - balance: Option, - autotopup: crate::views::credit_bar::AutoTopupFetch, - }, - GateRefreshed { - settings: Option, - }, - BillingError { - agent_id: AgentId, - error: String, - /// When true, swallow the error silently (background refresh). - silent: bool, + result: Result, String>, }, /// Debounce timer for shell suggestions expired. Routed by the arming /// `agent_id`. diff --git a/crates/codegen/kigi-tui/src/app/agent.rs b/crates/codegen/kigi-tui/src/app/agent.rs index 96a76eb..92ba9aa 100644 --- a/crates/codegen/kigi-tui/src/app/agent.rs +++ b/crates/codegen/kigi-tui/src/app/agent.rs @@ -653,16 +653,6 @@ pub struct AgentSession { /// fires, so the subsequent `TurnFailed` can be suppressed (the retry handler /// already displayed a user-friendly message). Cleared on `finish_turn`. pub model_incompatible: bool, - /// Set when a `RetryState::Failed` carries a 403 credit-limit error, so - /// the error message is suppressed in favour of the upsell modal. - /// Cleared on `finish_turn`. - pub credit_limit_blocked: bool, - /// Set when a rate-limit `RetryState::Exhausted` carries the - /// `subscription:free-usage-exhausted` code, so the PromptResponse - /// handler shows the free-usage paywall instead of the generic - /// rate-limit message. Always set together with [`Self::rate_limited`]. - /// Cleared on `finish_turn`. - pub free_usage_blocked: bool, pub(crate) tracker: AcpUpdateTracker, /// ACP-advertised slash commands. Seeded from `InitializeResponse.meta`, /// updated by `AvailableCommandsUpdate`. The prompt-side registry syncs @@ -796,8 +786,6 @@ impl AgentSession { self.state = AgentState::Idle; self.rate_limited = false; self.model_incompatible = false; - self.credit_limit_blocked = false; - self.free_usage_blocked = false; self.in_flight_prompt = None; self.current_prompt_id = None; } @@ -999,8 +987,6 @@ mod tests { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, diff --git a/crates/codegen/kigi-tui/src/app/agent_view/input.rs b/crates/codegen/kigi-tui/src/app/agent_view/input.rs index 6447cef..832dc3f 100644 --- a/crates/codegen/kigi-tui/src/app/agent_view/input.rs +++ b/crates/codegen/kigi-tui/src/app/agent_view/input.rs @@ -678,7 +678,6 @@ impl AgentView { .hit_plan_approval_status .update_hover(mouse.column, mouse.row); changed |= self.hit_context.update_hover(mouse.column, mouse.row); - changed |= self.hit_credits.update_hover(mouse.column, mouse.row); } MouseEventKind::Down(MouseButton::Left) => { if self.hit_plan_button.contains(mouse.column, mouse.row) { @@ -943,7 +942,7 @@ impl AgentView { || (key.code == KeyCode::Char('/') && key.modifiers.contains(KeyModifiers::SHIFT))) { self.active_modal = Some(crate::views::modal::ActiveModal::CommandPalette { - entries: crate::views::modal::default_palette_entries(self.sharing_enabled), + entries: crate::views::modal::default_palette_entries(), state: crate::views::picker::PickerState::input_active(), window: crate::views::modal_window::ModalWindowState::new(), }); @@ -1016,7 +1015,7 @@ impl AgentView { } ActionId::CommandPalette => { self.active_modal = Some(crate::views::modal::ActiveModal::CommandPalette { - entries: crate::views::modal::default_palette_entries(self.sharing_enabled), + entries: crate::views::modal::default_palette_entries(), state: crate::views::picker::PickerState::input_active(), window: crate::views::modal_window::ModalWindowState::new(), }); @@ -1462,7 +1461,7 @@ mod focus_gained_restore_tests { agent.session.state = AgentState::TurnRunning; with_permission(&mut agent); agent.active_modal = Some(ActiveModal::CommandPalette { - entries: crate::views::modal::default_palette_entries(false), + entries: crate::views::modal::default_palette_entries(), state: crate::views::picker::PickerState::input_active(), window: crate::views::modal_window::ModalWindowState::new(), }); diff --git a/crates/codegen/kigi-tui/src/app/agent_view/interactions.rs b/crates/codegen/kigi-tui/src/app/agent_view/interactions.rs index 1ec5a58..e0c4a1b 100644 --- a/crates/codegen/kigi-tui/src/app/agent_view/interactions.rs +++ b/crates/codegen/kigi-tui/src/app/agent_view/interactions.rs @@ -1261,8 +1261,6 @@ mod cancel_turn_mouse_tests { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, diff --git a/crates/codegen/kigi-tui/src/app/agent_view/mod.rs b/crates/codegen/kigi-tui/src/app/agent_view/mod.rs index b343299..2d2c745 100644 --- a/crates/codegen/kigi-tui/src/app/agent_view/mod.rs +++ b/crates/codegen/kigi-tui/src/app/agent_view/mod.rs @@ -727,10 +727,6 @@ pub struct AgentView { /// Stashed normal prompt state while editing a queued prompt. /// Restored when editing ends. pub stashed_prompt: Option, - /// Complete prompt stashed from a credit-limit-blocked turn. Used by - /// `CreditLimitRecheckComplete` to retry the prompt after a tier - /// upgrade instead of showing a stale upsell. - pub credit_limit_stashed_prompt: Option, /// Complete prompt stashed from a turn that failed because the login /// expired (401 / re-auth). Used by the `AuthComplete` handler to /// auto-resubmit the prompt after a successful mid-session re-auth so @@ -754,10 +750,6 @@ pub struct AgentView { /// Unlike `chat_kind`, stays `false` for a `/chat` one-shot session in /// a Build process, whose picker still lists local sessions. pub app_chat_mode: bool, - /// Mocked credit balance for the status bar indicator. - pub credit_balance: Option, - /// Auto top-up rule paired with `credit_balance` for the prompt warning. - pub auto_topup: Option, /// Current goal orchestration state. Set by `GoalUpdated` session /// notifications, cleared when a new session starts. pub goal_state: Option, @@ -947,7 +939,6 @@ pub struct AgentView { pub hovered_prompt: bool, pub hit_badge: HitArea, pub hit_context: HitArea, - pub hit_credits: HitArea, pub hit_todo_close: HitArea, pub hit_bg_close: HitArea, pub hit_subagent_close: HitArea, @@ -1235,8 +1226,6 @@ pub struct AgentView { /// Hit area for the [✗] close button in the subagent frame title bar. pub hit_subagent_frame_close: HitArea, /// Whether the `/share` slash command is available (mirrors - /// `AppView::sharing_enabled`). Used to gate palette entries. - pub sharing_enabled: bool, /// Input flight recorder — rolling buffer of recent key events. /// Dumped to file via Esc→d combo for debugging. pub(crate) input_log: crate::input_log::InputRingBuffer, @@ -1512,23 +1501,6 @@ fn translate_local_submit( persist_mode, }) } - LocalQuestionKind::CreditLimitUpsell => { - let q = qv.questions.first(); - let url = q - .and_then(|q| q.options.get(*idx)) - .and_then(|o| o.id.as_deref()) - .unwrap_or(super::dispatch::UPSELL_URL_PAYG); - InputOutcome::Action(Action::OpenUrl(url.to_string())) - } - LocalQuestionKind::FreeUsageUpsell => { - let url = qv - .questions - .first() - .and_then(|q| q.options.get(*idx)) - .and_then(|o| o.id.as_deref()) - .unwrap_or(super::dispatch::UPSELL_URL_UPGRADE); - InputOutcome::Action(Action::OpenUrl(url.to_string())) - } LocalQuestionKind::AgentTypeMismatch { model_id, effort } => { let start_new = *idx == 0; InputOutcome::Action(Action::AgentTypeMismatchAnswered { @@ -2217,8 +2189,6 @@ pub(super) mod test_fixtures { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, @@ -2278,8 +2248,6 @@ pub(super) mod test_fixtures { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, @@ -3010,8 +2978,6 @@ pub(crate) fn test_agent_view(session_id: Option<&str>, cwd: std::path::PathBuf) restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, diff --git a/crates/codegen/kigi-tui/src/app/agent_view/paste.rs b/crates/codegen/kigi-tui/src/app/agent_view/paste.rs index d730c88..2f18a4c 100644 --- a/crates/codegen/kigi-tui/src/app/agent_view/paste.rs +++ b/crates/codegen/kigi-tui/src/app/agent_view/paste.rs @@ -454,8 +454,6 @@ pub(super) mod paste_key_tests { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, diff --git a/crates/codegen/kigi-tui/src/app/agent_view/plan.rs b/crates/codegen/kigi-tui/src/app/agent_view/plan.rs index 7165346..d8fb796 100644 --- a/crates/codegen/kigi-tui/src/app/agent_view/plan.rs +++ b/crates/codegen/kigi-tui/src/app/agent_view/plan.rs @@ -703,8 +703,6 @@ mod plan_chip_tests { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, diff --git a/crates/codegen/kigi-tui/src/app/agent_view/render.rs b/crates/codegen/kigi-tui/src/app/agent_view/render.rs index 656d785..ec03ff1 100644 --- a/crates/codegen/kigi-tui/src/app/agent_view/render.rs +++ b/crates/codegen/kigi-tui/src/app/agent_view/render.rs @@ -329,7 +329,6 @@ impl AgentView { let thinking_label = self.scrollback.thinking_fold_label(); let selected_is_user_prompt = selected_entry.is_some_and(|e| e.block.is_user_prompt()); let selected_is_agent_message = selected_entry.is_some_and(|e| e.block.is_agent_message()); - let selected_is_credit_limit = selected_entry.is_some_and(|e| e.block.is_credit_limit()); let mut hints = agent::build_hints( self.active_pane, &self.prompt, @@ -355,7 +354,6 @@ impl AgentView { !self.visible_queue_is_empty(), selected_is_user_prompt, selected_is_agent_message, - selected_is_credit_limit, crate::terminal::terminal_context().shift_enter_unavailable(), self.scrollback_search.as_ref(), ); @@ -1300,7 +1298,6 @@ impl AgentView { self.hit_bg_status.rect = areas.get("bg_tasks").copied(); self.hit_goal_status.rect = areas.get("goal").copied(); self.hit_context.rect = areas.get("context").copied(); - self.hit_credits.rect = areas.get("credits").copied(); self.hit_plan_button.rect = areas.get("plan").copied(); self.hit_queue_badge.rect = areas.get("queue").copied(); self.hit_badge.rect = areas.get("badge").copied(); @@ -2041,23 +2038,6 @@ impl AgentView { } let mode_flags: &[PromptFlag] = &mode_flags_vec; let multiline = self.multiline_mode; - let usage_visible = self - .prompt - .slash_controller - .registry() - .get("usage") - .is_some(); - let warning = self.credit_balance.as_ref().and_then(|bal| { - crate::views::credit_bar::usage_warning_for_session( - bal, - self.auto_topup.as_ref(), - usage_visible, - self.chat_kind, - ) - }); - let usage_warning_text: Option = warning.as_ref().map(|(t, _)| t.clone()); - let usage_warning = usage_warning_text.as_deref(); - let usage_warning_critical = warning.is_some_and(|(_, critical)| critical); let model_label = match self.session.models.reasoning_effort { Some(eff) => format!("{model_id} ({eff})"), None => model_id, @@ -2067,8 +2047,6 @@ impl AgentView { model_name: &model_label, flags: mode_flags, multiline, - usage_warning, - usage_warning_critical, }, PromptMode::EditingQueued { id, .. } => { let pos = self.session.queue_position(*id).map(|i| i + 1).unwrap_or(1); @@ -2077,8 +2055,6 @@ impl AgentView { model_name: &editing_label, flags: mode_flags, multiline, - usage_warning, - usage_warning_critical, } } }; @@ -2087,8 +2063,6 @@ impl AgentView { model_name: label, flags: &[], multiline: false, - usage_warning, - usage_warning_critical, } } else { info diff --git a/crates/codegen/kigi-tui/src/app/agent_view/rewind.rs b/crates/codegen/kigi-tui/src/app/agent_view/rewind.rs index 5834d58..520d8b0 100644 --- a/crates/codegen/kigi-tui/src/app/agent_view/rewind.rs +++ b/crates/codegen/kigi-tui/src/app/agent_view/rewind.rs @@ -206,8 +206,6 @@ mod sync_rewind_anchor_to_picker_tests { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, diff --git a/crates/codegen/kigi-tui/src/app/agent_view/session.rs b/crates/codegen/kigi-tui/src/app/agent_view/session.rs index 7ba1ae9..19561f3 100644 --- a/crates/codegen/kigi-tui/src/app/agent_view/session.rs +++ b/crates/codegen/kigi-tui/src/app/agent_view/session.rs @@ -1,5 +1,5 @@ //! Session lifecycle: bind/reload/replay bookkeeping, turn activity -//! resolution, context/credit updates, and app-scoped gates. +//! resolution, context updates, and app-scoped gates. #[cfg(test)] use super::test_agent_view; use super::{ @@ -85,7 +85,6 @@ impl AgentView { bash_turn: false, cron_task_id: None, stashed_prompt: None, - credit_limit_stashed_prompt: None, reauth_stashed_prompt: None, active_modal: None, modal_buttons: Vec::new(), @@ -93,8 +92,6 @@ impl AgentView { context_state: None, chat_kind: false, app_chat_mode: false, - credit_balance: None, - auto_topup: None, goal_state: None, parked_wait_marker_for: None, end_work_announced: false, @@ -152,7 +149,6 @@ impl AgentView { hovered_prompt: false, hit_badge: Default::default(), hit_context: Default::default(), - hit_credits: Default::default(), hit_todo_close: Default::default(), hit_bg_close: Default::default(), hit_subagent_close: Default::default(), @@ -257,7 +253,6 @@ impl AgentView { active_subagent: None, is_subagent_view: false, hit_subagent_frame_close: Default::default(), - sharing_enabled: false, input_log: crate::input_log::InputRingBuffer::new(), esc_pressed_at: None, pending_first_prompt: None, @@ -740,21 +735,6 @@ impl AgentView { } } } - /// Apply Build coding-credit balance only for non-chat agents. - /// Gateway/chat-kind sessions keep credits unset so bars/warnings stay off. - pub fn apply_credit_balance( - &mut self, - balance: Option, - auto_topup: Option, - ) { - if self.chat_kind { - self.credit_balance = None; - self.auto_topup = None; - return; - } - self.credit_balance = balance; - self.auto_topup = auto_topup; - } /// Record a key event to the input flight recorder. /// /// Zero heap allocations — stores raw `Copy` types in the ring buffer. @@ -801,18 +781,6 @@ impl AgentView { textarea_changed: delta.textarea_changed, }); } - /// Set the sharing-enabled flag on this view and propagate it to the - /// slash-command registry so the `/share` entry stays hidden/visible in - /// lockstep with `AgentView::sharing_enabled`. Use this instead of - /// mutating `sharing_enabled` directly when a new agent is created or a - /// session is loaded, so the field and registry can't drift. - pub fn set_sharing_enabled(&mut self, enabled: bool) { - self.sharing_enabled = enabled; - self.prompt - .slash_controller - .registry_mut() - .set_share_visible(enabled); - } /// Show or hide the `/usage` slash command in this agent's registry. pub fn set_usage_visible(&mut self, visible: bool) { self.prompt @@ -839,13 +807,11 @@ impl AgentView { /// One place for the app-scoped gates a new/adopted session inherits so the session-creation sites cannot drift. pub(crate) fn apply_app_scoped_gates( &mut self, - sharing_enabled: bool, usage_visible: bool, chat_mode: bool, screen_mode: crate::app::ScreenMode, restricted_commands: &[String], ) { - self.set_sharing_enabled(sharing_enabled); self.set_usage_visible(usage_visible); self.app_chat_mode = chat_mode; self.prompt.set_screen_mode(screen_mode); diff --git a/crates/codegen/kigi-tui/src/app/app_view.rs b/crates/codegen/kigi-tui/src/app/app_view.rs index bffc134..7fef5b4 100644 --- a/crates/codegen/kigi-tui/src/app/app_view.rs +++ b/crates/codegen/kigi-tui/src/app/app_view.rs @@ -386,28 +386,6 @@ fn parse_esc_ttl(raw: Option) -> Duration { .map(|ms| Duration::from_millis(ms.min(ESC_DOUBLE_PRESS_TEST_MS))) .unwrap_or(PendingAction::ESC_DOUBLE_PRESS_TTL) } -/// Slash commands unavailable on the free and X Basic subscription tiers. -/// -/// To restrict another command for these tiers, add its canonical name -/// (no leading `/`) here — matching covers aliases automatically via -/// [`crate::slash::registry::CommandRegistry::set_restricted_commands`]. -/// -/// Current set: -/// - `usage` — coding credit / billing UI (alias: `/cost`) -/// - `imagine` — image generation entry point -/// - `imagine-video` — video generation entry point -pub(crate) const TIER_RESTRICTED_COMMANDS: &[&str] = &["usage", "imagine", "imagine-video"]; -/// Whether a subscription-tier display name is a tier with restricted -/// commands: the free tier (no subscription ⇒ `None`, or an explicit -/// "Free") and X Basic (CCP display name "X Basic"; JWT claim fallback -/// "x_basic"). Everything else — paid tiers and unknown future names — -/// is unrestricted (fail-open). -/// -/// Tier gating was an xAI concept; the Kimi Code subscription has no -/// client-visible tier, so nothing is ever restricted. -fn is_restricted_tier(_tier: Option<&str>) -> bool { - false -} /// True for API-key labels from shell/CCP: `"ApiKey"`, `"API Key"`, `"api_key"`. pub(crate) fn is_api_key_label(s: &str) -> bool { s.trim().to_ascii_lowercase().replace([' ', '_', '-'], "") == "apikey" @@ -519,19 +497,9 @@ pub struct AppView { pub tip: Option, /// Whether to show the resolved model ID in /session-info output. pub show_resolved_model: bool, - /// Whether the `/share` slash command is available. Gated by - /// `RemoteSettings.sharing_enabled`; defaults to `false` when remote - /// settings are unavailable or the field is absent. - pub sharing_enabled: bool, /// Whether the `/usage` slash command is available. Hidden for team /// (`team_name.is_some()`) and API-key auth. pub usage_visible: bool, - /// Slash commands denied for the current subscription tier - /// ([`TIER_RESTRICTED_COMMANDS`] when the user is on the free / X Basic - /// tier, empty otherwise). Recomputed by [`Self::apply_tier_restrictions`] - /// and fanned out to every slash registry (welcome prompt, agents, - /// dashboard); deny wins over all other visibility gates. - pub tier_restricted_commands: Vec, /// Whether the pager is connected via a leader (leader mode). The Agent /// Dashboard entry points (`/dashboard`, `Ctrl+\`, `grok dashboard`, the /// startup hook) are only meaningful when a leader is coordinating a @@ -539,13 +507,6 @@ pub struct AppView { /// `event_loop::run` from `connection.leader_status_rx.is_some()`; /// defaults to `false` (non-leader, dashboard hidden). pub leader_mode: bool, - /// App-level credit balance used to show the usage warning on the - /// welcome screen before any agent session exists. - pub credit_balance: Option, - /// App-level auto top-up rule paired with `credit_balance` for the warning. - pub auto_topup: Option, - /// Periodic billing poll requested (credits >= 99%). - pub billing_poll_wanted: bool, /// Leader-mode session roster (FleetView dashboard). Populated from /// `x.ai/sessions/list` polls and `x.ai/sessions/changed` broadcasts. /// Empty in non-leader mode, which naturally gates roster rendering. @@ -667,9 +628,7 @@ pub struct AppView { /// Hit-test rect for the "show full URL" fallback link. pub welcome_auth_fallback_rect: Option, /// Hit-test rect for the "[Refresh]" button on the paywall tier line. - pub welcome_refresh_rect: Option, /// Hit-test rect for the gate URL link on the paywall CTA. - pub welcome_gate_url_rect: Option, /// Hit-test rect for the clickable changelog info block (opens release notes). pub welcome_changelog_cta_rect: Option, /// Show the raw auth URL with mouse capture disabled for manual copy. @@ -833,16 +792,6 @@ pub struct AppView { pub auth_use_oauth: bool, /// Whether the last clipboard copy during auth succeeded. pub auth_clipboard_copied: bool, - /// Team principal UUID from auth (`None` for personal sessions). - pub team_id: Option, - /// Team name from auth (displayed in the shortcuts bar). - pub team_name: Option, - /// Whether the user's team has enterprise Zero Data Retention enabled. - pub is_zdr: bool, - /// Team role (e.g. "Admin", "Member", "Read Only") for access-control checks. - pub team_role: Option, - /// Whether the user has opted out of coding data retention. - pub coding_data_retention_opt_out: bool, /// Persisted `[cli].show_tips` mirror. `None` = no override (default `true`). pub show_tips: Option, /// Persisted `[cli].auto_update` mirror. `None` = no override (default `true`). @@ -851,31 +800,6 @@ pub struct AppView { /// from the effective TOML merge like `show_tips`. `None` = unset in TOML /// (default `true`); toggles write the user layer. pub ask_user_question_timeout_enabled: Option, - /// Whether ZDR users are allowed to use the product. - /// Server-controlled via RemoteSettings (remote settings). Default `false` (blocked) during beta. - pub zdr_access_enabled: bool, - /// When set, `/usage` shows a link to this URL instead of fetching billing - /// data from the backend. Server-controlled via RemoteSettings (remote settings - /// `grok_build_usage_redirect_url`, targeted at personal-team users). - /// `None` (default) fetches usage from the backend. - pub usage_billing_redirect_url: Option, - pub access_gate_shown_logged: bool, - /// Access gate from `grok_build_access_gate`. `Some` = blocked. - pub gate: Option, - /// User-friendly subscription tier name (e.g. "SuperGrok", "Free"). - pub subscription_tier: Option, - /// When the pager started auto-checking subscriptions (for 10-min timeout). - pub paywall_check_started: Option, - /// Debounce stamp for watch/focus subscription checks (see - /// [`super::subscription`]). - pub last_subscription_check_at: Option, - /// Server override (seconds) for the subscription-watch cadence. - pub subscription_watch_interval_secs: Option, - /// A stale-source gate held out of `gate` while a live check verifies - /// it (see [`super::subscription`]). - pub pending_gate_verification: Option, - /// Generation stamp of the current gate verification. - pub gate_verify_gen: u64, /// Whether a leader reconnect is in progress (blocks prompt submission). pub reconnect_pending: bool, /// Structured startup warnings collected from the terminal diagnostics @@ -923,17 +847,6 @@ pub struct AppView { pub(crate) keyboard_normalizer: KeyboardNormalizer, } impl AppView { - pub fn is_zdr_blocked(&self) -> bool { - self.is_zdr && !self.zdr_access_enabled - } - /// User is not gated (no gate from remote settings or subscription fallback). - pub fn has_access(&self) -> bool { - self.gate.is_none() - } - /// True when the user should not see the prompt (gate, subscription, or ZDR). - pub fn is_access_blocked(&self) -> bool { - !self.has_access() || self.is_zdr_blocked() - } /// Whether deferred session-startup actions may run: both auth AND folder /// trust must be resolved. Mirrors the auth gate at the session-creating /// startup sites; trust is gated AFTER auth so a pending trust question @@ -941,33 +854,12 @@ impl AppView { pub fn session_startup_allowed(&self) -> bool { matches!(self.auth_state, AuthState::Done) && matches!(self.trust_state, TrustState::Done) } - /// Extract `GateInfo` from `RemoteSettings`. - pub fn gate_from_settings( - rs: &kigi_shell::util::config::RemoteSettings, - ) -> Option { - let msg = rs.gate_message.as_ref()?; - if msg.is_empty() { - return None; - } - Some(kigi_shell::auth::GateInfo { - message: msg.clone(), - url: rs.gate_url.clone(), - label: rs.gate_label.clone(), - }) - } - /// Apply typed auth metadata from the shell. The Kimi auth model carries - /// no team/tier/gate info; those fields only ever come from remote - /// settings now. + /// Apply typed auth metadata from the shell: auth mode (drives the + /// API-key badge and `/usage` visibility) and the resolved-model display + /// preference. The Kimi auth model carries no team/tier/gate info. pub fn apply_auth_meta(&mut self, meta: &kigi_shell::auth::AuthMeta) { - self.pending_gate_verification = None; - let was_gated = self.gate.is_some(); - self.gate = None; - if was_gated { - self.paywall_check_started = None; - } self.is_api_key_auth = meta.auth_mode.as_deref().is_some_and(is_api_key_label); self.usage_visible = !self.is_api_key_auth; - self.apply_tier_restrictions(); if let Some(show) = meta.show_resolved_model { self.show_resolved_model = show; } @@ -1038,8 +930,6 @@ impl AppView { welcome_on_auth_url: false, welcome_on_changelog_cta: false, welcome_auth_fallback_rect: None, - welcome_refresh_rect: None, - welcome_gate_url_rect: None, welcome_changelog_cta_rect: None, auth_show_raw_url: false, auth_mouse_disabled: false, @@ -1098,24 +988,9 @@ impl AppView { deferred_startup: Default::default(), auth_use_oauth: false, auth_clipboard_copied: false, - team_id: None, - team_name: None, - is_zdr: false, - team_role: None, - coding_data_retention_opt_out: false, show_tips: None, auto_update: None, ask_user_question_timeout_enabled: None, - zdr_access_enabled: false, - usage_billing_redirect_url: None, - access_gate_shown_logged: false, - gate: None, - subscription_tier: None, - paywall_check_started: None, - last_subscription_check_at: None, - subscription_watch_interval_secs: None, - pending_gate_verification: None, - gate_verify_gen: 0, reconnect_pending: false, startup_warnings: Vec::new(), is_api_key_auth: false, @@ -1129,13 +1004,8 @@ impl AppView { welcome_doc_viewer: None, screen_mode: ScreenMode::Inline, show_resolved_model: true, - sharing_enabled: false, usage_visible: true, - tier_restricted_commands: Vec::new(), leader_mode: false, - credit_balance: None, - auto_topup: None, - billing_poll_wanted: false, leader_roster: Vec::new(), dashboard_local_sessions: Vec::new(), dashboard_sessions_loading: false, @@ -1178,35 +1048,6 @@ impl AppView { dashboard.set_auto_mode_available(available); } } - /// Recompute the tier-restricted slash commands from the current auth - /// state and sync the deny list into every slash surface (welcome - /// prompt, all agents, dashboard) so restricted commands hide/show in - /// lockstep. - /// - /// Called from [`Self::apply_auth_meta`] (startup / login) and from the - /// `x.ai/settings/update` handler when the subscription tier changes, so - /// a mid-session upgrade lifts the restrictions without a restart. - pub fn apply_tier_restrictions(&mut self) { - let restricted = self.team_name.is_none() - && !self.is_api_key_auth - && is_restricted_tier(self.subscription_tier.as_deref()); - let names: Vec = if restricted { - TIER_RESTRICTED_COMMANDS - .iter() - .map(|n| (*n).to_string()) - .collect() - } else { - Vec::new() - }; - for agent in self.agents.values_mut() { - agent.set_restricted_commands(&names); - } - self.welcome_prompt.set_restricted_commands(&names); - if let Some(dashboard) = self.dashboard.as_mut() { - dashboard.set_restricted_commands(&names); - } - self.tier_restricted_commands = names; - } /// Session ID of the active agent, if one exists and has an established session. pub fn active_session_id(&self) -> Option<&str> { match self.active_view { @@ -1667,8 +1508,6 @@ impl AppView { ); if is_mouse_action {} } - let zdr_blocked = self.is_zdr_blocked(); - let has_access = self.has_access(); let has_foreign_resume = self.foreign_resume_hint().is_some(); let outcome = match self.active_view { ActiveView::Welcome => handle_welcome_input( @@ -1684,27 +1523,20 @@ impl AppView { new_worktree_dialog: &mut self.new_worktree_dialog, menu_index: &mut self.welcome_menu_index, menu_rects: &self.welcome_menu_rects, - menu_count: if zdr_blocked { - 2 - } else { - 3 + if self.has_claude_import { 1 } else { 0 } - + if self.welcome_show_changelog_action { - 1 - } else { - 0 - } - }, + menu_count: 3 + + if self.has_claude_import { 1 } else { 0 } + + if self.welcome_show_changelog_action { + 1 + } else { + 0 + }, prompt_rect: self.welcome_prompt_rect.as_ref(), import_banner_rect: self.welcome_import_banner_rect.as_ref(), auth_url_rect: self.welcome_auth_url_rect.as_ref(), auth_fallback_rect: self.welcome_auth_fallback_rect.as_ref(), - refresh_rect: self.welcome_refresh_rect.as_ref(), - gate_url_rect: self.welcome_gate_url_rect.as_ref(), changelog_cta_rect: self.welcome_changelog_cta_rect.as_ref(), on_changelog_cta: &mut self.welcome_on_changelog_cta, show_raw_url: &mut self.auth_show_raw_url, - has_access, - is_zdr_blocked: zdr_blocked, sp_entries: &mut self.session_picker_entries, sp_state: &mut self.session_picker_state, sp_content_results: &self.session_picker_content_results, @@ -2228,15 +2060,11 @@ struct WelcomeInputCtx<'a> { import_banner_rect: Option<&'a ratatui::layout::Rect>, auth_url_rect: Option<&'a ratatui::layout::Rect>, auth_fallback_rect: Option<&'a ratatui::layout::Rect>, - refresh_rect: Option<&'a ratatui::layout::Rect>, - gate_url_rect: Option<&'a ratatui::layout::Rect>, /// Hit-test rect for the clickable changelog info block (opens release notes). changelog_cta_rect: Option<&'a ratatui::layout::Rect>, /// Sticky hover flag for the changelog block (redraw on enter/leave). on_changelog_cta: &'a mut bool, show_raw_url: &'a mut bool, - has_access: bool, - is_zdr_blocked: bool, sp_entries: &'a mut Option>, sp_state: &'a mut crate::views::picker::PickerState, sp_content_results: &'a Option>, @@ -2360,8 +2188,6 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco return InputOutcome::Unchanged; } if matches!(ctx.auth_state, AuthState::Done) - && ctx.has_access - && !ctx.is_zdr_blocked && matches!(ctx.trust_state, TrustState::Pending { .. }) { if let Event::Key(key) = ev { @@ -2603,22 +2429,6 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco if key.kind == KeyEventKind::Release { return InputOutcome::Unchanged; } - if ctx.is_zdr_blocked && matches!(ctx.auth_state, AuthState::Done) { - return handle_menu_shortcuts( - key, - ctx.menu_index, - &['l', 'q'], - dispatch_zdr_menu_action, - ); - } - if !ctx.has_access && matches!(ctx.auth_state, AuthState::Done) { - return handle_menu_shortcuts( - key, - ctx.menu_index, - &['g', 'l', 'q'], - dispatch_access_gate_menu_action, - ); - } if matches!(ctx.auth_state, AuthState::Done) && key!(Enter).matches(key) && key.modifiers.is_empty() @@ -2760,9 +2570,6 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco if let Event::Paste(text) = ev { match ctx.auth_state { AuthState::Done => { - if !ctx.has_access || ctx.is_zdr_blocked { - return InputOutcome::Unchanged; - } return InputOutcome::ActionThenForward(Action::NewSession); } AuthState::Authenticating { @@ -2792,12 +2599,6 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco if matches!(ctx.auth_state, AuthState::Pending { .. }) { return dispatch_pending_menu_action(i); } - if ctx.is_zdr_blocked { - return dispatch_zdr_menu_action(i); - } - if !ctx.has_access { - return dispatch_access_gate_menu_action(i); - } if ctx.has_claude_import && i == 0 && mouse.column >= rect.x + rect.width.saturating_sub(4) @@ -2813,16 +2614,6 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco ); } } - if let Some(rect) = ctx.refresh_rect - && rect.contains(ratatui::layout::Position::new(mouse.column, mouse.row)) - { - return InputOutcome::Action(Action::CheckSubscription); - } - if let Some(rect) = ctx.gate_url_rect - && rect.contains(ratatui::layout::Position::new(mouse.column, mouse.row)) - { - return InputOutcome::Action(Action::OpenSupergrokUrl); - } if let Some(rect) = ctx.changelog_cta_rect && rect.contains(ratatui::layout::Position::new(mouse.column, mouse.row)) && let Some(md) = ctx.changelog_markdown.as_deref() @@ -2901,32 +2692,6 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco InputOutcome::Unchanged } /// Handle Up/Down arrow key cycling through a menu of `count` items. -fn is_quit_signal(key: &crossterm::event::KeyEvent) -> bool { - key!('c', CONTROL).matches(key) || key!('d', CONTROL).matches(key) -} -/// `shortcuts[i]` triggers `menu_dispatch(i)`. -fn handle_menu_shortcuts( - key: &crossterm::event::KeyEvent, - menu_index: &mut Option, - shortcuts: &[char], - menu_dispatch: fn(usize) -> InputOutcome, -) -> InputOutcome { - if is_quit_signal(key) { - return InputOutcome::Action(Action::Quit); - } - for (i, &ch) in shortcuts.iter().enumerate() { - if key.code == KeyCode::Char(ch) { - return menu_dispatch(i); - } - } - if key!(Enter).matches(key) { - return menu_dispatch(menu_index.unwrap_or(0)); - } - if let Some(outcome) = handle_menu_nav(key, menu_index, shortcuts.len()) { - return outcome; - } - InputOutcome::Unchanged -} fn handle_menu_nav( key: &crossterm::event::KeyEvent, index: &mut Option, @@ -2959,25 +2724,6 @@ fn dispatch_pending_menu_action(index: usize) -> InputOutcome { _ => InputOutcome::Unchanged, } } -/// Dispatch an action for a welcome menu item when ZDR-blocked. -/// Menu layout: 0 = Switch account, 1 = Quit. -fn dispatch_zdr_menu_action(index: usize) -> InputOutcome { - match index { - 0 => InputOutcome::Action(Action::SwitchAccount), - 1 => InputOutcome::Action(Action::Quit), - _ => InputOutcome::Unchanged, - } -} -/// Menu actions when user is access-gated: 0 = Subscribe CTA, 1 = Logout, 2 = Quit. -/// "Refresh" (ctrl-r) is handled as a direct key shortcut, not a menu item. -fn dispatch_access_gate_menu_action(index: usize) -> InputOutcome { - match index { - 0 => InputOutcome::Action(Action::OpenSupergrokUrl), - 1 => InputOutcome::Action(Action::Logout), - 2 => InputOutcome::Action(Action::Quit), - _ => InputOutcome::Unchanged, - } -} /// Dispatch an action for a welcome menu item by index. /// /// Menu order: `[Import]`, New worktree, Resume session, `[Changelog]`, Quit. @@ -3242,8 +2988,6 @@ impl AppView { layout_cfg.eff_outer_vpad(compact), ) }; - let zdr_blocked_for_draw = self.is_zdr_blocked(); - let has_access = self.has_access(); let scroll_debug_panel = self.scroll_debug_panel(); let dev_fps_rows = self.dev_fps_rows(); let fps_overlay = self.fps_hud.overlay(dev_fps_rows); @@ -3337,11 +3081,8 @@ impl AppView { model_name: &model_name, flags: &flags_vec, selected: self.welcome_menu_index, - team_name: self.team_name.as_deref(), - has_access, has_claude_import: self.has_claude_import, mouse_pos: self.last_mouse_pos, - is_zdr_blocked: zdr_blocked_for_draw, session_picker: self.session_picker_entries.as_deref(), session_picker_loading: self.session_picker_entries.is_none() && (self.session_picker_loading @@ -3359,14 +3100,9 @@ impl AppView { .session_picker_entries_query .as_deref(), welcome_tick: self.welcome_tick, - gate: self.gate.as_ref(), - subscription_tier: self.subscription_tier.as_deref(), session_picker_grouped: self.session_picker_grouped, session_picker_source_filter: self.session_picker_source_filter, chat_mode: self.chat_mode, - credit_balance: self.credit_balance.as_ref(), - auto_topup: self.auto_topup.as_ref(), - usage_visible: self.usage_visible, is_api_key_auth: self.is_api_key_auth, changelog_bullets: &self.changelog_bullets, changelog_has_full_notes: self.changelog_markdown.is_some(), @@ -3384,8 +3120,6 @@ impl AppView { self.welcome_import_banner_rect = result.import_banner_rect; self.welcome_auth_url_rect = result.auth_url_rect; self.welcome_auth_fallback_rect = result.auth_fallback_rect; - self.welcome_refresh_rect = result.refresh_rect; - self.welcome_gate_url_rect = result.gate_url_rect; self.welcome_changelog_cta_rect = result.changelog_cta_rect; self.session_picker_state.hit_areas = result.session_picker_hit_areas; if let Some(modal) = self.import_claude_modal.as_mut() { @@ -3427,9 +3161,6 @@ impl AppView { &theme, ); } - if !has_access && !self.access_gate_shown_logged { - self.access_gate_shown_logged = true; - } if let Some(fps) = &fps_overlay { fps.render(full_area, f.buffer_mut()); } @@ -4474,24 +4205,9 @@ pub(crate) mod tests { deferred_startup: Default::default(), auth_use_oauth: false, auth_clipboard_copied: false, - team_id: None, - team_name: None, - is_zdr: false, - team_role: None, - coding_data_retention_opt_out: false, show_tips: None, auto_update: None, ask_user_question_timeout_enabled: None, - zdr_access_enabled: false, - usage_billing_redirect_url: None, - access_gate_shown_logged: false, - gate: None, - subscription_tier: None, - paywall_check_started: None, - last_subscription_check_at: None, - subscription_watch_interval_secs: None, - pending_gate_verification: None, - gate_verify_gen: 0, bundle_state: BundleState::default(), scroll_debug_hud: crate::views::scroll_debug_hud::ScrollDebugHud::new(), fps_hud: crate::views::fps_hud::FpsHud::new(), @@ -4513,8 +4229,6 @@ pub(crate) mod tests { welcome_on_auth_url: false, welcome_on_changelog_cta: false, welcome_auth_fallback_rect: None, - welcome_refresh_rect: None, - welcome_gate_url_rect: None, welcome_changelog_cta_rect: None, auth_show_raw_url: false, auth_mouse_disabled: false, @@ -4555,13 +4269,8 @@ pub(crate) mod tests { minimal_state: crate::minimal_api::MinimalState::default(), reconnect_pending: false, show_resolved_model: true, - sharing_enabled: false, usage_visible: true, - tier_restricted_commands: Vec::new(), leader_mode: true, - credit_balance: None, - auto_topup: None, - billing_poll_wanted: false, leader_roster: Vec::new(), dashboard_local_sessions: Vec::new(), dashboard_sessions_loading: false, @@ -4600,8 +4309,6 @@ pub(crate) mod tests { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, @@ -4791,8 +4498,6 @@ pub(crate) mod tests { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, @@ -5590,58 +5295,6 @@ pub(crate) mod tests { assert!(!app.is_api_key_auth); assert!(app.usage_visible); } - /// Make every tier-restricted command visible on the welcome prompt so the - /// present/absent assertions exercise the deny list, not incidental - /// fail-closed hiding: - /// - `/imagine`, `/imagine-video` are `required_tools()`-gated, so advertise - /// their tools (otherwise the registry fail-closes them). - fn advertise_media_tools(app: &mut AppView) { - app.welcome_prompt - .slash_controller - .registry_mut() - .set_available_tools( - ["image_gen", "image_to_video"] - .into_iter() - .map(str::to_string) - .collect(), - ); - } - fn assert_tier_restricted_commands_present(app: &AppView) { - let reg = app.welcome_prompt.slash_controller.registry(); - for name in TIER_RESTRICTED_COMMANDS { - assert!( - reg.get(name).is_some(), - "/{name} must be available when not tier-restricted (tools advertised)" - ); - } - } - #[test] - fn apply_auth_meta_never_restricts_tiers() { - let mut app = test_app(); - advertise_media_tools(&mut app); - app.apply_auth_meta(&kigi_shell::auth::AuthMeta::default()); - assert!(app.tier_restricted_commands.is_empty()); - assert_tier_restricted_commands_present(&app); - } - #[test] - fn is_restricted_tier_never_restricts() { - assert!(!is_restricted_tier(None)); - assert!(!is_restricted_tier(Some("Free"))); - assert!(!is_restricted_tier(Some("SomeFutureTier"))); - } - #[test] - fn apply_auth_meta_clears_gate_on_login() { - let mut app = test_app(); - app.gate = Some(kigi_shell::auth::GateInfo { - message: "Subscribe".into(), - url: None, - label: None, - }); - assert!(app.is_access_blocked()); - app.apply_auth_meta(&kigi_shell::auth::AuthMeta::default()); - assert!(app.gate.is_none()); - assert!(app.has_access()); - } #[test] fn welcome_ctrl_q_requires_confirmation() { let mut app = test_app(); diff --git a/crates/codegen/kigi-tui/src/app/cli.rs b/crates/codegen/kigi-tui/src/app/cli.rs index b033052..d8adbb9 100644 --- a/crates/codegen/kigi-tui/src/app/cli.rs +++ b/crates/codegen/kigi-tui/src/app/cli.rs @@ -38,9 +38,6 @@ pub enum Command { #[arg(long)] json: bool, }, - /// Share a session and print the share URL - #[command(hide = true)] - Share(crate::share_cmd::ShareArgs), /// Run any command with local clipboard support (OSC 52 → system clipboard). #[cfg_attr(not(any(unix, windows)), command(hide = true))] #[command(long_about = "\ @@ -268,8 +265,8 @@ pub struct AgentArgs { #[arg(long, conflicts_with = "leader")] pub no_leader: bool, /// Override the CLI chat proxy base URL. - #[arg(long = "cli-chat-proxy-base-url")] - pub cli_chat_proxy_base_url: Option, + #[arg(long = "coding-api-base-url")] + pub coding_api_base_url: Option, /// Override the public xAI API base URL. #[arg(long = "xai-api-base-url")] pub xai_api_base_url: Option, diff --git a/crates/codegen/kigi-tui/src/app/dispatch/auth.rs b/crates/codegen/kigi-tui/src/app/dispatch/auth.rs index ca4bdbd..4df3aa3 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/auth.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/auth.rs @@ -274,7 +274,7 @@ pub(super) fn handle_auth_complete( app.auth_state = AuthState::Done; app.auth_show_raw_url = false; - app.welcome_prompt_focused = !app.is_access_blocked(); + app.welcome_prompt_focused = true; app.auth_code_input.clear(); // Mid-session re-auth (`/login` or a 401 prompt): restore the @@ -312,9 +312,6 @@ pub(super) fn handle_auth_complete( } } let mut effects = dispatch(Action::RequestBundleStatus, app); - if app.usage_visible { - effects.push(Effect::FetchAppBilling); - } effects.extend(retry_effects); return effects; } @@ -322,27 +319,9 @@ pub(super) fn handle_auth_complete( // status only; shell auto-syncs post-auth let mut effects = dispatch(Action::RequestBundleStatus, app); - // Start auto-checking subscription if gated. - // Check immediately (don't wait 5s) then schedule the timer. - if !app.has_access() { - app.paywall_check_started = Some(std::time::Instant::now()); - effects.push(Effect::CheckSubscription { verify: None }); - effects.push(Effect::SchedulePaywallCheck); - } - // Fetch billing so the welcome screen can show a credit warning. - if app.usage_visible { - effects.push(Effect::FetchAppBilling); - } // Fetch changelog (mirrors startup path for interactive login). effects.push(Effect::FetchChangelog); - // ZDR-blocked users stay on the welcome screen — discard any - // deferred startup (they cannot start a session). - if app.is_zdr_blocked() { - clear_startup_actions(app); - return effects; - } - // Replay deferred session startup once BOTH gates are open. Auth // is now Done, so `session_startup_allowed()` here means "is trust // also resolved?" -- if trust is still Pending its question renders diff --git a/crates/codegen/kigi-tui/src/app/dispatch/billing.rs b/crates/codegen/kigi-tui/src/app/dispatch/billing.rs deleted file mode 100644 index aff6c81..0000000 --- a/crates/codegen/kigi-tui/src/app/dispatch/billing.rs +++ /dev/null @@ -1,532 +0,0 @@ -//! Subscription tier checks, credit-limit upsells, and auto-topup handling. - -use super::queue::maybe_drain_queue; -use crate::app::actions::Effect; -use crate::app::agent::AgentId; -use crate::app::agent_view::AgentView; -use crate::app::app_view::AppView; -use crate::scrollback::block::RenderBlock; -use std::time::Duration; - -/// How long the pager auto-checks subscription status before stopping. -/// After this, the user can still manually check via the [Refresh] button. -pub(super) const PAYWALL_AUTO_CHECK_TIMEOUT: Duration = Duration::from_secs(10 * 60); - -/// Whether the user is at the highest subscription tier (SuperGrok Heavy). -/// -/// Returns `true` only when `subscription_tier` **positively matches** a -/// known max-tier identifier. When the tier is unknown (`None`) or any -/// other value, returns `false` — the user gets the Q&A modal so lower- -/// tier users always see the upgrade option. -pub(super) fn is_max_tier(subscription_tier: Option<&str>) -> bool { - let Some(t) = subscription_tier else { - return false; // Unknown — default to Q&A. - }; - // Normalize: lowercase + spaces→underscores to match both JWT-derived - // keys ("supergrok_heavy") and CCP display names ("SuperGrok Heavy"). - t.to_ascii_lowercase().replace(' ', "_") == "supergrok_heavy" -} - -/// URL for upgrading the subscription tier. -pub(crate) const UPSELL_URL_UPGRADE: &str = "https://grok.com/supergrok?referrer=grok-build"; - -/// URL for managing pay-as-you-go / on-demand spending / purchasing credits. -pub(crate) const UPSELL_URL_PAYG: &str = "https://grok.com?_s=usage"; - -/// Billing mode for credit-limit upsell copy. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum CreditLimitUpsellMode { - /// Unified usage pool — suggest purchasing prepaid credits. - UnifiedCredits, - /// Legacy on-demand / PAYG (`enabled` = on-demand cap already active). - LegacyPayg { enabled: bool }, -} - -/// Resolve upsell copy mode from credits config. -/// -/// Prefers explicit `is_unified_billing_user` (`Option` — do not treat a -/// missing field as legacy). Positive `pay_as_you_go` (on-demand cap > 0) -/// only selects legacy when the unified flag is absent. Unknown defaults to -/// unified (buy credits) so pool users never get “enable on-demand” wrongly. -pub(super) fn credit_limit_upsell_mode( - balance: Option<&crate::views::credit_bar::CreditBalance>, -) -> CreditLimitUpsellMode { - match balance { - Some(b) if b.is_unified_billing_user == Some(true) => CreditLimitUpsellMode::UnifiedCredits, - Some(b) if b.is_unified_billing_user == Some(false) => CreditLimitUpsellMode::LegacyPayg { - enabled: b.pay_as_you_go, - }, - // Flag absent: only treat as legacy PAYG when we have a positive - // on-demand cap (pay_as_you_go is derived from cap > 0). - Some(b) if b.pay_as_you_go => CreditLimitUpsellMode::LegacyPayg { enabled: true }, - _ => CreditLimitUpsellMode::UnifiedCredits, - } -} - -/// Whether an API / retry error is a credit-limit / spend-block denial. -/// -/// - **402** Payment Required — always credit/spend block on this surface -/// (Build pool and IC spend blocks); no message filter. -/// - **403** — only when the body contains "run out of credits" (legacy IC -/// spend wording); other 403s (content-safety, ZDR, …) are excluded. -pub(crate) fn is_credit_limit_error(http_status: Option, message: &str) -> bool { - let m = message.to_ascii_lowercase(); - let legacy = m.contains("run out of credits"); - match http_status { - Some(402) => true, - Some(403) if legacy => true, - // Retry notifications embed "status 402" / "status 403" in the body - // without a separate status field. - None | Some(_) => m.contains("status 402") || (m.contains("status 403") && legacy), - } -} - -/// Well-known error code CCP returns (HTTP 429, flat body -/// `{"code": "...", "error": "..."}`) when a free-tier user exhausts the -/// free usage quota. Kept in sync with the shared well-known error code -/// `SUBSCRIPTION_FREE_USAGE_EXHAUSTED`. sampling-types' `parse_error_bytes` prepends the flat -/// `code` to the flattened message, so the code reaches the pager embedded -/// in `RetryState::Exhausted.reason` and the -32003 error's data string. -pub(crate) const FREE_USAGE_EXHAUSTED_ERROR_CODE: &str = "subscription:free-usage-exhausted"; - -/// Whether a rate-limit error is the free-usage-quota exhaustion (paywall) -/// rather than transient throttling. Text-sniff on the flattened message, -/// same precedent as [`is_credit_limit_error`]. -pub(crate) fn is_free_usage_exhausted_error(reason: &str) -> bool { - reason.contains(FREE_USAGE_EXHAUSTED_ERROR_CODE) -} - -/// Whether a rate-limited (-32003) ACP error is the free-usage exhaustion. -/// `data` may be a bare string or the `{message, promptUsage?}` object -/// `attach_prompt_usage` produces — always read via the shared detail helper. -pub(crate) fn acp_error_is_free_usage_exhausted(err: &agent_client_protocol::Error) -> bool { - err.data - .as_ref() - .and_then(kigi_shell::sampling::error::error_detail_from_data) - .as_deref() - .is_some_and(is_free_usage_exhausted_error) -} - -/// User-facing message for free-usage exhaustion. Shown by headless mode and -/// `format_acp_error` in place of auth-aware rate-limit copy. Deliberately -/// promises no reset duration — the quota window is backend-config-driven. -pub(crate) const FREE_USAGE_USER_MESSAGE: &str = "You\u{2019}ve reached your free Grok Build usage limit for now. Get SuperGrok for much higher limits, or try again later: https://grok.com/supergrok?referrer=grok-build"; - -/// Open the credit-limit upsell on the given agent. -/// -/// **`max_tier = false`** (default): shows the Q&A question modal with -/// two options ("Upgrade tier" + buy-credits or PAYG). Each option's `id` -/// carries the target URL so the submit handler is position-independent. -/// -/// **`max_tier = true`** (positively identified as SuperGrok Heavy): -/// pushes an inline scrollback card (`CreditLimitBlock`) with a single -/// continue action. No Q&A modal — the user can't upgrade further. -pub(super) fn open_credit_limit_upsell( - agent: &mut AgentView, - mode: CreditLimitUpsellMode, - max_tier: bool, -) { - use crate::scrollback::blocks::CreditLimitCardAction; - - let (heading, upgrade_tier_desc, secondary_label, secondary_desc, card_action): ( - &str, - &str, - &str, - &str, - CreditLimitCardAction, - ) = match mode { - CreditLimitUpsellMode::UnifiedCredits => ( - "You hit your weekly limit.", - "Upgrade to a higher tier for more usage", - "Buy more credits", - "Purchase credits to keep using Grok Build", - CreditLimitCardAction::PurchaseCredits, - ), - CreditLimitUpsellMode::LegacyPayg { enabled: true } => ( - "You\u{2019}ve hit your spending cap.", - "Upgrade to a higher tier for more credits", - "Increase limit", - "Raise your pay-as-you-go spending cap", - CreditLimitCardAction::IncreasePaygLimit, - ), - CreditLimitUpsellMode::LegacyPayg { enabled: false } => ( - "You\u{2019}ve hit the credit limit for your plan.", - "Upgrade to a higher tier for more credits", - "Pay as you go", - "Enable pay-as-you-go credits for on-demand usage", - CreditLimitCardAction::EnablePayg, - ), - }; - - // ── Max tier: inline scrollback card ───────────────────────── - if max_tier { - use crate::scrollback::block::RenderBlock; - agent.scrollback.push_block(RenderBlock::credit_limit_card( - heading, - card_action, - UPSELL_URL_PAYG, - )); - return; - } - - // ── Default: Q&A question modal with two options ──────────────── - use crate::views::question_view::{LocalQuestionKind, QuestionViewState}; - use kigi_tools::implementations::grok_build::ask_user_question::{Question, QuestionOption}; - - if agent.question_view.is_some() { - return; - } - - let question = Question { - question: heading.into(), - options: vec![ - QuestionOption { - label: "Upgrade tier".into(), - description: upgrade_tier_desc.into(), - preview: None, - id: Some(UPSELL_URL_UPGRADE.into()), - }, - QuestionOption { - label: secondary_label.into(), - description: secondary_desc.into(), - preview: None, - id: Some(UPSELL_URL_PAYG.into()), - }, - ], - multi_select: Some(false), - id: None, - }; - - let stashed = agent.prompt.stash(); - let state = QuestionViewState::new( - format!("credit-limit-upsell-{}", uuid::Uuid::new_v4()), - vec![question], - stashed, - ) - .with_local_kind(LocalQuestionKind::CreditLimitUpsell) - .with_no_freeform(); - agent.question_view = Some(state); - agent.prompt.set_text(""); -} - -/// Open the free-usage paywall on the given agent: a Q&A modal in the -/// [`open_credit_limit_upsell`] style with two upgrade options. Each -/// option's `id` carries its target URL so the submit handler is -/// position-independent. -/// -/// Driver-only by construction (called from the PromptResponse handler, -/// which viewers never receive). -pub(super) fn open_free_usage_upsell(agent: &mut AgentView) { - open_supergrok_upsell(agent, UpsellReason::FreeUsageLimit); -} - -/// Open the SuperGrok upsell for a tier-restricted slash command -/// (`/usage`, `/imagine`, …). Returns whether the modal opened (`false` -/// when another question modal is already up) so the caller can decide -/// whether to consume the input that triggered it. -pub(super) fn open_restricted_command_upsell(agent: &mut AgentView) -> bool { - open_supergrok_upsell(agent, UpsellReason::RestrictedCommand) -} - -/// Which situation opened the SuperGrok upsell modal. Controls the heading. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum UpsellReason { - /// Free-usage quota exhausted (429 paywall). - FreeUsageLimit, - /// A tier-restricted slash command was invoked. - RestrictedCommand, -} - -/// Shared builder behind [`open_free_usage_upsell`] / -/// [`open_restricted_command_upsell`]: a Q&A modal in the -/// [`open_credit_limit_upsell`] style. Upgrade options carry their target -/// URL in the option `id` (position-independent submit handling). -fn open_supergrok_upsell(agent: &mut AgentView, reason: UpsellReason) -> bool { - use crate::views::question_view::{LocalQuestionKind, QuestionViewState}; - use kigi_tools::implementations::grok_build::ask_user_question::{Question, QuestionOption}; - - // Never displace an already-open question modal. Callers that consume - // input on open must check this `false` and keep the input instead. - if agent.question_view.is_some() { - return false; - } - - let (heading, modal_id_prefix) = match reason { - UpsellReason::FreeUsageLimit => ("You hit your free usage limit.", "free-usage-upsell"), - UpsellReason::RestrictedCommand => ( - "Unlock all features with SuperGrok.", - "restricted-command-upsell", - ), - }; - - let options = vec![ - QuestionOption { - label: "Upgrade to SuperGrok".into(), - description: "For everyday coding and productivity tasks".into(), - preview: None, - id: Some(UPSELL_URL_UPGRADE.into()), - }, - QuestionOption { - label: "Upgrade to SuperGrok Heavy".into(), - description: "Get the most out of Grok Build. Highest usage limits.".into(), - preview: None, - // No Heavy-specific URL exists; the /supergrok page lists - // both plans, so both upgrade options land there. - id: Some(UPSELL_URL_UPGRADE.into()), - }, - ]; - let question = Question { - question: heading.into(), - options, - multi_select: Some(false), - id: None, - }; - - let stashed = agent.prompt.stash(); - let state = QuestionViewState::new( - format!("{modal_id_prefix}-{}", uuid::Uuid::new_v4()), - vec![question], - stashed, - ) - .with_local_kind(LocalQuestionKind::FreeUsageUpsell) - .with_no_freeform(); - agent.question_view = Some(state); - agent.prompt.set_text(""); - true -} - -/// Apply an [`AutoTopupFetch`] outcome to a cached `auto_topup` slot: `Resolved` -/// sets it, `Cleared` resets it to "unknown" (no credits), and `Unchanged` keeps -/// the last-known-good value (the fetch failed). -pub(super) fn apply_auto_topup( - slot: &mut Option, - fetch: &crate::views::credit_bar::AutoTopupFetch, -) { - use crate::views::credit_bar::AutoTopupFetch; - match fetch { - AutoTopupFetch::Resolved(rule) => *slot = Some(rule.clone()), - AutoTopupFetch::Cleared => *slot = None, - AutoTopupFetch::Unchanged => {} - } -} - -// TaskResult handlers. - -pub(super) fn handle_billing_fetched( - app: &mut AppView, - agent_id: AgentId, - balance: Option, - silent: bool, - subscription_tier: Option, - autotopup: crate::views::credit_bar::AutoTopupFetch, -) -> Vec { - // Parse/transport failures route to `BillingError`, so a `None` - // balance here means the response carried no billing config. Clear - // the cached balance + polling so the status bar agrees with the - // "No billing data available." message rather than showing a stale - // value. - app.credit_balance = balance.clone(); - // `Resolved` updates the cached rule, `Cleared` resets it to unknown - // (no credits), `Unchanged` keeps the last-known-good (fetch failed). - apply_auto_topup(&mut app.auto_topup, &autotopup); - app.billing_poll_wanted = balance - .as_ref() - .map(|b| b.usage_pct >= 99.0) - .unwrap_or(false); - if let Some(tier) = subscription_tier { - app.subscription_tier = Some(tier); - } - // Render the `/usage` summary from the now-current cached rule. - let summary_topup = app.auto_topup.clone(); - if let Some(agent) = app.agents.get_mut(&agent_id) { - // Gateway/chat-kind: do not attach Build coding credits. - let mut topup = agent.auto_topup.clone(); - apply_auto_topup(&mut topup, &autotopup); - agent.apply_credit_balance(balance.clone(), topup); - if !silent && !agent.chat_kind { - let msg = match &balance { - Some(bal) => { - crate::views::credit_bar::format_usage_summary(bal, summary_topup.as_ref()) - } - None => "No billing data available.".to_string(), - }; - agent.scrollback.push_block(RenderBlock::System( - crate::scrollback::blocks::SystemMessageBlock::new(msg), - )); - } - } - vec![] -} - -pub(super) fn handle_gate_refreshed( - app: &mut AppView, - settings: Option, -) -> Vec { - let Some(rs) = settings else { - return vec![]; - }; - app.usage_billing_redirect_url = rs.usage_billing_redirect_url.clone(); - if let Some(secs) = rs.subscription_watch_interval_secs { - app.subscription_watch_interval_secs = Some(secs); - } - match AppView::gate_from_settings(&rs) { - Some(gate) => app.impose_gate(gate), - None => app.lift_gate(), - } -} - -/// `x.ai/auth/check_subscription` completed. Meta is authoritative -/// (`apply_auth_meta` also drops any deferred gate). A failed check only -/// promotes the deferred gate it was verifying (`verify` generation); -/// generic watch/focus/paywall-chain failures never touch it. -pub(super) fn handle_check_subscription_complete( - app: &mut AppView, - verify: Option, - meta: Option, -) -> Vec { - let was_blocked = !app.has_access(); - let applied = match meta { - Some(meta_val) => { - match serde_json::from_value::(meta_val) { - Ok(auth_meta) => { - app.apply_auth_meta(&auth_meta); - true - } - Err(e) => { - // Shell sent meta we can't decode — a protocol bug, not - // a transient failure. The check result is lost, so a - // verify deferral falls through to promotion below. - crate::unified_log::error( - "subscription.check.meta_parse_failed", - None, - Some(serde_json::json!({ - "verify": verify, - "error": e.to_string(), - })), - ); - false - } - } - } - // meta: None = shell reports "not authenticated" or the check RPC - // failed (already logged as subscription.check.rpc_failed). - None => false, - }; - if !applied && let Some(generation) = verify { - app.promote_deferred_gate(generation, "check_failed"); - } - crate::unified_log::info( - "subscription.check.complete", - None, - Some(serde_json::json!({ - "verify": verify, - "meta_applied": applied, - "was_blocked": was_blocked, - "gated": !app.has_access(), - "tier": app.subscription_tier, - })), - ); - maybe_start_paywall_chain(app, was_blocked) -} - -/// Safety net for a hung verification check: show the still-pending -/// deferred gate (err on blocking). -pub(super) fn handle_gate_verify_timeout(app: &mut AppView, generation: u64) -> Vec { - let was_blocked = !app.has_access(); - app.promote_deferred_gate(generation, "verify_timeout"); - maybe_start_paywall_chain(app, was_blocked) -} - -/// Arm the 5s paywall auto-check chain on an ungated→gated transition, so a -/// paywall shown by verify-before-paywall self-lifts exactly like the -/// login-path one. Guarded so steady-state paywall-poller responses and -/// repeated checks can't fan out extra timers. -fn maybe_start_paywall_chain(app: &mut AppView, was_blocked: bool) -> Vec { - if !was_blocked && !app.has_access() && app.paywall_check_started.is_none() { - app.paywall_check_started = Some(std::time::Instant::now()); - return vec![Effect::SchedulePaywallCheck]; - } - vec![] -} - -pub(super) fn handle_credit_limit_recheck_complete( - app: &mut AppView, - agent_id: AgentId, - meta: Option, -) -> Vec { - if let Some(meta_val) = meta - && let Ok(auth_meta) = serde_json::from_value::(meta_val) - { - app.apply_auth_meta(&auth_meta); - } - - let Some(agent) = app.agents.get_mut(&agent_id) else { - return vec![]; - }; - - // If the user already submitted another prompt while the - // recheck was in flight, don't show the upsell — they've moved on. - let user_moved_on = !agent.session.state.is_idle() || !agent.session.pending_prompts.is_empty(); - - if !user_moved_on { - let balance = agent - .credit_balance - .as_ref() - .or(app.credit_balance.as_ref()); - let mode = credit_limit_upsell_mode(balance); - let max_tier = is_max_tier(app.subscription_tier.as_deref()); - open_credit_limit_upsell(agent, mode, max_tier); - } - // Either way, drop the stashed prompt. - agent.credit_limit_stashed_prompt = None; - - let mut effects = maybe_drain_queue(agent); - effects.push(Effect::FetchBilling { - agent_id, - silent: true, - }); - effects -} - -// Action handlers. - -pub(super) fn dispatch_open_supergrok_url(app: &mut AppView) -> Vec { - let url = app - .gate - .as_ref() - .and_then(|g| g.url.as_deref()) - .unwrap_or("https://grok.com/supergrok?referrer=grok-build"); - // Funnel attribution: tag CLI-originated SuperGrok upsell clicks - // with `referrer=grok-build`, matching the OAuth consent flow and - // x.ai/cli marketing links. Applied even when the URL came from - // remote settings's `gate_url`, so we don't depend on the remote flag - // being correctly configured. If the URL already specifies a - // referrer it's left alone. - let url = crate::app::link_opener::ensure_query_param(url, "referrer", "grok-build"); - crate::app::link_opener::open_url(&url); - vec![] -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn free_usage_dual_read_string_and_wrapped_object_data() { - let free = "subscription:free-usage-exhausted quota hit"; - let string_err = agent_client_protocol::Error::new(-32003, "Rate limited").data(free); - assert!(acp_error_is_free_usage_exhausted(&string_err)); - - // attach_prompt_usage wraps string data as {"message": ..., "promptUsage": ...}. - let wrapped = - agent_client_protocol::Error::new(-32003, "Rate limited").data(serde_json::json!({ - "message": free, - "promptUsage": { "inputTokens": 1, "outputTokens": 0, "numTurns": 1 } - })); - assert!(acp_error_is_free_usage_exhausted(&wrapped)); - assert!(!wrapped.data.as_ref().unwrap().is_string()); - - let other = agent_client_protocol::Error::new(-32003, "Rate limited").data("throttled"); - assert!(!acp_error_is_free_usage_exhausted(&other)); - } -} diff --git a/crates/codegen/kigi-tui/src/app/dispatch/dashboard.rs b/crates/codegen/kigi-tui/src/app/dispatch/dashboard.rs index 727e1e7..a7c95dd 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/dashboard.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/dashboard.rs @@ -48,7 +48,7 @@ pub(super) fn ensure_dashboard_state(app: &mut AppView) { state.adopt_slash_mru(app.slash_mru.clone()); state.set_screen_mode(app.screen_mode); state.set_recap_visible(app.session_recap_available); - state.set_restricted_commands(&app.tier_restricted_commands); + state.set_restricted_commands(&[]); app.dashboard = Some(state); } @@ -144,7 +144,7 @@ pub(super) fn dispatch_open_dashboard(app: &mut AppView) -> Vec { // Subsequent reopen — just gc dead ids; in-memory state stays. d.gc_stale_refs(&dashboard_alive_fn(&app.agents)); d.set_recap_visible(app.session_recap_available); - d.set_restricted_commands(&app.tier_restricted_commands); + d.set_restricted_commands(&[]); } // Refresh each local agent's git context (branch / worktree / label) // from disk so the row subtitles show the LATEST branch and worktree @@ -1207,7 +1207,6 @@ pub(super) fn dispatch_dashboard_dispatch_slash(app: &mut AppView, text: String) return vec![]; } - let coding_data_sharing_opt_out_from_app = app.coding_data_retention_opt_out; let show_tips_from_app = app.show_tips; let auto_update_from_app = app.auto_update; let respect_manual_folds_from_app = app.appearance.scrollback.scroll.respect_manual_folds; @@ -1231,23 +1230,6 @@ pub(super) fn dispatch_dashboard_dispatch_slash(app: &mut AppView, text: String) }; let reg = dashboard.dispatch.slash_controller.registry(); - // Tier-restricted commands stay visible for discoverability but must - // not execute — and must not fall through to the unknown-command - // path below (which would spawn a session with the raw slash text as - // its first prompt). The dashboard has no question-modal surface, so - // upsell via the feedback toast. - if reg.is_restricted(invocation.token) { - let token = invocation.token.to_string(); - if let Some(d) = app.dashboard.as_mut() { - d.dispatch.set_text(""); - d.set_error_toast(&format!( - "/{token} requires SuperGrok — upgrade at {}", - super::billing::UPSELL_URL_UPGRADE - )); - } - return vec![]; - } - let Some(command) = reg.get(invocation.token).cloned() else { // Unknown command. Fall back to the regular dispatch // path so the text becomes a new session's prompt. @@ -1295,7 +1277,6 @@ pub(super) fn dispatch_dashboard_dispatch_slash(app: &mut AppView, text: String) .iter() .map(|(id, info)| (info.name.clone(), id.clone())) .collect(), - coding_data_sharing_opt_out: coding_data_sharing_opt_out_from_app, plan_mode_active: false, show_tips: show_tips_from_app, auto_update: auto_update_from_app, diff --git a/crates/codegen/kigi-tui/src/app/dispatch/mod.rs b/crates/codegen/kigi-tui/src/app/dispatch/mod.rs index 4415527..e5f9472 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/mod.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/mod.rs @@ -13,7 +13,6 @@ //! otherwise); tests/ shares a fixture prelude via `use super::*;`. mod auth; -mod billing; mod ctx; mod dashboard; mod import_claude; @@ -33,10 +32,6 @@ mod task_result; mod transcript; mod turn; -pub(crate) use billing::{ - FREE_USAGE_USER_MESSAGE, UPSELL_URL_PAYG, UPSELL_URL_UPGRADE, - acp_error_is_free_usage_exhausted, is_credit_limit_error, is_free_usage_exhausted_error, -}; pub(crate) use modes::{downgrade_displayed_auto_if_gated, effective_auto}; pub(crate) use notes::{recap_unavailable_toast, scrollback_has_user_messages}; pub(crate) use permissions::resolve_permission_queue_transition; diff --git a/crates/codegen/kigi-tui/src/app/dispatch/prompt.rs b/crates/codegen/kigi-tui/src/app/dispatch/prompt.rs index 55c0332..1f77487 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/prompt.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/prompt.rs @@ -1,7 +1,6 @@ //! Prompt and bash-command submission dispatchers and reload-window helpers. use super::auth::{scrollback_has_recent_context_too_large, scrollback_has_recent_reauth_prompt}; -use super::billing::is_credit_limit_error; use super::ctx::with_active_agent; use super::interject; use super::permissions::drain_permission_queue; @@ -292,7 +291,6 @@ pub(super) fn dispatch_send_prompt_inner( return vec![]; }; // Capture app-level fields before the mut-borrow on `agent`. - let coding_data_sharing_opt_out_from_app = app.coding_data_retention_opt_out; let show_tips_from_app = app.show_tips; let auto_update_from_app = app.auto_update; let respect_manual_folds_from_app = app.appearance.scrollback.scroll.respect_manual_folds; @@ -324,37 +322,6 @@ pub(super) fn dispatch_send_prompt_inner( let mut effects = Vec::new(); - // ── Tier-restricted command upsell ───────────────────────────── - // Restricted commands (`/usage`, `/imagine`, …) are hidden from the - // registry's `get()`, so a typed invocation would otherwise fall - // through the unknown-command path below and leak to the model as a - // raw prompt. Upsell instead; genuinely unknown commands still pass - // through (shell/ACP commands depend on that). - if !literal - && trimmed.starts_with('/') - && let Some(invocation) = crate::slash::parse_invocation(trimmed) - && agent - .prompt - .slash_controller - .registry() - .is_restricted(invocation.token) - { - // Only consume the composer when the upsell can actually open: with - // another question modal already up, `open_supergrok_upsell` would - // no-op and wiping the composer here would silently drop the typed - // text. Keep it instead so the user can resubmit after closing the - // modal — and never fall through to passthrough for restricted - // commands. - if agent.question_view.is_none() { - if consume_input { - agent.prompt.set_text(""); - } - let opened = super::billing::open_restricted_command_upsell(agent); - debug_assert!(opened, "no modal was open, so the upsell must open"); - } - return vec![]; - } - // ── Registry-based slash command execution ───────────────────── // If the text starts with `/`, run it through the slash registry. // The registry resolves builtins, ACP-advertised commands, and @@ -384,7 +351,6 @@ pub(super) fn dispatch_send_prompt_inner( .iter() .map(|(id, info)| (info.name.clone(), id.clone())) .collect(), - coding_data_sharing_opt_out: coding_data_sharing_opt_out_from_app, // Prefer optimistic pending over confirmed active. plan_mode_active: agent.plan_mode_pending.unwrap_or(agent.plan_mode_active), show_tips: show_tips_from_app, @@ -1004,29 +970,11 @@ pub(super) fn handle_prompt_response( None => expected_send_now.is_some(), }; let rate_limited = agent.session.rate_limited; - // Fallback mirroring the credit-limit race guard below: if the retry - // notification lost the race with (or never reached) this - // PromptResponse, detect the free-usage code from the prompt error - // itself — the flattened 429 body embeds it. - let free_usage_blocked = agent.session.free_usage_blocked - || result - .as_ref() - .err() - .is_some_and(|e| super::billing::is_free_usage_exhausted_error(e)); let model_incompatible = agent.session.model_incompatible; // Context overflow: the RetryState handler already pushed the actionable // block, so the generic TurnFailed + error toast are redundant. Derived // from the scrollback (mirrors reauth), not a session flag. let context_overflow = scrollback_has_recent_context_too_large(&agent.scrollback); - // Fallback: if the retry notification didn't set the flag, - // detect credit-limit denials (legacy 403 or pool 402) from - // the PromptResponse error + HTTP status. Covers races where - // the retry notification arrives after the PromptResponse. - let credit_limit_blocked = agent.session.credit_limit_blocked - || result - .as_ref() - .err() - .is_some_and(|e| is_credit_limit_error(http_status, e)); // A 401/auth failure already surfaced an actionable // `ReAuthRequired` prompt via the RetryState handler (which // runs before this PromptResponse). Suppress the redundant @@ -1055,12 +1003,7 @@ pub(super) fn handle_prompt_response( ); } - // Stash the complete in-flight prompt before finish_turn clears it. - // Used by CreditLimitRecheckComplete to retry after a tier upgrade. - if credit_limit_blocked { - agent.credit_limit_stashed_prompt = agent.session.in_flight_prompt.clone(); - } - // Likewise, stash the prompt from a turn that failed on an + // Stash the prompt from a turn that failed on an // expired login (401 / re-auth). The AuthComplete handler // auto-resubmits it after a successful mid-session re-auth. // A non-rewindable turn (None) must not clobber an earlier stash. @@ -1110,17 +1053,11 @@ pub(super) fn handle_prompt_response( elapsed: Some(elapsed.unwrap_or_default()), }), (Err(_), _) - if rate_limited - || free_usage_blocked - || model_incompatible - || credit_limit_blocked - || reauth_prompted - || context_overflow => + if rate_limited || model_incompatible || reauth_prompted || context_overflow => { // Skip TurnFailed when a dedicated prompt/modal shows instead - // (rate limit, free-usage paywall, model incompatibility, - // credit 403, 401 re-auth, or a terminal context-window - // overflow). + // (rate limit, model incompatibility, 401 re-auth, or a + // terminal context-window overflow). None } (Err(err), _) => Some(SessionEvent::TurnFailed { @@ -1147,9 +1084,7 @@ pub(super) fn handle_prompt_response( } (Err(err), _) if !rate_limited - && !free_usage_blocked && !model_incompatible - && !credit_limit_blocked && !reauth_prompted && !context_overflow => { @@ -1256,7 +1191,7 @@ pub(super) fn handle_prompt_response( // Predicted-next-prompt (tab autocomplete): wipe any stale suggestion // at every turn boundary. This must run before the reconnect / - // credit-limit early returns below, which skip the fetch gate + // paywall early returns below, which skip the fetch gate // entirely — a prior ghost would otherwise survive those paths. agent.prompt.prompt_suggestion.clear(); @@ -1271,60 +1206,6 @@ pub(super) fn handle_prompt_response( return vec![]; } - // Credit-limit (403 legacy / 402 pool): strip stale error - // blocks, then do a one-shot subscription re-check. If the - // tier changed (user upgraded mid-session), the stashed - // prompt is retried automatically; otherwise the upsell - // is shown. - if credit_limit_blocked { - // Strip stale "Retry failed" / "Turn failed" error blocks - // that were pushed before the credit-limit was detected. - // Walk backwards from the end and remove matching events. - let mut to_remove = Vec::new(); - for idx in (0..agent.scrollback.len()).rev() { - match agent.scrollback.entry(idx).map(|e| &e.block) { - Some(crate::scrollback::block::RenderBlock::SessionEvent(ev)) - if matches!( - &ev.event, - SessionEvent::RetryFailed { .. } | SessionEvent::TurnFailed { .. } - ) => - { - to_remove.push(idx); - } - // Stop at the first non-error block. - Some( - crate::scrollback::block::RenderBlock::SessionEvent(_) - | crate::scrollback::block::RenderBlock::System(_), - ) => continue, - _ => break, - } - } - for idx in to_remove { - agent.scrollback.remove_from(idx); - } - - // Defer the upsell until the subscription re-check - // completes. Queue drain + billing fetch happen in the - // CreditLimitRecheckComplete handler. - if let Some(p) = pending_adoption { - agent.discard_pending_adoption_updates(&p.prompt_id); - } - return vec![Effect::CreditLimitRecheck { agent_id }]; - } - - // Free-usage paywall (429 + subscription:free-usage-exhausted): the - // RetryState handler set the flag and suppressed the generic - // rate-limit block; show the upsell modal. Driver-only by - // construction — viewers never receive a PromptResponse. No queue - // drain: queued prompts would fail on the same exhausted quota. - if free_usage_blocked { - super::billing::open_free_usage_upsell(agent); - if let Some(p) = pending_adoption { - agent.discard_pending_adoption_updates(&p.prompt_id); - } - return vec![]; - } - // FIFO handoff: if a server-authoritative prompt drained // into the running slot during this turn's teardown, adopt it // now (finish_turn cleared current_prompt_id) and run the @@ -1372,10 +1253,6 @@ pub(super) fn handle_prompt_response( }); } - effects.push(Effect::FetchBilling { - agent_id, - silent: true, - }); return effects; } vec![] diff --git a/crates/codegen/kigi-tui/src/app/dispatch/queue.rs b/crates/codegen/kigi-tui/src/app/dispatch/queue.rs index f522cef..b93dacb 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/queue.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/queue.rs @@ -876,13 +876,9 @@ mod tests { kind: crate::app::agent::QueueEntryKind::Prompt, }; - // Turn ends → should NOT drain "second" (user is editing it), only FetchBilling. + // Turn ends → should NOT drain "second" (user is editing it). let effects = dispatch(end_turn(), &mut app); - assert_eq!(effects.len(), 1); - assert!(matches!( - &effects[0], - Effect::FetchBilling { silent: true, .. } - )); + assert!(effects.is_empty(), "drain should be blocked: {effects:?}"); assert!(app.agents[&id].session.state.is_idle()); // "second" should still be in the queue. assert_eq!(app.agents[&id].session.queue_len(), 2); @@ -908,14 +904,10 @@ mod tests { kind: crate::app::agent::QueueEntryKind::Prompt, }; - // Turn ends → should drain "second" (front, not being edited) + FetchBilling. + // Turn ends → should drain "second" (front, not being edited). let effects = dispatch(end_turn(), &mut app); - assert_eq!(effects.len(), 2); + assert_eq!(effects.len(), 1); assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "second")); - assert!(matches!( - &effects[1], - Effect::FetchBilling { silent: true, .. } - )); // "third" should still be in queue. assert_eq!(app.agents[&id].session.queue_len(), 1); assert_eq!(app.agents[&id].session.pending_prompts[0].text, "third"); @@ -1809,13 +1801,9 @@ mod tests { kind: crate::app::agent::QueueEntryKind::Prompt, }; - // End turn for p2 → should NOT drain p3 (being edited), only FetchBilling. + // End turn for p2 → should NOT drain p3 (being edited). let effects = dispatch(end_turn(), &mut app); - assert_eq!(effects.len(), 1); - assert!( - matches!(&effects[0], Effect::FetchBilling { silent: true, .. }), - "drain should be blocked, only billing refresh" - ); + assert!(effects.is_empty(), "drain should be blocked: {effects:?}"); assert_eq!(app.agents[&id].session.queue_len(), 2); // p3, p4 // Simulate user saving edited text. diff --git a/crates/codegen/kigi-tui/src/app/dispatch/router.rs b/crates/codegen/kigi-tui/src/app/dispatch/router.rs index e8d4528..4737192 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/router.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/router.rs @@ -3,7 +3,6 @@ use super::auth::{ dispatch_cancel_login, dispatch_login, dispatch_logout, dispatch_submit_auth_code, dispatch_switch_account, }; -use super::billing::dispatch_open_supergrok_url; use super::ctx::{ active_agent_session_id, get_active_agent_mut, navigate_clearing_selection, sync_sleep_inhibitor, with_active_agent, with_scrollback, @@ -93,10 +92,9 @@ use super::settings::ui::{ dispatch_toggle_vim_mode, }; use super::status::{ - dispatch_copy_session_id, dispatch_open_gboom, dispatch_share_session, - dispatch_show_context_info, dispatch_show_privacy_info, dispatch_show_queue, + dispatch_copy_session_id, dispatch_open_gboom, dispatch_show_context_info, dispatch_show_queue, dispatch_show_release_notes, dispatch_show_session_info, dispatch_show_tasks, - dispatch_show_usage, set_coding_data_sharing, + dispatch_show_usage, }; use super::task_result::{dispatch_task_result, unregister_all_active_sessions}; use super::transcript::{ @@ -551,23 +549,7 @@ pub(crate) fn dispatch(action: Action, app: &mut AppView) -> Vec { if group_toggled { return vec![]; } - let mut credit_card: Option = None; - with_scrollback(app, |s| { - if let Some(idx) = s.selected() - && let Some(entry) = s.entry(idx) - && let crate::scrollback::block::RenderBlock::CreditLimit(ref blk) = entry.block - { - credit_card = Some(blk.url.clone()); - } - }); - if let Some(url) = credit_card { - crate::app::link_opener::open_url_if_safe( - &url, - crate::terminal::hyperlinks::SchemeFilter::Standard, - ); - } else { - dispatch_open_block_viewer(app); - } + dispatch_open_block_viewer(app); vec![] } Action::OpenExtensionsModal { tab } => { @@ -789,7 +771,6 @@ pub(crate) fn dispatch(action: Action, app: &mut AppView) -> Vec { vec![Effect::FetchCatalogEntry { kind, name }] } Action::CycleMode => dispatch_cycle_mode(app), - Action::ShareSession => dispatch_share_session(app), Action::ShowSessionInfo => dispatch_show_session_info(app), Action::ShowReleaseNotes { title, content } => { dispatch_show_release_notes(app, title, content) @@ -809,8 +790,6 @@ pub(crate) fn dispatch(action: Action, app: &mut AppView) -> Vec { Action::SaveRememberNoteFromModal => dispatch_save_remember_note_from_modal(app), Action::SendBtw(question) => dispatch_send_btw(app, question), Action::SendRecap { auto } => dispatch_send_recap(app, auto), - Action::ShowPrivacyInfo => dispatch_show_privacy_info(app), - Action::SetCodingDataSharing { opted_in } => set_coding_data_sharing(app, opted_in), Action::ToggleYolo => dispatch_toggle_yolo(app), Action::ToggleMultiline => dispatch_toggle_multiline(app), Action::ToggleCompactMode => dispatch_toggle_compact_mode(app), @@ -873,8 +852,6 @@ pub(crate) fn dispatch(action: Action, app: &mut AppView) -> Vec { Action::PermissionCancel => dispatch_permission_cancel(app), Action::Logout => dispatch_logout(app), Action::SwitchAccount => dispatch_switch_account(app), - Action::CheckSubscription => vec![Effect::CheckSubscription { verify: None }], - Action::OpenSupergrokUrl => dispatch_open_supergrok_url(app), Action::OpenUrl(url) => { use crate::terminal::hyperlinks::SchemeFilter; if url.starts_with("file://") { @@ -894,7 +871,7 @@ pub(crate) fn dispatch(action: Action, app: &mut AppView) -> Vec { } Action::OpenManagedConnectors => { use crate::terminal::hyperlinks::SchemeFilter; - let url = crate::views::mcps_modal::managed_connectors_url(app.team_id.as_deref()); + let url = crate::views::mcps_modal::managed_connectors_url(None); crate::app::link_opener::open_url_if_safe(&url, SchemeFilter::Standard); vec![] } diff --git a/crates/codegen/kigi-tui/src/app/dispatch/session/foreign.rs b/crates/codegen/kigi-tui/src/app/dispatch/session/foreign.rs index ab697a3..e67f9d3 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/session/foreign.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/session/foreign.rs @@ -1,7 +1,6 @@ use crate::app::actions::Effect; use crate::app::app_view::{AppView, SessionPickerEntry}; use crate::app::dispatch::ctx::get_active_agent_mut; -use crate::app::effects::ConversationsPartial; use crate::views::modal::ActiveModal; use crate::views::picker::PickerState; use crate::views::session_picker::{ @@ -61,7 +60,6 @@ impl PickerSurface<'_> { query: Option, chat_mode: bool, empty_notice: String, - partial_notice: Option<&'static str>, ) -> Option { let anchor = self.capture_selection(); let is_search = query.is_some(); @@ -88,11 +86,7 @@ impl PickerSurface<'_> { } } else { self.lanes.pending_notice = None; - if chat_mode { - partial_notice.map(str::to_owned) - } else { - None - } + None }; self.restore_selection(anchor); notice @@ -188,7 +182,6 @@ pub(in crate::app::dispatch) fn dispatch_fetch_session_list(app: &mut AppView) - pub(in crate::app::dispatch) fn handle_session_list_loaded( app: &mut AppView, sessions: Vec, - partial: Option, seq: u64, query: Option, ) -> Vec { @@ -196,18 +189,7 @@ pub(in crate::app::dispatch) fn handle_session_list_loaded( return vec![]; } app.session_picker_detail_generation += 1; - if let Some(partial) = partial { - crate::unified_log::warn( - "session.list.partial", - None, - Some(serde_json::json!({ "reason": format!("{partial:?}") })), - ); - } - let empty_notice = partial.map_or_else( - || "No sessions found for this directory".to_owned(), - |partial| partial.picker_notice().to_owned(), - ); - let partial_notice = partial.map(ConversationsPartial::picker_notice); + let empty_notice = "No sessions found for this directory".to_owned(); let chat_mode = app.chat_mode; let mut sessions = Some(sessions); let mut notice = None; @@ -242,7 +224,6 @@ pub(in crate::app::dispatch) fn handle_session_list_loaded( query.clone(), chat_mode, empty_notice.clone(), - partial_notice, ); } } @@ -260,7 +241,7 @@ pub(in crate::app::dispatch) fn handle_session_list_loaded( grouped: app.session_picker_grouped, current_repo, } - .native_loaded(sessions, query, chat_mode, empty_notice, partial_notice); + .native_loaded(sessions, query, chat_mode, empty_notice); } if let Some(notice) = notice { app.show_toast(¬ice); diff --git a/crates/codegen/kigi-tui/src/app/dispatch/session/fork.rs b/crates/codegen/kigi-tui/src/app/dispatch/session/fork.rs index 9e2bbb3..8d9059a 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/session/fork.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/session/fork.rs @@ -206,15 +206,8 @@ pub(in crate::app::dispatch) fn dispatch_fork_resolved( .prompt .set_contextual_hints(app.contextual_hints.undo, app.contextual_hints.plan_mode); agent.set_session_recap_available(app.session_recap_available); - agent.apply_app_scoped_gates( - app.sharing_enabled, - app.usage_visible, - app.chat_mode, - app.screen_mode, - &app.tier_restricted_commands, - ); + agent.apply_app_scoped_gates(app.usage_visible, app.chat_mode, app.screen_mode, &[]); agent.chat_kind = parent_chat_kind; - agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent .prompt .slash_controller @@ -353,7 +346,6 @@ pub(in crate::app::dispatch) fn dispatch_project_selected( let chat_kind = consume_chat_kind(app); if let Some(agent) = app.agents.get_mut(&id) { agent.chat_kind = chat_kind; - agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); } effects.push(Effect::CreateSession { agent_id: id, @@ -398,8 +390,6 @@ fn build_fork_placeholder( restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: app.bootstrap_acp_commands.clone(), available_commands_generation: 1, available_tools: None, @@ -543,7 +533,6 @@ pub(in crate::app::dispatch) fn handle_worktree_forked( } let effective_chat = conversation_entry || app.chat_mode; agent.chat_kind = effective_chat; - agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); return vec![Effect::LoadSession { agent_id, session_id: session_id_str, diff --git a/crates/codegen/kigi-tui/src/app/dispatch/session/lifecycle.rs b/crates/codegen/kigi-tui/src/app/dispatch/session/lifecycle.rs index 1226dc0..6241c5c 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/session/lifecycle.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/session/lifecycle.rs @@ -304,8 +304,6 @@ pub(in crate::app::dispatch) fn dispatch_new_session_inner_with_id( restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: app.bootstrap_acp_commands.clone(), available_commands_generation: 1, available_tools: None, @@ -330,14 +328,7 @@ pub(in crate::app::dispatch) fn dispatch_new_session_inner_with_id( .prompt .set_contextual_hints(app.contextual_hints.undo, app.contextual_hints.plan_mode); agent.set_session_recap_available(app.session_recap_available); - agent.apply_app_scoped_gates( - app.sharing_enabled, - app.usage_visible, - app.chat_mode, - app.screen_mode, - &app.tier_restricted_commands, - ); - agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); + agent.apply_app_scoped_gates(app.usage_visible, app.chat_mode, app.screen_mode, &[]); agent .prompt .slash_controller @@ -353,7 +344,6 @@ pub(in crate::app::dispatch) fn dispatch_new_session_inner_with_id( let chat_kind = consume_chat_kind(app); if let Some(agent) = app.agents.get_mut(&agent_id) { agent.chat_kind = chat_kind; - agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent.mcp_init_progress = Some(McpInitProgress { total: 0, connected: 0, @@ -402,7 +392,7 @@ pub(in crate::app::dispatch) fn dispatch_trust_folder(app: &mut AppView) -> Vec< /// `AuthComplete` uses, so whichever gate resolves last drains exactly once. pub(in crate::app::dispatch) fn finish_trust(app: &mut AppView) -> Vec { app.trust_state = TrustState::Done; - app.welcome_prompt_focused = !app.is_access_blocked(); + app.welcome_prompt_focused = true; if app.session_startup_allowed() { drain_startup_actions(app) } else { @@ -629,8 +619,6 @@ pub(in crate::app::dispatch) fn dispatch_new_worktree_session( restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: app.bootstrap_acp_commands.clone(), available_commands_generation: 1, available_tools: None, @@ -667,15 +655,8 @@ pub(in crate::app::dispatch) fn dispatch_new_worktree_session( .prompt .set_contextual_hints(app.contextual_hints.undo, app.contextual_hints.plan_mode); agent.set_session_recap_available(app.session_recap_available); - agent.apply_app_scoped_gates( - app.sharing_enabled, - app.usage_visible, - app.chat_mode, - app.screen_mode, - &app.tier_restricted_commands, - ); + agent.apply_app_scoped_gates(app.usage_visible, app.chat_mode, app.screen_mode, &[]); agent.chat_kind = chat_kind; - agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent .prompt .slash_controller @@ -764,7 +745,6 @@ pub(in crate::app::dispatch) fn skip_picker_and_create_session( let chat_kind = consume_chat_kind(app); if let Some(agent) = app.agents.get_mut(&agent_id) { agent.chat_kind = chat_kind; - agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent.mcp_init_progress = Some(McpInitProgress { total: 0, connected: 0, @@ -836,10 +816,6 @@ pub(in crate::app::dispatch) fn handle_session_created( session_id: session_id_clone.clone(), }); effects.push(Effect::RefreshAvailableCommands { agent_id, cwd }); - effects.push(Effect::FetchBilling { - agent_id, - silent: true, - }); if let Some((model_id, effort)) = deferred { effects.push(Effect::SwitchModel { agent_id, @@ -916,10 +892,6 @@ pub(in crate::app::dispatch) fn handle_worktree_session_created( session_id: session_id_clone.clone(), }); effects.push(Effect::RefreshAvailableCommands { agent_id, cwd }); - effects.push(Effect::FetchBilling { - agent_id, - silent: true, - }); if let Some((model_id, effort)) = deferred { effects.push(Effect::SwitchModel { agent_id, diff --git a/crates/codegen/kigi-tui/src/app/dispatch/session/load.rs b/crates/codegen/kigi-tui/src/app/dispatch/session/load.rs index d89e973..3fad218 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/session/load.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/session/load.rs @@ -162,8 +162,6 @@ fn dispatch_load_session_ungated( restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: app.bootstrap_acp_commands.clone(), available_commands_generation: 1, available_tools: None, @@ -195,15 +193,8 @@ fn dispatch_load_session_ungated( agent_mut.session.start_command(AgentCommand::RestoreCode); agent_mut.turn_started_at = Some(std::time::Instant::now()); } - agent_mut.apply_app_scoped_gates( - app.sharing_enabled, - app.usage_visible, - app.chat_mode, - app.screen_mode, - &app.tier_restricted_commands, - ); + agent_mut.apply_app_scoped_gates(app.usage_visible, app.chat_mode, app.screen_mode, &[]); agent_mut.chat_kind = chat_kind || app.chat_mode; - agent_mut.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent_mut .prompt .slash_controller @@ -808,8 +799,6 @@ pub(in crate::app::dispatch) fn dispatch_load_session_with_restore( restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: app.bootstrap_acp_commands.clone(), available_commands_generation: 1, available_tools: None, @@ -836,15 +825,8 @@ pub(in crate::app::dispatch) fn dispatch_load_session_with_restore( .prompt .set_contextual_hints(app.contextual_hints.undo, app.contextual_hints.plan_mode); agent.set_session_recap_available(app.session_recap_available); - agent.apply_app_scoped_gates( - app.sharing_enabled, - app.usage_visible, - app.chat_mode, - app.screen_mode, - &app.tier_restricted_commands, - ); + agent.apply_app_scoped_gates(app.usage_visible, app.chat_mode, app.screen_mode, &[]); agent.chat_kind = app.chat_mode; - agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent .prompt .slash_controller @@ -958,10 +940,6 @@ pub(in crate::app::dispatch) fn handle_session_loaded( agent_id, session_id: hydrate_sid.clone(), }); - effects.push(Effect::FetchBilling { - agent_id, - silent: true, - }); if let Some((model_id, effort)) = deferred { agent.session.model_switch_pending = true; effects.push(Effect::SwitchModel { @@ -1110,7 +1088,6 @@ pub(in crate::app::dispatch) fn handle_session_restored( supersede_open_reload_window(agent, agent_id, "SessionRestored"); agent.bind_session_id(sid); agent.chat_kind = app.chat_mode; - agent.apply_credit_balance(app.credit_balance.clone(), app.auto_topup.clone()); agent.scrollback.push_block(RenderBlock::system(format!( "Session restored. Loading {local_session_id}..." ))); diff --git a/crates/codegen/kigi-tui/src/app/dispatch/settings/ui.rs b/crates/codegen/kigi-tui/src/app/dispatch/settings/ui.rs index b4fc240..1247619 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/settings/ui.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/settings/ui.rs @@ -44,7 +44,6 @@ pub(crate) fn refresh_open_settings_modals(app: &mut AppView) { } let ui_snapshot = app.current_ui.clone(); // Capture app-level fields before the mut-borrow loop. - let coding_data_sharing_opt_out_from_app = app.coding_data_retention_opt_out; let show_tips_from_app = app.show_tips; let auto_update_from_app = app.auto_update; let respect_manual_folds_from_app = app.appearance.scrollback.scroll.respect_manual_folds; @@ -76,7 +75,6 @@ pub(crate) fn refresh_open_settings_modals(app: &mut AppView) { .iter() .map(|(id, info)| (info.name.clone(), id.clone())) .collect(), - coding_data_sharing_opt_out: coding_data_sharing_opt_out_from_app, // Prefer optimistic pending over confirmed active. plan_mode_active: agent.plan_mode_pending.unwrap_or(agent.plan_mode_active), show_tips: show_tips_from_app, @@ -110,7 +108,7 @@ pub(in crate::app::dispatch) fn dispatch_open_command_palette(app: &mut AppView) return vec![]; } agent.active_modal = Some(ActiveModal::CommandPalette { - entries: crate::views::modal::default_palette_entries(agent.sharing_enabled), + entries: crate::views::modal::default_palette_entries(), // Type-to-find: open in input mode (matches Ctrl+P). state: crate::views::picker::PickerState::input_active(), window: crate::views::modal_window::ModalWindowState::new(), @@ -149,7 +147,6 @@ pub(in crate::app::dispatch) fn dispatch_open_settings(app: &mut AppView) -> Vec let registry = app.settings_registry.clone(); let ui_snapshot = app.current_ui.clone(); // Capture app-level fields before the mut-borrow on the agent. - let coding_data_sharing_opt_out_from_app = app.coding_data_retention_opt_out; let show_tips_from_app = app.show_tips; let auto_update_from_app = app.auto_update; let respect_manual_folds_from_app = app.appearance.scrollback.scroll.respect_manual_folds; @@ -187,7 +184,6 @@ pub(in crate::app::dispatch) fn dispatch_open_settings(app: &mut AppView) -> Vec .iter() .map(|(id, info)| (info.name.clone(), id.clone())) .collect(), - coding_data_sharing_opt_out: coding_data_sharing_opt_out_from_app, // Prefer optimistic pending over confirmed active. plan_mode_active: agent.plan_mode_pending.unwrap_or(agent.plan_mode_active), show_tips: show_tips_from_app, @@ -657,7 +653,6 @@ pub(crate) fn build_pager_snapshot(app: &AppView) -> crate::settings::PagerLocal auto_mode: agent_auto_mode(app), current_model_name: agent_current_model_name(app), available_models: agent_available_models(app), - coding_data_sharing_opt_out: app.coding_data_retention_opt_out, plan_mode_active: agent_plan_mode(app), show_tips: app.show_tips, auto_update: app.auto_update, @@ -793,14 +788,6 @@ pub(in crate::app::dispatch) fn action_for_reset( } // max_thoughts_width: direct round-trip. ("max_thoughts_width", SettingValue::Int(i)) => Some(Action::SetMaxThoughtsWidth(*i)), - // coding_data_sharing: "opt-in" / "opt-out" → bool. - // "opt-out" arm is a skew guard (default is "opt-in"). - ("coding_data_sharing", SettingValue::Enum("opt-in")) => { - Some(Action::SetCodingDataSharing { opted_in: true }) - } - ("coding_data_sharing", SettingValue::Enum("opt-out")) => { - Some(Action::SetCodingDataSharing { opted_in: false }) - } // plan_mode: "on" / "off" → PlanModeKind. // "on" arm is a skew guard (default is "off"). ("plan_mode", SettingValue::Enum("off")) => { diff --git a/crates/codegen/kigi-tui/src/app/dispatch/status.rs b/crates/codegen/kigi-tui/src/app/dispatch/status.rs index 0d65d76..a48a3ac 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/status.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/status.rs @@ -1,7 +1,6 @@ -//! Session status, sharing, privacy, usage, and info dispatchers. +//! Session status, privacy, usage, and info dispatchers. use super::ctx::get_active_agent; -use super::settings::ui::refresh_open_settings_modals; use crate::app::actions::Effect; use crate::app::agent::AgentId; use crate::app::agent_view::AgentView; @@ -9,39 +8,6 @@ use crate::app::app_view::{ActiveView, AppView}; use crate::notifications::{NotificationEvent, NotificationEventKind}; use crate::scrollback::block::RenderBlock; -/// Toggle YOLO mode (auto-approve all permissions). -/// -/// When turning ON: auto-approve all currently queued permissions and -/// restore the stashed prompt. Future incoming permissions will be -/// auto-approved in `handle_permission_request`. -/// -/// Share the current session via a public URL. -/// -/// Produces Effect::ShareSession which spawns an async ACP ext request. -/// On completion, TaskResult::ShareSessionComplete shows the URL in scrollback. -pub(super) fn dispatch_share_session(app: &mut AppView) -> Vec { - if !app.sharing_enabled { - app.show_toast("Sharing is disabled"); - return vec![]; - } - let ActiveView::Agent(id) = app.active_view else { - return vec![]; - }; - let Some(agent) = app.agents.get_mut(&id) else { - return vec![]; - }; - let Some(session_id) = agent.session.session_id.clone() else { - // No active session — error should have been caught by slash command, - // but guard here just in case. - return vec![]; - }; - - vec![Effect::ShareSession { - agent_id: id, - session_id, - }] -} - /// Show session info: fetch via x.ai/session/info and display in scrollback. /// /// Produces Effect::ShowSessionInfo which spawns an async ACP ext request. @@ -66,131 +32,6 @@ pub(super) fn dispatch_show_session_info(app: &mut AppView) -> Vec { }] } -/// Show privacy and data retention status as a system message in scrollback. -/// -/// Three-state display: Enterprise ZDR, coding data sharing opted out, -/// or opted in. Labels align with `CODING_DATA_SHARING_CHOICES` in -/// `settings/defs.rs` and the `coding_data_sharing_toast` format. -pub(super) fn dispatch_show_privacy_info(app: &mut AppView) -> Vec { - let mut lines = Vec::new(); - - if app.is_zdr { - // Enterprise ZDR -- the team has disabled retention entirely. - lines.push(" Zero Data Retention: enabled"); - lines.push(" Your data is not retained or used for training (ZDR enabled)."); - } else if app.coding_data_retention_opt_out { - // Coding data sharing opted out -- matches desktop's "Privacy mode" state. - lines.push(" Privacy: privacy mode"); - lines.push(" Your code data will not be trained on or used to improve the product."); - lines.push(""); - lines.push(" Use /privacy opt-in to share data and help improve the product."); - } else { - // Coding data sharing opted in -- matches desktop's "Share data" state. - lines.push(" Privacy: share data"); - lines.push(" Usage and code data may be used by SpaceXAI to improve the product."); - lines.push(""); - lines.push(" Use /privacy opt-out to enable privacy mode."); - } - - lines.push(""); - lines.push(" Learn more: https://x.ai/legal"); - let text = lines.join("\n"); - push_system_to_any_agent(app, &text); - vec![] -} - -/// State-only mutation for `coding_data_sharing`. SHELL-owned. -pub(super) fn set_coding_data_sharing_inner(app: &mut AppView, opted_in: bool) { - app.coding_data_retention_opt_out = !opted_in; -} - -/// Set coding-data-sharing preference. SHELL-owned, auth-metadata-backed -/// (persists via ACP ext-request, NOT `~/.kigi/config.toml`). -pub(super) fn set_coding_data_sharing(app: &mut AppView, opted_in: bool) -> Vec { - // ── Guard 1: Enterprise ZDR ────────────────────────────────────── - if app.is_zdr { - app.show_toast("\u{2717} Cannot change: Zero Data Retention enabled"); - return vec![]; - } - // ── Guard 2: Non-admin team member ─────────────────────────────── - if app.team_name.is_some() { - let is_admin = app - .team_role - .as_deref() - .is_some_and(|r| r.eq_ignore_ascii_case("admin")); - if !is_admin { - app.show_toast("\u{2717} Data sharing is controlled by your team admin"); - return vec![]; - } - } - // ── Guard 3: an agent must exist to thread the ACP call through ── - let agent_id = match app.active_view { - crate::app::app_view::ActiveView::Agent(id) => id, - _ => match app.agents.keys().next().copied() { - Some(id) => id, - None => { - tracing::warn!( - target: "settings", - key = "coding_data_sharing", - opted_in, - "set_coding_data_sharing called with no agents — unreachable in \ - practice; returning empty (no toast: app.show_toast would no-op)", - ); - return vec![]; - } - }, - }; - - let prev = !app.coding_data_retention_opt_out; - - // ── Idempotent path: toast but skip the ACP round-trip. ────────── - if prev == opted_in { - app.show_toast(&coding_data_sharing_toast(opted_in)); - return vec![]; - } - - // ── Optimistic mutation: state, then UI feedback, then effect. ─── - set_coding_data_sharing_inner(app, opted_in); - refresh_open_settings_modals(app); - app.show_toast(&coding_data_sharing_toast(opted_in)); - - tracing::info!( - target: "settings", - key = "coding_data_sharing", - opted_in, - "setting changed", - ); - - vec![Effect::SetCodingDataSharing { - agent_id, - opted_in, - rollback_to_opted_in: prev, - }] -} - -/// Format the `Coding data sharing` toast. Asymmetric: opt-in -/// (privacy-degrading) uses ⚠ + consequence text; opt-out (safe -/// default) uses ✓. Uses display names from the registry catalog. -pub(super) fn coding_data_sharing_toast(opted_in: bool) -> String { - let display = display_for_coding_data_sharing_canonical(opted_in); - if opted_in { - // Privacy-degrading: warn glyph + spelled-out consequence. - format!( - "\u{26A0} Coding data sharing: {display} \u{2014} code samples may be retained \ - for training" - ) - } else { - // Safe default — uniform ✓ glyph. - format!("\u{2713} Coding data sharing: {display}") - } -} - -/// Display string for the canonical bool. Keep aligned with -/// `CODING_DATA_SHARING_CHOICES` in `settings/defs.rs`. -fn display_for_coding_data_sharing_canonical(opted_in: bool) -> &'static str { - if opted_in { "Opt in" } else { "Opt out" } -} - /// Scrub an untrusted error string for toast display. Substitutes a /// generic placeholder when the input exceeds 120 chars or contains /// control / bidi-override characters (prevents escape-sequence @@ -208,21 +49,6 @@ pub(super) fn scrub_error_for_toast(error: &str) -> String { } } -/// Push a system message to the active agent's scrollback, or to any available -/// agent if on the welcome screen. -fn push_system_to_any_agent(app: &mut AppView, msg: &str) { - let block = crate::scrollback::block::RenderBlock::system(msg.to_string()); - if let ActiveView::Agent(id) = app.active_view - && let Some(agent) = app.agents.get_mut(&id) - { - agent.scrollback.push_block(block); - return; - } - if let Some(agent) = app.agents.values_mut().next() { - agent.scrollback.push_block(block); - } -} - /// Show context info: fetch via x.ai/session/info and display rich breakdown. /// /// Produces Effect::ShowContextInfo which spawns an async ACP ext request. @@ -244,32 +70,73 @@ pub(super) fn dispatch_show_context_info(app: &mut AppView) -> Vec { }] } -/// Show credit usage: fetch billing data and display inline. +/// `/usage` — fetch Kimi usage/quota rows and display them inline. /// -/// When the remote settings `grok_build_usage_redirect_url` flag is set (delivered via -/// RemoteSettings, targeted at personal-team users), skip the backend fetch and -/// just point the user at that URL instead. This is a kill switch for the -/// personal-team billing path while it is unreliable. +/// Produces [`Effect::FetchUsage`], which asks the shell's `x.ai/billing` +/// extension (`GET {base}/usages`); [`handle_usage_fetched`] renders the +/// rows as a system block in scrollback. pub(super) fn dispatch_show_usage(app: &mut AppView) -> Vec { let ActiveView::Agent(id) = app.active_view else { return vec![]; }; - if let Some(url) = app.usage_billing_redirect_url.clone() { - if let Some(agent) = app.agents.get_mut(&id) { - agent.scrollback.push_block(RenderBlock::System( - crate::scrollback::blocks::SystemMessageBlock::new(format!( - "Please check your usage on {url}" - )), - )); - } - return vec![]; + vec![Effect::FetchUsage { agent_id: id }] +} + +/// Render the `/usage` result: the fetched quota rows (kimi-cli +/// `usage.py` semantics — label, remaining-quota bar, percent left, reset +/// hint), "No usage data available." for an empty list, or the error. +pub(super) fn handle_usage_fetched( + app: &mut AppView, + agent_id: AgentId, + result: Result, String>, +) -> Vec { + let msg = match &result { + Ok(rows) if rows.is_empty() => "No usage data available.".to_string(), + Ok(rows) => format_usage_rows(rows), + Err(e) => format!("Couldn't fetch usage: {e}"), + }; + if let Some(agent) = app.agents.get_mut(&agent_id) { + agent.scrollback.push_block(RenderBlock::system(msg)); } - // Non-silent fetch: the effect also pulls the auto top-up rule so the - // summary can render usage, prepaid credits, and auto top-up together. - vec![Effect::FetchBilling { - agent_id: id, - silent: false, - }] + vec![] +} + +/// Width of the remaining-quota bar, matching kimi-cli's usage panel. +const USAGE_BAR_WIDTH: usize = 20; + +/// Format usage rows as aligned text lines (kimi-cli `_format_row` +/// parity): `label [bar] N% left (reset hint)`. The percentage is +/// derived from `used`/`limit` only — a row without a positive limit +/// renders as 0% left with an empty bar, exactly like kimi-cli. +fn format_usage_rows(rows: &[kigi_shell::extensions::billing::UsageRow]) -> String { + let label_width = rows + .iter() + .map(|r| r.label.chars().count()) + .max() + .unwrap_or(0) + .max(6); + let mut lines = vec!["API Usage".to_string()]; + for row in rows { + let ratio = if row.limit <= 0 { + 0.0 + } else { + (row.limit - row.used).clamp(0, row.limit) as f64 / row.limit as f64 + }; + let filled = (ratio * USAGE_BAR_WIDTH as f64).round() as usize; + let filled = filled.min(USAGE_BAR_WIDTH); + let bar: String = "\u{2588}".repeat(filled) + &"\u{2591}".repeat(USAGE_BAR_WIDTH - filled); + let mut line = format!( + " {: Vec { - // Re-anchor mirror to server-confirmed value (defense-in- - // depth against server reshaping the boolean). `agent_id` - // discarded — privacy is app-level, not per-agent. - set_coding_data_sharing_inner(app, opted_in); - refresh_open_settings_modals(app); - // Re-toast on confirmation. Without this, a slow ACP - // round-trip would leave the user with only the - // optimistic toast (already faded) and no - // server-confirmed feedback. - app.show_toast(&coding_data_sharing_toast(opted_in)); - tracing::info!( - target: "settings", - key = "coding_data_sharing", - ?agent_id, - opted_in, - "ACP update confirmed; mirror re-anchored", - ); - vec![] -} - -pub(super) fn handle_coding_data_sharing_failed( - app: &mut AppView, - agent_id: AgentId, - error: String, - rollback_to_opted_in: bool, -) -> Vec { - // Revert optimistic mutation: inner → refresh → toast. - // - // `agent_id` discarded — privacy is global. - set_coding_data_sharing_inner(app, rollback_to_opted_in); - refresh_open_settings_modals(app); - // Scrub long/unsafe error strings before toasting. - let scrubbed = scrub_error_for_toast(&error); - app.show_toast(&format!( - "\u{2717} Couldn't update coding data sharing: {scrubbed}" - )); - tracing::warn!( - target: "settings", - key = "coding_data_sharing", - ?agent_id, - rollback_to_opted_in, - %error, - "ACP update failed; reverted optimistic mutation", - ); - vec![] -} - pub(super) fn handle_context_info_complete( app: &mut AppView, agent_id: AgentId, diff --git a/crates/codegen/kigi-tui/src/app/dispatch/task_result.rs b/crates/codegen/kigi-tui/src/app/dispatch/task_result.rs index d79a7b6..0a6cd05 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/task_result.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/task_result.rs @@ -2,11 +2,6 @@ use super::auth::{ ensure_login_method, handle_auth_complete, handle_auth_url_ready, handle_mcp_auth_trigger_done, }; -use super::billing::{ - PAYWALL_AUTO_CHECK_TIMEOUT, apply_auto_topup, handle_billing_fetched, - handle_check_subscription_complete, handle_credit_limit_recheck_complete, - handle_gate_refreshed, handle_gate_verify_timeout, -}; use super::ctx::{find_agent_by_session_id, get_active_agent_mut}; use super::notes::{handle_btw_response, handle_memory_note_saved}; use super::prompt::{ @@ -34,10 +29,7 @@ use super::session::load::{ handle_session_search_debounce_expired, remove_session_from_pickers, }; use super::settings::ui::apply_setting_rollback; -use super::status::{ - handle_coding_data_sharing_failed, handle_coding_data_sharing_updated, - handle_context_info_complete, scrub_error_for_toast, -}; +use super::status::{handle_context_info_complete, handle_usage_fetched, scrub_error_for_toast}; use super::transcript::{ handle_hooks_list_loaded, handle_mcp_toggle_done, handle_plugins_list_loaded, handle_skills_toggle_done, @@ -224,33 +216,9 @@ pub(super) fn dispatch_task_result(result: TaskResult, app: &mut AppView) -> Vec TaskResult::ForkSessionFailed { agent_id, error } => { handle_fork_session_failed(app, agent_id, error) } - TaskResult::BillingFetched { - agent_id, - balance, - silent, - subscription_tier, - autotopup, - } => handle_billing_fetched(app, agent_id, balance, silent, subscription_tier, autotopup), - TaskResult::BillingError { - agent_id, - error, - silent, - } => { - if !silent && let Some(agent) = app.agents.get_mut(&agent_id) { - agent.scrollback.push_block(RenderBlock::System( - crate::scrollback::blocks::SystemMessageBlock::new(format!( - "Billing error: {error}" - )), - )); - } - vec![] + TaskResult::UsageFetched { agent_id, result } => { + handle_usage_fetched(app, agent_id, result) } - TaskResult::AppBillingFetched { balance, autotopup } => { - app.credit_balance = balance; - apply_auto_topup(&mut app.auto_topup, &autotopup); - vec![] - } - TaskResult::GateRefreshed { settings } => handle_gate_refreshed(app, settings), TaskResult::SessionLoaded { agent_id, session_id, @@ -287,10 +255,9 @@ pub(super) fn dispatch_task_result(result: TaskResult, app: &mut AppView) -> Vec } => handle_session_load_failed(app, agent_id, session_id, error), TaskResult::SessionListLoaded { sessions, - partial, seq, query, - } => handle_session_list_loaded(app, sessions, partial, seq, query), + } => handle_session_list_loaded(app, sessions, seq, query), TaskResult::ForeignSessionsScanned { entries, seq } => { handle_foreign_sessions_scanned(app, entries, seq) } @@ -606,29 +573,6 @@ pub(super) fn dispatch_task_result(result: TaskResult, app: &mut AppView) -> Vec TaskResult::SkillsToggleDone { agent_id, result } => { handle_skills_toggle_done(app, agent_id, result) } - TaskResult::ShareSessionComplete { - agent_id, - share_url, - } => { - if let Some(agent) = app.agents.get_mut(&agent_id) { - agent - .scrollback - .push_block(crate::scrollback::block::RenderBlock::system(format!( - "Session shared: {share_url}" - ))); - } - vec![] - } - TaskResult::ShareSessionFailed { agent_id, error } => { - if let Some(agent) = app.agents.get_mut(&agent_id) { - agent - .scrollback - .push_block(crate::scrollback::block::RenderBlock::system(format!( - "Couldn't share session: {error}" - ))); - } - vec![] - } TaskResult::SessionAgentNameResolved { agent_id, agent_name, @@ -668,14 +612,6 @@ pub(super) fn dispatch_task_result(result: TaskResult, app: &mut AppView) -> Vec } vec![] } - TaskResult::CodingDataSharingUpdated { agent_id, opted_in } => { - handle_coding_data_sharing_updated(app, agent_id, opted_in) - } - TaskResult::CodingDataSharingFailed { - agent_id, - error, - rollback_to_opted_in, - } => handle_coding_data_sharing_failed(app, agent_id, error, rollback_to_opted_in), TaskResult::RenameSessionComplete { agent_id, title } => { if let Some(agent) = app.agents.get_mut(&agent_id) { let safe = crate::views::session_title::sanitize_display_text(&title); @@ -873,32 +809,8 @@ pub(super) fn dispatch_task_result(result: TaskResult, app: &mut AppView) -> Vec app.auth_clipboard_copied = false; vec![] } - TaskResult::PaywallCheckTick => { - let timed_out = app - .paywall_check_started - .is_some_and(|t| t.elapsed() >= PAYWALL_AUTO_CHECK_TIMEOUT); - if !app.has_access() && !timed_out { - vec![ - Effect::CheckSubscription { verify: None }, - Effect::SchedulePaywallCheck, - ] - } else { - vec![] - } - } - TaskResult::CheckSubscriptionComplete { verify, meta } => { - handle_check_subscription_complete(app, verify, meta) - } - TaskResult::GateVerifyTimeout { generation } => handle_gate_verify_timeout(app, generation), - TaskResult::CreditLimitRecheckComplete { agent_id, meta } => { - handle_credit_limit_recheck_complete(app, agent_id, meta) - } TaskResult::LogoutComplete => { app.auth_state = AuthState::Pending { error: None }; - app.access_gate_shown_logged = false; - app.gate = None; - app.pending_gate_verification = None; - app.last_subscription_check_at = None; app.login_method_id = None; ensure_login_method(app); app.auth_clipboard_copied = false; diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/billing.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/billing.rs index 2e347ab..4c305f4 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/billing.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/billing.rs @@ -1,983 +1,123 @@ -//! Tests for credit-limit upsells, paywall gating, and auto-topup. +//! Tests for the `/usage` view. use super::*; -// ── Credit-limit upsell / max-tier tests ─────────────────────────── - -/// Open the non-max-tier Q&A upsell modal. Panics if the modal was not created. -fn open_upsell_qa(app: &mut AppView, mode: CreditLimitUpsellMode) { - let agent = app.agents.get_mut(&AgentId(0)).unwrap(); - open_credit_limit_upsell(agent, mode, false); -} - -/// Open the max-tier inline scrollback card upsell. -fn open_upsell_max_card(app: &mut AppView, mode: CreditLimitUpsellMode) { - let agent = app.agents.get_mut(&AgentId(0)).unwrap(); - open_credit_limit_upsell(agent, mode, true); -} - -/// Return the `QuestionViewState` from agent 0. Panics if absent. -fn agent_qv(app: &AppView) -> &crate::views::question_view::QuestionViewState { - app.agents - .get(&AgentId(0)) - .unwrap() - .question_view - .as_ref() - .unwrap() -} - -/// Extract the last-pushed `CreditLimitBlock` from agent 0 scrollback. -fn last_credit_limit_block( - app: &AppView, - idx: usize, -) -> &crate::scrollback::blocks::CreditLimitBlock { +/// Text of the last System block in agent 0's scrollback. Panics if the +/// last block is not a System message. +fn last_system_text(app: &AppView) -> String { let agent = app.agents.get(&AgentId(0)).unwrap(); - if let crate::scrollback::block::RenderBlock::CreditLimit(ref blk) = - agent.scrollback.entry(idx).unwrap().block - { - blk - } else { - panic!("expected CreditLimit block at index {idx}"); + let idx = agent.scrollback.len() - 1; + match &agent.scrollback.entry(idx).unwrap().block { + crate::scrollback::block::RenderBlock::System(b) => b.text.clone(), + other => panic!("expected System block, got {other:?}"), } } -/// Dispatch a `BillingFetched` task result with sensible defaults. -fn dispatch_billing( - app: &mut AppView, - balance: Option, - silent: bool, - subscription_tier: Option, -) { - dispatch( - Action::TaskComplete(TaskResult::BillingFetched { - agent_id: AgentId(0), - balance, - silent, - subscription_tier, - autotopup: crate::views::credit_bar::AutoTopupFetch::Unchanged, - }), - app, - ); -} +// ── /usage dispatch tests ─────────────────────────────────── #[test] -fn is_max_tier_positive_match() { - assert!(is_max_tier(Some("supergrok_heavy"))); - assert!(is_max_tier(Some("SuperGrok Heavy"))); - assert!(is_max_tier(Some("SUPERGROK_HEAVY"))); -} - -#[test] -fn is_max_tier_non_max_and_unknown() { - assert!(!is_max_tier(Some("supergrok"))); - assert!(!is_max_tier(Some("premium"))); - assert!(!is_max_tier(Some("free"))); - // Unknown defaults to non-max → Q&A shown. - assert!(!is_max_tier(None)); -} - -#[test] -fn is_max_tier_handles_mixed_case_and_whitespace() { - assert!(is_max_tier(Some("SuperGrok_Heavy"))); - assert!(is_max_tier(Some("supergrok heavy"))); - assert!(is_max_tier(Some("SUPERGROK HEAVY"))); -} - -#[test] -fn is_max_tier_rejects_partial_matches() { - assert!(!is_max_tier(Some("supergrok_heav"))); - assert!(!is_max_tier(Some("supergrok_heavy_plus"))); - assert!(!is_max_tier(Some(""))); -} - -#[test] -fn upsell_non_max_shows_qa_with_two_options() { - let mut app = test_app_with_agent(); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - let q = &agent_qv(&app).questions[0]; - assert_eq!(q.options.len(), 2); - assert_eq!(q.options[0].label, "Upgrade tier"); - assert_eq!(q.options[0].id.as_deref(), Some(UPSELL_URL_UPGRADE)); - assert_eq!(q.options[1].label, "Pay as you go"); - assert_eq!(q.options[1].id.as_deref(), Some(UPSELL_URL_PAYG)); -} - -#[test] -fn upsell_non_max_payg_on_shows_increase_label() { - let mut app = test_app_with_agent(); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: true }, - ); - let q = &agent_qv(&app).questions[0]; - assert_eq!(q.options.len(), 2); - assert_eq!(q.options[1].label, "Increase limit"); -} - -#[test] -fn upsell_non_max_qa_heading_is_credit_limit_when_payg_off() { - let mut app = test_app_with_agent(); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - let heading = &agent_qv(&app).questions[0].question; - assert!( - heading.contains("credit limit"), - "expected 'credit limit' in heading, got: {heading}" - ); -} - -#[test] -fn upsell_non_max_qa_heading_is_spending_cap_when_payg_on() { - let mut app = test_app_with_agent(); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: true }, - ); - let heading = &agent_qv(&app).questions[0].question; - assert!( - heading.contains("spending cap"), - "expected 'spending cap' in heading, got: {heading}" - ); -} - -#[test] -fn upsell_non_max_upgrade_url_is_supergrok() { - let mut app = test_app_with_agent(); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - let url = agent_qv(&app).questions[0].options[0] - .id - .as_deref() - .unwrap(); - assert!(url.contains("supergrok"), "got: {url}"); - assert!(url.contains("referrer=grok-build"), "got: {url}"); -} - -#[test] -fn upsell_non_max_payg_url_is_usage() { - let mut app = test_app_with_agent(); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - let url = agent_qv(&app).questions[0].options[1] - .id - .as_deref() - .unwrap(); - assert!(url.contains("_s=usage"), "got: {url}"); -} - -#[test] -fn upsell_non_max_payg_on_description_mentions_spending_cap() { - let mut app = test_app_with_agent(); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: true }, - ); - assert_eq!( - agent_qv(&app).questions[0].options[1].description, - "Raise your pay-as-you-go spending cap" - ); -} - -#[test] -fn upsell_non_max_payg_off_description_mentions_on_demand() { - let mut app = test_app_with_agent(); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - assert_eq!( - agent_qv(&app).questions[0].options[1].description, - "Enable pay-as-you-go credits for on-demand usage" - ); -} - -#[test] -fn upsell_non_max_unified_shows_buy_credits() { - let mut app = test_app_with_agent(); - open_upsell_qa(&mut app, CreditLimitUpsellMode::UnifiedCredits); - let q = &agent_qv(&app).questions[0]; - assert!(q.question.contains("weekly limit")); - assert_eq!( - q.options[0].description, - "Upgrade to a higher tier for more usage" - ); - assert_eq!(q.options[1].label, "Buy more credits"); - assert_eq!( - q.options[1].description, - "Purchase credits to keep using Grok Build" - ); -} - -#[test] -fn upsell_max_unified_card_mentions_purchasing() { - let mut app = test_app_with_agent(); - let before = agent_scrollback_len(&app); - open_upsell_max_card(&mut app, CreditLimitUpsellMode::UnifiedCredits); - let blk = last_credit_limit_block(&app, before); - assert_eq!( - blk.action, - crate::scrollback::blocks::CreditLimitCardAction::PurchaseCredits - ); - assert!(blk.heading.contains("weekly limit")); -} - -#[test] -fn credit_limit_upsell_mode_prefers_unified_flag() { - let mut bal = test_bal(100.0); - bal.is_unified_billing_user = Some(true); - bal.pay_as_you_go = true; // must not override explicit unified - assert_eq!( - credit_limit_upsell_mode(Some(&bal)), - CreditLimitUpsellMode::UnifiedCredits - ); - bal.is_unified_billing_user = Some(false); - bal.pay_as_you_go = false; - assert_eq!( - credit_limit_upsell_mode(Some(&bal)), - CreditLimitUpsellMode::LegacyPayg { enabled: false } - ); - bal.is_unified_billing_user = None; - bal.pay_as_you_go = true; - assert_eq!( - credit_limit_upsell_mode(Some(&bal)), - CreditLimitUpsellMode::LegacyPayg { enabled: true } - ); - bal.pay_as_you_go = false; - assert_eq!( - credit_limit_upsell_mode(Some(&bal)), - CreditLimitUpsellMode::UnifiedCredits - ); - assert_eq!( - credit_limit_upsell_mode(None), - CreditLimitUpsellMode::UnifiedCredits - ); -} - -#[test] -fn is_credit_limit_error_matches_legacy_403_and_pool_402() { - assert!(is_credit_limit_error( - Some(403), - "status 403: run out of credits" - )); - // 402 Payment Required is always credit/spend on this surface. - assert!(is_credit_limit_error(Some(402), "anything")); - assert!(is_credit_limit_error( - None, - "API error (status 402 Payment Required): Grok Build usage balance exhausted" - )); - assert!(is_credit_limit_error( - None, - "status 403: run out of credits" - )); - assert!(!is_credit_limit_error(Some(403), "content safety blocked")); - assert!(!is_credit_limit_error(Some(500), "internal server error")); - // Pool phrases alone without 402/403 status do not match. - assert!(!is_credit_limit_error( - None, - "usage balance exhausted without status" - )); -} - -#[test] -fn upsell_non_max_sets_no_freeform() { - let mut app = test_app_with_agent(); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - assert!( - agent_qv(&app).no_freeform, - "upsell Q&A should disable freeform input" - ); -} - -#[test] -fn upsell_non_max_qa_has_single_select() { - let mut app = test_app_with_agent(); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - assert_eq!( - agent_qv(&app).questions[0].multi_select, - Some(false), - "upsell should be single-select" - ); -} - -#[test] -fn upsell_non_max_does_not_push_scrollback_block() { - let mut app = test_app_with_agent(); - let before = agent_scrollback_len(&app); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - assert_eq!( - agent_scrollback_len(&app), - before, - "non-max-tier upsell should NOT push a scrollback block" - ); -} - -#[test] -fn upsell_non_max_idempotent_when_question_view_already_open() { - let mut app = test_app_with_agent(); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - assert!(app.agents.get(&AgentId(0)).unwrap().question_view.is_some()); - - // Second call should be a no-op (guard at line 2070). - let before = agent_scrollback_len(&app); - open_upsell_qa( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - assert_eq!( - agent_scrollback_len(&app), - before, - "second call should not push a block" - ); -} - -#[test] -fn upsell_max_tier_pushes_scrollback_card_payg_off() { - let mut app = test_app_with_agent(); - let before = agent_scrollback_len(&app); - open_upsell_max_card( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - assert!( - app.agents.get(&AgentId(0)).unwrap().question_view.is_none(), - "max-tier should NOT open the question modal" - ); - assert_eq!(agent_scrollback_len(&app), before + 1); - let blk = last_credit_limit_block(&app, before); - assert!(blk.heading.contains("credit limit")); - assert_eq!( - blk.action, - crate::scrollback::blocks::CreditLimitCardAction::EnablePayg - ); - assert_eq!(blk.url, UPSELL_URL_PAYG); -} - -#[test] -fn upsell_max_tier_pushes_scrollback_card_payg_on() { - let mut app = test_app_with_agent(); - let before = agent_scrollback_len(&app); - open_upsell_max_card( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: true }, - ); - assert!(app.agents.get(&AgentId(0)).unwrap().question_view.is_none()); - assert_eq!(agent_scrollback_len(&app), before + 1); - let blk = last_credit_limit_block(&app, before); - assert!(blk.heading.contains("spending cap")); - assert_eq!( - blk.action, - crate::scrollback::blocks::CreditLimitCardAction::IncreasePaygLimit - ); - assert_eq!(blk.url, UPSELL_URL_PAYG); -} - -#[test] -fn upsell_max_tier_does_not_open_question_view() { - let mut app = test_app_with_agent(); - open_upsell_max_card( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - assert!( - app.agents.get(&AgentId(0)).unwrap().question_view.is_none(), - "max-tier should use inline card, not question modal" - ); -} - -#[test] -fn upsell_max_tier_scrollback_card_url_is_payg() { - let mut app = test_app_with_agent(); - let before = agent_scrollback_len(&app); - open_upsell_max_card( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - assert_eq!(last_credit_limit_block(&app, before).url, UPSELL_URL_PAYG); -} - -#[test] -fn upsell_max_tier_not_idempotent_pushes_multiple_cards() { - let mut app = test_app_with_agent(); - let before = agent_scrollback_len(&app); - // Max-tier path doesn't guard against duplicates — each call - // pushes a new inline card. - open_upsell_max_card( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - open_upsell_max_card( - &mut app, - CreditLimitUpsellMode::LegacyPayg { enabled: false }, - ); - assert_eq!( - agent_scrollback_len(&app), - before + 2, - "max-tier path pushes a card on every call" - ); -} - -// ── ShowUsage dispatch tests ──────────────────────────────────────── - -#[test] -fn show_usage_returns_fetch_billing_effect() { +fn show_usage_returns_fetch_usage_effect() { let mut app = test_app_with_agent(); let effects = dispatch(Action::ShowUsage, &mut app); - // One non-silent FetchBilling — the effect pulls billing + auto-topup - // together and renders a single summary. assert_eq!(effects.len(), 1, "got: {effects:?}"); assert!( - matches!(&effects[0], Effect::FetchBilling { agent_id, silent } if *agent_id == AgentId(0) && !*silent), - "effect should be a non-silent FetchBilling, got: {effects:?}" + matches!(&effects[0], Effect::FetchUsage { agent_id } if *agent_id == AgentId(0)), + "effect should be FetchUsage for the active agent, got: {effects:?}" ); } -// ── BillingFetched dispatch tests ─────────────────────────────────── - #[test] -fn billing_fetched_updates_app_credit_balance() { - let mut app = test_app_with_agent(); - dispatch_billing(&mut app, Some(test_bal(42.0)), true, None); - assert!(app.credit_balance.is_some()); - assert_eq!(app.credit_balance.as_ref().unwrap().usage_pct, 42.0); -} - -#[test] -fn billing_fetched_updates_subscription_tier() { - let mut app = test_app_with_agent(); - dispatch_billing(&mut app, None, true, Some("supergrok_heavy".into())); - assert_eq!(app.subscription_tier.as_deref(), Some("supergrok_heavy")); -} - -#[test] -fn billing_fetched_silent_does_not_push_scrollback() { +fn usage_fetched_renders_rows_with_percent_left_and_reset_hint() { + use kigi_shell::extensions::billing::UsageRow; let mut app = test_app_with_agent(); let before = agent_scrollback_len(&app); - dispatch_billing(&mut app, Some(test_bal(50.0)), true, None); - assert_eq!( - agent_scrollback_len(&app), - before, - "silent billing fetch should not push a scrollback message" - ); -} - -#[test] -fn billing_fetched_non_silent_pushes_scrollback_message() { - let mut app = test_app_with_agent(); - let before = agent_scrollback_len(&app); - let bal = crate::views::credit_bar::CreditBalance { - pay_as_you_go: true, - on_demand_cap_cents: Some(1000), - on_demand_used_cents: Some(350), - period_end_display: Some("Jul 1, 00:00".into()), - ..test_bal(75.5) - }; - dispatch_billing(&mut app, Some(bal), false, None); - assert_eq!( - agent_scrollback_len(&app), - before + 1, - "non-silent billing fetch should push a scrollback message" - ); -} - -#[test] -fn billing_fetched_none_balance_shows_no_data_message() { - let mut app = test_app_with_agent(); - let before = agent_scrollback_len(&app); - dispatch_billing(&mut app, None, false, None); - assert_eq!(agent_scrollback_len(&app), before + 1); -} - -#[test] -fn billing_fetched_none_balance_clears_cached() { - let mut app = test_app_with_agent(); - // Seed a known balance + polling, as a prior successful fetch would. - dispatch_billing(&mut app, Some(test_bal(80.0)), true, None); - app.billing_poll_wanted = true; - // A response carrying no billing config clears the cached balance and - // polling so the status bar agrees with the "No billing data" message - // (parse/transport failures route to BillingError, not here). - dispatch_billing(&mut app, None, false, None); - assert!( - app.credit_balance.is_none(), - "None balance should clear the cached credit balance" - ); - assert!( - !app.billing_poll_wanted, - "None balance should disable billing polling" - ); -} - -#[test] -fn billing_fetched_high_usage_enables_poll() { - let mut app = test_app_with_agent(); - assert!(!app.billing_poll_wanted); - dispatch_billing(&mut app, Some(test_bal(99.5)), true, None); - assert!( - app.billing_poll_wanted, - "usage >= 99% should enable billing polling" - ); -} - -#[test] -fn billing_fetched_low_usage_disables_poll() { - let mut app = test_app_with_agent(); - app.billing_poll_wanted = true; - dispatch_billing(&mut app, Some(test_bal(50.0)), true, None); - assert!( - !app.billing_poll_wanted, - "usage < 99% should disable billing polling" - ); -} - -#[test] -fn billing_fetched_propagates_balance_to_agent() { - let mut app = test_app_with_agent(); - let bal = crate::views::credit_bar::CreditBalance { - effective_usage_pct: 60.0, - pay_as_you_go: true, - on_demand_cap_cents: Some(5000), - on_demand_used_cents: Some(1200), - period_end_display: Some("Aug 15, 00:00".into()), - ..test_bal(88.0) - }; - dispatch_billing(&mut app, Some(bal), true, None); - let agent_bal = app - .agents - .get(&AgentId(0)) - .unwrap() - .credit_balance - .as_ref() - .unwrap(); - assert_eq!(agent_bal.usage_pct, 88.0); - assert_eq!(agent_bal.effective_usage_pct, 60.0); - assert!(agent_bal.pay_as_you_go); - assert_eq!(agent_bal.on_demand_cap_cents, Some(5000)); - assert_eq!(agent_bal.on_demand_used_cents, Some(1200)); -} - -#[test] -fn billing_fetched_stores_autotopup_on_app_and_agent() { - let mut app = test_app_with_agent(); - let bal = crate::views::credit_bar::CreditBalance { - prepaid_balance_cents: Some(1500), - ..test_bal(100.0) - }; - let autotopup = crate::views::credit_bar::AutoTopupInfo { - enabled: true, - topup_amount_cents: Some(2000), - max_amount_cents: Some(10000), - }; dispatch( - Action::TaskComplete(TaskResult::BillingFetched { + Action::TaskComplete(TaskResult::UsageFetched { agent_id: AgentId(0), - balance: Some(bal), - silent: true, - subscription_tier: None, - autotopup: crate::views::credit_bar::AutoTopupFetch::Resolved(autotopup), - }), - &mut app, - ); - assert!(app.auto_topup.as_ref().is_some_and(|at| at.enabled)); - let agent_at = app.agents.get(&AgentId(0)).unwrap().auto_topup.as_ref(); - assert_eq!(agent_at.and_then(|at| at.max_amount_cents), Some(10000)); -} - -#[test] -fn billing_fetched_unchanged_autotopup_keeps_cached_rule() { - let mut app = test_app_with_agent(); - let bal = || crate::views::credit_bar::CreditBalance { - prepaid_balance_cents: Some(1500), - ..test_bal(100.0) - }; - let resolved = crate::views::credit_bar::AutoTopupFetch::Resolved( - crate::views::credit_bar::AutoTopupInfo { - enabled: true, - topup_amount_cents: Some(2000), - max_amount_cents: None, - }, - ); - dispatch( - Action::TaskComplete(TaskResult::BillingFetched { - agent_id: AgentId(0), - balance: Some(bal()), - silent: true, - subscription_tier: None, - autotopup: resolved, - }), - &mut app, - ); - // A later refresh whose auto-topup fetch failed must not clear the rule. - dispatch( - Action::TaskComplete(TaskResult::BillingFetched { - agent_id: AgentId(0), - balance: Some(bal()), - silent: true, - subscription_tier: None, - autotopup: crate::views::credit_bar::AutoTopupFetch::Unchanged, - }), - &mut app, - ); - assert!(app.auto_topup.as_ref().is_some_and(|at| at.enabled)); - let agent_at = app.agents.get(&AgentId(0)).unwrap().auto_topup.as_ref(); - assert!(agent_at.is_some_and(|at| at.enabled)); -} - -#[test] -fn billing_fetched_cleared_autotopup_resets_cache() { - let mut app = test_app_with_agent(); - // Seed a known rule while credits exist. - dispatch( - Action::TaskComplete(TaskResult::BillingFetched { - agent_id: AgentId(0), - balance: Some(crate::views::credit_bar::CreditBalance { - prepaid_balance_cents: Some(1500), - ..test_bal(100.0) - }), - silent: true, - subscription_tier: None, - autotopup: crate::views::credit_bar::AutoTopupFetch::Resolved( - crate::views::credit_bar::AutoTopupInfo { - enabled: true, - topup_amount_cents: Some(2000), - max_amount_cents: None, + result: Ok(vec![ + UsageRow { + label: "Weekly limit".into(), + used: 250, + limit: 1000, + reset_hint: Some("resets in 2d 1h".into()), }, - ), + UsageRow { + label: "5h limit".into(), + used: 20, + limit: 50, + reset_hint: None, + }, + ]), }), &mut app, ); - // Credits gone → `Cleared` resets the cached rule to "unknown" so a later - // credits period can't read a stale rule. - dispatch( - Action::TaskComplete(TaskResult::BillingFetched { - agent_id: AgentId(0), - balance: Some(test_bal(50.0)), - silent: true, - subscription_tier: None, - autotopup: crate::views::credit_bar::AutoTopupFetch::Cleared, - }), - &mut app, + assert_eq!(agent_scrollback_len(&app), before + 1); + let text = last_system_text(&app); + assert!(text.contains("API Usage"), "got: {text}"); + assert!( + text.contains("Weekly limit") && text.contains("75% left"), + "summary row shows remaining percent: {text}" + ); + assert!( + text.contains("(resets in 2d 1h)"), + "reset hint kept: {text}" + ); + assert!( + text.contains("5h limit") && text.contains("60% left"), + "limit row shows remaining percent: {text}" ); - assert!(app.auto_topup.is_none()); - assert!(app.agents.get(&AgentId(0)).unwrap().auto_topup.is_none()); } #[test] -fn app_billing_fetched_stores_autotopup() { +fn usage_fetched_zero_limit_row_renders_zero_percent_left() { + use kigi_shell::extensions::billing::UsageRow; let mut app = test_app_with_agent(); - let bal = crate::views::credit_bar::CreditBalance { - prepaid_balance_cents: Some(500), - ..test_bal(0.0) - }; dispatch( - Action::TaskComplete(TaskResult::AppBillingFetched { - balance: Some(bal), - autotopup: crate::views::credit_bar::AutoTopupFetch::Resolved( - crate::views::credit_bar::AutoTopupInfo::disabled(), - ), + Action::TaskComplete(TaskResult::UsageFetched { + agent_id: AgentId(0), + result: Ok(vec![UsageRow { + label: "RPM".into(), + used: 12, + limit: 0, + reset_hint: None, + }]), }), &mut app, ); - assert_eq!( - app.credit_balance.and_then(|b| b.prepaid_balance_cents), - Some(500) - ); - assert!(app.auto_topup.is_some_and(|at| !at.enabled)); + let text = last_system_text(&app); + assert!(text.contains("0% left"), "no invented percentage: {text}"); } -// ── BillingError dispatch tests ───────────────────────────────────── - #[test] -fn billing_error_silent_does_not_push_scrollback() { +fn usage_fetched_empty_rows_shows_no_data_message() { let mut app = test_app_with_agent(); let before = agent_scrollback_len(&app); dispatch( - Action::TaskComplete(TaskResult::BillingError { + Action::TaskComplete(TaskResult::UsageFetched { agent_id: AgentId(0), - error: "network timeout".into(), - silent: true, + result: Ok(vec![]), }), &mut app, ); - assert_eq!( - agent_scrollback_len(&app), - before, - "silent billing error should not push a scrollback message" - ); + assert_eq!(agent_scrollback_len(&app), before + 1); + assert_eq!(last_system_text(&app), "No usage data available."); } #[test] -fn billing_error_non_silent_pushes_error_message() { +fn usage_fetched_error_pushes_error_message() { let mut app = test_app_with_agent(); let before = agent_scrollback_len(&app); dispatch( - Action::TaskComplete(TaskResult::BillingError { + Action::TaskComplete(TaskResult::UsageFetched { agent_id: AgentId(0), - error: "service unavailable".into(), - silent: false, + result: Err("Usage endpoint not available. Try Kimi for Coding.".into()), }), &mut app, ); + assert_eq!(agent_scrollback_len(&app), before + 1); assert_eq!( - agent_scrollback_len(&app), - before + 1, - "non-silent billing error should push an error message" - ); -} - -// ── Free-usage paywall tests ──────────────────────────────────────── - -#[test] -fn free_usage_error_detected_by_embedded_code() { - // parse_error_bytes flattens the 429 body to ": ". - assert!(is_free_usage_exhausted_error( - "API error (status 429 Too Many Requests): \ - subscription:free-usage-exhausted: You have used all your free usage." - )); - // Generic rate limits and other WKE codes must not match. - assert!(!is_free_usage_exhausted_error( - "API error (status 429 Too Many Requests): Rate limit exceeded" - )); - assert!(!is_free_usage_exhausted_error( - "unauthorized:missing-acl: nope" - )); -} - -#[test] -fn free_usage_upsell_shows_two_options_with_exact_labels() { - let mut app = test_app_with_agent(); - let agent = app.agents.get_mut(&AgentId(0)).unwrap(); - open_free_usage_upsell(agent); - - let qv = agent_qv(&app); - assert!(matches!( - qv.local_kind, - Some(crate::views::question_view::LocalQuestionKind::FreeUsageUpsell) - )); - let q = &qv.questions[0]; - assert_eq!(q.question, "You hit your free usage limit."); - let expected = [ - ( - "Upgrade to SuperGrok", - "For everyday coding and productivity tasks", - Some(UPSELL_URL_UPGRADE), - ), - ( - "Upgrade to SuperGrok Heavy", - "Get the most out of Grok Build. Highest usage limits.", - Some(UPSELL_URL_UPGRADE), - ), - ]; - assert_eq!(q.options.len(), expected.len()); - for (opt, (label, desc, id)) in q.options.iter().zip(expected) { - assert_eq!(opt.label, label); - assert_eq!(opt.description, desc); - assert_eq!(opt.id.as_deref(), id); - } -} - -/// Replay the REAL free-usage sequence — send → `RetryState::Retrying` → -/// `Exhausted` → PromptResponse error — through the production handlers: -/// the paywall modal must open on the turn-end error. -#[test] -fn free_usage_failure_opens_paywall_modal() { - use crate::app::acp_handler::apply_session_event_for_test; - use kigi_shell::extensions::notification::{RetryState, SessionUpdate}; - - let mut app = test_app_with_agent(); - let id = AgentId(0); - - // 1. Real send. - let effects = dispatch(Action::SendPrompt("draw me a cat".into()), &mut app); - assert!( - matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "draw me a cat"), - "send must dispatch: {effects:?}" - ); - let prompt_id = app.agents[&id].session.current_prompt_id.clone(); - assert!(prompt_id.is_some(), "send must mint a prompt id"); - - // 2+3. Real notification sequence through the production handler. - { - let agent = app.agents.get_mut(&id).unwrap(); - apply_session_event_for_test( - &SessionUpdate::RetryState(RetryState::Retrying { - attempt: 1, - max_retries: 2, - reason: "429 Too Many Requests".into(), - }), - &mut agent.session, - &mut agent.scrollback, - ); - apply_session_event_for_test( - &SessionUpdate::RetryState(RetryState::Exhausted { - attempts: 2, - reason: "API error (status 429 Too Many Requests): \ - subscription:free-usage-exhausted: You have used all your free usage." - .into(), - is_rate_limited: true, - }), - &mut agent.session, - &mut agent.scrollback, - ); - assert!(agent.session.free_usage_blocked); - } - - // 4. Turn-end RPC error opens the upsell modal. - let _ = dispatch( - Action::TaskComplete(TaskResult::PromptResponse { - agent_id: id, - result: Err("rate limited".into()), - http_status: Some(429), - prompt_id, - }), - &mut app, - ); - assert!( - app.agents[&id].question_view.is_some(), - "paywall modal must open" - ); -} - -/// Answer translation: both upgrade options open their URL. -#[test] -fn free_usage_translate_local_submit_maps_options() { - use crate::app::agent_view::translate_local_submit_for_test; - use crate::app::app_view::InputOutcome; - use crate::views::question_view::{LocalQuestionKind, QuestionSelection}; - - let mut app = test_app_with_agent(); - let agent = app.agents.get_mut(&AgentId(0)).unwrap(); - open_free_usage_upsell(agent); - let mut qv = agent.question_view.take().unwrap(); - let kind = || LocalQuestionKind::FreeUsageUpsell; - - for idx in [0, 1] { - qv.selections[0] = QuestionSelection::Single(Some(idx)); - match translate_local_submit_for_test(&qv, kind(), false) { - InputOutcome::Action(Action::OpenUrl(url)) => assert_eq!(url, UPSELL_URL_UPGRADE), - other => panic!("expected OpenUrl for option {idx}, got {other:?}"), - } - } -} - -// ── Restricted-command upsell tests ───────────────────────────────── - -/// Submitting a tier-restricted command opens the two-option SuperGrok -/// upsell and neither runs the command nor leaks the text to the model. -#[test] -fn restricted_command_submit_opens_two_option_upsell() { - let mut app = test_app_with_agent(); - let id = AgentId(0); - app.agents - .get_mut(&id) - .unwrap() - .set_restricted_commands(&["imagine".to_string()]); - - let effects = dispatch(Action::SendPrompt("/imagine a sunset".into()), &mut app); - - assert!( - effects.is_empty(), - "restricted command must not produce a SendPrompt: {effects:?}" - ); - let agent = &app.agents[&id]; - assert!( - agent.session.pending_prompts.is_empty(), - "restricted command must not be enqueued" - ); - assert!(agent.prompt.text().is_empty(), "composer consumed"); - - let qv = agent_qv(&app); - assert!(matches!( - qv.local_kind, - Some(crate::views::question_view::LocalQuestionKind::FreeUsageUpsell) - )); - let q = &qv.questions[0]; - assert_eq!(q.question, "Unlock all features with SuperGrok."); - assert_eq!(q.options.len(), 2); - assert_eq!(q.options[0].label, "Upgrade to SuperGrok"); - assert_eq!(q.options[0].id.as_deref(), Some(UPSELL_URL_UPGRADE)); - assert_eq!(q.options[1].label, "Upgrade to SuperGrok Heavy"); - assert_eq!(q.options[1].id.as_deref(), Some(UPSELL_URL_UPGRADE)); -} - -/// Aliases of a restricted command hit the same upsell (deny-list -/// matching covers aliases via the registry). -#[test] -fn restricted_command_alias_also_upsells() { - let mut app = test_app_with_agent(); - let id = AgentId(0); - app.agents - .get_mut(&id) - .unwrap() - .set_restricted_commands(&["usage".to_string()]); - - let effects = dispatch(Action::SendPrompt("/cost".into()), &mut app); - - assert!(effects.is_empty()); - assert!(app.agents[&id].question_view.is_some(), "upsell must open"); -} - -/// A restricted submit while ANOTHER question modal is already open -/// must not silently drop the typed text — the upsell can't open (the guard -/// never displaces a modal), so the composer keeps the text for a resubmit -/// after the modal closes. No passthrough, nothing enqueued, and the -/// existing modal survives untouched. -#[test] -fn restricted_command_with_open_modal_keeps_composer_text() { - let mut app = test_app_with_agent(); - let id = AgentId(0); - { - let agent = app.agents.get_mut(&id).unwrap(); - agent.set_restricted_commands(&["imagine".to_string()]); - // A question modal is already up (credit-limit upsell). - open_credit_limit_upsell(agent, CreditLimitUpsellMode::UnifiedCredits, false); - assert!(agent.question_view.is_some()); - // The user typed the restricted command into the composer. - agent.prompt.set_text("/imagine a sunset"); - } - - let effects = dispatch(Action::SendPrompt("/imagine a sunset".into()), &mut app); - - assert!(effects.is_empty(), "no passthrough / send: {effects:?}"); - let agent = &app.agents[&id]; - assert_eq!( - agent.prompt.text(), - "/imagine a sunset", - "composer text must be preserved for a later resubmit" - ); - assert!( - matches!( - agent - .question_view - .as_ref() - .and_then(|qv| qv.local_kind.as_ref()), - Some(crate::views::question_view::LocalQuestionKind::CreditLimitUpsell) - ), - "the pre-existing modal must survive (no second modal)" - ); - assert!( - agent.session.pending_prompts.is_empty(), - "nothing may be enqueued" + last_system_text(&app), + "Couldn't fetch usage: Usage endpoint not available. Try Kimi for Coding." ); } diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/dashboard.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/dashboard.rs index 837437e..3ac6e3e 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/dashboard.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/dashboard.rs @@ -1270,37 +1270,6 @@ fn dashboard_slash_model_stages_pending_model() { ); } -/// A tier-restricted command typed into the dashboard dispatch input must -/// upsell via the feedback toast — not execute, and (crucially) not fall -/// through the unknown-command path, which would spawn a session whose -/// first prompt is the raw slash text. -#[serial_test::serial(KIGI_AGENT_DASHBOARD)] -#[test] -fn dashboard_slash_restricted_command_upsells_via_toast() { - let mut app = test_app(); - app.tier_restricted_commands = vec!["imagine".to_string()]; - open_dashboard(&mut app); - - let effects = dispatch_dashboard_dispatch_slash(&mut app, "/imagine a sunset".into()); - - assert!(effects.is_empty(), "restricted command must not dispatch"); - assert!( - app.agents.is_empty(), - "no session may be spawned for the raw slash text" - ); - let toast = app - .dashboard - .as_ref() - .unwrap() - .error_toast - .as_deref() - .expect("restricted command must set the upsell toast"); - assert!( - toast.contains("/imagine") && toast.contains("SuperGrok"), - "toast must carry the upsell: {toast}" - ); -} - /// A slash command that fails (`CommandResult::Error`) surfaces on /// the dashboard with the `✗` error prefix — command error strings /// carry no glyph of their own, and the feedback badge paints the diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/mod.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/mod.rs index 358a9ee..e63fa7a 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/mod.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/mod.rs @@ -15,10 +15,6 @@ mod status; mod task_result; mod transcript; mod turn; -use super::billing::{ - CreditLimitUpsellMode, credit_limit_upsell_mode, is_max_tier, open_credit_limit_upsell, - open_free_usage_upsell, -}; use super::ctx::{find_agent_by_session_id, get_active_agent, get_active_agent_mut}; use super::dashboard::{ apply_pending_dispatch_config, dispatch_dashboard_attach, dispatch_dashboard_begin_rename, @@ -133,24 +129,9 @@ fn test_app() -> AppView { deferred_startup: Default::default(), auth_use_oauth: false, auth_clipboard_copied: false, - team_id: None, - team_name: None, - is_zdr: false, - team_role: None, - coding_data_retention_opt_out: false, show_tips: None, auto_update: None, ask_user_question_timeout_enabled: None, - zdr_access_enabled: false, - usage_billing_redirect_url: None, - access_gate_shown_logged: false, - gate: None, - subscription_tier: None, - paywall_check_started: None, - last_subscription_check_at: None, - subscription_watch_interval_secs: None, - pending_gate_verification: None, - gate_verify_gen: 0, bundle_state: crate::app::bundle::BundleState::default(), scroll_debug_hud: crate::views::scroll_debug_hud::ScrollDebugHud::new(), fps_hud: crate::views::fps_hud::FpsHud::new(), @@ -172,8 +153,6 @@ fn test_app() -> AppView { welcome_on_auth_url: false, welcome_on_changelog_cta: false, welcome_auth_fallback_rect: None, - welcome_refresh_rect: None, - welcome_gate_url_rect: None, welcome_changelog_cta_rect: None, auth_show_raw_url: false, auth_mouse_disabled: false, @@ -213,13 +192,8 @@ fn test_app() -> AppView { minimal_state: crate::minimal_api::MinimalState::default(), reconnect_pending: false, show_resolved_model: true, - sharing_enabled: false, usage_visible: true, - tier_restricted_commands: Vec::new(), leader_mode: true, - credit_balance: None, - auto_topup: None, - billing_poll_wanted: false, leader_roster: Vec::new(), dashboard_local_sessions: Vec::new(), dashboard_sessions_loading: false, @@ -262,8 +236,6 @@ fn make_test_agent_session(app: &AppView, id: AgentId, sid: &str) -> AgentSessio restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, @@ -446,8 +418,6 @@ fn insert_placeholder_agent(app: &mut AppView, id: AgentId) { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, @@ -591,8 +561,6 @@ fn two_agent_app_with_bg_task() -> AppView { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, @@ -895,17 +863,3 @@ fn reset_mouse_capture_enabled(on: bool) { fn mouse_capture_is_enabled() -> bool { crate::app::MOUSE_CAPTURE_ENABLED.load(std::sync::atomic::Ordering::Acquire) } -/// Build a minimal `CreditBalance` for billing dispatch tests. -fn test_bal(usage_pct: f64) -> crate::views::credit_bar::CreditBalance { - crate::views::credit_bar::CreditBalance { - usage_pct, - effective_usage_pct: usage_pct, - period_end_display: None, - pay_as_you_go: false, - on_demand_cap_cents: None, - on_demand_used_cents: None, - prepaid_balance_cents: None, - period_type: None, - is_unified_billing_user: None, - } -} diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/prompt.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/prompt.rs index 6a467ca..5d78bca 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/prompt.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/prompt.rs @@ -733,13 +733,8 @@ fn turn_end_drains_next_queued_prompt() { &mut app, ); - // No re-send (the prompt was already sent at enqueue time): only the - // billing refresh effect. - assert_eq!(effects.len(), 1); - assert!(matches!( - &effects[0], - Effect::FetchBilling { silent: true, .. } - )); + // No re-send (the prompt was already sent at enqueue time). + assert!(effects.is_empty(), "no effects expected: {effects:?}"); assert!(app.agents[&id].session.state.is_turn_running()); // current_prompt_id was handed off to the second prompt for correlation. assert_eq!( @@ -772,12 +767,8 @@ fn turn_end_with_empty_queue_stays_idle() { &mut app, ); - // Silent billing refresh after turn completion. - assert_eq!(effects.len(), 1); - assert!(matches!( - &effects[0], - Effect::FetchBilling { silent: true, .. } - )); + // Turn completion produces no follow-up effects. + assert!(effects.is_empty(), "no effects expected: {effects:?}"); assert!(app.agents[&id].session.state.is_idle()); // Session event "Worked for" added. assert_eq!(app.agents[&id].scrollback.len(), 1); @@ -806,31 +797,19 @@ fn multiple_queued_prompts_drain_one_per_turn() { }) }; - // Turn end → drain "b" + FetchBilling. + // Turn end → drain "b". let effects = dispatch(end_turn(), &mut app); assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "b")); - assert!(matches!( - &effects[1], - Effect::FetchBilling { silent: true, .. } - )); assert_eq!(app.agents[&id].session.queue_len(), 1); - // Turn end → drain "c" + FetchBilling. + // Turn end → drain "c". let effects = dispatch(end_turn(), &mut app); assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "c")); - assert!(matches!( - &effects[1], - Effect::FetchBilling { silent: true, .. } - )); assert_eq!(app.agents[&id].session.queue_len(), 0); - // Turn end → FetchBilling only. + // Turn end with an empty queue → nothing to drain. let effects = dispatch(end_turn(), &mut app); - assert_eq!(effects.len(), 1); - assert!(matches!( - &effects[0], - Effect::FetchBilling { silent: true, .. } - )); + assert!(effects.is_empty(), "no effects expected: {effects:?}"); assert!(app.agents[&id].session.state.is_idle()); } @@ -852,12 +831,8 @@ fn prompt_response_resets_turn_state() { }), &mut app, ); - // Silent billing refresh after turn completion. - assert_eq!(effects.len(), 1); - assert!(matches!( - &effects[0], - Effect::FetchBilling { silent: true, .. } - )); + // Turn completion produces no follow-up effects. + assert!(effects.is_empty(), "no effects expected: {effects:?}"); assert!(app.agents[&id].session.state.is_idle()); assert!(app.agents[&id].turn_started_at.is_none()); // mark_turn_finished must stamp the activity anchor used by the @@ -892,7 +867,7 @@ fn turn_end_fetches_prompt_suggestion_when_enabled() { &mut app, ); - assert_eq!(effects.len(), 2, "suggestion fetch + billing: {effects:?}"); + assert_eq!(effects.len(), 1, "suggestion fetch: {effects:?}"); let Effect::FetchPromptSuggestion { agent_id, generation, @@ -1249,12 +1224,8 @@ fn turn_complete_notification_suppressed_when_queue_non_empty() { }), &mut app, ); - // No re-send; only billing refresh. The second prompt is adopted. - assert_eq!(effects.len(), 1); - assert!(matches!( - &effects[0], - Effect::FetchBilling { silent: true, .. } - )); + // No re-send. The second prompt is adopted. + assert!(effects.is_empty(), "no effects expected: {effects:?}"); assert!(app.agents[&id].session.state.is_turn_running()); assert!( app.deferred_notification.is_none(), @@ -1536,12 +1507,8 @@ fn prompt_response_resets_cancelling_to_idle() { }), &mut app, ); - // Silent billing refresh after turn completion. - assert_eq!(effects.len(), 1); - assert!(matches!( - &effects[0], - Effect::FetchBilling { silent: true, .. } - )); + // Turn completion produces no follow-up effects. + assert!(effects.is_empty(), "no effects expected: {effects:?}"); assert!(app.agents[&id].session.state.is_idle()); // Cancellation produces a "Turn cancelled" session event. assert_eq!(app.agents[&id].scrollback.len(), 1); @@ -1578,12 +1545,8 @@ fn cancel_with_queued_prompt_drains_on_completion() { &mut app, ); - assert_eq!(effects.len(), 2); + assert_eq!(effects.len(), 1); assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "queued")); - assert!(matches!( - &effects[1], - Effect::FetchBilling { silent: true, .. } - )); assert!(app.agents[&id].session.state.is_turn_running()); assert_eq!(app.agents[&id].session.queue_len(), 0); } @@ -1605,12 +1568,8 @@ fn cancel_with_empty_queue_stays_idle() { }), &mut app, ); - // Silent billing refresh after turn completion. - assert_eq!(effects.len(), 1); - assert!(matches!( - &effects[0], - Effect::FetchBilling { silent: true, .. } - )); + // Turn completion produces no follow-up effects. + assert!(effects.is_empty(), "no effects expected: {effects:?}"); assert!(app.agents[&id].session.state.is_idle()); } @@ -1660,12 +1619,8 @@ fn cancel_with_multiple_queued_prompts_drains_only_front_prompt() { &mut app, ); - assert_eq!(effects.len(), 2); + assert_eq!(effects.len(), 1); assert!(matches!(&effects[0], Effect::SendPrompt { text, .. } if text == "queued-1")); - assert!(matches!( - &effects[1], - Effect::FetchBilling { silent: true, .. } - )); assert!(app.agents[&id].session.state.is_turn_running()); assert_eq!(app.agents[&id].session.queue_len(), 1); assert_eq!(app.agents[&id].session.pending_prompts[0].text, "queued-2"); @@ -1706,12 +1661,8 @@ fn cancel_drain_is_blocked_when_editing_front_prompt() { &mut app, ); - // Drain blocked but billing refresh still happens. - assert_eq!(effects.len(), 1); - assert!(matches!( - &effects[0], - Effect::FetchBilling { silent: true, .. } - )); + // Drain blocked — no effects. + assert!(effects.is_empty(), "drain should be blocked: {effects:?}"); assert!(app.agents[&id].session.state.is_idle()); assert_eq!(app.agents[&id].session.queue_len(), 2); assert_eq!(app.agents[&id].session.pending_prompts[0].text, "queued-1"); diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/session/foreign.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/session/foreign.rs index cb74c86..1647fb2 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/session/foreign.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/session/foreign.rs @@ -220,7 +220,6 @@ fn native_empty_waits_for_foreign_and_foreign_only_rows_survive() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![], - partial: None, seq: 1, query: None, }), @@ -267,7 +266,6 @@ fn foreign_empty_then_native_empty_finishes_once_without_resurrecting() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![], - partial: None, seq: 3, query: None, }), @@ -331,7 +329,6 @@ fn modal_empty_notice_waits_until_both_lanes_are_empty() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![], - partial: None, seq: 9, query: None, }), @@ -435,7 +432,6 @@ fn modal_selection_survives_native_and_foreign_completion_races() { at(make_picker_entry("a", "/repo"), 20), at(make_picker_entry("b", "/repo"), 10), ], - partial: None, seq: 2, query: None, }), diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/session/fork.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/session/fork.rs index d946ea9..6031683 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/session/fork.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/session/fork.rs @@ -723,27 +723,14 @@ fn dispatch_fork_stashes_directive_in_pending_first_prompt() { } #[test] -fn dispatch_fork_inherits_appearance_sharing_and_plugin_visibility() { +fn dispatch_fork_inherits_appearance_and_plugin_visibility() { let mut app = fork_test_app(); // Tweak app-level state so we can verify the sweep applied it. app.appearance.prompt.compact = true; - app.sharing_enabled = false; app.usage_visible = false; app.appearance.disable_plugins = true; - // Cached billing state must be inherited so the credits warning is - // correct from the first frame (not just after a billing fetch). - app.credit_balance = Some(crate::views::credit_bar::CreditBalance { - prepaid_balance_cents: Some(1500), - ..test_bal(50.0) - }); - app.auto_topup = Some(crate::views::credit_bar::AutoTopupInfo { - enabled: true, - topup_amount_cents: Some(2000), - max_amount_cents: None, - }); dispatch(Action::Fork(fork_args(Some(false), None)), &mut app); let new_agent = app.agents.get(&AgentId(1)).unwrap(); - assert!(!new_agent.sharing_enabled); assert!( new_agent .prompt @@ -752,14 +739,6 @@ fn dispatch_fork_inherits_appearance_sharing_and_plugin_visibility() { .get("usage") .is_none() ); - assert_eq!( - new_agent - .credit_balance - .as_ref() - .and_then(|b| b.prepaid_balance_cents), - Some(1500) - ); - assert!(new_agent.auto_topup.as_ref().is_some_and(|at| at.enabled)); } #[test] diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/session/lifecycle.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/session/lifecycle.rs index 825f219..1b7419f 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/session/lifecycle.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/session/lifecycle.rs @@ -50,7 +50,7 @@ fn session_created_sets_session_id() { }), &mut app, ); - assert_eq!(effects.len(), 5); + assert_eq!(effects.len(), 4); assert!( matches!(& effects[0], Effect::FetchPromptHistory { session_id, .. } if session_id == "new-session-123") @@ -60,11 +60,7 @@ fn session_created_sets_session_id() { &effects[2], Effect::RefreshAvailableCommands { .. } )); - assert!(matches!( - &effects[3], - Effect::FetchBilling { silent: true, .. } - )); - assert!(matches!(&effects[4], Effect::RegisterActiveSession { .. })); + assert!(matches!(&effects[3], Effect::RegisterActiveSession { .. })); assert_eq!( app.agents[&id] .session @@ -198,11 +194,6 @@ fn worktree_session_created_sets_session_and_cwd() { .iter() .any(|e| matches!(e, Effect::FetchSessionAgentName { .. })) ); - assert!( - effects - .iter() - .any(|e| matches!(e, Effect::FetchBilling { silent: true, .. })) - ); assert!( effects .iter() diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/session/load.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/session/load.rs index d05181d..cf4bd93 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/session/load.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/session/load.rs @@ -1836,7 +1836,6 @@ fn stale_session_list_responses_are_dropped() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![make_conversation_entry("conv-stale-1")], - partial: None, seq: 1, query: None, }), @@ -1861,7 +1860,6 @@ fn stale_session_list_responses_are_dropped() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![make_conversation_entry("conv-fresh-2")], - partial: None, seq: 2, query: Some("abcd".into()), }), @@ -1906,7 +1904,6 @@ fn modal_search_response_lands_and_stale_is_dropped() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![make_conversation_entry("conv-hit-1")], - partial: None, seq: 1, query: Some("hit".into()), }), @@ -1944,7 +1941,6 @@ fn modal_search_response_lands_and_stale_is_dropped() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![make_conversation_entry("conv-stale-m")], - partial: None, seq: 1, query: Some("hit".into()), }), @@ -1999,7 +1995,6 @@ fn modal_close_drops_in_flight_search_response() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![make_conversation_entry("conv-late-1")], - partial: None, seq, query: Some("hit".into()), }), @@ -2045,7 +2040,6 @@ fn modal_pick_drops_in_flight_search_response() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![make_conversation_entry("conv-late-p")], - partial: None, seq, query: Some("hit".into()), }), @@ -2089,7 +2083,6 @@ fn welcome_esc_drops_in_flight_fetch_response() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![make_conversation_entry("conv-late-w")], - partial: None, seq, query: None, }), @@ -2119,7 +2112,6 @@ fn build_mode_modal_close_does_not_invalidate_plain_fetch() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![make_picker_entry("build-late-1", "/tmp/repo")], - partial: None, seq, query: None, }), @@ -2142,7 +2134,6 @@ fn zero_hit_search_shows_empty_list_without_toast() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![], - partial: None, seq: 1, query: Some("zzz".into()), }), @@ -2163,7 +2154,6 @@ fn zero_hit_search_shows_empty_list_without_toast() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![], - partial: None, seq: 2, query: None, }), @@ -2347,7 +2337,6 @@ fn build_mode_list_response_preserves_deep_search_spinner() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![make_picker_entry("local-1", "/r")], - partial: None, seq: app.session_picker_list_seq, query: None, }), @@ -2400,7 +2389,6 @@ fn build_mode_rapid_plain_fetches_keep_last_write_wins() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![make_picker_entry("build-first", "/r")], - partial: None, seq: 0, query: None, }), @@ -2416,7 +2404,6 @@ fn build_mode_rapid_plain_fetches_keep_last_write_wins() { let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![make_picker_entry("build-second", "/r")], - partial: None, seq: 0, query: None, }), diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/settings.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/settings.rs index c5eb077..a0ee8d7 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/settings.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/settings.rs @@ -1209,9 +1209,6 @@ fn move_setting_away_from_default(app: &mut AppView, key: crate::settings::Setti "max_thoughts_width" => { let _ = dispatch(Action::SetMaxThoughtsWidth(200), app); } - "coding_data_sharing" => { - let _ = dispatch(Action::SetCodingDataSharing { opted_in: false }, app); - } "plan_mode" => { let _ = dispatch( Action::SetPlanMode(crate::app::actions::PlanModeKind::On), @@ -1388,8 +1385,6 @@ fn set_simple_mode_propagates_to_every_agent() { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/status.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/status.rs index 09995aa..34f1398 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/status.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/status.rs @@ -1,4 +1,4 @@ -//! Tests for session status, sharing, privacy, and coding-data-sharing dispatchers. +//! Tests for session status and sharing dispatchers. use super::*; @@ -60,659 +60,6 @@ fn send_while_idle_with_nonempty_shared_queue_routes_to_server() { assert_eq!(q.last().map(|e| e.text.as_str()), Some("c")); } -#[test] -fn show_privacy_info_zdr() { - let mut app = test_app_with_agent(); - app.is_zdr = true; - let effects = dispatch(Action::ShowPrivacyInfo, &mut app); - assert!(effects.is_empty()); - let text = last_system_text(&app, AgentId(0)); - assert!(text.contains("Zero Data Retention")); -} - -/// `/privacy` info-print uses the desktop-aligned "privacy mode" / -/// "share data" labels from the user's intentional rewrite. -#[test] -fn show_privacy_info_opted_out() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = true; - let effects = dispatch(Action::ShowPrivacyInfo, &mut app); - assert!(effects.is_empty()); - let text = last_system_text(&app, AgentId(0)); - assert!( - text.contains("Privacy: privacy mode"), - "info-print must use 'Privacy: privacy mode' (desktop-aligned label): {text}", - ); - assert!(text.contains("/privacy opt-in")); -} - -#[test] -fn show_privacy_info_opted_in() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = false; - let effects = dispatch(Action::ShowPrivacyInfo, &mut app); - assert!(effects.is_empty()); - let text = last_system_text(&app, AgentId(0)); - assert!( - text.contains("Privacy: share data"), - "info-print must use 'Privacy: share data' (desktop-aligned label): {text}", - ); - assert!(text.contains("/privacy opt-out")); -} - -/// The info-print uses desktop-aligned labels ("privacy mode" / -/// "share data"). This test pins those labels to catch accidental -/// regressions to the registry's "Opt in" / "Opt out" display -/// strings. -#[test] -fn show_privacy_info_does_not_use_old_desktop_labels() { - // opted-out → "Privacy: privacy mode" - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = true; - let _ = dispatch(Action::ShowPrivacyInfo, &mut app); - let text = last_system_text(&app, AgentId(0)); - assert!( - text.contains("privacy mode"), - "[opted-out] info-print must contain 'privacy mode': {text:?}", - ); - - // opted-in → "Privacy: share data" - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = false; - let _ = dispatch(Action::ShowPrivacyInfo, &mut app); - let text = last_system_text(&app, AgentId(0)); - assert!( - text.contains("share data"), - "[opted-in] info-print must contain 'share data': {text:?}", - ); -} - -// ── coding_data_sharing dispatch tests ─── -// -// The dispatcher uses **optimistic + rollback + toast**, matching the -// `set_yolo_mode` pattern. These tests pin the contract: -// - Guards (ZDR, non-admin team) toast and short-circuit. -// - Idempotent dispatch toasts but emits no Effect. -// - Optimistic mutation flips `app.coding_data_retention_opt_out` -// BEFORE the Effect is emitted. -// - `Effect::SetCodingDataSharing` carries -// `rollback_to_opted_in = previous_value`. -// - `TaskResult::CodingDataSharingFailed` reverts the optimistic -// mutation; `TaskResult::CodingDataSharingUpdated` re-anchors -// to the server-confirmed value. - -/// Idempotent re-dispatch when already opted-in toasts but emits -/// no Effect (avoids a wasted ACP round-trip). -/// -/// Toast uses the **display name** ("Opt in", not the -/// snake-case canonical "opt-in") AND the **destructive `⚠` -/// glyph** on the opt-in direction (privacy-degrading). -#[test] -fn set_coding_data_sharing_idempotent_opt_in() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = false; // currently opted-in - let effects = dispatch(Action::SetCodingDataSharing { opted_in: true }, &mut app); - assert!( - effects.is_empty(), - "idempotent re-dispatch must NOT emit Effect" - ); - let toast = read_toast(&app); - assert!( - toast.contains("Opt in"), - "toast must show display name 'Opt in' (PR 9 R1, General-3 Issue 6): {toast}", - ); - assert!( - !toast.contains("opt-in"), - "toast must NOT use snake-case canonical 'opt-in' — display name only: {toast}", - ); - assert!( - toast.contains('\u{26A0}'), - "idempotent opt-in toast uses ⚠ destructive-warning glyph (PR 9 R1, \ - General-3 Issue 5): {toast}", - ); - // State unchanged. - assert!( - !app.coding_data_retention_opt_out, - "idempotent path must not mutate state", - ); -} - -/// Idempotent re-dispatch when already opted-out toasts but emits -/// no Effect. -/// -/// Opt-out direction uses the **uniform `✓` glyph** -/// (restoring the safe default) and the display name "Opt out". -#[test] -fn set_coding_data_sharing_idempotent_opt_out() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = true; // currently opted-out - let effects = dispatch(Action::SetCodingDataSharing { opted_in: false }, &mut app); - assert!( - effects.is_empty(), - "idempotent re-dispatch must NOT emit Effect" - ); - let toast = read_toast(&app); - assert!( - toast.contains("Opt out"), - "toast must show display name 'Opt out': {toast}", - ); - assert!( - toast.contains('\u{2713}'), - "idempotent opt-out toast uses ✓ safe-default glyph: {toast}", - ); - assert!( - !toast.contains('\u{26A0}'), - "opt-out is the safe direction — must NOT use ⚠: {toast}", - ); - // State unchanged. - assert!( - app.coding_data_retention_opt_out, - "idempotent path must not mutate state", - ); -} - -/// ZDR teams are blocked from toggling. The blocked path -/// toasts (not scrollback) and short-circuits with no Effect. -#[test] -fn set_coding_data_sharing_blocked_by_zdr() { - let mut app = test_app_with_agent(); - app.is_zdr = true; - app.coding_data_retention_opt_out = false; - let effects = dispatch(Action::SetCodingDataSharing { opted_in: false }, &mut app); - assert!(effects.is_empty(), "ZDR block must NOT emit Effect"); - let toast = read_toast(&app); - assert!( - toast.contains("Zero Data Retention"), - "ZDR toast must surface the policy: {toast}", - ); - assert!( - toast.contains('\u{2717}'), - "blocked toast uses ✗ glyph: {toast}" - ); - // State unchanged — the user was blocked, the optimistic - // mutation never happened. - assert!( - !app.coding_data_retention_opt_out, - "ZDR block must not mutate state", - ); -} - -/// ZDR block fires even when the toggle would be a no-op -/// (defense-in-depth: don't quietly accept a same-value toggle -/// from a user the policy says shouldn't be touching this). -#[test] -fn set_coding_data_sharing_blocked_by_zdr_even_if_idempotent() { - let mut app = test_app_with_agent(); - app.is_zdr = true; - app.coding_data_retention_opt_out = false; - let effects = dispatch(Action::SetCodingDataSharing { opted_in: true }, &mut app); - assert!(effects.is_empty()); - assert!(read_toast(&app).contains("Zero Data Retention")); -} - -/// Non-admin team members are blocked from toggling (matches -/// desktop). The blocked path toasts and short-circuits. -#[test] -fn set_coding_data_sharing_blocked_non_admin() { - let mut app = test_app_with_agent(); - app.team_name = Some("Acme".into()); - app.team_role = Some("Member".into()); - app.coding_data_retention_opt_out = false; - let effects = dispatch(Action::SetCodingDataSharing { opted_in: false }, &mut app); - assert!(effects.is_empty()); - let toast = read_toast(&app); - assert!( - toast.contains("team admin"), - "non-admin toast must mention team admin: {toast}", - ); -} - -/// Admin team members CAN toggle. The admin-allowed path produces -/// an Effect carrying the rollback value. -#[test] -fn set_coding_data_sharing_allowed_for_admin() { - let mut app = test_app_with_agent(); - app.team_name = Some("Acme".into()); - app.team_role = Some("Admin".into()); - app.coding_data_retention_opt_out = false; // currently opted-in - let effects = dispatch(Action::SetCodingDataSharing { opted_in: false }, &mut app); - assert_eq!(effects.len(), 1); - match &effects[0] { - Effect::SetCodingDataSharing { - opted_in, - rollback_to_opted_in, - .. - } => { - assert!(!*opted_in, "Effect must carry opted_in=false"); - assert!( - *rollback_to_opted_in, - "rollback_to_opted_in must capture pre-toggle opt-in=true", - ); - } - other => panic!("expected SetCodingDataSharing Effect, got {other:?}"), - } - // Optimistic mutation already applied. - assert!( - app.coding_data_retention_opt_out, - "admin-allowed dispatch must optimistically flip state", - ); -} - -/// Non-idempotent dispatch emits one Effect AND mutates state -/// optimistically AND toasts. -#[test] -fn set_coding_data_sharing_produces_effect_and_optimistic_mutation() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = false; // currently opted-in - let effects = dispatch(Action::SetCodingDataSharing { opted_in: false }, &mut app); - assert_eq!(effects.len(), 1, "non-idempotent dispatch emits one Effect"); - match &effects[0] { - Effect::SetCodingDataSharing { - agent_id, - opted_in, - rollback_to_opted_in, - } => { - assert_eq!(*agent_id, AgentId(0)); - assert!(!*opted_in); - assert!( - *rollback_to_opted_in, - "rollback_to_opted_in must be pre-toggle value (true == opted-in)", - ); - } - other => panic!("expected SetCodingDataSharing Effect, got {other:?}"), - } - // Optimistic mutation applied. - assert!( - app.coding_data_retention_opt_out, - "dispatch must optimistically mutate state", - ); - // Toast on every dispatch (SHELL setter contract). - assert!(app.agents[&AgentId(0)].toast.is_some()); -} - -/// `TaskResult::CodingDataSharingUpdated` re-anchors state to the -/// server-confirmed value (defense-in-depth) and re-toasts. -#[test] -fn coding_data_sharing_updated_re_anchors_state_and_re_toasts() { - let mut app = test_app_with_agent(); - // Simulate post-optimistic state: opted-out. - app.coding_data_retention_opt_out = true; - let id = AgentId(0); - // Server confirms opt-out (same as optimistic). - let effects = dispatch( - Action::TaskComplete(TaskResult::CodingDataSharingUpdated { - agent_id: id, - opted_in: false, - }), - &mut app, - ); - assert!(effects.is_empty(), "TaskResult arm must NOT emit Effect"); - // State re-anchored (was already true, stays true). - assert!(app.coding_data_retention_opt_out); - // Re-toast on confirmation uses display name + ✓. - let toast = read_toast(&app); - assert!( - toast.contains("Opt out"), - "confirmation toast must use display name 'Opt out': {toast}", - ); - assert!( - toast.contains('\u{2713}'), - "opt-out confirmation toast uses ✓: {toast}", - ); -} - -/// `TaskResult::CodingDataSharingUpdated` corrects the in-memory -/// state if the server reshapes the boolean (e.g. policy -/// override). Pins the defense-in-depth re-anchor contract. -#[test] -fn coding_data_sharing_updated_corrects_state_if_server_disagrees() { - let mut app = test_app_with_agent(); - // Optimistic mutation said "opt-out" — but the server - // overrides to "opt-in" (e.g. policy that prevents opt-out). - app.coding_data_retention_opt_out = true; - let id = AgentId(0); - let effects = dispatch( - Action::TaskComplete(TaskResult::CodingDataSharingUpdated { - agent_id: id, - opted_in: true, // server says opted-in - }), - &mut app, - ); - assert!(effects.is_empty()); - // State corrected to match server. - assert!( - !app.coding_data_retention_opt_out, - "server-confirmed opt-in must overwrite optimistic opt-out", - ); - // Server-correction toast uses the destructive ⚠ - // pattern for the opt-in direction (the privacy-degrading - // override deserves the warning glyph even if the SERVER, not - // the user, made the call). - let toast = read_toast(&app); - assert!( - toast.contains("Opt in"), - "post-correction toast uses display name 'Opt in': {toast}", - ); - assert!( - toast.contains('\u{26A0}'), - "opt-in direction always uses ⚠ glyph, even on server-correction path: {toast}", - ); -} - -/// `TaskResult::CodingDataSharingFailed` REVERTS the optimistic -/// mutation and surfaces a failure toast. Pins the rollback -/// contract. -/// -/// Failure toast uses the standardised "coding data sharing" -/// wording. -#[test] -fn coding_data_sharing_failed_rolls_back_and_toasts_error() { - let mut app = test_app_with_agent(); - // Simulate post-optimistic state: user picked opt-out, state - // was flipped, then the ACP call failed. The pre-toggle value - // was opt-in (true), so `rollback_to_opted_in = true`. - app.coding_data_retention_opt_out = true; - let id = AgentId(0); - let effects = dispatch( - Action::TaskComplete(TaskResult::CodingDataSharingFailed { - agent_id: id, - error: "server error".into(), - rollback_to_opted_in: true, - }), - &mut app, - ); - assert!(effects.is_empty(), "rollback path must NOT emit Effect"); - // State reverted to pre-toggle (opted-in). - assert!( - !app.coding_data_retention_opt_out, - "rollback must revert optimistic mutation", - ); - // Failure toast surfaces the error using full label. - let toast = read_toast(&app); - assert!( - toast.contains("coding data sharing"), - "PR 9 R1: failure toast wording standardised to include 'coding data sharing' \ - (G2 Issue 2): {toast}", - ); - assert!(toast.contains("server error"), "error in toast: {toast}"); - assert!(toast.contains('\u{2717}'), "failure toast uses ✗: {toast}"); -} - -/// `TaskResult::CodingDataSharingFailed` reverts in the OTHER -/// direction too (the pre-toggle state could have been either). -#[test] -fn coding_data_sharing_failed_rolls_back_to_opt_out() { - let mut app = test_app_with_agent(); - // Post-optimistic: opted-in (user picked opt-in, server - // failed, pre-toggle was opt-out). - app.coding_data_retention_opt_out = false; - let id = AgentId(0); - let effects = dispatch( - Action::TaskComplete(TaskResult::CodingDataSharingFailed { - agent_id: id, - error: "network timeout".into(), - rollback_to_opted_in: false, - }), - &mut app, - ); - assert!(effects.is_empty()); - // Reverted to pre-toggle opt-out. - assert!( - app.coding_data_retention_opt_out, - "rollback to opt-out must set state=true", - ); -} - -/// Optimistic mutation refreshes any open settings modal. -/// Without this refresh, the modal indicator would stay at the -/// pre-toggle value until manual re-render. -#[test] -fn set_coding_data_sharing_refreshes_open_modal_snapshot() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = false; - // Open a settings modal (capture initial snapshot). - let _ = dispatch(Action::OpenSettings, &mut app); - // Verify snapshot reads opted-in. - let agent_id = AgentId(0); - { - let state = match &app.agents[&agent_id].active_modal { - Some(crate::views::modal::ActiveModal::Settings { state }) => state, - _ => panic!("expected Settings modal open after OpenSettings dispatch"), - }; - assert!( - !state.pager_snapshot.coding_data_sharing_opt_out, - "initial snapshot must read opt_out=false (opted-in)", - ); - } - // Dispatch the toggle. - let _ = dispatch(Action::SetCodingDataSharing { opted_in: false }, &mut app); - // Snapshot now reflects the optimistic mutation. - let state = match &app.agents[&agent_id].active_modal { - Some(crate::views::modal::ActiveModal::Settings { state }) => state, - _ => panic!("Settings modal must still be open after SetCodingDataSharing dispatch"), - }; - assert!( - state.pager_snapshot.coding_data_sharing_opt_out, - "snapshot must refresh to reflect opt_out=true (opted-out) after dispatch", - ); -} - -/// Rollback also refreshes the modal — the user sees the -/// reverted value, not the stale optimistic one. -#[test] -fn coding_data_sharing_failed_refreshes_open_modal_snapshot() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = false; - let _ = dispatch(Action::OpenSettings, &mut app); - // Optimistic flip. - let _ = dispatch(Action::SetCodingDataSharing { opted_in: false }, &mut app); - // ACP failure. - let _ = dispatch( - Action::TaskComplete(TaskResult::CodingDataSharingFailed { - agent_id: AgentId(0), - error: "x".into(), - rollback_to_opted_in: true, - }), - &mut app, - ); - let state = match &app.agents[&AgentId(0)].active_modal { - Some(crate::views::modal::ActiveModal::Settings { state }) => state, - _ => panic!("Settings modal must still be open after rollback TaskResult"), - }; - assert!( - !state.pager_snapshot.coding_data_sharing_opt_out, - "rollback must refresh snapshot back to opt_out=false (opted-in)", - ); -} - -// ── coding_data_sharing toast tests ───────────── - -/// The opt-in transition -/// uses the **`⚠` destructive-warning glyph** + spelled-out -/// consequence text — mirroring `yolo_toast`'s -/// "Always-approve ON: all tool actions auto-run" pattern. The -/// consequence text is verbatim-pinned because the toast is the -/// only post-commit feedback for a privacy-degrading transition; -/// a future PR that softens the wording silently degrades the -/// safety affordance. -#[test] -fn set_coding_data_sharing_opt_in_renders_destructive_warning_toast() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = true; // currently opted-out - let effects = dispatch(Action::SetCodingDataSharing { opted_in: true }, &mut app); - assert_eq!(effects.len(), 1, "non-idempotent opt-in must emit Effect"); - let toast = read_toast(&app); - assert!( - toast.contains('\u{26A0}'), - "opt-in toast MUST use ⚠ glyph (PR 9 R1, General-3 Issue 5 — \ - privacy-degrading transition deserves destructive-warning glyph): {toast}", - ); - assert!( - !toast.contains('\u{2713}'), - "opt-in toast MUST NOT use the uniform ✓ glyph — that's the \ - safe-default toast for opt-out: {toast}", - ); - assert!( - toast.contains("Opt in"), - "destructive toast still uses display name 'Opt in': {toast}", - ); - // Consequence text pinned: a future PR softening this loses - // the safety affordance. - assert!( - toast.contains("code samples"), - "destructive toast must spell out the consequence \ - (mention 'code samples'): {toast}", - ); - assert!( - toast.contains("training"), - "destructive toast must spell out the consequence \ - (mention 'training'): {toast}", - ); -} - -/// The opt-out transition uses the -/// uniform `✓` glyph (safe default), NOT the destructive `⚠`. -/// Mirrors `yolo_toast(false)` precedent — restoring the safe -/// default doesn't warrant the heavier visual. -#[test] -fn set_coding_data_sharing_opt_out_renders_safe_default_toast() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = false; // currently opted-in - let _ = dispatch(Action::SetCodingDataSharing { opted_in: false }, &mut app); - let toast = read_toast(&app); - assert!( - toast.contains('\u{2713}'), - "opt-out toast uses ✓ safe-default glyph: {toast}", - ); - assert!( - !toast.contains('\u{26A0}'), - "opt-out toast MUST NOT use ⚠ — that's reserved for the privacy-degrading \ - direction (PR 9 R1): {toast}", - ); - assert!(toast.contains("Opt out")); -} - -/// The toast renders -/// the registered `EnumChoice.display` ("Opt in" / "Opt out"), -/// NOT the persisted canonical ("opt-in" / "opt-out"). Mirrors -/// the `set_theme_toast_format_uses_display_name` contract. -/// The display strings here are pinned by the -/// `coding_data_sharing_choices_use_canonical_strings` e2e test -/// (registry side) AND -/// `pr9_coding_data_sharing_choices_use_canonical_strings` (which -/// also pins the display labels via the same EnumChoice -/// entries). -#[test] -fn coding_data_sharing_toast_format_uses_display_name() { - let mut app = test_app_with_agent(); - // Opt-in direction. - app.coding_data_retention_opt_out = true; - let _ = dispatch(Action::SetCodingDataSharing { opted_in: true }, &mut app); - let opt_in_toast = read_toast(&app); - assert!( - opt_in_toast.contains("Opt in"), - "opt-in toast uses display 'Opt in', not canonical 'opt-in': {opt_in_toast}", - ); - // Clear and test opt-out direction. - app.agents.get_mut(&AgentId(0)).unwrap().toast = None; - app.coding_data_retention_opt_out = false; - let _ = dispatch(Action::SetCodingDataSharing { opted_in: false }, &mut app); - let opt_out_toast = read_toast(&app); - assert!( - opt_out_toast.contains("Opt out"), - "opt-out toast uses display 'Opt out', not canonical 'opt-out': {opt_out_toast}", - ); -} - -/// The failure toast -/// substitutes a generic placeholder when the error string is -/// too long OR contains control characters / newlines. Pins the -/// scrub contract. -#[test] -fn coding_data_sharing_failed_scrubs_long_error_messages() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = true; - let id = AgentId(0); - // ~500-char error simulating a stack trace / HTML 502 page. - let huge_error = "a".repeat(500); - let _ = dispatch( - Action::TaskComplete(TaskResult::CodingDataSharingFailed { - agent_id: id, - error: huge_error.clone(), - rollback_to_opted_in: false, - }), - &mut app, - ); - let toast = read_toast(&app); - assert!( - !toast.contains(&huge_error), - "long error MUST be scrubbed from the toast: {} chars", - toast.len(), - ); - assert!( - toast.contains("see logs"), - "scrubbed toast must point at the log for full details: {toast}", - ); -} - -/// Control characters (CR/LF/NUL) -/// in the error trigger the scrub path even on short strings — -/// preserves the toast's single-line layout. -#[test] -fn coding_data_sharing_failed_scrubs_control_chars_in_error() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = true; - let id = AgentId(0); - // Short message with embedded newlines. - let multiline = "line1\nline2\nline3".to_string(); - let _ = dispatch( - Action::TaskComplete(TaskResult::CodingDataSharingFailed { - agent_id: id, - error: multiline.clone(), - rollback_to_opted_in: false, - }), - &mut app, - ); - let toast = read_toast(&app); - assert!( - !toast.contains('\n'), - "newlines MUST be scrubbed from the toast (would break single-line layout): \ - {toast:?}", - ); - assert!( - toast.contains("see logs"), - "control-char-scrubbed toast points at logs: {toast}", - ); -} - -/// The scrub path preserves short, -/// sanitised error messages verbatim — the typical happy-path -/// shell-side error string stays unscrubbed. -#[test] -fn coding_data_sharing_failed_preserves_short_clean_error_message() { - let mut app = test_app_with_agent(); - app.coding_data_retention_opt_out = true; - let id = AgentId(0); - let short_clean = "network timeout".to_string(); - let _ = dispatch( - Action::TaskComplete(TaskResult::CodingDataSharingFailed { - agent_id: id, - error: short_clean.clone(), - rollback_to_opted_in: false, - }), - &mut app, - ); - let toast = read_toast(&app); - assert!( - toast.contains(&short_clean), - "short clean error must appear verbatim in the toast: {toast}", - ); - assert!( - !toast.contains("see logs"), - "short clean error must NOT trigger the scrub fallback: {toast}", - ); -} - /// Direct unit test of the `scrub_error_for_toast` helper — /// pins the threshold and the fallback string against drift. #[test] @@ -766,31 +113,6 @@ fn scrub_error_for_toast_unit() { ); } -/// The no-agent path -/// returns empty cleanly — no toast (the show_toast call would -/// no-op anyway), no panic, no Effect emitted. A "✗ No active -/// session" toast would be dead UX (no agent = no toast surface -/// to render on), so this path emits a tracing::warn! instead. -#[test] -fn set_coding_data_sharing_no_agents_returns_empty_without_panic() { - let mut app = test_app_with_agent(); - // Remove every agent so the dispatcher hits the no-agent path. - app.agents.clear(); - // Force the view off Agent so the dispatcher falls through to - // app.agents.keys().next() which is now empty. - app.active_view = ActiveView::Welcome; - let effects = dispatch(Action::SetCodingDataSharing { opted_in: false }, &mut app); - assert!( - effects.is_empty(), - "no-agent path must return empty (no Effect to fire)", - ); - // State unchanged (we never reach the optimistic mutation). - assert!( - !app.coding_data_retention_opt_out, - "no-agent path must NOT mutate state", - ); -} - #[test] fn dispatch_rename_session_updates_display_name_locally() { let mut app = test_app_with_agent(); @@ -898,27 +220,6 @@ fn show_usage_on_welcome_screen_is_noop() { ); } -#[test] -fn show_usage_with_redirect_url_shows_link_and_skips_fetch() { - let mut app = test_app_with_agent(); - app.usage_billing_redirect_url = Some("https://billing.example.com/me".to_string()); - let before = agent_scrollback_len(&app); - let effects = dispatch(Action::ShowUsage, &mut app); - assert!( - effects.is_empty(), - "with a redirect URL set, ShowUsage should not fetch (billing or auto-topup), got: {effects:?}" - ); - assert_eq!( - agent_scrollback_len(&app), - before + 1, - "redirect path should push one system message with the billing link" - ); - assert!( - last_system_text(&app, AgentId(0)).contains("https://billing.example.com/me"), - "redirect message should use the remote settings-provided URL" - ); -} - // ── Minimal update-notice tests ────────────────────────────────────── #[test] diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs index c399650..495b35e 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs @@ -1494,391 +1494,6 @@ fn rename_session_failed_keeps_local_display_name_and_pushes_system_block() { ); } -// ── GateRefreshed subscription flow ───────────────────────────── - -/// Regression: when the 30s gate poll detects the subscription gate has -/// been lifted, it must emit `CheckSubscription` so the shell refreshes -/// the JWT. Without this the auth token still lacks the subscription -/// claim and all API calls return 403. -#[test] -fn gate_refreshed_emits_check_subscription_on_gate_lift() { - let mut app = test_app(); - // User starts gated (no subscription). - app.gate = Some(kigi_shell::auth::GateInfo { - message: "SuperGrok subscription required".into(), - url: Some("https://grok.com/supergrok".into()), - label: Some("Subscribe".into()), - }); - assert!(!app.has_access()); - - // Server-side settings now show no gate (user purchased subscription). - let settings = kigi_shell::util::config::RemoteSettings::default(); - let effects = dispatch_task_result( - TaskResult::GateRefreshed { - settings: Some(settings), - }, - &mut app, - ); - - // Gate must be lifted. - assert!(app.has_access(), "gate should be lifted"); - assert!(app.welcome_prompt_focused, "prompt should be focused"); - - // Must emit CheckSubscription to trigger shell-side JWT refresh. - assert!( - effects - .iter() - .any(|e| matches!(e, Effect::CheckSubscription { verify: None })), - "must emit CheckSubscription to refresh JWT; got: {effects:?}" - ); -} - -/// When the gate poll returns settings that still have a gate, no -/// effects should be emitted and the user stays blocked. -#[test] -fn gate_refreshed_no_effect_when_still_gated() { - let mut app = test_app(); - app.gate = Some(kigi_shell::auth::GateInfo { - message: "Subscribe".into(), - url: None, - label: None, - }); - - let settings = kigi_shell::util::config::RemoteSettings { - gate_message: Some("Subscribe".into()), - ..Default::default() - }; - let effects = dispatch_task_result( - TaskResult::GateRefreshed { - settings: Some(settings), - }, - &mut app, - ); - - assert!(!app.has_access(), "gate should remain"); - assert!(effects.is_empty(), "no effects when still gated"); -} - -/// When the user was never gated, GateRefreshed is a no-op. -#[test] -fn gate_refreshed_no_effect_when_already_unblocked() { - let mut app = test_app(); - assert!(app.has_access()); // no gate - - let settings = kigi_shell::util::config::RemoteSettings::default(); - let effects = dispatch_task_result( - TaskResult::GateRefreshed { - settings: Some(settings), - }, - &mut app, - ); - - assert!(effects.is_empty(), "no effects when already unblocked"); -} - -/// A gate newly imposed by the 30s settings poll (possibly stale) must be -/// deferred for live verification instead of painting the paywall directly: -/// the gate is held out of `app.gate` and a `CheckSubscription` + -/// verify-timeout pair is emitted. -#[test] -fn gate_refreshed_newly_blocked_defers_gate_for_verification() { - let mut app = test_app(); - assert!(app.has_access()); // ungated - - let settings = kigi_shell::util::config::RemoteSettings { - gate_message: Some("Subscribe".into()), - ..Default::default() - }; - let effects = dispatch_task_result( - TaskResult::GateRefreshed { - settings: Some(settings), - }, - &mut app, - ); - - assert!( - app.has_access(), - "deferred gate must not show as paywall before verification" - ); - assert!(app.pending_gate_verification.is_some()); - assert!( - effects - .iter() - .any(|e| matches!(e, Effect::CheckSubscription { verify: Some(_) })), - "must live-check before showing the paywall; got: {effects:?}" - ); - assert!( - effects - .iter() - .any(|e| matches!(e, Effect::ScheduleGateVerifyTimeout { .. })), - "must arm the verification timeout; got: {effects:?}" - ); -} - -// ── Stale-gate verification resolution ────────────────────────── - -fn test_gate() -> kigi_shell::auth::GateInfo { - kigi_shell::auth::GateInfo { - message: "Subscribe".into(), - url: None, - label: None, - } -} - -/// The live check confirmed access (meta without a gate): the deferred -/// stale gate is dropped and the paywall never shows. -#[test] -fn verify_check_with_meta_resolves_pending_gate() { - let mut app = test_app(); - let _effs = app.impose_gate(test_gate()); - assert!(app.has_access()); - - let meta = serde_json::to_value(kigi_shell::auth::AuthMeta::default()).unwrap(); - dispatch_task_result( - TaskResult::CheckSubscriptionComplete { - verify: Some(app.gate_verify_gen), - meta: Some(meta), - }, - &mut app, - ); - - assert!(app.has_access(), "live check says subscribed — no paywall"); - assert!(app.pending_gate_verification.is_none()); -} - -/// The verification's own check failed (meta None) while its stale gate -/// was deferred: err on blocking — the deferred gate is promoted. -#[test] -fn verify_check_failure_promotes_pending_gate() { - let mut app = test_app(); - let _effs = app.impose_gate(test_gate()); - - let effects = dispatch_task_result( - TaskResult::CheckSubscriptionComplete { - verify: Some(app.gate_verify_gen), - meta: None, - }, - &mut app, - ); - - assert!(!app.has_access(), "check failed — deferred gate must show"); - assert!(app.pending_gate_verification.is_none()); - assert!( - effects - .iter() - .any(|e| matches!(e, Effect::SchedulePaywallCheck)), - "freshly shown gate must arm the 5s auto-lift chain; got: {effects:?}" - ); -} - -/// A failed GENERIC check (watch / focus / paywall chain — no generation) -/// must never promote a deferred gate: only the deferral's own -/// generation-scoped check or timeout may (a superseded or unrelated check -/// failing is not evidence about the current verification). -#[test] -fn check_subscription_complete_failure_leaves_pending_gate_untouched() { - let mut app = test_app(); - let _effs = app.impose_gate(test_gate()); - - let effects = dispatch_task_result( - TaskResult::CheckSubscriptionComplete { - verify: None, - meta: None, - }, - &mut app, - ); - - assert!(effects.is_empty()); - assert!( - app.has_access(), - "generic check failure must not promote the deferred gate" - ); - assert!( - app.pending_gate_verification.is_some(), - "verification must stay in flight" - ); -} - -/// A failed verification check from a SUPERSEDED deferral (older -/// generation) must not promote the newer pending gate. -#[test] -fn verify_check_stale_generation_failure_is_ignored() { - let mut app = test_app(); - let _effs = app.impose_gate(test_gate()); - let stale_gen = app.gate_verify_gen; - // Second deferral supersedes the first (its check is in flight). - let _effs = app.impose_gate(test_gate()); - - let effects = dispatch_task_result( - TaskResult::CheckSubscriptionComplete { - verify: Some(stale_gen), - meta: None, - }, - &mut app, - ); - - assert!(effects.is_empty()); - assert!( - app.has_access(), - "superseded verification failure must not promote the newer gate" - ); - assert!(app.pending_gate_verification.is_some()); -} - -/// A check failure with no deferred gate (the plain paywall-poller path) -/// must not invent a gate. -#[test] -fn check_subscription_complete_failure_without_pending_gate_is_noop() { - let mut app = test_app(); - dispatch_task_result( - TaskResult::CheckSubscriptionComplete { - verify: None, - meta: None, - }, - &mut app, - ); - assert!(app.has_access()); -} - -/// The verification window expired before the live check resolved: -/// err on blocking — the deferred gate is promoted, and the freshly shown -/// paywall gets the 5s auto-lift chain. -#[test] -fn gate_verify_timeout_promotes_pending_gate() { - let mut app = test_app(); - let _effs = app.impose_gate(test_gate()); - assert!(app.has_access()); - - let effects = dispatch_task_result( - TaskResult::GateVerifyTimeout { - generation: app.gate_verify_gen, - }, - &mut app, - ); - - assert!(!app.has_access(), "timeout — deferred gate must show"); - assert!(app.pending_gate_verification.is_none()); - assert!( - app.paywall_check_started.is_some(), - "promoted gate must arm the paywall auto-check chain" - ); - assert!( - effects - .iter() - .any(|e| matches!(e, Effect::SchedulePaywallCheck)), - "promoted gate must schedule the 5s chain; got: {effects:?}" - ); -} - -/// The timeout fires after the check already resolved the gate: no-op. -#[test] -fn gate_verify_timeout_noop_when_already_resolved() { - let mut app = test_app(); - let _effs = app.impose_gate(test_gate()); - let generation = app.gate_verify_gen; - // Live check resolved first (access confirmed). - let meta = serde_json::to_value(kigi_shell::auth::AuthMeta::default()).unwrap(); - dispatch_task_result( - TaskResult::CheckSubscriptionComplete { - verify: None, - meta: Some(meta), - }, - &mut app, - ); - - dispatch_task_result(TaskResult::GateVerifyTimeout { generation }, &mut app); - assert!( - app.has_access(), - "stale timeout must not re-impose the gate" - ); -} - -/// A timeout from a SUPERSEDED verification (older generation) must not -/// promote a newer deferred gate whose own live check is still in flight. -#[test] -fn gate_verify_timeout_stale_generation_is_ignored() { - let mut app = test_app(); - // First deferral resolves (access confirmed) ... - let _effs = app.impose_gate(test_gate()); - let stale_gen = app.gate_verify_gen; - let meta = serde_json::to_value(kigi_shell::auth::AuthMeta::default()).unwrap(); - dispatch_task_result( - TaskResult::CheckSubscriptionComplete { - verify: None, - meta: Some(meta), - }, - &mut app, - ); - // ... then a SECOND gate is deferred (check in flight). - let _effs = app.impose_gate(test_gate()); - assert!(app.has_access()); - - // The FIRST deferral's timer fires now — it must not promote the - // second deferral's pending gate. - let effects = dispatch_task_result( - TaskResult::GateVerifyTimeout { - generation: stale_gen, - }, - &mut app, - ); - - assert!(effects.is_empty()); - assert!( - app.has_access(), - "stale-generation timer must not promote the newer pending gate" - ); - assert!( - app.pending_gate_verification.is_some(), - "the newer verification must stay in flight" - ); -} - -/// `GateRefreshed` with gate-free settings while a deferred gate awaits -/// verification must drop the pending copy — the fresh settings are newer -/// than the stale snapshot that produced it — and still run the lift -/// bookkeeping (`CheckSubscription` for the JWT refresh), since the pending -/// deferral means the user was conceptually blocked. -#[test] -fn gate_refreshed_without_gate_clears_pending_verification() { - let mut app = test_app(); - let _effs = app.impose_gate(test_gate()); - let generation = app.gate_verify_gen; - - let settings = kigi_shell::util::config::RemoteSettings::default(); - let effects = dispatch_task_result( - TaskResult::GateRefreshed { - settings: Some(settings), - }, - &mut app, - ); - - assert!(app.pending_gate_verification.is_none()); - assert!( - effects - .iter() - .any(|e| matches!(e, Effect::CheckSubscription { verify: None })), - "settings-confirmed lift of a pending gate must refresh the JWT; got: {effects:?}" - ); - // The still-armed timer must find nothing to promote. - dispatch_task_result(TaskResult::GateVerifyTimeout { generation }, &mut app); - assert!( - app.has_access(), - "cleared pending gate must not resurface via the timer" - ); -} - -/// Logout clears any deferred gate and the check debounce. -#[test] -fn logout_clears_pending_gate_verification() { - let mut app = test_app(); - let _effs = app.impose_gate(test_gate()); - - dispatch_task_result(TaskResult::LogoutComplete, &mut app); - - assert!(app.pending_gate_verification.is_none()); - assert!(app.last_subscription_check_at.is_none()); -} - /// `apply_setting_rollback` on a known key reverts the in-memory /// cache without emitting any new effects. #[test] @@ -2002,38 +1617,16 @@ fn rollback_to_always_approve_blocked_by_policy_pin() { assert!(!app.default_yolo); } -// -- Degraded conversations lane (SessionListLoaded.partial) ---------- +// -- SessionListLoaded ------------------------------------------------ -/// A degraded conversations lane surfaces an actionable notice instead of -/// the misleading "No sessions found" toast. +/// Canary: an empty list surfaces the generic "no sessions" toast. #[test] -fn session_list_partial_no_oauth_surfaces_login_hint() { +fn session_list_empty_shows_generic_toast() { let mut app = test_app_with_agent(); open_session_picker_with(&mut app, vec![]); let _ = dispatch( Action::TaskComplete(TaskResult::SessionListLoaded { sessions: vec![], - partial: Some(crate::app::effects::ConversationsPartial::NoOauth), - seq: 0, - query: None, - }), - &mut app, - ); - assert!( - read_toast(&app).contains("/login"), - "no_oauth must point at /login" - ); -} - -/// Canary: an empty list without a degraded lane keeps the generic toast. -#[test] -fn session_list_empty_without_partial_keeps_generic_toast() { - let mut app = test_app_with_agent(); - open_session_picker_with(&mut app, vec![]); - let _ = dispatch( - Action::TaskComplete(TaskResult::SessionListLoaded { - sessions: vec![], - partial: None, seq: 0, query: None, }), @@ -2041,95 +1634,3 @@ fn session_list_empty_without_partial_keeps_generic_toast() { ); assert!(read_toast(&app).contains("No sessions found")); } - -/// Non-empty degraded list under chat mode (welcome-fallback branch): -/// entries land AND the retry notice surfaces; Build mode stays silent. -#[test] -fn session_list_nonempty_partial_toasts_retry_in_chat_mode_only() { - let mut app = test_app_with_agent(); - app.chat_mode = true; - let _ = dispatch( - Action::TaskComplete(TaskResult::SessionListLoaded { - sessions: vec![make_conversation_entry("conv-part-1")], - partial: Some(crate::app::effects::ConversationsPartial::Timeout), - seq: 0, - query: None, - }), - &mut app, - ); - assert!( - app.session_picker_entries.is_some(), - "entries must still land on a degraded lane" - ); - assert!( - read_toast(&app).contains("retry"), - "timeout must surface the retry notice" - ); - - // Build-mode canary: stays silent on a degraded lane. - let mut app = test_app_with_agent(); - let _ = dispatch( - Action::TaskComplete(TaskResult::SessionListLoaded { - sessions: vec![make_picker_entry("local-part-1", "/r")], - partial: Some(crate::app::effects::ConversationsPartial::Timeout), - seq: 0, - query: None, - }), - &mut app, - ); - assert!( - app.agents[&AgentId(0)].toast.is_none(), - "Build-mode non-empty degraded list stays silent" - ); -} - -/// Modal variant of the non-empty degraded-lane notice: same chat-mode-only -/// gating as the welcome-fallback branch. -#[test] -fn session_list_nonempty_partial_modal_toasts_in_chat_mode_only() { - use crate::views::modal::ActiveModal; - let mut app = test_app_with_agent(); - app.chat_mode = true; - open_session_picker_with(&mut app, vec![]); - let _ = dispatch( - Action::TaskComplete(TaskResult::SessionListLoaded { - sessions: vec![make_conversation_entry("conv-part-m1")], - partial: Some(crate::app::effects::ConversationsPartial::Timeout), - seq: 0, - query: None, - }), - &mut app, - ); - let agent = get_active_agent(&app).expect("active agent"); - assert!( - matches!( - agent.active_modal.as_ref(), - Some(ActiveModal::SessionPicker { - entries: Some(list), - .. - }) if list.len() == 1 - ), - "entries must land in the open modal on a degraded lane" - ); - assert!( - read_toast(&app).contains("retry"), - "chat-mode modal must surface the retry notice" - ); - - // Build-mode canary: the open modal stays silent. - let mut app = test_app_with_agent(); - open_session_picker_with(&mut app, vec![]); - let _ = dispatch( - Action::TaskComplete(TaskResult::SessionListLoaded { - sessions: vec![make_picker_entry("local-part-m1", "/r")], - partial: Some(crate::app::effects::ConversationsPartial::Timeout), - seq: 0, - query: None, - }), - &mut app, - ); - assert!( - app.agents[&AgentId(0)].toast.is_none(), - "Build-mode modal non-empty degraded list stays silent" - ); -} diff --git a/crates/codegen/kigi-tui/src/app/dispatch/transcript.rs b/crates/codegen/kigi-tui/src/app/dispatch/transcript.rs index 714ae1e..7307175 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/transcript.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/transcript.rs @@ -372,8 +372,7 @@ pub(super) fn dispatch_open_extensions_modal( // Mutual exclusivity: close agents modal when opening extensions. agent.agents_modal = None; - let mut modal = ExtensionsModalState::new(tab); - modal.session_team_id = app.team_id.clone(); + let modal = ExtensionsModalState::new(tab); agent.extensions_modal = Some(modal); let Some(session_id) = agent.session.session_id.clone() else { diff --git a/crates/codegen/kigi-tui/src/app/effects/helpers.rs b/crates/codegen/kigi-tui/src/app/effects/helpers.rs index f561d16..5a25510 100644 --- a/crates/codegen/kigi-tui/src/app/effects/helpers.rs +++ b/crates/codegen/kigi-tui/src/app/effects/helpers.rs @@ -43,9 +43,6 @@ pub(super) const SESSION_SEARCH_DEBOUNCE_MS: u64 = 250; /// All other errors are sanitized to remove internal service names and jargon. pub(super) fn format_acp_error(err: &acp::Error, is_api_key_auth: bool) -> String { if i32::from(err.code) == RATE_LIMITED_ERROR_CODE { - if super::dispatch::acp_error_is_free_usage_exhausted(err) { - return super::dispatch::FREE_USAGE_USER_MESSAGE.into(); - } return rate_limited_user_message(is_api_key_auth).into(); } if err.code == acp::ErrorCode::InvalidParams && let Some(data) = &err.data @@ -132,8 +129,6 @@ pub(crate) fn sanitize_user_error(raw: &str) -> String { return "Out of disk space.".to_string(); } static REPLACEMENTS: &[(&str, &str)] = &[ - ("cli-chat-proxy", "server"), - ("cli_chat_proxy", "server"), ("inference-api", "server"), ("inference_api", "server"), ("research-api", "server"), @@ -365,41 +360,6 @@ pub(super) fn count_chat_history_stats(history_path: &Path) -> (usize, usize) { } (turn_count, tool_call_count) } -/// Degraded conversations lane on `x.ai/session/list`, parsed from the -/// response's `_meta["x.ai/partial"]` envelope. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ConversationsPartial { - NoOauth, - Timeout, - Error, -} -impl ConversationsPartial { - /// Actionable picker notice for a degraded conversations lane. - pub(crate) fn picker_notice(self) -> &'static str { - match self { - Self::NoOauth => "Couldn't load your chats \u{2014} log in with /login", - Self::Timeout | Self::Error => "Couldn't load conversations \u{2014} retry", - } - } -} -/// Read `_meta["x.ai/partial"]` from a session-list payload. `None` when the -/// conversations lane completed (or was skipped); unknown reasons degrade to -/// [`ConversationsPartial::Error`]. -pub(super) fn parse_session_list_partial( - payload: &serde_json::Value, -) -> Option { - let partial = payload.get("_meta")?.get("x.ai/partial")?; - if partial.get("conversations").and_then(|v| v.as_bool()) != Some(true) { - return None; - } - Some( - match partial.get("reason").and_then(|v| v.as_str()) { - Some("no_oauth") => ConversationsPartial::NoOauth, - Some("timeout") => ConversationsPartial::Timeout, - _ => ConversationsPartial::Error, - }, - ) -} /// Parse the `x.ai/session/list` response payload (the unwrapped /// `{ "sessions": [...] }` object) into [`SessionPickerEntry`] rows. /// @@ -597,73 +557,6 @@ pub(super) async fn send_logout(tx: &AcpAgentTx) { tracing::warn!(error = % e, "logout failed"); } } -pub(super) async fn send_check_subscription( - tx: &AcpAgentTx, - verify: Option, -) -> TaskResult { - let req = acp::ExtRequest::new( - "x.ai/auth/check_subscription", - serde_json::value::to_raw_value(&serde_json::json!({})) - .expect("serialize check_subscription params") - .into(), - ); - match acp_send(req, tx).await { - Ok(resp) => { - let meta = serde_json::from_str::(resp.0.get()) - .ok() - .and_then(|v| v.get("meta").cloned()); - TaskResult::CheckSubscriptionComplete { - verify, - meta, - } - } - Err(e) => { - tracing::warn!(error = % e, "check_subscription failed"); - crate::unified_log::warn( - "subscription.check.rpc_failed", - None, - Some(serde_json::json!({ "verify" : verify, "error" : e.to_string(), })), - ); - TaskResult::CheckSubscriptionComplete { - verify, - meta: None, - } - } - } -} -/// One-shot subscription re-check for the credit-limit retry flow. -/// Same ACP call as `send_check_subscription` but returns a -/// `CreditLimitRecheckComplete` so the dispatch layer can decide -/// whether to retry the stashed prompt or show the upsell. -pub(super) async fn send_credit_limit_recheck( - tx: &AcpAgentTx, - agent_id: AgentId, -) -> TaskResult { - let req = acp::ExtRequest::new( - "x.ai/auth/check_subscription", - serde_json::value::to_raw_value(&serde_json::json!({})) - .expect("serialize check_subscription params") - .into(), - ); - match acp_send(req, tx).await { - Ok(resp) => { - let meta = serde_json::from_str::(resp.0.get()) - .ok() - .and_then(|v| v.get("meta").cloned()); - TaskResult::CreditLimitRecheckComplete { - agent_id, - meta, - } - } - Err(e) => { - tracing::warn!(error = % e, "credit_limit_recheck failed"); - TaskResult::CreditLimitRecheckComplete { - agent_id, - meta: None, - } - } - } -} pub(super) async fn send_authenticate( tx: &AcpAgentTx, request_seq: u64, @@ -1199,125 +1092,16 @@ pub(super) fn persist_hint( TaskResult::CancelComplete }); } -/// Map a billing config into a [`CreditBalance`]. -/// -/// Prefers the newer credits-config fields (`credit_usage_percent`, -/// `current_period`) and falls back to the deprecated -/// `monthly_limit`/`used`/`billing_period_end`. Shared by `Effect::FetchBilling` -/// and `Effect::FetchAppBilling` so every pager UI path derives identical usage -/// values from the same config. -pub(super) fn credit_balance_from_config( - c: kigi_shell::extensions::billing::BillingConfig, -) -> crate::views::credit_bar::CreditBalance { - let limit = c.monthly_limit.map(|v| v.val).unwrap_or(0); - let used = c.used.map(|v| v.val).unwrap_or(0); - let has_credit_pct = c.credit_usage_percent.is_some(); - let usage_pct = match c.credit_usage_percent { - Some(pct) => pct.clamp(0.0, 100.0), - None if limit > 0 => (used as f64 / limit as f64 * 100.0).min(100.0), - None => 0.0, - }; - let period_end_display = c - .current_period - .as_ref() - .and_then(|p| p.end.clone()) - .or(c.billing_period_end) - .and_then(|s| { - chrono::DateTime::parse_from_rfc3339(&s) - .ok() - .map(|dt| { - dt.with_timezone(&chrono::Local).format("%B %-d, %H:%M").to_string() - }) - }); - let on_demand_val = c.on_demand_cap.map(|v| v.val).unwrap_or(0); - let pay_as_you_go = on_demand_val > 0; - let on_demand_cap_cents = if on_demand_val > 0 { Some(on_demand_val) } else { None }; - let on_demand_used_cents = c - .on_demand_used - .map(|v| v.val) - .unwrap_or_else(|| (used - limit).max(0)); - let effective_usage_pct = if on_demand_val > 0 { - if usage_pct >= 100.0 { - (on_demand_used_cents as f64 / on_demand_val as f64 * 100.0).min(100.0) - } else if has_credit_pct { - usage_pct - } else { - let total_budget = limit + on_demand_val; - if total_budget > 0 { - (used as f64 / total_budget as f64 * 100.0).min(100.0) - } else { - 0.0 - } - } - } else { - usage_pct - }; - let period_type = c.current_period.as_ref().and_then(|p| p.period_type.clone()); - crate::views::credit_bar::CreditBalance { - usage_pct, - effective_usage_pct, - period_end_display, - pay_as_you_go, - on_demand_cap_cents, - on_demand_used_cents: Some(on_demand_used_cents), - prepaid_balance_cents: c.prepaid_balance.map(|v| v.val), - period_type, - is_unified_billing_user: c.is_unified_billing_user, - } -} -/// Whether the balance carries a non-zero prepaid credit balance (signed cents). -pub(super) fn has_prepaid_credits( - balance: Option<&crate::views::credit_bar::CreditBalance>, -) -> bool { - balance.and_then(|b| b.prepaid_balance_cents).map(i64::abs).is_some_and(|c| c > 0) -} -/// Fetch the user's auto top-up rule via the `x.ai/auto-topup-rule` extension. -/// A transport failure yields [`AutoTopupFetch::Unchanged`] so the caller keeps -/// any cached rule rather than treating the blip as "no auto top-up". -pub(super) async fn fetch_auto_topup_info( - tx: &kigi_acp_lib::AcpAgentTx, -) -> crate::views::credit_bar::AutoTopupFetch { - use crate::views::credit_bar::AutoTopupFetch; - let req = acp::ExtRequest::new( - "x.ai/auto-topup-rule", - serde_json::value::to_raw_value(&serde_json::json!({})) - .expect("serialize auto-topup params") - .into(), - ); - let Ok(resp) = acp_send(req, tx).await else { - return AutoTopupFetch::Unchanged; - }; - let wrapper: serde_json::Value = serde_json::from_str(resp.0.get()) - .unwrap_or_default(); - let result = wrapper.get("result").unwrap_or(&wrapper); - parse_auto_topup_response(result) -} -/// Map an `x.ai/auto-topup-rule` payload to an [`AutoTopupFetch`]. A body that -/// fails to deserialize is a fetch error (→ `Unchanged`, keep the cached rule), -/// not a definitive "no rule", so a malformed response can't silently flip the -/// credits warning. -pub(super) fn parse_auto_topup_response( +/// Parse an `x.ai/billing` ext response body (the unwrapped `result` +/// payload) into Kimi usage rows. A body that fails to deserialize is an +/// error, not an empty quota list, so a malformed response can't render +/// as "no usage data". +pub(super) fn parse_usage_response( result: &serde_json::Value, -) -> crate::views::credit_bar::AutoTopupFetch { - use crate::views::credit_bar::{AutoTopupFetch, AutoTopupInfo}; - use kigi_shell::extensions::billing::GetAutoTopupRuleResponse; - match serde_json::from_value::(result.clone()) { - Ok(parsed) => { - AutoTopupFetch::Resolved( - parsed - .rule - .map_or_else( - AutoTopupInfo::disabled, - |rule| AutoTopupInfo { - enabled: rule.enabled, - topup_amount_cents: rule.topup_amount.map(|c| c.val), - max_amount_cents: rule.max_amount_per_month.map(|c| c.val), - }, - ), - ) - } - Err(_) => AutoTopupFetch::Unchanged, - } +) -> Result, String> { + serde_json::from_value::(result.clone()) + .map(|usage| usage.rows) + .map_err(|e| format!("Parse error: {e}")) } /// A blocking flock on the shared, possibly-network `~/.kigi` lock must never /// stall the event-loop thread (and would hang exit on `/quit`); the registry diff --git a/crates/codegen/kigi-tui/src/app/effects/mod.rs b/crates/codegen/kigi-tui/src/app/effects/mod.rs index 7890374..ae006db 100644 --- a/crates/codegen/kigi-tui/src/app/effects/mod.rs +++ b/crates/codegen/kigi-tui/src/app/effects/mod.rs @@ -9,7 +9,6 @@ mod helpers; use super::actions; #[allow(unused_imports)] use super::{agent, dispatch}; -pub use helpers::ConversationsPartial; pub(super) use helpers::parse_session_load_running_prompt_id; pub(crate) use helpers::{ EffectMeta, RestoreProgressMsg, SessionFlags, persist_permission_mode_and_notify, @@ -79,31 +78,6 @@ pub(crate) fn execute( TaskResult::LogoutComplete }); } - Effect::CheckSubscription { verify } => { - let tx = acp_tx.clone(); - tasks.spawn(async move { send_check_subscription(&tx, verify).await }); - } - Effect::CreditLimitRecheck { agent_id } => { - let tx = acp_tx.clone(); - tasks.spawn(async move { send_credit_limit_recheck(&tx, agent_id).await }); - } - Effect::SchedulePaywallCheck => { - tasks - .spawn(async move { - tokio::time::sleep(std::time::Duration::from_secs(5)).await; - TaskResult::PaywallCheckTick - }); - } - Effect::ScheduleGateVerifyTimeout { generation } => { - tasks - .spawn(async move { - tokio::time::sleep(crate::app::subscription::GATE_VERIFY_TIMEOUT) - .await; - TaskResult::GateVerifyTimeout { - generation, - } - }); - } Effect::SwitchAccount { request_seq, method_id, use_oauth } => { let tx = acp_tx.clone(); let abort_handle = tasks @@ -717,10 +691,8 @@ pub(crate) fn execute( } let payload = wrapper.get("result").unwrap_or(&wrapper); let sessions = parse_session_picker_entries(payload); - let partial = parse_session_list_partial(payload); TaskResult::SessionListLoaded { sessions, - partial, seq, query, } @@ -2414,66 +2386,6 @@ pub(crate) fn execute( } }); } - Effect::ShareSession { agent_id, session_id } => { - use kigi_shell::session::{ShareSessionRequest, ShareSessionResponse}; - let tx = acp_tx.clone(); - tasks - .spawn(async move { - let request = acp::ExtRequest::new( - "x.ai/share_session", - serde_json::value::to_raw_value( - &ShareSessionRequest { - session_id: session_id.0.to_string(), - }, - ) - .expect("serialize share session params") - .into(), - ); - match acp_send(request, &tx).await { - Ok(resp) => { - let wrapper: serde_json::Value = serde_json::from_str( - resp.0.get(), - ) - .unwrap_or_default(); - if let Some(err) = wrapper.get("error") { - let msg = err - .as_str() - .map(String::from) - .unwrap_or_else(|| "unknown error".to_string()); - return TaskResult::ShareSessionFailed { - agent_id, - error: msg, - }; - } - let inner = wrapper.get("result").unwrap_or(&wrapper); - match serde_json::from_value::< - ShareSessionResponse, - >(inner.clone()) { - Ok(share_resp) => { - TaskResult::ShareSessionComplete { - agent_id, - share_url: share_resp.share_url, - } - } - Err(_) => { - TaskResult::ShareSessionFailed { - agent_id, - error: "couldn't share session".to_string(), - } - } - } - } - Err(e) => { - TaskResult::ShareSessionFailed { - agent_id, - error: sanitize_user_error( - &format!("couldn't share session: {e}"), - ), - } - } - } - }); - } Effect::FetchSessionAgentName { agent_id, session_id } => { let tx = acp_tx.clone(); tasks @@ -2638,68 +2550,6 @@ pub(crate) fn execute( } }); } - Effect::SetCodingDataSharing { agent_id, opted_in, rollback_to_opted_in } => { - let tx = acp_tx.clone(); - tasks - .spawn(async move { - let request = acp::ExtRequest::new( - "x.ai/privacy/setCodingDataRetention", - serde_json::value::to_raw_value( - &serde_json::json!( - { "codingDataRetentionOptOut" : ! opted_in } - ), - ) - .expect("serialize params") - .into(), - ); - match acp_send(request, &tx).await { - Ok(resp) => { - let wrapper: serde_json::Value = match serde_json::from_str( - resp.0.get(), - ) { - Ok(v) => v, - Err(e) => { - return TaskResult::CodingDataSharingFailed { - agent_id, - error: format!("malformed response: {e}"), - rollback_to_opted_in, - }; - } - }; - if let Some(err) = wrapper - .get("error") - .filter(|v| !v.is_null()) - { - let msg = err - .as_str() - .map(String::from) - .unwrap_or_else(|| err.to_string()); - return TaskResult::CodingDataSharingFailed { - agent_id, - error: msg, - rollback_to_opted_in, - }; - } - let confirmed_opted_in = wrapper - .get("codingDataRetentionOptOut") - .and_then(|v| v.as_bool()) - .map(|opt_out| !opt_out) - .unwrap_or(opted_in); - TaskResult::CodingDataSharingUpdated { - agent_id, - opted_in: confirmed_opted_in, - } - } - Err(e) => { - TaskResult::CodingDataSharingFailed { - agent_id, - error: format!("{e}"), - rollback_to_opted_in, - } - } - } - }); - } Effect::ShowContextInfo { agent_id, session_id } => { let tx = acp_tx.clone(); tasks @@ -3453,150 +3303,30 @@ pub(crate) fn execute( } }); } - Effect::FetchBilling { agent_id, silent } => { + Effect::FetchUsage { agent_id } => { let tx = acp_tx.clone(); tasks .spawn(async move { - use kigi_shell::extensions::billing::BillingConfigResponse; let req = acp::ExtRequest::new( "x.ai/billing", serde_json::value::to_raw_value(&serde_json::json!({})) - .expect("serialize billing params") + .expect("serialize usage params") .into(), ); - let parsed = match acp_send(req, &tx).await { + let result = match acp_send(req, &tx).await { Ok(resp) => { let wrapper: serde_json::Value = serde_json::from_str( resp.0.get(), ) .unwrap_or_default(); - let result = wrapper.get("result").unwrap_or(&wrapper); - serde_json::from_value::< - BillingConfigResponse, - >(result.clone()) - } - Err(e) => { - return TaskResult::BillingError { - agent_id, - error: sanitize_user_error(&format!("{e}")), - silent, - }; + let payload = wrapper.get("result").unwrap_or(&wrapper); + parse_usage_response(payload) } + Err(e) => Err(sanitize_user_error(&format!("{e}"))), }; - let billing = match parsed { - Ok(billing) => billing, - Err(e) => { - return TaskResult::BillingError { - agent_id, - error: format!("Parse error: {e}"), - silent, - }; - } - }; - let subscription_tier = billing.subscription_tier; - let balance = billing.config.map(credit_balance_from_config); - let autotopup = if has_prepaid_credits(balance.as_ref()) { - fetch_auto_topup_info(&tx).await - } else { - crate::views::credit_bar::AutoTopupFetch::Cleared - }; - TaskResult::BillingFetched { + TaskResult::UsageFetched { agent_id, - balance, - silent, - subscription_tier, - autotopup, - } - }); - } - Effect::RefreshGate => { - tasks - .spawn(async move { - let settings = tokio::task::spawn_blocking(|| { - if !kigi_shell::util::config::resolve_remote_fetch_enabled() { - return None; - } - let kigi_home = kigi_shell::util::kigi_home::kigi_home(); - let store = kigi_shell::auth::read_auth_json( - &kigi_home.join("auth.json"), - ) - .ok()?; - let scope = kigi_shell::auth::KimiCodeConfig::default() - .auth_scope(); - let auth = kigi_shell::auth::lookup_auth( - &store, - &scope, - )?; - let proxy_base = std::env::var( - "KIGI_CLI_CHAT_PROXY_BASE_URL", - ) - .unwrap_or_else(|_| kigi_env::coding_api_base_url()); - kigi_shell::remote::fetch_settings_blocking( - &proxy_base, - &auth, - None, - ) - }) - .await - .ok() - .flatten(); - TaskResult::GateRefreshed { - settings, - } - }); - } - Effect::FetchAppBilling => { - let tx = acp_tx.clone(); - tasks - .spawn(async move { - use kigi_shell::extensions::billing::BillingConfigResponse; - let req = acp::ExtRequest::new( - "x.ai/billing", - serde_json::value::to_raw_value(&serde_json::json!({})) - .expect("serialize billing params") - .into(), - ); - match acp_send(req, &tx).await { - Ok(resp) => { - let wrapper: serde_json::Value = serde_json::from_str( - resp.0.get(), - ) - .unwrap_or_default(); - let result = wrapper.get("result").unwrap_or(&wrapper); - match serde_json::from_value::< - BillingConfigResponse, - >(result.clone()) { - Ok(billing) => { - let balance = billing - .config - .map(|c| crate::views::credit_bar::CreditBalance { - period_end_display: None, - ..credit_balance_from_config(c) - }); - let autotopup = if has_prepaid_credits(balance.as_ref()) { - fetch_auto_topup_info(&tx).await - } else { - crate::views::credit_bar::AutoTopupFetch::Cleared - }; - TaskResult::AppBillingFetched { - balance, - autotopup, - } - } - Err(_) => { - TaskResult::AppBillingFetched { - balance: None, - autotopup: crate::views::credit_bar::AutoTopupFetch::Unchanged, - } - } - } - } - Err(_) => { - TaskResult::AppBillingFetched { - balance: None, - autotopup: crate::views::credit_bar::AutoTopupFetch::Unchanged, - } - } + result, } }); } diff --git a/crates/codegen/kigi-tui/src/app/effects/tests.rs b/crates/codegen/kigi-tui/src/app/effects/tests.rs index 990dbc3..393ef94 100644 --- a/crates/codegen/kigi-tui/src/app/effects/tests.rs +++ b/crates/codegen/kigi-tui/src/app/effects/tests.rs @@ -1,6 +1,5 @@ #![cfg_attr(rustfmt, rustfmt::skip)] use super::*; -use kigi_shell::extensions::billing::{BillingConfig, Cent, UsagePeriod}; /// The invalid-params server detail survives `attach_prompt_usage` /// wrapping `error.data` as `{message, promptUsage}`. #[test] @@ -23,7 +22,7 @@ fn format_acp_error_rate_limit_is_auth_aware() { RATE_LIMITED_USER_MESSAGE_OAUTH, }; let err = acp::Error::new(RATE_LIMITED_ERROR_CODE, "Rate limited").data("slow down"); - assert_eq!(format_acp_error(& err, false), RATE_LIMITED_USER_MESSAGE_OAUTH); + assert_eq!(format_acp_error(&err, false), RATE_LIMITED_USER_MESSAGE_OAUTH.as_str()); assert_eq!(format_acp_error(& err, true), RATE_LIMITED_USER_MESSAGE_API_KEY); } /// Non-empty token ranges ride the wire block meta as `skillTokenRanges` @@ -138,43 +137,6 @@ fn picker_still_drops_build_row_with_empty_summary() { assert!(entries.is_empty(), "empty-summary Build rows stay dropped"); } #[test] -fn session_list_partial_parses_reasons() { - let payload = |reason: &str| { - serde_json::json!( - { "sessions" : [], "_meta" : { "x.ai/partial" : { "conversations" : true, - "reason" : reason } } } - ) - }; - assert_eq!( - parse_session_list_partial(& payload("no_oauth")), - Some(ConversationsPartial::NoOauth) - ); - assert_eq!( - parse_session_list_partial(& payload("timeout")), - Some(ConversationsPartial::Timeout) - ); - assert_eq!( - parse_session_list_partial(& payload("error")), Some(ConversationsPartial::Error) - ); - assert_eq!( - parse_session_list_partial(& payload("something_new")), - Some(ConversationsPartial::Error) - ); -} -#[test] -fn session_list_partial_absent_for_healthy_or_meta_less_responses() { - let healthy = serde_json::json!( - { "sessions" : [], "_meta" : { "x.ai/partial" : { "conversations" : false } } } - ); - assert_eq!(parse_session_list_partial(& healthy), None); - let legacy = serde_json::json!({ "sessions" : [] }); - assert_eq!(parse_session_list_partial(& legacy), None); -} -/// The agent serializes `ExtMethodResult`: the outcome -/// lives at `result.outcome`. Probing the top level (the pre-fix code) -/// was why the tasks-pane ✗ never removed stale (`not_found`) rows after -/// a session resume. -#[test] fn parse_kill_outcome_reads_result_envelope() { use kigi_tools::types::KillOutcome; let resp = r#"{"result":{"taskId":"t-1","outcome":"not_found"}}"#; @@ -315,250 +277,61 @@ fn interject_params_carry_content_when_blocks_present() { assert_eq!(content.len(), 1); assert_eq!(content[0] ["text"], "look at [Image #1]"); } -/// A billing config with every field unset, for use as a base in -/// `credit_balance_from_config` tests via struct-update syntax. -fn empty_billing_config() -> BillingConfig { - BillingConfig { - credit_usage_percent: None, - current_period: None, - monthly_limit: None, - used: None, - on_demand_cap: None, - on_demand_used: None, - prepaid_balance: None, - is_unified_billing_user: None, - billing_period_start: None, - billing_period_end: None, - history: vec![], - } -} +/// `x.ai/billing` ext result (a serialized shell `UsageResponse`) parses +/// into typed rows; unknown labels/reset hints survive the round trip. #[test] -fn credit_balance_prefers_credit_usage_percent_over_limit_used() { - let c = BillingConfig { - credit_usage_percent: Some(42.0), - monthly_limit: Some(Cent { val: 10_000 }), - used: Some(Cent { val: 9_000 }), - ..empty_billing_config() - }; - assert_eq!(credit_balance_from_config(c).usage_pct, 42.0); +fn parse_usage_response_reads_rows_from_fixture() { + let fixture = serde_json::json!({ + "rows": [ + { + "label": "Weekly limit", + "used": 250, + "limit": 1000, + "resetHint": "resets in 2d 1h" + }, + { "label": "5h limit", "used": 20, "limit": 50 } + ] + }); + let rows = parse_usage_response(&fixture).expect("fixture must parse"); + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].label, "Weekly limit"); + assert_eq!(rows[0].used, 250); + assert_eq!(rows[0].limit, 1000); + assert_eq!(rows[0].reset_hint.as_deref(), Some("resets in 2d 1h")); + assert_eq!(rows[1].label, "5h limit"); + assert!(rows[1].reset_hint.is_none()); } +/// Round-trip through the shell's own serializer: what the billing +/// extension emits must parse back to the same typed rows. #[test] -fn credit_balance_forwards_is_unified_billing_user() { - let c = BillingConfig { - is_unified_billing_user: Some(true), - ..empty_billing_config() - }; - assert_eq!(credit_balance_from_config(c).is_unified_billing_user, Some(true)); - assert_eq!( - credit_balance_from_config(empty_billing_config()).is_unified_billing_user, None - ); +fn parse_usage_response_round_trips_shell_serialization() { + use kigi_shell::extensions::billing::{UsageResponse, UsageRow}; + let wire = serde_json::to_value(&UsageResponse { + rows: vec![UsageRow { + label: "RPM".into(), + used: 12, + limit: 60, + reset_hint: None, + }], + }) + .expect("serialize"); + let rows = parse_usage_response(&wire).expect("round trip"); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].label, "RPM"); + assert_eq!(rows[0].used, 12); + assert_eq!(rows[0].limit, 60); } +/// Empty rows parse to an empty list (→ "No usage data available." in the +/// dispatch layer); a malformed body is an error, not an empty quota list. #[test] -fn credit_balance_falls_back_to_limit_used_when_percent_absent() { - let c = BillingConfig { - monthly_limit: Some(Cent { val: 10_000 }), - used: Some(Cent { val: 2_500 }), - ..empty_billing_config() - }; - assert_eq!(credit_balance_from_config(c).usage_pct, 25.0); -} -/// Match production: RFC 3339 → user's local wall-clock (no zone label). -fn expected_period_end_display(rfc3339: &str) -> String { - chrono::DateTime::parse_from_rfc3339(rfc3339) - .expect("test fixture is valid RFC 3339") - .with_timezone(&chrono::Local) - .format("%B %-d, %H:%M") - .to_string() -} -#[test] -fn credit_balance_prefers_current_period_end_over_billing_period_end() { - let end = "2026-06-08T20:00:00Z"; - let c = BillingConfig { - credit_usage_percent: Some(10.0), - current_period: Some(UsagePeriod { - period_type: Some("USAGE_PERIOD_TYPE_WEEKLY".into()), - start: Some("2026-06-01T00:00:00Z".into()), - end: Some(end.into()), - }), - billing_period_end: Some("2026-07-01T20:00:00Z".into()), - ..empty_billing_config() - }; - assert_eq!( - credit_balance_from_config(c).period_end_display.as_deref(), - Some(expected_period_end_display(end).as_str()) - ); -} -#[test] -fn credit_balance_period_end_uses_local_timezone() { - let winter = "2026-01-15T20:00:00Z"; - let summer = "2026-07-15T20:00:00Z"; - let winter_cfg = BillingConfig { - billing_period_end: Some(winter.into()), - ..empty_billing_config() - }; - let summer_cfg = BillingConfig { - billing_period_end: Some(summer.into()), - ..empty_billing_config() - }; - assert_eq!( - credit_balance_from_config(winter_cfg).period_end_display.as_deref(), - Some(expected_period_end_display(winter).as_str()) - ); - assert_eq!( - credit_balance_from_config(summer_cfg).period_end_display.as_deref(), - Some(expected_period_end_display(summer).as_str()) - ); - assert_ne!(expected_period_end_display(winter), expected_period_end_display(summer)); -} -#[test] -fn credit_balance_falls_back_to_billing_period_end() { - let end = "2026-07-01T20:00:00Z"; - let c = BillingConfig { - billing_period_end: Some(end.into()), - ..empty_billing_config() - }; - assert_eq!( - credit_balance_from_config(c).period_end_display.as_deref(), - Some(expected_period_end_display(end).as_str()) - ); -} -#[test] -fn credit_balance_period_end_falls_back_when_current_period_has_no_end() { - let end = "2026-07-01T20:00:00Z"; - let c = BillingConfig { - current_period: Some(UsagePeriod { - period_type: None, - start: Some("2026-06-01T00:00:00Z".into()), - end: None, - }), - billing_period_end: Some(end.into()), - ..empty_billing_config() - }; - assert_eq!( - credit_balance_from_config(c).period_end_display.as_deref(), - Some(expected_period_end_display(end).as_str()) - ); -} -#[test] -fn credit_balance_period_end_none_when_unavailable() { +fn parse_usage_response_empty_and_malformed() { assert!( - credit_balance_from_config(empty_billing_config()).period_end_display.is_none() + parse_usage_response(&serde_json::json!({ "rows": [] })) + .expect("empty rows parse") + .is_empty() ); -} -#[test] -fn credit_balance_clamps_new_percent_above_100() { - let c = BillingConfig { - credit_usage_percent: Some(150.0), - ..empty_billing_config() - }; - assert_eq!(credit_balance_from_config(c).usage_pct, 100.0); -} -#[test] -fn credit_balance_clamps_legacy_used_above_limit() { - let c = BillingConfig { - monthly_limit: Some(Cent { val: 1_000 }), - used: Some(Cent { val: 2_500 }), - ..empty_billing_config() - }; - assert_eq!(credit_balance_from_config(c).usage_pct, 100.0); -} -#[test] -fn credit_balance_effective_equals_usage_when_no_on_demand() { - let c = BillingConfig { - credit_usage_percent: Some(40.0), - ..empty_billing_config() - }; - let bal = credit_balance_from_config(c); - assert!(! bal.pay_as_you_go); - assert_eq!(bal.on_demand_cap_cents, None); - assert_eq!(bal.effective_usage_pct, 40.0); -} -#[test] -fn credit_balance_effective_uses_on_demand_ratio_when_included_exhausted() { - let c = BillingConfig { - credit_usage_percent: Some(100.0), - on_demand_cap: Some(Cent { val: 5_000 }), - on_demand_used: Some(Cent { val: 1_000 }), - ..empty_billing_config() - }; - let bal = credit_balance_from_config(c); - assert!(bal.pay_as_you_go); - assert_eq!(bal.usage_pct, 100.0); - assert_eq!(bal.effective_usage_pct, 20.0); - assert_eq!(bal.on_demand_cap_cents, Some(5_000)); - assert_eq!(bal.on_demand_used_cents, Some(1_000)); -} -#[test] -fn parse_auto_topup_present_rule_resolves() { - let v = serde_json::json!( - { "rule" : { "enabled" : true, "topupAmount" : { "val" : 2000 }, - "maxAmountPerMonth" : { "val" : 10000 } } } - ); - match parse_auto_topup_response(&v) { - crate::views::credit_bar::AutoTopupFetch::Resolved(at) => { - assert!(at.enabled); - assert_eq!(at.topup_amount_cents, Some(2000)); - assert_eq!(at.max_amount_cents, Some(10000)); - } - other => panic!("expected Resolved, got {other:?}"), - } -} -#[test] -fn parse_auto_topup_empty_body_resolves_to_disabled() { - for v in [serde_json::json!({}), serde_json::json!({ "rule" : null })] { - match parse_auto_topup_response(&v) { - crate::views::credit_bar::AutoTopupFetch::Resolved(at) => { - assert!(! at.enabled); - } - other => panic!("expected Resolved(disabled), got {other:?}"), - } - } -} -#[test] -fn parse_auto_topup_rule_without_enabled_is_disabled() { - let v = serde_json::json!({ "rule" : { "topupAmount" : { "val" : 500 } } }); - match parse_auto_topup_response(&v) { - crate::views::credit_bar::AutoTopupFetch::Resolved(at) => { - assert!(! at.enabled); - assert_eq!(at.topup_amount_cents, Some(500)); - } - other => panic!("expected Resolved(disabled), got {other:?}"), - } -} -#[test] -fn parse_auto_topup_malformed_body_is_unchanged() { - for v in [serde_json::json!(null), serde_json::json!(42)] { - match parse_auto_topup_response(&v) { - crate::views::credit_bar::AutoTopupFetch::Unchanged => {} - other => panic!("expected Unchanged, got {other:?}"), - } - } -} -#[test] -fn credit_balance_effective_tracks_included_for_new_shape_under_100() { - let c = BillingConfig { - credit_usage_percent: Some(95.0), - on_demand_cap: Some(Cent { val: 5_000 }), - on_demand_used: Some(Cent { val: 0 }), - ..empty_billing_config() - }; - let bal = credit_balance_from_config(c); - assert!(bal.pay_as_you_go); - assert_eq!(bal.effective_usage_pct, 95.0); -} -#[test] -fn credit_balance_effective_blends_budget_for_legacy_shape_under_100() { - let c = BillingConfig { - monthly_limit: Some(Cent { val: 10_000 }), - used: Some(Cent { val: 5_000 }), - on_demand_cap: Some(Cent { val: 10_000 }), - on_demand_used: Some(Cent { val: 0 }), - ..empty_billing_config() - }; - let bal = credit_balance_from_config(c); - assert!(bal.pay_as_you_go); - assert_eq!(bal.usage_pct, 50.0); - assert_eq!(bal.effective_usage_pct, 25.0); + assert!(parse_usage_response(&serde_json::json!({ "bogus": 1 })).is_err()); + assert!(parse_usage_response(&serde_json::json!(null)).is_err()); } #[test] fn parse_worktree_restore_payload_full() { diff --git a/crates/codegen/kigi-tui/src/app/event_loop.rs b/crates/codegen/kigi-tui/src/app/event_loop.rs index 151442d..abcee71 100644 --- a/crates/codegen/kigi-tui/src/app/event_loop.rs +++ b/crates/codegen/kigi-tui/src/app/event_loop.rs @@ -595,10 +595,6 @@ pub(crate) async fn run( .as_ref() .and_then(|s| s.show_resolved_model) .unwrap_or(true); - app.sharing_enabled = remote_settings - .as_ref() - .and_then(|s| s.sharing_enabled) - .unwrap_or(false); app.session_picker_grouped = std::env::var("KIGI_SESSION_PICKER_GROUPED") .ok() .and_then(|v| match v.as_str() { @@ -681,7 +677,7 @@ pub(crate) async fn run( // else: auth_state defaults to Done (already authenticated eagerly) // Effects stashed until after the initial render, so the user sees the // welcome/auth UI right away. - let mut post_render_effects = if needs_interactive_login { + let post_render_effects = if needs_interactive_login { if connection.auth_methods.is_empty() { app.auth_state = super::app_view::AuthState::Pending { error: Some("No login method available".to_string()), @@ -711,21 +707,6 @@ pub(crate) async fn run( } } - // Fallback: prefetch may have gate info the shell's AuthMeta missed. - // Errs on the side of blocking if stale. - if app.gate.is_none() - && let Some(rs) = remote_settings.as_ref() - { - app.gate = AppView::gate_from_settings(rs); - } - - // Re-impose the startup gate through the chokepoint: cached auth meta - // and the settings prefetch are both possibly stale, so a consumer - // session's gate is deferred for live verification before first paint. - if let Some(gate) = app.gate.take() { - post_render_effects.extend(app.impose_gate(gate)); - } - // Load config layers once, resolve tips and feature flags. let requirements = kigi_shell::config::load_merged_requirements(); let user_config = kigi_shell::config::load_from_disk().ok(); @@ -756,17 +737,6 @@ pub(crate) async fn run( ); } - app.zdr_access_enabled = kigi_shell::util::config::resolve_zdr_access_enabled( - requirements.as_ref(), - user_config.as_ref(), - managed_config.as_ref(), - remote_settings.as_ref(), - ); - - app.subscription_watch_interval_secs = remote_settings - .as_ref() - .and_then(|rs| rs.subscription_watch_interval_secs); - // Full layered resolve (env/requirements/remote may beat plain `[ui]`). crate::appearance::cache::set_show_thinking_blocks( kigi_shell::util::config::resolve_show_thinking_blocks( @@ -796,14 +766,6 @@ pub(crate) async fn run( .value, ); - app.usage_billing_redirect_url = remote_settings - .as_ref() - .and_then(|s| s.usage_billing_redirect_url.clone()); - - if app.is_access_blocked() { - app.welcome_prompt_focused = false; - } - { use kigi_shell::util::config::resolve_tips; @@ -1124,20 +1086,6 @@ pub(crate) async fn run( // iteration so it is popped on every close path. let mut gboom_keyboard_pushed = false; - const BILLING_POLL_INTERVAL: Duration = Duration::from_secs(30); - let mut billing_poll_at: Option = None; - - const GATE_POLL_INTERVAL: Duration = Duration::from_secs(30); - let mut gate_poll_at: Option = None; - - // Free→paid subscription watch (see `app::subscription`). - let mut subscription_watch_at: Option = if app.subscription_watch_wanted() { - app.subscription_watch_interval() - .map(|iv| Instant::now() + iv) - } else { - None - }; - // Leader-mode roster poll (FleetView dashboard). Only fires while the // dashboard is open AND we're connected via a leader. Armed to fire // immediately at loop start so an already-open dashboard refreshes @@ -1168,22 +1116,12 @@ pub(crate) async fn run( if process_effects(effs, &mut tasks, &mut app, &progress_tx) { return Ok(make_run_result(&app)); } - // Fetch billing early so the welcome screen can show a credit warning. - if app.usage_visible { - let effs = vec![super::actions::Effect::FetchAppBilling]; - if process_effects(effs, &mut tasks, &mut app, &progress_tx) { - return Ok(make_run_result(&app)); - } - } // Fetch changelog off the render path so the welcome screen // can display bullets and /release-notes uses the cached result. let effs = vec![super::actions::Effect::FetchChangelog]; if process_effects(effs, &mut tasks, &mut app, &progress_tx) { return Ok(make_run_result(&app)); } - if !app.has_access() { - gate_poll_at = Some(Instant::now() + GATE_POLL_INTERVAL); - } } if !post_render_effects.is_empty() @@ -1276,16 +1214,14 @@ pub(crate) async fn run( app.draw(terminal); } - // Initial prompt from the CLI positional (`grok "fix the bug"`). When + // Initial prompt from the CLI positional (`kigi "fix the bug"`). When // already authenticated, hand it to the shared dispatcher helper (same - // `NewSession`/`SendPrompt` path the welcome screen uses). ZDR-blocked - // accounts cannot start a session, so drop the prompt — this mirrors the - // deferred post-login path, which clears the startup prompt for ZDR-blocked - // accounts. When not yet authenticated, stash it for `AuthComplete`. + // `NewSession`/`SendPrompt` path the welcome screen uses). When not yet + // authenticated, stash it for `AuthComplete`. if let Some(initial_prompt) = args.initial_prompt() { if !app.session_startup_allowed() { app.deferred_startup.prompt = Some(initial_prompt.to_string()); - } else if !app.is_zdr_blocked() { + } else { let effs = dispatch::dispatch_initial_prompt(&mut app, initial_prompt.to_string()); if process_effects(effs, &mut tasks, &mut app, &progress_tx) { return Ok(make_run_result(&app)); @@ -1322,10 +1258,7 @@ pub(crate) async fn run( // empty one so the user lands directly at the prompt. Unauthenticated / // ZDR-blocked startup stays on Welcome, where `crate::minimal::live` shows // a sign-in hint instead of a blank region. - if term_state.screen_mode.is_minimal() - && matches!(app.active_view, ActiveView::Welcome) - && !app.is_zdr_blocked() - { + if term_state.screen_mode.is_minimal() && matches!(app.active_view, ActiveView::Welcome) { if app.session_startup_allowed() { // Already authenticated + trusted: open the empty session now so the // user lands directly at the prompt. @@ -1456,15 +1389,6 @@ pub(crate) async fn run( roster_poll_at = Some(Instant::now()); } - // (Re-)arm the subscription watch on the dormant→wanted transition - // and after each fired tick. - if subscription_watch_at.is_none() - && app.subscription_watch_wanted() - && let Some(iv) = app.subscription_watch_interval() - { - subscription_watch_at = Some(Instant::now() + iv); - } - // Future that sleeps until the next animation tick, or waits forever if none. let animation_tick = async { match animation_tick_at { @@ -1511,27 +1435,6 @@ pub(crate) async fn run( } }; - let billing_poll = async { - match billing_poll_at { - Some(at) => sleep_until(at).await, - None => std::future::pending().await, - } - }; - - let gate_poll = async { - match gate_poll_at { - Some(at) => sleep_until(at).await, - None => std::future::pending().await, - } - }; - - let subscription_watch = async { - match subscription_watch_at { - Some(at) => sleep_until(at).await, - None => std::future::pending().await, - } - }; - let roster_poll = async { match roster_poll_at { Some(at) => sleep_until(at).await, @@ -1629,18 +1532,6 @@ pub(crate) async fn run( schedule_tick(&mut animation_tick_at, &app, tick_interval); resize_debounce_at = None; - // Schedule/clear poll timers. - if app.billing_poll_wanted && billing_poll_at.is_none() { - billing_poll_at = Some(Instant::now() + BILLING_POLL_INTERVAL); - } else if !app.billing_poll_wanted { - billing_poll_at = None; - } - if !app.has_access() && gate_poll_at.is_none() { - gate_poll_at = Some(Instant::now() + GATE_POLL_INTERVAL); - } else if app.has_access() { - gate_poll_at = None; - } - app.draw(terminal); last_draw_at = Instant::now(); draw_scheduled_at = None; @@ -1810,41 +1701,6 @@ pub(crate) async fn run( schedule_tick(&mut animation_tick_at, &app, tick_interval); } - _ = billing_poll => { - billing_poll_at = None; - if let ActiveView::Agent(id) = app.active_view { - let effs = vec![Effect::FetchBilling { - agent_id: id, - silent: true, - }]; - if process_effects(effs, &mut tasks, &mut app, &progress_tx) { - break; - } - } - if app.billing_poll_wanted { - billing_poll_at = Some(Instant::now() + BILLING_POLL_INTERVAL); - } - } - - _ = gate_poll => { - gate_poll_at = None; - let effs = vec![Effect::RefreshGate]; - if process_effects(effs, &mut tasks, &mut app, &progress_tx) { - break; - } - if !app.has_access() { - gate_poll_at = Some(Instant::now() + GATE_POLL_INTERVAL); - } - } - - _ = subscription_watch => { - subscription_watch_at = None; - let effs = app.fire_subscription_check("watch"); - if process_effects(effs, &mut tasks, &mut app, &progress_tx) { - break; - } - } - _ = roster_poll => { roster_poll_at = None; // Only poll while the dashboard is open. When it is not active @@ -2519,12 +2375,6 @@ async fn drain_and_process( { crate::clipboard::prewarm_image_probe(); } - // The user may have just subscribed in the browser and - // tabbed back. - let effs = app.fire_subscription_check("focus"); - if process_effects(effs, tasks, app, progress_tx) { - return true; - } // Restore Prompt on refocus: needs-input overlay always, else idle non-vim. match app.active_view { ActiveView::Agent(id) => { diff --git a/crates/codegen/kigi-tui/src/app/foreign_sessions.rs b/crates/codegen/kigi-tui/src/app/foreign_sessions.rs index 0afb6d4..2a3c186 100644 --- a/crates/codegen/kigi-tui/src/app/foreign_sessions.rs +++ b/crates/codegen/kigi-tui/src/app/foreign_sessions.rs @@ -34,7 +34,6 @@ impl AppView { && self.agents.is_empty() && self.next_agent_id == 0 && !self.chat_mode - && !self.is_zdr_blocked() && self.pending_update_version.is_none() } diff --git a/crates/codegen/kigi-tui/src/app/leader_cluster/mod.rs b/crates/codegen/kigi-tui/src/app/leader_cluster/mod.rs index b79ff50..a1a9066 100644 --- a/crates/codegen/kigi-tui/src/app/leader_cluster/mod.rs +++ b/crates/codegen/kigi-tui/src/app/leader_cluster/mod.rs @@ -317,7 +317,7 @@ impl PagerLeaderCluster { let env = vec![ crate::test_util::EnvVarGuard::set("KIGI_SHARE_DIR", kigi_home.path()), - crate::test_util::EnvVarGuard::set("KIGI_CLI_CHAT_PROXY_BASE_URL", server.url()), + crate::test_util::EnvVarGuard::set("KIGI_CODE_BASE_URL", server.url()), crate::test_util::EnvVarGuard::set("KIGI_XAI_API_BASE_URL", server.url()), crate::test_util::EnvVarGuard::set("XAI_API_KEY", "test-key-for-ci"), crate::test_util::EnvVarGuard::set("KIGI_TELEMETRY_ENABLED", "false"), diff --git a/crates/codegen/kigi-tui/src/app/mod.rs b/crates/codegen/kigi-tui/src/app/mod.rs index 7822e50..8f06a7f 100644 --- a/crates/codegen/kigi-tui/src/app/mod.rs +++ b/crates/codegen/kigi-tui/src/app/mod.rs @@ -41,7 +41,6 @@ pub mod session_startup; mod signal_handler; pub mod status_blocks; pub mod subagent; -pub mod subscription; mod turn_completion; mod xt_filter; pub(crate) use crate::terminal::kitty_flags_pushed; @@ -56,7 +55,6 @@ use crossterm::execute; use crossterm::terminal::{ self, Clear, ClearType, EnterAlternateScreen, LeaveAlternateScreen, SetTitle, }; -pub(crate) use dispatch::{FREE_USAGE_USER_MESSAGE, acp_error_is_free_usage_exhausted}; pub use foreign_sessions::ForeignScanCoordinator; pub(crate) use foreign_sessions::{ badge_for_picker_source, foreign_tool_display_label, is_foreign_picker_source, @@ -272,7 +270,6 @@ pub fn resolve_use_leader( leader_flag: bool, no_leader_flag: bool, raw_config: &toml::Value, - _remote_settings: Option<&kigi_shell::util::config::RemoteSettings>, eligible: bool, ) -> (bool, Option<&'static str>) { if no_leader_flag { @@ -287,35 +284,8 @@ pub fn resolve_use_leader( if let Some(v) = config::use_leader_from_toml_opt(raw_config) { return (v, (!v).then_some("config")); } - #[cfg(feature = "release-dist")] - if let Some(remote_val) = _remote_settings.and_then(|s| s.leader_mode) { - return (remote_val, (!remote_val).then_some("remote")); - } (false, None) } -/// Join early prefetch to get remote settings (with timeout). -/// -/// Remote settings come from the product settings API and contain `leader_mode`, -/// feature gates, etc. Waits up to 2 s for the background thread. -pub fn join_early_prefetch( - handle: Option, -) -> Option { - let handle = handle?; - if handle.is_finished() { - return match handle.join() { - Ok(r) => r.settings, - Err(_) => None, - }; - } - let (tx, rx) = std::sync::mpsc::channel(); - std::thread::spawn(move || { - let _ = tx.send(handle.join()); - }); - match rx.recv_timeout(std::time::Duration::from_secs(2)) { - Ok(Ok(r)) => r.settings, - _ => None, - } -} /// First non-blank of CLI > env > config (precedence + blank-skip). `None` → /// nothing set; `acp::initialize` canonicalizes and applies the default. fn resolve_hunk_tracker_mode( @@ -360,27 +330,21 @@ pub async fn run( } }; let refreshed_auth = kigi_shell::auth::try_ensure_fresh_auth(&kimi_code_config).await; - let early_prefetch = kigi_shell::agent::models::start_early_prefetch_with_auth(refreshed_auth); + // Fire-and-forget model-catalog warmup; nothing joins the handle now that + // the xAI settings fetch it used to carry is gone. + drop(kigi_shell::agent::models::start_early_prefetch_with_auth( + refreshed_auth, + )); kigi_shell::agent::mvp_agent::warm_async_http_client(); tokio::task::spawn_blocking(|| {}); if let Ok(cwd) = std::env::current_dir() { crate::git_info::populate_from_cwd_async(cwd); } - let remote_settings = join_early_prefetch(early_prefetch); - kigi_shell::util::config::cache_remote_auto_mode( - remote_settings.as_ref().and_then(|s| s.auto_mode.clone()), - ); - kigi_shell::util::config::set_remote_campaigns_from_settings(remote_settings.as_ref()); let raw_config = kigi_shell::config::load_effective_config() .map_err(|e| anyhow::anyhow!("Failed to load config: {e}"))?; let prefetch_elapsed = startup_start.elapsed(); - let (use_leader, policy_disable_reason) = resolve_use_leader( - args.leader, - args.no_leader, - &raw_config, - remote_settings.as_ref(), - true, - ); + let (use_leader, policy_disable_reason) = + resolve_use_leader(args.leader, args.no_leader, &raw_config, true); tracing::info!( use_leader, ?policy_disable_reason, @@ -473,9 +437,7 @@ pub async fn run( env_hunk_tracker_mode.as_deref(), config_hunk_tracker_mode, ); - let remote_permission_mode = remote_settings - .as_ref() - .and_then(|s| s.permission_mode.as_deref()); + let remote_permission_mode = None; let launch_yolo = kigi_shell::util::config::effective_yolo_for_launch( args.yolo, args.permission_mode_flag.as_deref(), @@ -500,7 +462,7 @@ pub async fn run( fs_read: args.fs_read, fs_write: args.fs_write, installer: args.installer.clone(), - remote_settings: remote_settings.clone(), + remote_settings: None, system_prompt_override: args.system_prompt_override.clone(), rules: args.rules.clone(), reasoning_effort_override: args @@ -613,7 +575,7 @@ pub async fn run( &mut config_watcher, &effective_args, session_cwd, - remote_settings, + None, term_state, materialized, bg_update_rx, @@ -1305,47 +1267,47 @@ mod tests { #[test] fn no_leader_flag_wins_over_leader_flag_and_config() { let cfg = config_with_leader(true); - let (use_leader, reason) = resolve_use_leader(true, true, &cfg, None, true); + let (use_leader, reason) = resolve_use_leader(true, true, &cfg, true); assert!(!use_leader); assert_eq!(reason, None); } #[test] fn leader_flag_enables() { - let (use_leader, reason) = resolve_use_leader(true, false, &empty_config(), None, true); + let (use_leader, reason) = resolve_use_leader(true, false, &empty_config(), true); assert!(use_leader); assert_eq!(reason, None); } #[test] fn not_eligible_returns_false() { let cfg = config_with_leader(true); - let (use_leader, reason) = resolve_use_leader(false, false, &cfg, None, false); + let (use_leader, reason) = resolve_use_leader(false, false, &cfg, false); assert!(!use_leader); assert_eq!(reason, None); } #[test] fn config_toml_enables() { let cfg = config_with_leader(true); - let (use_leader, reason) = resolve_use_leader(false, false, &cfg, None, true); + let (use_leader, reason) = resolve_use_leader(false, false, &cfg, true); assert!(use_leader); assert_eq!(reason, None); } #[test] fn config_toml_disables() { let cfg = config_with_leader(false); - let (use_leader, reason) = resolve_use_leader(false, false, &cfg, None, true); + let (use_leader, reason) = resolve_use_leader(false, false, &cfg, true); assert!(!use_leader); assert_eq!(reason, Some("config")); } #[test] fn default_is_false() { - let (use_leader, reason) = resolve_use_leader(false, false, &empty_config(), None, true); + let (use_leader, reason) = resolve_use_leader(false, false, &empty_config(), true); assert!(!use_leader); assert_eq!(reason, None); } #[test] fn cli_flag_overrides_config() { let cfg = config_with_leader(false); - let (use_leader, reason) = resolve_use_leader(true, false, &cfg, None, true); + let (use_leader, reason) = resolve_use_leader(true, false, &cfg, true); assert!(use_leader); assert_eq!(reason, None); } @@ -1379,7 +1341,7 @@ mod tests { #[test] fn no_leader_flag_overrides_config_for_tui_fallback() { let cfg = config_with_leader(true); - let (use_leader, reason) = resolve_use_leader(false, true, &cfg, None, true); + let (use_leader, reason) = resolve_use_leader(false, true, &cfg, true); assert!(!use_leader); assert_eq!(reason, None); } @@ -1405,59 +1367,11 @@ mod tests { assert!(try_parse_pager(&["grok-pager", "agent"]).is_err()); } #[test] - fn remote_settings_none_falls_through_to_default() { - let (use_leader, reason) = resolve_use_leader(false, false, &empty_config(), None, true); + fn leader_defaults_off_without_config() { + let (use_leader, reason) = resolve_use_leader(false, false, &empty_config(), true); assert!(!use_leader); assert_eq!(reason, None); } - #[cfg(feature = "release-dist")] - #[test] - fn remote_settings_leader_mode_true_enables_leader() { - let rs = kigi_shell::util::config::RemoteSettings { - leader_mode: Some(true), - ..Default::default() - }; - let (use_leader, reason) = - resolve_use_leader(false, false, &empty_config(), Some(&rs), true); - assert!(use_leader); - assert_eq!(reason, None); - } - #[cfg(feature = "release-dist")] - #[test] - fn remote_settings_leader_mode_false_disables_leader() { - let rs = kigi_shell::util::config::RemoteSettings { - leader_mode: Some(false), - ..Default::default() - }; - let (use_leader, reason) = - resolve_use_leader(false, false, &empty_config(), Some(&rs), true); - assert!(!use_leader); - assert_eq!(reason, Some("remote")); - } - #[cfg(feature = "release-dist")] - #[test] - fn remote_settings_unknown_leader_mode_is_not_policy_disable() { - let rs = kigi_shell::util::config::RemoteSettings { - leader_mode: None, - ..Default::default() - }; - let (use_leader, reason) = - resolve_use_leader(false, false, &empty_config(), Some(&rs), true); - assert!(!use_leader); - assert_eq!(reason, None); - } - #[cfg(feature = "release-dist")] - #[test] - fn config_toml_overrides_remote_settings() { - let rs = kigi_shell::util::config::RemoteSettings { - leader_mode: Some(true), - ..Default::default() - }; - let cfg = config_with_leader(false); - let (use_leader, reason) = resolve_use_leader(false, false, &cfg, Some(&rs), true); - assert!(!use_leader); - assert_eq!(reason, Some("config")); - } #[test] fn cli_resume_parses_session_id() { let args = try_parse_pager(&["grok-pager", "--resume", "abc-123"]).unwrap(); diff --git a/crates/codegen/kigi-tui/src/app/modals.rs b/crates/codegen/kigi-tui/src/app/modals.rs index b2f709e..8cb728b 100644 --- a/crates/codegen/kigi-tui/src/app/modals.rs +++ b/crates/codegen/kigi-tui/src/app/modals.rs @@ -635,8 +635,7 @@ impl AgentView { entries: _, state, .. } => { // Build filtered entries for count and non-selectable indices. - let filtered = - crate::views::modal::filter_palette_entries(&state.query, self.sharing_enabled); + let filtered = crate::views::modal::filter_palette_entries(&state.query); let non_sel: Vec = filtered .iter() .map(|e| matches!(e.command, PaletteCommand::SectionHeader(_))) @@ -840,14 +839,10 @@ impl AgentView { } PickerOutcome::Changed => { // Re-filter entries based on updated query. - let sharing_enabled = self.sharing_enabled; if let Some(ActiveModal::CommandPalette { entries, state, .. }) = self.active_modal.as_mut() { - *entries = crate::views::modal::filter_palette_entries( - &state.query, - sharing_enabled, - ); + *entries = crate::views::modal::filter_palette_entries(&state.query); state.selected = state.selected.min(entries.len().saturating_sub(1)); } InputOutcome::Changed @@ -1623,7 +1618,7 @@ impl AgentView { } = active_modal { // Command palette: ModalWindow chrome + picker content. - let filtered = modal::filter_palette_entries(&state.query, self.sharing_enabled); + let filtered = modal::filter_palette_entries(&state.query); let non_sel: Vec = filtered .iter() .map(|e| matches!(e.command, modal::PaletteCommand::SectionHeader(_))) @@ -2614,7 +2609,7 @@ mod command_palette_vim_input_tests { // INPUT mode (`input_active`) over the full palette entries. fn open_command_palette(agent: &mut AgentView) { agent.active_modal = Some(ActiveModal::CommandPalette { - entries: crate::views::modal::default_palette_entries(agent.sharing_enabled), + entries: crate::views::modal::default_palette_entries(), state: PickerState::input_active(), window: crate::views::modal_window::ModalWindowState::new(), }); diff --git a/crates/codegen/kigi-tui/src/app/mouse.rs b/crates/codegen/kigi-tui/src/app/mouse.rs index 24d2e3e..dcaa095 100644 --- a/crates/codegen/kigi-tui/src/app/mouse.rs +++ b/crates/codegen/kigi-tui/src/app/mouse.rs @@ -698,31 +698,6 @@ impl AgentView { .scrollback .entry_index_at_screen_row(click_row, self.pane_areas.scrollback); if let Some(idx) = hit_idx { - let credit_click = self.scrollback.entry(idx).and_then(|entry| { - if let crate::scrollback::block::RenderBlock::CreditLimit(ref blk) = - entry.block - { - Some(blk.url.clone()) - } else { - None - } - }); - if let Some(url) = credit_click - && let Some((area, _, _)) = self - .scrollback - .entry_screen_area(idx, self.pane_areas.scrollback) - { - let url_row = area.y + area.height.saturating_sub(2); - if click_row >= url_row { - self.scrollback.set_selected(Some(idx)); - crate::app::link_opener::open_url_if_safe( - &url, - crate::terminal::hyperlinks::SchemeFilter::Standard, - ); - self.last_click = None; - return InputOutcome::Changed; - } - } let selectable = self .scrollback .get(idx) @@ -878,7 +853,6 @@ impl AgentView { .set_hovered_follow_up_chip(self.follow_up_chip_at(mouse.column, mouse.row)); changed |= self.hit_badge.update_hover(mouse.column, mouse.row); changed |= self.hit_context.update_hover(mouse.column, mouse.row); - changed |= self.hit_credits.update_hover(mouse.column, mouse.row); changed |= self.hit_todo_close.update_hover(mouse.column, mouse.row); changed |= self.hit_queue_close.update_hover(mouse.column, mouse.row); changed |= self.hit_queue_badge.update_hover(mouse.column, mouse.row); diff --git a/crates/codegen/kigi-tui/src/app/queue_edit.rs b/crates/codegen/kigi-tui/src/app/queue_edit.rs index a80da47..674f8f6 100644 --- a/crates/codegen/kigi-tui/src/app/queue_edit.rs +++ b/crates/codegen/kigi-tui/src/app/queue_edit.rs @@ -1150,7 +1150,7 @@ mod tests { // early-return) is what leaves the palette alone. agent.prompt_mode = editing_lone_local(); agent.active_modal = Some(ActiveModal::CommandPalette { - entries: crate::views::modal::default_palette_entries(agent.sharing_enabled), + entries: crate::views::modal::default_palette_entries(), state: crate::views::picker::PickerState::input_active(), window: crate::views::modal_window::ModalWindowState::new(), }); diff --git a/crates/codegen/kigi-tui/src/app/subagent.rs b/crates/codegen/kigi-tui/src/app/subagent.rs index 1039f82..dd8f604 100644 --- a/crates/codegen/kigi-tui/src/app/subagent.rs +++ b/crates/codegen/kigi-tui/src/app/subagent.rs @@ -561,8 +561,6 @@ mod tests { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, diff --git a/crates/codegen/kigi-tui/src/app/subscription.rs b/crates/codegen/kigi-tui/src/app/subscription.rs deleted file mode 100644 index 1a171cb..0000000 --- a/crates/codegen/kigi-tui/src/app/subscription.rs +++ /dev/null @@ -1,487 +0,0 @@ -//! Free→paid subscription detection and gate imposition/lift. -//! -//! All gate transitions go through [`AppView::impose_gate`] / -//! [`AppView::lift_gate`] so the defer-vs-show decision and the lift -//! bookkeeping (focus, telemetry, JWT-refresh check) live in one place. -//! -//! Design constraints that are not obvious from the code: -//! - Gates arriving from cached auth meta, prefetched settings, or settings -//! pushes can be stale: the user may have subscribed since the snapshot -//! was computed. Painting such a gate directly flashes a paywall at a -//! paying user, so it is held in `pending_gate_verification` while a live -//! check runs. On check failure or timeout we err on blocking. -//! - Timer effects have no cancellation, so verifications are stamped with -//! `gate_verify_gen`; results and timeouts from superseded deferrals are -//! ignored by generation mismatch. - -use super::actions::Effect; -use super::app_view::{AppView, AuthState}; - -/// Default watch cadence. Overridable via the remote settings -/// `grok_build_settings.subscription_watch_interval_secs` field. -pub(crate) const SUBSCRIPTION_WATCH_INTERVAL: std::time::Duration = - std::time::Duration::from_secs(60); - -/// Floor for the server-supplied cadence: a fat-fingered remote settings value -/// must not turn the fleet into a hot-poller. `0` means "disabled" and is -/// special-cased before this clamp. -pub(crate) const SUBSCRIPTION_WATCH_MIN_INTERVAL_SECS: u64 = 30; - -/// Floor for the `KIGI_SUBSCRIPTION_WATCH_INTERVAL_SECS` env override -/// (test seam / power user — deliberately below the server floor). -const SUBSCRIPTION_WATCH_ENV_MIN_SECS: u64 = 1; - -/// Cap on the spacing between watch/focus-triggered checks. -pub(crate) const SUBSCRIPTION_CHECK_DEBOUNCE: std::time::Duration = - std::time::Duration::from_secs(30); - -/// How long a deferred gate is held before being shown anyway. This is a -/// safety net for a hung ACP round-trip only — a completed check (even a -/// failed one) resolves the deferral immediately. Generous on purpose: the -/// check can chain a `/user` fetch, a JWT refresh, and a settings re-fetch; -/// 5s was observed timing out in CI under full-suite contention. -pub(crate) const GATE_VERIFY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); - -impl AppView { - /// Consumer xAI session auth: not an API key, not an enterprise team. - /// Subscription gates and the watch only apply to these sessions. - fn is_consumer_session(&self) -> bool { - matches!(self.auth_state, AuthState::Done) - && !self.is_api_key_auth - && self.team_name.is_none() - } - - /// `None` tier counts as potentially-free so detection works before the - /// first auth meta lands. Not a confirmed-free signal. - pub fn may_be_free_tier(&self) -> bool { - match self.subscription_tier.as_deref() { - Some(t) => t.trim().eq_ignore_ascii_case("free"), - None => true, - } - } - - /// Effective watch cadence; `None` = disabled. Precedence: env override - /// (`0` disables), server override (`0` disables, floor-clamped), - /// default. - pub fn subscription_watch_interval(&self) -> Option { - if let Ok(v) = std::env::var("KIGI_SUBSCRIPTION_WATCH_INTERVAL_SECS") - && let Ok(secs) = v.trim().parse::() - { - return match secs { - 0 => None, - s => Some(std::time::Duration::from_secs( - s.max(SUBSCRIPTION_WATCH_ENV_MIN_SECS), - )), - }; - } - match self.subscription_watch_interval_secs { - Some(0) => None, - Some(secs) => Some(std::time::Duration::from_secs( - secs.max(SUBSCRIPTION_WATCH_MIN_INTERVAL_SECS), - )), - None => Some(SUBSCRIPTION_WATCH_INTERVAL), - } - } - - /// Whether the watch (and the refocus check) should run: enabled, - /// consumer session, and gated or possibly-free. - pub fn subscription_watch_wanted(&self) -> bool { - self.subscription_watch_interval().is_some() - && self.is_consumer_session() - && (self.gate.is_some() || self.may_be_free_tier()) - } - - /// Half the effective interval, capped at [`SUBSCRIPTION_CHECK_DEBOUNCE`] - /// — scaling keeps the debounce from swallowing watch ticks when the - /// cadence is tightened. - fn subscription_check_allowed(&self) -> bool { - let debounce = self - .subscription_watch_interval() - .map(|iv| (iv / 2).min(SUBSCRIPTION_CHECK_DEBOUNCE)) - .unwrap_or(SUBSCRIPTION_CHECK_DEBOUNCE); - self.last_subscription_check_at - .is_none_or(|t| t.elapsed() >= debounce) - } - - fn note_subscription_check(&mut self) { - self.last_subscription_check_at = Some(std::time::Instant::now()); - } - - /// Single guard-and-fire for the watch tick and the terminal-refocus - /// trigger. Empty when unwanted or debounced. The 5s paywall chain - /// deliberately bypasses this. `trigger` tags the unified-log entry - /// (`"watch"` / `"focus"`) so the check cadence is reconstructable - /// from logs. - #[must_use] - pub fn fire_subscription_check(&mut self, trigger: &'static str) -> Vec { - if self.subscription_watch_wanted() && self.subscription_check_allowed() { - self.note_subscription_check(); - crate::unified_log::info( - "subscription.check.fired", - None, - Some(serde_json::json!({ - "trigger": trigger, - "interval_secs": self - .subscription_watch_interval() - .map(|iv| iv.as_secs()), - "gated": self.gate.is_some(), - "tier": self.subscription_tier, - })), - ); - vec![Effect::CheckSubscription { verify: None }] - } else { - vec![] - } - } - - /// Chokepoint for showing a gate. Already gated → update the copy. - /// Consumer session with access → defer for live verification (the gate - /// source may be stale). Otherwise → show directly. - #[must_use] - pub fn impose_gate(&mut self, gate: kigi_shell::auth::GateInfo) -> Vec { - if self.gate.is_some() { - self.gate = Some(gate); - return vec![]; - } - if self.is_consumer_session() { - return self.defer_gate_for_verification(gate); - } - crate::unified_log::info( - "subscription.gate.imposed", - None, - Some(serde_json::json!({ "deferred": false })), - ); - self.gate = Some(gate); - vec![] - } - - /// Chokepoint for a settings-confirmed gate lift. Clears the visible - /// gate and any pending deferral; when either existed, runs the lift - /// bookkeeping and returns the JWT-refresh check (the tier claim is - /// baked into the JWT, so the shell must re-mint it). - #[must_use] - pub fn lift_gate(&mut self) -> Vec { - let was_blocked = self.gate.is_some() || self.pending_gate_verification.is_some(); - self.gate = None; - self.pending_gate_verification = None; - if !was_blocked { - return vec![]; - } - self.welcome_prompt_focused = true; - self.paywall_check_started = None; - crate::unified_log::info( - "subscription.gate.lifted", - None, - Some(serde_json::json!({ "tier": self.subscription_tier })), - ); - vec![Effect::CheckSubscription { verify: None }] - } - - /// Hold `gate` out of `self.gate` while a generation-stamped live check - /// verifies it. Resolution: authoritative meta via `apply_auth_meta` - /// (drops the deferral), or promotion on same-generation check failure / - /// timeout via [`Self::promote_deferred_gate`]. - #[must_use] - fn defer_gate_for_verification(&mut self, gate: kigi_shell::auth::GateInfo) -> Vec { - self.pending_gate_verification = Some(gate); - self.gate_verify_gen = self.gate_verify_gen.wrapping_add(1); - self.note_subscription_check(); - crate::unified_log::info( - "subscription.gate.deferred", - None, - Some(serde_json::json!({ - "generation": self.gate_verify_gen, - "tier": self.subscription_tier, - })), - ); - vec![ - Effect::CheckSubscription { - verify: Some(self.gate_verify_gen), - }, - Effect::ScheduleGateVerifyTimeout { - generation: self.gate_verify_gen, - }, - ] - } - - /// Show a deferred gate (err on blocking) — no-op unless `generation` - /// is the current verification and nothing resolved it meanwhile. - /// `reason` tags the unified-log entry (`"check_failed"` / - /// `"verify_timeout"`). - pub(crate) fn promote_deferred_gate(&mut self, generation: u64, reason: &'static str) { - if generation == self.gate_verify_gen - && let Some(gate) = self.pending_gate_verification.take() - && self.gate.is_none() - { - // Warn: the verification did not confirm access, so the user is - // now blocked. If this is wrong (paying user paywalled), this - // entry plus the preceding check.fired/check.complete lines - // show which path failed. - crate::unified_log::warn( - "subscription.gate.promoted", - None, - Some(serde_json::json!({ - "generation": generation, - "reason": reason, - "tier": self.subscription_tier, - })), - ); - self.gate = Some(gate); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::app::app_view::tests::test_app; - - fn watch_gate() -> kigi_shell::auth::GateInfo { - kigi_shell::auth::GateInfo { - message: "Subscribe".into(), - url: None, - label: None, - } - } - - #[test] - fn may_be_free_tier_matrix() { - let mut app = test_app(); - app.subscription_tier = None; - assert!(app.may_be_free_tier(), "unknown tier is potentially free"); - app.subscription_tier = Some("Free".into()); - assert!(app.may_be_free_tier()); - app.subscription_tier = Some(" FREE ".into()); - assert!(app.may_be_free_tier(), "case/whitespace-insensitive"); - app.subscription_tier = Some("SuperGrok Heavy".into()); - assert!(!app.may_be_free_tier()); - app.subscription_tier = Some("X Premium".into()); - assert!(!app.may_be_free_tier()); - } - - #[test] - fn subscription_watch_wanted_matrix() { - let mut app = test_app(); // AuthState::Done, consumer, tier unknown - assert!( - app.subscription_watch_wanted(), - "unknown-tier consumer session watches" - ); - - app.subscription_tier = Some("Free".into()); - assert!(app.subscription_watch_wanted(), "free tier watches"); - - app.subscription_tier = Some("SuperGrok".into()); - assert!(!app.subscription_watch_wanted(), "paid tier is dormant"); - - // Gated — watches regardless of the (stale) tier string. - app.gate = Some(watch_gate()); - assert!(app.subscription_watch_wanted(), "gated session watches"); - app.gate = None; - - app.subscription_tier = Some("Free".into()); - app.is_api_key_auth = true; - assert!( - !app.subscription_watch_wanted(), - "API-key auth never watches" - ); - app.is_api_key_auth = false; - app.team_name = Some("Acme Corp".into()); - assert!( - !app.subscription_watch_wanted(), - "team session never watches" - ); - app.team_name = None; - - app.auth_state = AuthState::Pending { error: None }; - assert!(!app.subscription_watch_wanted(), "pre-auth never watches"); - } - - #[test] - fn subscription_watch_interval_override_clamp_and_disable() { - let mut app = test_app(); - assert_eq!( - app.subscription_watch_interval(), - Some(SUBSCRIPTION_WATCH_INTERVAL) - ); - app.subscription_watch_interval_secs = Some(120); - assert_eq!( - app.subscription_watch_interval(), - Some(std::time::Duration::from_secs(120)) - ); - app.subscription_watch_interval_secs = Some(1); - assert_eq!( - app.subscription_watch_interval(), - Some(std::time::Duration::from_secs( - SUBSCRIPTION_WATCH_MIN_INTERVAL_SECS - )), - "sub-floor values are clamped" - ); - app.subscription_watch_interval_secs = Some(0); - assert_eq!(app.subscription_watch_interval(), None); - app.subscription_tier = Some("Free".into()); - assert!( - !app.subscription_watch_wanted(), - "interval 0 must disable the watch even on the free tier" - ); - } - - #[test] - fn subscription_check_debounce() { - let mut app = test_app(); - assert!(app.subscription_check_allowed(), "no prior check — allowed"); - - app.note_subscription_check(); - assert!( - !app.subscription_check_allowed(), - "right after a check — debounced" - ); - - app.last_subscription_check_at = - Some(std::time::Instant::now() - SUBSCRIPTION_CHECK_DEBOUNCE); - assert!(app.subscription_check_allowed()); - } - - #[test] - fn fire_subscription_check_guards_and_debounces() { - let mut app = test_app(); - let effs = app.fire_subscription_check("watch"); - assert!(matches!( - effs.as_slice(), - [Effect::CheckSubscription { verify: None }] - )); - assert!( - app.fire_subscription_check("watch").is_empty(), - "second fire inside the debounce window must be empty" - ); - - let mut paid = test_app(); - paid.subscription_tier = Some("SuperGrok".into()); - assert!( - paid.fire_subscription_check("watch").is_empty(), - "paid tier never fires" - ); - } - - #[test] - fn impose_gate_defers_for_consumer_session() { - let mut app = test_app(); - let effs = app.impose_gate(watch_gate()); - - assert!( - app.has_access(), - "deferred gate must not render as a paywall" - ); - assert!(app.pending_gate_verification.is_some()); - assert!( - !app.subscription_check_allowed(), - "the deferral's own check counts for the debounce" - ); - assert!(matches!( - effs.as_slice(), - [ - Effect::CheckSubscription { - verify: Some(check_gen) - }, - Effect::ScheduleGateVerifyTimeout { - generation: timeout_gen - } - ] if *check_gen == app.gate_verify_gen && *timeout_gen == app.gate_verify_gen - )); - } - - #[test] - fn impose_gate_direct_for_non_consumer_and_already_gated() { - // Team session: no live verification possible — show directly. - let mut app = test_app(); - app.team_name = Some("Acme Corp".into()); - assert!(app.impose_gate(watch_gate()).is_empty()); - assert!(!app.has_access()); - assert!(app.pending_gate_verification.is_none()); - - // Already gated: update the copy only. - let mut gated = test_app(); - gated.gate = Some(watch_gate()); - let new_copy = kigi_shell::auth::GateInfo { - message: "New copy".into(), - url: None, - label: None, - }; - assert!(gated.impose_gate(new_copy).is_empty()); - assert_eq!(gated.gate.as_ref().unwrap().message, "New copy"); - } - - #[test] - fn impose_gate_bumps_generation_each_time() { - let mut app = test_app(); - let _ = app.impose_gate(watch_gate()); - let first = app.gate_verify_gen; - app.pending_gate_verification = None; // simulate resolution - let _ = app.impose_gate(watch_gate()); - assert_eq!(app.gate_verify_gen, first + 1, "each deferral re-stamps"); - } - - #[test] - fn lift_gate_runs_bookkeeping_once() { - let mut app = test_app(); - app.gate = Some(watch_gate()); - app.paywall_check_started = Some(std::time::Instant::now()); - - let effs = app.lift_gate(); - assert!(app.has_access()); - assert!(app.welcome_prompt_focused); - assert!(app.paywall_check_started.is_none()); - assert!(matches!( - effs.as_slice(), - [Effect::CheckSubscription { verify: None }] - )); - - assert!( - app.lift_gate().is_empty(), - "lift without a gate or deferral is a no-op" - ); - } - - #[test] - fn lift_gate_counts_pending_deferral_as_blocked() { - let mut app = test_app(); - let _ = app.impose_gate(watch_gate()); - - let effs = app.lift_gate(); - assert!(app.pending_gate_verification.is_none()); - assert!( - matches!( - effs.as_slice(), - [Effect::CheckSubscription { verify: None }] - ), - "a confirmed lift of a pending gate must still refresh the JWT" - ); - } - - #[test] - fn promote_deferred_gate_is_generation_scoped() { - let mut app = test_app(); - let _ = app.impose_gate(watch_gate()); - let stale_gen = app.gate_verify_gen; - let _ = app.impose_gate(watch_gate()); - - app.promote_deferred_gate(stale_gen, "verify_timeout"); - assert!( - app.has_access(), - "stale generation must not promote the newer deferral" - ); - - app.promote_deferred_gate(app.gate_verify_gen, "verify_timeout"); - assert!(!app.has_access(), "current generation promotes"); - } - - #[test] - fn apply_auth_meta_drops_pending_gate_verification() { - let mut app = test_app(); - let _effs = app.impose_gate(watch_gate()); - - app.apply_auth_meta(&kigi_shell::auth::AuthMeta::default()); - - assert!(app.pending_gate_verification.is_none()); - assert!(app.has_access()); - } -} diff --git a/crates/codegen/kigi-tui/src/client_identity.rs b/crates/codegen/kigi-tui/src/client_identity.rs index d989fbb..d12a845 100644 --- a/crates/codegen/kigi-tui/src/client_identity.rs +++ b/crates/codegen/kigi-tui/src/client_identity.rs @@ -1,4 +1,4 @@ pub const PAGER_CLIENT_TYPE: &str = "grok-pager"; -pub const HEADLESS_CLIENT_TYPE: &str = "grok-shell"; +pub const HEADLESS_CLIENT_TYPE: &str = "kigi"; pub const PAGER_CLIENT_VERSION: &str = kigi_version::VERSION; diff --git a/crates/codegen/kigi-tui/src/headless.rs b/crates/codegen/kigi-tui/src/headless.rs index 11eb331..1018c4d 100644 --- a/crates/codegen/kigi-tui/src/headless.rs +++ b/crates/codegen/kigi-tui/src/headless.rs @@ -838,7 +838,6 @@ pub async fn run_single_turn( ) -> Result<()> { // Stamp proxy requests as headless before the agent spawns and issues // its first request (auth enrichment, model list, etc.). - kigi_shell::http::set_process_client_mode_headless(); let cwd = match options.cwd { None => std::env::current_dir()?, @@ -1318,13 +1317,7 @@ pub async fn run_single_turn( } Some(Err(err)) => { let msg = if i32::from(err.code) == RATE_LIMITED_ERROR_CODE { - // The -32003 data is the flattened server message; a - // free-usage 429 carries the well-known code inline there. - if crate::app::acp_error_is_free_usage_exhausted(&err) { - crate::app::FREE_USAGE_USER_MESSAGE.to_string() - } else { - rate_limited_user_message(is_api_key_auth).to_string() - } + rate_limited_user_message(is_api_key_auth).to_string() } else { err.to_string() }; diff --git a/crates/codegen/kigi-tui/src/lib.rs b/crates/codegen/kigi-tui/src/lib.rs index dc85170..a91701a 100644 --- a/crates/codegen/kigi-tui/src/lib.rs +++ b/crates/codegen/kigi-tui/src/lib.rs @@ -47,7 +47,6 @@ pub mod scrollback; pub mod search; pub mod sessions_cmd; pub mod settings; -pub mod share_cmd; pub mod slash; pub mod startup; pub mod tips; diff --git a/crates/codegen/kigi-tui/src/scrollback/block.rs b/crates/codegen/kigi-tui/src/scrollback/block.rs index 1e55659..d11edc3 100644 --- a/crates/codegen/kigi-tui/src/scrollback/block.rs +++ b/crates/codegen/kigi-tui/src/scrollback/block.rs @@ -13,11 +13,10 @@ use crate::prompt_images::{InlineMediaInfo, ScrollbackImageRef, ScrollbackVideoR use super::blocks::mermaid_content::DiagramAffordance; use super::blocks::{ - AgentMessageBlock, BgTaskBlock, BtwBlock, ContextInfoBlock, CreditLimitBlock, - EditToolCallBlock, ExecuteToolCallBlock, LineRange, ListDirToolCallBlock, OtherToolCallBlock, - ReadToolCallBlock, SearchFileMatch, SearchToolCallBlock, SessionEvent, SessionEventBlock, - SubagentBlock, SubagentBlockKind, SystemMessageBlock, ThinkingBlock, ToolCallBlock, - UserPromptBlock, + AgentMessageBlock, BgTaskBlock, BtwBlock, ContextInfoBlock, EditToolCallBlock, + ExecuteToolCallBlock, LineRange, ListDirToolCallBlock, OtherToolCallBlock, ReadToolCallBlock, + SearchFileMatch, SearchToolCallBlock, SessionEvent, SessionEventBlock, SubagentBlock, + SubagentBlockKind, SystemMessageBlock, ThinkingBlock, ToolCallBlock, UserPromptBlock, }; use super::types::{ AccentStyle, BlockBackground, BlockContext, BlockOutput, DisplayMode, RenderedBlockOutput, @@ -383,8 +382,6 @@ pub enum RenderBlock { Btw(BtwBlock), /// `/context` snapshot with categorical bar + breakdown. ContextInfo(ContextInfoBlock), - /// Credit-limit card for max-tier users (red accent, single action). - CreditLimit(CreditLimitBlock), } /// Delegate a method call to the inner block variant. @@ -402,7 +399,6 @@ macro_rules! delegate_block { RenderBlock::Subagent(b) => b.$method($($arg),*), RenderBlock::Btw(b) => b.$method($($arg),*), RenderBlock::ContextInfo(b) => b.$method($($arg),*), - RenderBlock::CreditLimit(b) => b.$method($($arg),*), } }; } @@ -778,15 +774,6 @@ impl RenderBlock { RenderBlock::SessionEvent(SessionEventBlock::new(event)) } - /// Create a credit-limit card (inline scrollback block for max-tier users). - pub fn credit_limit_card( - heading: impl Into, - action: crate::scrollback::blocks::CreditLimitCardAction, - url: impl Into, - ) -> Self { - RenderBlock::CreditLimit(CreditLimitBlock::new(heading, action, url)) - } - /// Create a "Task started" background task block. pub fn bg_task(command: impl Into, task_id: impl Into) -> Self { RenderBlock::BgTask(BgTaskBlock::started(command, task_id)) @@ -912,11 +899,6 @@ impl RenderBlock { matches!(self, RenderBlock::AgentMessage(_)) } - /// Check if this block is a CreditLimit card. - pub fn is_credit_limit(&self) -> bool { - matches!(self, RenderBlock::CreditLimit(_)) - } - /// Check if this block is a plan mode tool call (enter or exit). /// /// Exact-matches the canonical tool-name set rather than substring-matching @@ -1014,10 +996,9 @@ impl RenderBlock { None } } - RenderBlock::System(_) - | RenderBlock::SessionEvent(_) - | RenderBlock::ContextInfo(_) - | RenderBlock::CreditLimit(_) => None, + RenderBlock::System(_) | RenderBlock::SessionEvent(_) | RenderBlock::ContextInfo(_) => { + None + } RenderBlock::Btw(_) => Some(theme.accent_plan), RenderBlock::Stub(block) => Some(block.accent_color), } @@ -1161,9 +1142,6 @@ impl RenderBlock { Some(b.content().rendered_plain_text()), ]), RenderBlock::ContextInfo(b) => join_searchable([Some(b.model.clone())]), - RenderBlock::CreditLimit(b) => { - join_searchable([Some(b.heading.clone()), Some(b.url.clone())]) - } RenderBlock::ToolCall(tc) => tc.searchable_text(), } } @@ -1566,18 +1544,6 @@ mod searchable_text_tests { assert_eq!(block.searchable_text().as_deref(), Some("grok-4.5")); } - #[test] - fn credit_limit_indexes_heading_and_url() { - let block = RenderBlock::credit_limit_card( - "credit limit reached", - crate::scrollback::blocks::CreditLimitCardAction::EnablePayg, - "https://grok.com?_s=usage", - ); - let text = block.searchable_text().expect("credit limit text"); - assert!(text.contains("credit limit reached"), "got: {text:?}"); - assert!(text.contains("https://grok.com?_s=usage"), "got: {text:?}"); - } - #[test] fn search_tool_indexes_pattern_and_match_line() { let block = RenderBlock::search( diff --git a/crates/codegen/kigi-tui/src/scrollback/blocks/credit_limit.rs b/crates/codegen/kigi-tui/src/scrollback/blocks/credit_limit.rs deleted file mode 100644 index 9ffc4e4..0000000 --- a/crates/codegen/kigi-tui/src/scrollback/blocks/credit_limit.rs +++ /dev/null @@ -1,253 +0,0 @@ -//! CreditLimitBlock — scrollback card shown when a max-tier user exhausts credits. -//! -//! Replaces the Q&A question modal for users already at the highest tier -//! (SuperGrok Heavy). Instead of offering "Upgrade tier" + PAYG / buy-credits -//! options in the question overlay, this block renders an inline card with a -//! descriptive message and a link to the usage/billing page. - -use ratatui::style::{Modifier, Style}; -use ratatui::text::{Line, Span}; - -use crate::scrollback::block::BlockContent; -use crate::scrollback::types::{AccentStyle, BlockContext, BlockLine, BlockOutput, DisplayMode}; -use crate::theme::Theme; - -/// Which continue-path the max-tier credit-limit card recommends. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum CreditLimitCardAction { - /// Legacy on-demand: PAYG not enabled yet. - EnablePayg, - /// Legacy on-demand: PAYG on but at spending cap. - IncreasePaygLimit, - /// Unified usage billing: purchase prepaid credits. - PurchaseCredits, -} - -/// Inline scrollback card for credit-limit exhaustion on max-tier accounts. -#[derive(Debug, Clone)] -pub struct CreditLimitBlock { - /// Card heading (e.g. "You've hit your free credits limit."). - pub heading: String, - /// Continue-path body copy selector. - pub action: CreditLimitCardAction, - /// URL to the usage/billing page. - pub url: String, -} - -impl CreditLimitBlock { - /// Create a new credit-limit card. - pub fn new( - heading: impl Into, - action: CreditLimitCardAction, - url: impl Into, - ) -> Self { - Self { - heading: heading.into(), - action, - url: url.into(), - } - } -} - -impl BlockContent for CreditLimitBlock { - fn output(&self, _ctx: &BlockContext) -> BlockOutput { - let theme = Theme::current(); - - // Heading in bold warning color (amber/yellow). - let heading_style = Style::default() - .fg(theme.warning) - .add_modifier(Modifier::BOLD); - let heading = Line::from(Span::styled(self.heading.clone(), heading_style)); - - // Body copy — contextual message based on billing mode. - let muted = theme.muted(); - let body = match self.action { - CreditLimitCardAction::IncreasePaygLimit => { - "You can continue by increasing your spending limit." - } - CreditLimitCardAction::EnablePayg => { - "You can continue by enabling pay-as-you-go usage." - } - CreditLimitCardAction::PurchaseCredits => { - "You can continue by purchasing more credits." - } - }; - let body_line = Line::from(Span::styled(body.to_string(), muted)); - - // Clickable link styled as a button. - let link_style = theme.link_style(); - let link_line = Line::from(vec![Span::styled(self.url.clone(), link_style)]); - - BlockOutput { - lines: vec![ - BlockLine::styled(heading).with_selection_range(Some(0)), - BlockLine::separator(Line::from("")), - BlockLine::styled(body_line).with_selection_range(Some(0)), - BlockLine::styled(link_line).with_selection_range(Some(0)), - ], - } - } - - fn accent(&self, _ctx: &BlockContext) -> Option { - let theme = Theme::current(); - Some(AccentStyle::static_color(theme.warning)) - } - - fn has_vpad(&self, _ctx: &BlockContext) -> bool { - true - } - - fn has_raw_mode(&self) -> bool { - false - } - - fn is_foldable(&self) -> bool { - false - } - - fn default_display_mode(&self) -> DisplayMode { - DisplayMode::Expanded - } - - fn is_selectable(&self) -> bool { - true - } - - fn is_groupable(&self) -> bool { - false - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::appearance::AppearanceConfig; - - fn ctx() -> BlockContext { - BlockContext { - mode: DisplayMode::Expanded, - is_running: false, - width: 80, - raw: false, - max_lines: None, - appearance: AppearanceConfig::default(), - is_selected: false, - cwd: None, - } - } - - #[test] - fn output_payg_off_mentions_enabling() { - let block = CreditLimitBlock::new( - "You\u{2019}ve hit your credit limit.", - CreditLimitCardAction::EnablePayg, - "https://grok.com?_s=usage", - ); - let output = block.output(&ctx()); - let all_text: String = output - .lines - .iter() - .flat_map(|l| l.content.spans.iter().map(|s| s.content.as_ref())) - .collect(); - assert!(all_text.contains("credit limit")); - assert!(all_text.contains("enabling pay-as-you-go")); - assert!(all_text.contains("grok.com?_s=usage")); - } - - #[test] - fn output_payg_on_mentions_increasing() { - let block = CreditLimitBlock::new( - "You\u{2019}ve hit your spending cap.", - CreditLimitCardAction::IncreasePaygLimit, - "https://grok.com?_s=usage", - ); - let output = block.output(&ctx()); - let all_text: String = output - .lines - .iter() - .flat_map(|l| l.content.spans.iter().map(|s| s.content.as_ref())) - .collect(); - assert!(all_text.contains("spending cap")); - assert!(all_text.contains("increasing your spending limit")); - assert!(all_text.contains("grok.com?_s=usage")); - } - - #[test] - fn output_unified_mentions_purchasing_credits() { - let block = CreditLimitBlock::new( - "You hit your weekly limit.", - CreditLimitCardAction::PurchaseCredits, - "https://grok.com?_s=usage", - ); - let output = block.output(&ctx()); - let all_text: String = output - .lines - .iter() - .flat_map(|l| l.content.spans.iter().map(|s| s.content.as_ref())) - .collect(); - assert!(all_text.contains("purchasing more credits")); - assert!(all_text.contains("grok.com?_s=usage")); - } - - #[test] - fn has_warning_accent() { - let block = CreditLimitBlock::new("heading", CreditLimitCardAction::EnablePayg, "url"); - let accent = block.accent(&ctx()); - let theme = Theme::current(); - assert!(accent.is_some()); - assert_eq!(accent.unwrap().color, theme.warning); - } - - #[test] - fn block_content_contract() { - let block = CreditLimitBlock::new("heading", CreditLimitCardAction::EnablePayg, "url"); - let c = ctx(); - assert!(!block.is_foldable()); - assert!(block.is_selectable()); - assert!(!block.is_groupable()); - assert!(matches!( - block.default_display_mode(), - DisplayMode::Expanded - )); - assert!(block.has_vpad(&c)); - assert!(!block.has_raw_mode()); - } - - #[test] - fn output_structure_and_content() { - let url = "https://grok.com?_s=usage"; - let block = CreditLimitBlock::new("Test heading", CreditLimitCardAction::EnablePayg, url); - let output = block.output(&ctx()); - - // heading, separator, body, link = 4 lines - assert_eq!(output.lines.len(), 4); - - let all_text: String = output - .lines - .iter() - .flat_map(|l| l.content.spans.iter().map(|s| s.content.as_ref())) - .collect(); - assert!(all_text.contains(url)); - - // Heading uses bold modifier. - assert!( - output.lines[0] - .content - .spans - .iter() - .any(|s| s.style.add_modifier.contains(Modifier::BOLD)) - ); - } - - #[test] - fn new_stores_fields_correctly() { - let block = CreditLimitBlock::new( - "my heading", - CreditLimitCardAction::IncreasePaygLimit, - "https://example.com", - ); - assert_eq!(block.heading, "my heading"); - assert_eq!(block.action, CreditLimitCardAction::IncreasePaygLimit); - assert_eq!(block.url, "https://example.com"); - } -} diff --git a/crates/codegen/kigi-tui/src/scrollback/blocks/mod.rs b/crates/codegen/kigi-tui/src/scrollback/blocks/mod.rs index 3682b30..46c5792 100644 --- a/crates/codegen/kigi-tui/src/scrollback/blocks/mod.rs +++ b/crates/codegen/kigi-tui/src/scrollback/blocks/mod.rs @@ -6,7 +6,6 @@ mod agent; mod bg_task; mod btw; mod context_info; -mod credit_limit; pub mod markdown_content; pub mod mermaid_content; mod quote_bar; @@ -21,7 +20,6 @@ pub use agent::AgentMessageBlock; pub use bg_task::{BgTaskBlock, BgTaskKind}; pub use btw::BtwBlock; pub use context_info::ContextInfoBlock; -pub use credit_limit::{CreditLimitBlock, CreditLimitCardAction}; pub use session_event::{EndWork, SessionEvent, SessionEventBlock}; pub use subagent::{SubagentBlock, SubagentBlockKind}; pub use system::SystemMessageBlock; diff --git a/crates/codegen/kigi-tui/src/scrollback/export.rs b/crates/codegen/kigi-tui/src/scrollback/export.rs index 29db5de..b7ea327 100644 --- a/crates/codegen/kigi-tui/src/scrollback/export.rs +++ b/crates/codegen/kigi-tui/src/scrollback/export.rs @@ -60,7 +60,7 @@ pub fn render_blocks_to_markdown<'a>(blocks: impl IntoIterator {} } diff --git a/crates/codegen/kigi-tui/src/scrollback/wrappers/entry_renderer.rs b/crates/codegen/kigi-tui/src/scrollback/wrappers/entry_renderer.rs index 0e182f4..40539ca 100644 --- a/crates/codegen/kigi-tui/src/scrollback/wrappers/entry_renderer.rs +++ b/crates/codegen/kigi-tui/src/scrollback/wrappers/entry_renderer.rs @@ -526,7 +526,7 @@ impl<'a> EntryRenderer<'a> { /// /// EXACT for blocks whose searchable text mirrors their selectable rendered /// lines (plain source blocks, markdown/thinking bodies); a best-effort - /// estimate for field-joined source (Subagent/BgTask/CreditLimit), kept on + /// estimate for field-joined source (Subagent/BgTask/etc.), kept on /// screen by the caller's entry-height clamp. Past the last logical line, /// clamps to the final content row. pub fn rendered_row_of_logical_line(&self, width: u16, logical_line: usize) -> u16 { diff --git a/crates/codegen/kigi-tui/src/sessions_cmd.rs b/crates/codegen/kigi-tui/src/sessions_cmd.rs index 15134da..15a36bf 100644 --- a/crates/codegen/kigi-tui/src/sessions_cmd.rs +++ b/crates/codegen/kigi-tui/src/sessions_cmd.rs @@ -1,7 +1,5 @@ use anyhow::Result; use clap::Subcommand; -use kigi_shell::agent::config::Config as AgentConfig; -use kigi_shell::auth::{AuthManager, try_ensure_fresh_auth}; use kigi_shell::session::merge::MergedSession; use kigi_shell::util::kigi_home::kigi_home; #[derive(Debug, clap::Args, Clone)] @@ -33,41 +31,17 @@ enum SessionsCommand { }, } -pub async fn run(args: SessionsArgs, agent_config: &AgentConfig) -> Result<()> { - // Best-effort only. Do not force an interactive public login for enterprise - // deployments that only configure a deployment_key + custom xai_api_base_url. - // If the user has previously run the interactive `grok` TUI (which succeeds - // for these setups), any cached credential will be used. Otherwise we still - // proceed so the SessionRegistryClient can use the deployment_key when - // talking to the custom proxy. - let auth = try_ensure_fresh_auth(&agent_config.kimi_code_config).await; - - let auth_manager = std::sync::Arc::new(AuthManager::new( - &kigi_home(), - agent_config.kimi_code_config.clone(), - )); - - let client = kigi_shell::agent::session_registry_client::SessionRegistryClient::new( - agent_config.endpoints.proxy_url(), - String::new(), - ) - .with_deployment_key(agent_config.endpoints.deployment_key.clone()) - .with_alpha_test_key(agent_config.endpoints.alpha_test_key.clone()) - .with_auth(auth_manager.clone()); - +pub async fn run(args: SessionsArgs) -> Result<()> { let cwd = std::env::current_dir().unwrap_or_else(|_| ".".into()); match args.command { SessionsCommand::List { limit } => { let sessions = - kigi_shell::session::merge::fetch_merged(Some(&client), cwd.to_str(), None, limit) - .await; + kigi_shell::session::merge::fetch_merged(None, cwd.to_str(), None, limit).await; print_sessions_grouped(&sessions); } SessionsCommand::Search { query, limit } => { - use kigi_shell::session::merge::REMOTE_TIMEOUT; use kigi_shell::session::storage::search::{SessionSearchRequest, execute_search}; - use std::collections::HashSet; let req = SessionSearchRequest { query, @@ -78,28 +52,7 @@ pub async fn run(args: SessionsArgs, agent_config: &AgentConfig) -> Result<()> { }; let root = kigi_home(); - let remote_limit = (limit * 3).max(100) as i64; - let (local_resp, remote_results) = tokio::join!(execute_search(&root, &req), async { - tokio::time::timeout( - REMOTE_TIMEOUT, - client.search(Some(&req.query), remote_limit), - ) - .await - .unwrap_or_else(|_| { - eprintln!( - "warning: remote session search timed out, showing local results only" - ); - Ok(Vec::new()) - }) - .unwrap_or_else(|e| { - eprintln!("warning: remote session search failed: {e}"); - Vec::new() - }) - }); - - let resp = local_resp?; - let local_ids: HashSet<&str> = - resp.results.iter().map(|r| r.session_id.as_str()).collect(); + let resp = execute_search(&root, &req).await?; for hit in &resp.results { let title = if hit.title.is_empty() { @@ -124,63 +77,14 @@ pub async fn run(args: SessionsArgs, agent_config: &AgentConfig) -> Result<()> { ); } - let remaining = limit.saturating_sub(resp.results.len()); - let mut remote_shown = 0usize; - for r in &remote_results { - if remote_shown >= remaining { - break; - } - if local_ids.contains(r.session_id.as_str()) { - continue; - } - let title = if r.summary.is_empty() { - "(untitled)" - } else { - &r.summary - }; - let time = chrono::DateTime::parse_from_rfc3339(&r.updated_at) - .map(|dt| { - dt.with_timezone(&chrono::Local) - .format("%b %d, %l:%M%P") - .to_string() - }) - .unwrap_or_default(); - let snippet: String = r - .first_prompt - .as_deref() - .unwrap_or("") - .chars() - .take(80) - .collect(); - println!( - "{} (remote) {}\n {}\n {}", - r.session_id, time, title, snippet - ); - remote_shown += 1; - } - - println!("\nTotal: {}", resp.results.len() + remote_shown); + println!("\nTotal: {}", resp.results.len()); } SessionsCommand::Delete { id } => { - // Always attempt the remote delete when authenticated and not - // ZDR — `list` / `search` likewise query remote unconditionally - // rather than gating on storage mode (which the CLI cannot - // resolve here: it builds config without remote settings). The - // backend delete is idempotent (a `404` is treated as success), - // so this is safe for local-only sessions with no remote copy. - // ZDR teams never upload, so there is nothing remote to delete. - let needs_remote = auth.is_some(); - // Pass `cwd = None` so the session is found by id regardless of // which workspace it was created in; the local delete still uses // the resolved per-session cwd. - let deletion = kigi_shell::session::persistence::delete_session_history( - &id, - None, - needs_remote, - auth_manager.clone(), - ) - .await?; + let deletion = + kigi_shell::session::persistence::delete_session_history(&id, None).await?; if deletion.any_removed() { println!("Deleted session {id}"); diff --git a/crates/codegen/kigi-tui/src/settings/defs.rs b/crates/codegen/kigi-tui/src/settings/defs.rs index 2f30049..a5f3cc7 100644 --- a/crates/codegen/kigi-tui/src/settings/defs.rs +++ b/crates/codegen/kigi-tui/src/settings/defs.rs @@ -119,30 +119,6 @@ const PERMISSION_MODE_CHOICES: &[EnumChoice] = &[ }, ]; -// --------------------------------------------------------------------------- -// Coding-data-sharing catalog. -// -// Persisted in auth metadata (`AuthEntry::coding_data_retention_opt_out`), -// NOT config.toml. Two choices only — the pager has no `Option`/`Unset` -// representation for this field. -// -// `supports_preview: false` — toggling fires an async ACP call that -// can fail. Commit on Enter only. -// --------------------------------------------------------------------------- - -const CODING_DATA_SHARING_CHOICES: &[EnumChoice] = &[ - EnumChoice { - canonical: "opt-in", - display: "Opt in", - description: "Allow SpaceXAI to retain and use coding session data for training and product improvement.", - }, - EnumChoice { - canonical: "opt-out", - display: "Opt out", - description: "Do not retain coding session data. Code requests will not be used for training.", - }, -]; - // --------------------------------------------------------------------------- // Plan-mode catalog. // @@ -947,35 +923,6 @@ pub fn default_settings() -> Vec { restart_required: false, hidden_in_minimal: false, }, - // SHELL-owned. Persisted in auth metadata (not config.toml). - // Reads from `PagerLocalSnapshot.coding_data_sharing_opt_out`. - // Default "opt-in" matches `AuthEntry::coding_data_retention_opt_out = false`. - // ZDR / non-admin guards are enforced at dispatch time. - SettingMeta { - key: "coding_data_sharing", - category: SettingCategory::Privacy, - owner: SettingOwner::Shell, - label: "Coding data sharing", - description: "Controls whether SpaceXAI may retain and train on coding session data.", - keywords: &[ - "privacy", - "data", - "sharing", - "coding", - "retention", - "telemetry", - "training", - "opt-in", - "opt-out", - ], - kind: SettingKind::Enum { - default: "opt-in", - choices: CODING_DATA_SHARING_CHOICES, - supports_preview: false, - }, - restart_required: false, - hidden_in_minimal: false, - }, // SHELL-owned, persisted to `[ui].default_selected_permission` in // config.toml. Read by the pager via `appearance::permission_cursor`. // Canonical `always_allow_all_sessions` (the effective default) lands diff --git a/crates/codegen/kigi-tui/src/settings/registry.rs b/crates/codegen/kigi-tui/src/settings/registry.rs index b630ceb..f50e45b 100644 --- a/crates/codegen/kigi-tui/src/settings/registry.rs +++ b/crates/codegen/kigi-tui/src/settings/registry.rs @@ -38,7 +38,6 @@ pub enum SettingCategory { Mouse, Editor, Agent, - Privacy, Models, Session, Advanced, @@ -51,7 +50,6 @@ impl SettingCategory { Self::Mouse, Self::Editor, Self::Agent, - Self::Privacy, Self::Models, Self::Session, Self::Advanced, @@ -64,7 +62,6 @@ impl SettingCategory { Self::Mouse => "Mouse", Self::Editor => "Editor & Input", Self::Agent => "Agent & Approval", - Self::Privacy => "Privacy", Self::Models => "Models", Self::Session => "Session", Self::Advanced => "Advanced", @@ -247,10 +244,6 @@ pub struct PagerLocalSnapshot { /// Cloned into the snapshot so the modal's validator/resolver is /// self-contained (the modal outlives the borrow on `app.agents`). pub available_models: Vec<(String, acp::ModelId)>, - /// Whether the user has opted OUT of coding data sharing. - /// Lives in auth metadata (no `UiConfig` field). Inverted mapping: - /// `opt_out == false` → canonical "opt-in". - pub coding_data_sharing_opt_out: bool, /// Whether plan mode is active. Uses effective state /// (`pending.unwrap_or(active)`) so rapid toggles don't double-send. /// Refreshed on all mutation paths including ACP `CurrentModeUpdate`. @@ -285,7 +278,6 @@ impl Default for PagerLocalSnapshot { auto_mode: false, current_model_name: None, available_models: Vec::new(), - coding_data_sharing_opt_out: false, plan_mode_active: false, show_tips: None, auto_update: None, @@ -602,12 +594,6 @@ pub fn current_value_for( )), // max_thoughts_width: `u16` widened to `i64`. "max_thoughts_width" => Some(SettingValue::Int(ui.max_thoughts_width as i64)), - // coding_data_sharing: inverts the `_opt_out` bool. - "coding_data_sharing" => Some(SettingValue::Enum(if pager.coding_data_sharing_opt_out { - "opt-out" - } else { - "opt-in" - })), // plan_mode: canonical via `PlanModeKind::from_bool().as_canonical()`. "plan_mode" => Some(SettingValue::Enum( crate::app::actions::PlanModeKind::from_bool(pager.plan_mode_active).as_canonical(), @@ -816,17 +802,6 @@ mod tests { "max_thoughts_width default drifts from UiConfig::default()", ); } - // coding_data_sharing: no UiConfig field; default pinned - // against auth metadata (opt_out=false → "opt-in"). - ("coding_data_sharing", SettingKind::Enum { default, .. }) => { - let expected = "opt-in"; - assert_eq!( - *default, expected, - "coding_data_sharing registry default must be 'opt-in' — \ - the on-disk source of truth is `AuthEntry::coding_data_retention_opt_out: \ - bool` (defaults to `false`, i.e. user has NOT opted out)", - ); - } // CLI batch: fields live on CliConfig, not UiConfig. // Defaults pinned literally. ("show_tips", SettingKind::Bool { default }) => { diff --git a/crates/codegen/kigi-tui/src/share_cmd.rs b/crates/codegen/kigi-tui/src/share_cmd.rs deleted file mode 100644 index ea105d6..0000000 --- a/crates/codegen/kigi-tui/src/share_cmd.rs +++ /dev/null @@ -1,49 +0,0 @@ -use anyhow::Result; -use kigi_shell::agent::config::Config as AgentConfig; -use kigi_shell::session::share::{ShareSessionRequest, ShareSessionResponse}; -use tokio_util::sync::CancellationToken; - -use agent_client_protocol as acp; -use kigi_acp_lib::acp_send; - -#[derive(Debug, clap::Args, Clone)] -pub struct ShareArgs { - /// Session ID to share - pub session_id: String, -} - -pub async fn run(args: &ShareArgs, agent_config: &AgentConfig) -> Result<()> { - let cancel = CancellationToken::new(); - let spawned = crate::acp::spawn::spawn_grok_shell(agent_config.clone(), &cancel, None).await?; - - let _init: acp::InitializeResponse = acp_send( - acp::InitializeRequest::new(acp::ProtocolVersion::V1) - .client_capabilities( - acp::ClientCapabilities::new() - .fs(acp::FileSystemCapabilities::new()) - .terminal(false), - ) - .meta( - serde_json::json!({ - "clientType": crate::client_identity::HEADLESS_CLIENT_TYPE, - "clientVersion": crate::client_identity::PAGER_CLIENT_VERSION - }) - .as_object() - .cloned(), - ), - &spawned.channel.tx, - ) - .await?; - - let params = serde_json::value::to_raw_value(&ShareSessionRequest { - session_id: args.session_id.clone(), - })?; - let ext_req = acp::ExtRequest::new("x.ai/share_session", params.into()); - - let ext_resp: acp::ExtResponse = acp_send(ext_req, &spawned.channel.tx).await?; - let response: ShareSessionResponse = serde_json::from_str(ext_resp.0.get())?; - - println!("{}", response.share_url); - cancel.cancel(); - Ok(()) -} diff --git a/crates/codegen/kigi-tui/src/slash/commands/mod.rs b/crates/codegen/kigi-tui/src/slash/commands/mod.rs index fc729f0..f6c2052 100644 --- a/crates/codegen/kigi-tui/src/slash/commands/mod.rs +++ b/crates/codegen/kigi-tui/src/slash/commands/mod.rs @@ -41,7 +41,6 @@ pub mod new; pub mod personas; pub mod plan; pub mod plugin; -pub mod privacy; pub mod queue; pub mod recap; pub mod release_notes; @@ -53,7 +52,6 @@ pub mod screen_mode_switch; pub mod scroll_debug; pub mod session_info; pub mod settings_cmd; -pub mod share; pub mod tasks; pub mod terminal_setup; pub mod theme; @@ -98,7 +96,6 @@ pub fn builtin_commands() -> Vec> { Arc::new(plugin::HooksCommand), Arc::new(plugin::PluginsCommand), Arc::new(plugin::SkillsCommand), - Arc::new(share::ShareCommand), Arc::new(session_info::SessionInfoCommand), Arc::new(rename::RenameCommand), Arc::new(dashboard::DashboardCommand), @@ -120,7 +117,6 @@ pub fn builtin_commands() -> Vec> { Arc::new(timeline::TimelineCommand), Arc::new(toggle_mouse_reporting::ToggleMouseReportingCommand), Arc::new(settings_cmd::SettingsCommand), - Arc::new(privacy::PrivacyCommand), Arc::new(rewind::RewindCommand), Arc::new(jump::JumpCommand), Arc::new(login::LoginCommand), @@ -455,83 +451,24 @@ mod tests { usage::UsageCommand.run(&mut ctx, args) } #[test] - fn usage_no_args_returns_show_usage() { + fn usage_returns_show_usage() { assert!(matches!( run_usage(""), CommandResult::Action(Action::ShowUsage) )); } #[test] - fn usage_show_returns_show_usage() { + fn usage_ignores_stray_args() { assert!(matches!( - run_usage("show"), + run_usage(" anything "), CommandResult::Action(Action::ShowUsage) )); } #[test] - fn usage_manage_returns_open_url() { - match run_usage("manage") { - CommandResult::Action(Action::OpenUrl(url)) => { - assert_eq!(url, "https://grok.com/?_s=usage"); - } - other => panic!("expected Action(OpenUrl), got {other:?}"), - } - } - #[test] - fn usage_invalid_arg_returns_error() { - match run_usage("delete") { - CommandResult::Error(msg) => { - assert!(msg.contains("delete"), "got: {msg}"); - } - other => panic!("expected Error, got {other:?}"), - } - } - #[test] - fn usage_whitespace_only_treated_as_no_args() { - assert!(matches!( - run_usage(" "), - CommandResult::Action(Action::ShowUsage) - )); - } - #[test] - fn usage_show_with_leading_whitespace() { - assert!(matches!( - run_usage(" show "), - CommandResult::Action(Action::ShowUsage) - )); - } - #[test] - fn usage_manage_with_leading_whitespace() { - match run_usage(" manage ") { - CommandResult::Action(Action::OpenUrl(url)) => { - assert_eq!(url, "https://grok.com/?_s=usage"); - } - other => panic!("expected Action(OpenUrl), got {other:?}"), - } - } - #[test] - fn usage_suggest_args_returns_show_and_manage() { - let models = ModelState::default(); - let ctx = crate::slash::command::AppCtx { - models: &models, - cwd: std::path::Path::new("."), - screen_mode: crate::app::ScreenMode::Fullscreen, - }; - let items = usage::UsageCommand - .suggest_args(&ctx, "") - .expect("should have suggestions"); - assert_eq!(items.len(), 2); - assert_eq!(items[0].display, "show"); - assert_eq!(items[0].insert_text, "show"); - assert_eq!(items[1].display, "manage"); - assert_eq!(items[1].insert_text, "manage"); - } - #[test] fn usage_metadata() { let cmd = usage::UsageCommand; assert_eq!(cmd.name(), "usage"); - assert!(cmd.takes_args()); - assert_eq!(cmd.arg_placeholder(), Some("show | manage")); + assert!(!cmd.takes_args()); assert!(!cmd.description().is_empty()); assert!(!cmd.usage().is_empty()); } diff --git a/crates/codegen/kigi-tui/src/slash/commands/privacy.rs b/crates/codegen/kigi-tui/src/slash/commands/privacy.rs deleted file mode 100644 index d0dceb3..0000000 --- a/crates/codegen/kigi-tui/src/slash/commands/privacy.rs +++ /dev/null @@ -1,212 +0,0 @@ -//! `/privacy` -- show or toggle privacy and data retention status. - -use crate::app::actions::Action; -use crate::slash::command::{CommandExecCtx, CommandResult, SlashCommand}; - -/// Show or toggle privacy and data retention status. -/// -/// Usage: -/// - `/privacy` show current status -/// - `/privacy opt-in` opt in to coding data sharing -/// - `/privacy opt-out` opt out of coding data sharing -/// -/// Case-insensitive. Only unambiguous aliases are accepted (e.g. `in`, -/// `share`, `out`, `private`) — generic toggles like `on`/`off` are -/// rejected because they're ambiguous in privacy context. -pub struct PrivacyCommand; - -impl SlashCommand for PrivacyCommand { - fn name(&self) -> &str { - "privacy" - } - - fn description(&self) -> &str { - "Show or toggle privacy & data retention status" - } - - fn usage(&self) -> &str { - "/privacy [opt-in|opt-out]" - } - - fn takes_args(&self) -> bool { - true - } - - fn run(&self, _ctx: &mut CommandExecCtx, args: &str) -> CommandResult { - let arg = args.trim(); - if arg.is_empty() { - return CommandResult::Action(Action::ShowPrivacyInfo); - } - match parse_privacy_arg(arg) { - Some(opted_in) => CommandResult::Action(Action::SetCodingDataSharing { opted_in }), - None => CommandResult::Error(format!( - "Unknown argument `{arg}`. Valid options: opt-in (aliases: in, share) | \ - opt-out (aliases: out, private)." - )), - } - } -} - -/// Parse `/privacy ` into `Some(true)` (opt-in), `Some(false)` -/// (opt-out), or `None` (unknown). Case-insensitive ASCII matching. -#[doc(hidden)] -pub fn parse_privacy_arg(arg: &str) -> Option { - const OPT_IN_ALIASES: &[&str] = &["opt-in", "in", "share"]; - const OPT_OUT_ALIASES: &[&str] = &["opt-out", "out", "private"]; - - if OPT_IN_ALIASES.iter().any(|a| arg.eq_ignore_ascii_case(a)) { - return Some(true); - } - if OPT_OUT_ALIASES.iter().any(|a| arg.eq_ignore_ascii_case(a)) { - return Some(false); - } - None -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn parse_opt_in_canonical() { - assert_eq!(parse_privacy_arg("opt-in"), Some(true)); - } - - #[test] - fn parse_opt_out_canonical() { - assert_eq!(parse_privacy_arg("opt-out"), Some(false)); - } - - /// Case-insensitive matching. - #[test] - fn parse_case_insensitive() { - for variant in &["OPT-IN", "Opt-In", "opt-IN", "OpT-iN"] { - assert_eq!( - parse_privacy_arg(variant), - Some(true), - "case-insensitive parse must accept `{variant}` as opt-in", - ); - } - for variant in &["OPT-OUT", "Opt-Out", "opt-OUT", "OpT-oUt"] { - assert_eq!( - parse_privacy_arg(variant), - Some(false), - "case-insensitive parse must accept `{variant}` as opt-out", - ); - } - } - - /// Pins the accepted alias catalog. - #[test] - fn parse_opt_in_aliases() { - for alias in &["in", "share"] { - assert_eq!( - parse_privacy_arg(alias), - Some(true), - "alias `{alias}` must map to opt-in", - ); - } - } - - #[test] - fn parse_opt_out_aliases() { - for alias in &["out", "private"] { - assert_eq!( - parse_privacy_arg(alias), - Some(false), - "alias `{alias}` must map to opt-out", - ); - } - } - - /// Ambiguous generic-toggle aliases must be rejected — `/privacy on` - /// is ambiguous (could mean opt-in or opt-out). - #[test] - fn parse_rejects_ambiguous_generic_aliases() { - for ambiguous in &[ - "on", "off", "true", "false", "enable", "enabled", "disable", "disabled", - ] { - assert_eq!( - parse_privacy_arg(ambiguous), - None, - "ambiguous alias `{ambiguous}` MUST be rejected — it would let a user typing \ - `/privacy {ambiguous}` get the OPPOSITE of their intent in privacy context. \ - See Security Issue 10 in PR 9 R1.", - ); - } - } - - /// Unknown arguments return None → the command surfaces an error - /// listing valid options. Pins the "no silent fallback" contract. - #[test] - fn parse_unknown_returns_none() { - for unknown in &["yes", "no", "maybe", "opt-maybe", "", " ", "1", "0"] { - assert_eq!( - parse_privacy_arg(unknown), - None, - "unknown arg `{unknown}` must NOT parse", - ); - } - } - - /// Alias families must not overlap. - #[test] - fn alias_families_disjoint() { - let opt_in_results: Vec = ["opt-in", "in", "share"] - .iter() - .map(|a| parse_privacy_arg(a).unwrap()) - .collect(); - assert!( - opt_in_results.iter().all(|b| *b), - "every opt-in alias must parse to true", - ); - let opt_out_results: Vec = ["opt-out", "out", "private"] - .iter() - .map(|a| parse_privacy_arg(a).unwrap()) - .collect(); - assert!( - opt_out_results.iter().all(|b| !*b), - "every opt-out alias must parse to false", - ); - } - - /// Error message must list every accepted alias. - #[test] - fn error_message_lists_all_accepted_aliases() { - use crate::acp::model_state::ModelState; - use crate::app::bundle::BundleState; - - let cmd = PrivacyCommand; - let models = ModelState::default(); - let bundle = BundleState::default(); - let mut ctx = CommandExecCtx { - models: &models, - session_id: None, - bundle_state: &bundle, - screen_mode: crate::app::ScreenMode::Inline, - pager_state: crate::settings::PagerLocalSnapshot::default(), - }; - let result = cmd.run(&mut ctx, "garbage-input"); - match result { - CommandResult::Error(msg) => { - // Every accepted alias appears in the error message. - for alias in &["opt-in", "in", "share", "opt-out", "out", "private"] { - assert!( - msg.contains(alias), - "error message must mention alias `{alias}` so the user knows \ - what to type; msg = {msg:?}", - ); - } - // Dropped ambiguous aliases must not appear. - for dropped in &["off", "true", "false", "enable", "disable"] { - assert!( - !msg.contains(dropped), - "dropped alias `{dropped}` must NOT appear in error message \ - (would suggest it's still accepted); msg = {msg:?}", - ); - } - } - other => panic!("expected Error result for unknown arg, got {other:?}"), - } - } -} diff --git a/crates/codegen/kigi-tui/src/slash/commands/share.rs b/crates/codegen/kigi-tui/src/slash/commands/share.rs deleted file mode 100644 index 247907e..0000000 --- a/crates/codegen/kigi-tui/src/slash/commands/share.rs +++ /dev/null @@ -1,34 +0,0 @@ -//! `/share` -- share current session via URL. - -use crate::app::actions::Action; -use crate::slash::command::{CommandExecCtx, CommandResult, SlashCommand}; - -/// Share the current session via a public URL. -pub struct ShareCommand; - -impl SlashCommand for ShareCommand { - fn name(&self) -> &str { - "share" - } - - fn description(&self) -> &str { - "Share this session via URL" - } - - fn session_scoped(&self) -> bool { - true - } - - fn usage(&self) -> &str { - "/share" - } - - fn run(&self, ctx: &mut CommandExecCtx, _args: &str) -> CommandResult { - // Check if we have an active session - if ctx.session_id.is_none() { - return CommandResult::Error("No active session to share".to_string()); - } - - CommandResult::Action(Action::ShareSession) - } -} diff --git a/crates/codegen/kigi-tui/src/slash/commands/usage.rs b/crates/codegen/kigi-tui/src/slash/commands/usage.rs index 35d1196..d32db0a 100644 --- a/crates/codegen/kigi-tui/src/slash/commands/usage.rs +++ b/crates/codegen/kigi-tui/src/slash/commands/usage.rs @@ -1,13 +1,9 @@ -//! `/usage` -- show credit usage or open billing management page. +//! `/usage` -- display Kimi API usage and quota information. use crate::app::actions::Action; -use crate::slash::command::{AppCtx, ArgItem, CommandExecCtx, CommandResult, SlashCommand}; +use crate::slash::command::{CommandExecCtx, CommandResult, SlashCommand}; -/// Show coding credit usage or manage billing. -/// -/// `/usage` -- show current credit usage -/// `/usage show` -- same as above -/// `/usage manage` -- open billing management page in browser +/// Display API usage and quota information. pub struct UsageCommand; impl SlashCommand for UsageCommand { @@ -15,56 +11,22 @@ impl SlashCommand for UsageCommand { "usage" } - /// `/cost` is the minimal-mode name for the same credit-usage summary: - /// it commits a usage/cost system block rather than opening a - /// pane, so it's an alias rather than a separate command. + /// `/cost` is the minimal-mode name for the same usage summary: it + /// commits a usage system block rather than opening a pane, so it's + /// an alias rather than a separate command. fn aliases(&self) -> &[&str] { &["cost"] } fn description(&self) -> &str { - "View credit usage or manage billing" + "Display API usage and quota information" } fn usage(&self) -> &str { - "/usage [show|manage]" + "/usage" } - fn takes_args(&self) -> bool { - true - } - - fn arg_placeholder(&self) -> Option<&str> { - Some("show | manage") - } - - fn suggest_args(&self, _ctx: &AppCtx, _args_query: &str) -> Option> { - Some(vec![ - ArgItem { - display: "show".to_string(), - match_text: "show".to_string(), - insert_text: "show".to_string(), - description: "View credit usage".to_string(), - }, - ArgItem { - display: "manage".to_string(), - match_text: "manage".to_string(), - insert_text: "manage".to_string(), - description: "Open billing management page".to_string(), - }, - ]) - } - - fn run(&self, _ctx: &mut CommandExecCtx, args: &str) -> CommandResult { - let arg = args.trim(); - match arg { - "" | "show" => CommandResult::Action(Action::ShowUsage), - "manage" => { - CommandResult::Action(Action::OpenUrl("https://grok.com/?_s=usage".to_string())) - } - _ => CommandResult::Error(format!( - "Unknown argument: {arg}. Use /usage show or /usage manage" - )), - } + fn run(&self, _ctx: &mut CommandExecCtx, _args: &str) -> CommandResult { + CommandResult::Action(Action::ShowUsage) } } diff --git a/crates/codegen/kigi-tui/src/slash/mod.rs b/crates/codegen/kigi-tui/src/slash/mod.rs index 1c0c59d..8c8cdc9 100644 --- a/crates/codegen/kigi-tui/src/slash/mod.rs +++ b/crates/codegen/kigi-tui/src/slash/mod.rs @@ -1828,7 +1828,6 @@ mod tests { "/compact", "/fork", "/rewind", - "/share", "/context", "/copy", "/export", diff --git a/crates/codegen/kigi-tui/src/slash/registry.rs b/crates/codegen/kigi-tui/src/slash/registry.rs index ef940d0..1a3c9bd 100644 --- a/crates/codegen/kigi-tui/src/slash/registry.rs +++ b/crates/codegen/kigi-tui/src/slash/registry.rs @@ -351,12 +351,6 @@ impl CommandRegistry { self.available_tools = Some(tools); } - /// Show or hide the /share command. - /// When hidden, it won't appear in the dropdown or be executable. - pub fn set_share_visible(&mut self, visible: bool) { - self.set_command_visible("share", visible); - } - /// Show or hide the /usage command. /// When hidden, it won't appear in the dropdown or be executable. pub fn set_usage_visible(&mut self, visible: bool) { @@ -699,35 +693,6 @@ mod tests { assert!(registry.get("flush").is_none()); } - #[test] - fn set_share_visible_hides_and_restores_share_command() { - let share: Arc = Arc::new(DummyCommand { - name: "share", - aliases: &[], - }); - let other: Arc = Arc::new(DummyCommand { - name: "exit", - aliases: &[], - }); - let mut registry = CommandRegistry::new(vec![share, other]); - - // Default: /share is visible. - assert!(registry.get("share").is_some()); - assert!(registry.triggers().iter().any(|t| t.canonical == "share")); - - // Hiding /share removes it from lookup and triggers. - registry.set_share_visible(false); - assert!(registry.get("share").is_none()); - assert!(!registry.triggers().iter().any(|t| t.canonical == "share")); - // Other commands are unaffected. - assert!(registry.get("exit").is_some()); - - // Re-enabling restores it. - registry.set_share_visible(true); - assert!(registry.get("share").is_some()); - assert!(registry.triggers().iter().any(|t| t.canonical == "share")); - } - #[test] fn set_usage_visible_hides_and_restores_usage_command() { let usage: Arc = Arc::new(DummyCommand { @@ -1110,9 +1075,9 @@ mod tests { /// unresolvable for dispatch, exactly like `get()`. #[test] fn get_for_dispatch_respects_hard_gates() { - // Hard-hidden by name (e.g. /dashboard default, /share toggle). - let share: Arc = Arc::new(DummyCommand { - name: "share", + // Hard-hidden by name (e.g. /dashboard default). + let dashboard: Arc = Arc::new(DummyCommand { + name: "dashboard", aliases: &[], }); // Tier-restricted. @@ -1125,11 +1090,14 @@ mod tests { name: "loop", required: &["scheduler_create"], }); - let mut reg = CommandRegistry::new(vec![share, usage, gated]); - reg.set_share_visible(false); + let mut reg = CommandRegistry::new(vec![dashboard, usage, gated]); + reg.set_dashboard_visible(false); reg.set_restricted_commands(&["usage".to_string()]); - assert!(reg.get_for_dispatch("share").is_none(), "hidden stays hard"); + assert!( + reg.get_for_dispatch("dashboard").is_none(), + "hidden stays hard" + ); assert!( reg.get_for_dispatch("usage").is_none(), "restricted stays blocked (upsell path owns it)" diff --git a/crates/codegen/kigi-tui/src/test_util.rs b/crates/codegen/kigi-tui/src/test_util.rs index fdf6206..46c0751 100644 --- a/crates/codegen/kigi-tui/src/test_util.rs +++ b/crates/codegen/kigi-tui/src/test_util.rs @@ -26,8 +26,6 @@ pub fn make_agent_view(session_id: Option<&str>, cwd: &str) -> crate::app::agent restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, diff --git a/crates/codegen/kigi-tui/src/views/agent.rs b/crates/codegen/kigi-tui/src/views/agent.rs index 7d2b942..c8da88a 100644 --- a/crates/codegen/kigi-tui/src/views/agent.rs +++ b/crates/codegen/kigi-tui/src/views/agent.rs @@ -889,7 +889,6 @@ pub fn build_hints( has_queued_follow_up: bool, selected_is_user_prompt: bool, selected_is_agent_message: bool, - selected_is_credit_limit: bool, shift_enter_unavailable: bool, scrollback_search: Option<&ScrollbackSearchState>, ) -> Vec { @@ -1029,19 +1028,12 @@ pub fn build_hints( let mut hints = Vec::new(); let nothing_special = !selected_is_agent_message && !selected_is_user_prompt - && !selected_is_credit_limit && fold_label.is_none() && group_header_label.is_none() && !selected_supports_fullscreen; if nothing_special { hints.push(space_prompt_hint()); } - if selected_is_credit_limit { - if let Some(key) = registry.key_for(ActionId::OpenBlockViewer) { - hints.push(HintItem::new(key, "open")); - } - hints.push(space_prompt_hint()); - } if selected_is_agent_message { if vim_mode && selected_supports_copy @@ -1217,7 +1209,6 @@ mod tests { selected_is_user_prompt, selected_is_agent_message, false, - false, None, ) } @@ -1249,7 +1240,6 @@ mod tests { false, false, false, - false, None, ); let labels: Vec<&str> = hints.iter().map(|h| h.label.as_ref()).collect(); @@ -1413,7 +1403,6 @@ mod tests { false, false, false, - false, Some(&search), ) } @@ -1516,7 +1505,6 @@ mod tests { false, false, false, - false, None, ); assert!( @@ -1559,7 +1547,6 @@ mod tests { false, false, false, - false, shift_enter_unavailable, None, ) @@ -1619,7 +1606,6 @@ mod tests { false, false, false, - false, None, ); let labels: Vec<&str> = hints.iter().map(|h| h.label.as_ref()).collect(); diff --git a/crates/codegen/kigi-tui/src/views/credit_bar.rs b/crates/codegen/kigi-tui/src/views/credit_bar.rs deleted file mode 100644 index 1d869c5..0000000 --- a/crates/codegen/kigi-tui/src/views/credit_bar.rs +++ /dev/null @@ -1,817 +0,0 @@ -//! Credit balance indicator for the agent status bar. -//! -//! Shows the user's coding credit usage as a compact status bar item. -//! Fetches real data from the `x.ai/billing` agent extension. - -use ratatui::style::Style; -use ratatui::text::{Line, Span}; - -use crate::theme::Theme; - -/// Credit balance state from the billing API. -#[derive(Debug, Clone)] -pub struct CreditBalance { - /// Usage as a percentage of the allowance (0.0–100.0). - pub usage_pct: f64, - /// Usage as a percentage of total budget (free + on-demand when enabled). - pub effective_usage_pct: f64, - /// Billing period end as a formatted local wall-clock string (no zone - /// label), e.g. "Mar 31, 12:00". - pub period_end_display: Option, - /// Whether pay-as-you-go (on-demand) billing is enabled. - pub pay_as_you_go: bool, - /// On-demand spending cap in USD cents (e.g. 500 = $5.00). - pub on_demand_cap_cents: Option, - /// On-demand usage this period in USD cents. - pub on_demand_used_cents: Option, - /// Remaining prepaid ("bought") credit balance in USD cents. - pub prepaid_balance_cents: Option, - /// Usage period type from the billing response (the proto enum name, e.g. - /// `USAGE_PERIOD_TYPE_WEEKLY`). Drives the "Weekly/Monthly limit" label. - pub period_type: Option, - /// From credits config `is_unified_billing_user` (`None` if absent). - /// `Some(true)` = unified pool / buy-credits UX; `Some(false)` = legacy - /// on-demand / PAYG UX. - pub is_unified_billing_user: Option, -} - -impl CreditBalance { - /// Label for the percentage allowance, chosen from the period type: - /// "Weekly limit" / "Monthly limit", falling back to "Usage" when unknown. - pub fn usage_label(&self) -> &'static str { - match self.period_type.as_deref() { - Some(t) if t.contains("WEEKLY") => "Weekly limit", - Some(t) if t.contains("MONTHLY") => "Monthly limit", - _ => "Usage", - } - } -} - -/// Auto top-up rule data used by the `/usage` summary. -#[derive(Debug, Clone)] -pub struct AutoTopupInfo { - /// Whether auto top-up is enabled. - pub enabled: bool, - /// Per-trigger top-up amount in USD cents. - pub topup_amount_cents: Option, - /// Optional maximum monthly top-up amount in USD cents. - pub max_amount_cents: Option, -} - -impl AutoTopupInfo { - /// A known "no / disabled auto top-up" state — distinct from an unresolved - /// `None`, which means the rule hasn't been fetched yet. - pub fn disabled() -> Self { - Self { - enabled: false, - topup_amount_cents: None, - max_amount_cents: None, - } - } -} - -/// Outcome of an auto top-up rule fetch, so a transient failure doesn't clear a -/// previously cached rule. -#[derive(Debug, Clone)] -pub enum AutoTopupFetch { - /// A definitive rule state (a real rule, or [`AutoTopupInfo::disabled`] when - /// the backend reports none). Stored as the *known* auto top-up state. - Resolved(AutoTopupInfo), - /// Fetch failed — keep the cached value (last-known-good). A stored `None` - /// therefore means "not yet known", not "no auto top-up". - Unchanged, - /// The rule is not applicable (no prepaid credits) — reset the cache to - /// "unknown" so a later credits period doesn't read a stale rule. - Cleared, -} - -/// Format `cents` as a dollar string: whole dollars as `$N`, otherwise `$N.NN`. -fn fmt_dollars(cents: i64) -> String { - let dollars = cents as f64 / 100.0; - if dollars.fract() == 0.0 { - format!("${dollars:.0}") - } else { - format!("${dollars:.2}") - } -} - -/// Build the `/usage` summary block shown in scrollback. -/// -/// Always shows usage % and (when known) the next reset time. The credits -/// block is rendered only when the user has a positive prepaid balance: -/// - no prepaid balance → credits block omitted entirely -/// - auto top-up off/unknown → `Auto topup: disabled` (no max line) -/// - auto top-up on, no max → `Auto topup: $N` -/// - auto top-up on, max set → `Auto topup: $N` + `Max monthly topup: $M` -pub fn format_usage_summary(balance: &CreditBalance, autotopup: Option<&AutoTopupInfo>) -> String { - // Floor to match the backend SpendingLimiter's `as u8` truncation - // (99.994% → 99%, never 100% until truly exhausted). - let mut lines = vec![format!( - "{}: {}%", - balance.usage_label(), - balance.usage_pct.floor() as i64 - )]; - if let Some(reset) = &balance.period_end_display { - lines.push(format!("Next reset: {reset}")); - } - - // Billing stores credit / top-up amounts as negative cents (accounting - // convention); display the absolute USD value, matching the web clients. - if let Some(prepaid) = balance - .prepaid_balance_cents - .map(i64::abs) - .filter(|c| *c > 0) - { - lines.push(String::new()); - lines.push(format!("Credits: {}", fmt_dollars(prepaid))); - match autotopup { - Some(at) if at.enabled && at.topup_amount_cents.is_some() => { - lines.push(format!( - "Auto topup: {}", - fmt_dollars(at.topup_amount_cents.unwrap().abs()) - )); - if let Some(max) = at.max_amount_cents { - lines.push(format!("Max monthly topup: {}", fmt_dollars(max.abs()))); - } - } - _ => lines.push("Auto topup: disabled".to_string()), - } - } - - // Legacy on-demand (pay-as-you-go) billing — shown only when enabled, for - // users on the older monthly + on-demand model. Amounts always carry cents - // (e.g. `$50.00`), matching the web client. - if balance.pay_as_you_go { - let used = balance.on_demand_used_cents.unwrap_or(0).abs() as f64 / 100.0; - let cap = balance.on_demand_cap_cents.unwrap_or(0).abs() as f64 / 100.0; - lines.push(String::new()); - lines.push(format!("Pay-as-you-go: ${used:.2} used of ${cap:.2} limit")); - } - - lines.join("\n") -} - -/// Low-balance ($10) and pay-as-you-go critical ($5) warning thresholds, in cents. -const LOW_BALANCE_CENTS: i64 = 1000; -const PAY_AS_YOU_GO_CRITICAL_CENTS: i64 = 500; - -/// The prompt's usage/credits warning as `(text, critical)`, or `None` -/// (`critical` = yellow, else grey; team users with `usage_visible = false` -/// never warn). Behaviour splits by billing model — prepaid credits, -/// pay-as-you-go on-demand, or the included-allowance percentage — with exact -/// thresholds and copy pinned by the unit tests. -/// -/// Gateway light-frontend (`kind: "chat"`) sessions must not surface Build -/// coding-credit warnings — use [`usage_warning_for_session`] with -/// `gateway_chat = true` so the prompt shows no fake local sampler telemetry. -pub fn usage_warning( - balance: &CreditBalance, - autotopup: Option<&AutoTopupInfo>, - usage_visible: bool, -) -> Option<(String, bool)> { - usage_warning_for_session(balance, autotopup, usage_visible, false) -} - -/// Like [`usage_warning`], but suppresses output for gateway/chat-kind sessions. -pub fn usage_warning_for_session( - balance: &CreditBalance, - autotopup: Option<&AutoTopupInfo>, - usage_visible: bool, - gateway_chat: bool, -) -> Option<(String, bool)> { - if gateway_chat || !usage_visible { - return None; - } - - // A non-zero prepaid balance (stored as signed cents) means the credits model. - let credits = balance - .prepaid_balance_cents - .map(i64::abs) - .filter(|c| *c > 0); - - let Some(credits_cents) = credits else { - // Pay-as-you-go (legacy on-demand): warn on dollars left in the cap once - // the included allowance is spent. - if balance.pay_as_you_go { - if balance.usage_pct >= 100.0 { - let cap = balance.on_demand_cap_cents.unwrap_or(0).abs(); - let used = balance.on_demand_used_cents.unwrap_or(0).abs(); - let remaining = (cap - used).max(0); - if remaining <= LOW_BALANCE_CENTS { - let text = format!("Pay-as-you-go limit left: {}", fmt_dollars(remaining)); - return Some((text, remaining <= PAY_AS_YOU_GO_CRITICAL_CENTS)); - } - } - return None; - } - - let pct = balance.effective_usage_pct; - if pct > 90.0 { - // "Left" = complement of floored usage, so it agrees with the - // floored summary (99.994% → "1% left", not "0%"). - let remaining = (100 - pct.floor() as i64).max(0); - let label = balance.usage_label(); - return Some((format!("{label} left: {remaining}%"), pct > 95.0)); - } - return None; - }; - - // Credits are only drawn down at 100% usage; don't warn before then. - if balance.usage_pct < 100.0 { - return None; - } - - let credits_warning = || { - ( - format!("Credits left: {}", fmt_dollars(credits_cents)), - true, - ) - }; - - // Auto top-up gates the warning: unknown → silent; disabled → warn when low; - // enabled w/o max → never; enabled w/ max → warn below one top-up amount. - match autotopup { - None => None, - Some(at) if !at.enabled => (credits_cents <= LOW_BALANCE_CENTS).then(credits_warning), - Some(at) if at.max_amount_cents.is_none() => None, - Some(at) => at - .topup_amount_cents - .map(i64::abs) - .and_then(|amt| (credits_cents < amt).then(credits_warning)), - } -} - -/// Build the credit balance indicator as a `Line<'static>`. -/// -/// Shows `Credits used: XX%` in the status bar. -/// -/// Gateway light-frontend (`kind: "chat"`) sessions must not show Build coding -/// credits — use [`credit_bar_line_for_session`] with `gateway_chat = true` -/// (returns `None`). remote settings / managed opt-in for chat entry can share the -/// same gate later; for now it only zeros/suppresses misleading local telemetry. -pub fn credit_bar_line(balance: &CreditBalance, hovered: bool, theme: &Theme) -> Line<'static> { - credit_bar_line_for_session(balance, hovered, theme, false) - .expect("non-chat credit_bar_line always renders") -} - -/// Like [`credit_bar_line`], but returns `None` for gateway/chat-kind sessions -/// so the status bar never implies Build sampler / coding-credit usage. -pub fn credit_bar_line_for_session( - balance: &CreditBalance, - _hovered: bool, - theme: &Theme, - gateway_chat: bool, -) -> Option> { - if gateway_chat { - return None; - } - let pct = balance.usage_pct; - let color = if pct >= 100.0 { - theme.accent_error - } else if pct >= 80.0 { - theme.warning - } else { - theme.accent_success - }; - - let text = format!("Credits used: {pct:.0}%"); - - let style = Style::default().fg(color).bg(theme.bg_base); - Some(Line::from(Span::styled(text, style))) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn bal(pct: f64) -> CreditBalance { - CreditBalance { - usage_pct: pct, - effective_usage_pct: pct, - period_end_display: None, - pay_as_you_go: false, - on_demand_cap_cents: None, - on_demand_used_cents: None, - prepaid_balance_cents: None, - period_type: None, - is_unified_billing_user: None, - } - } - - fn topup(enabled: bool, amount: Option, max: Option) -> AutoTopupInfo { - AutoTopupInfo { - enabled, - topup_amount_cents: amount, - max_amount_cents: max, - } - } - - #[test] - fn summary_no_credits_omits_credits_block() { - let b = CreditBalance { - period_end_display: Some("June 14, 16:00".into()), - prepaid_balance_cents: Some(0), - ..bal(25.0) - }; - // Even with an auto-topup rule present, zero prepaid → no credits block. - let out = format_usage_summary(&b, Some(&topup(true, Some(2000), Some(10000)))); - assert_eq!(out, "Usage: 25%\nNext reset: June 14, 16:00"); - } - - #[test] - fn summary_credits_without_autotopup_shows_disabled() { - let b = CreditBalance { - prepaid_balance_cents: Some(10000), - ..bal(25.0) - }; - assert_eq!( - format_usage_summary(&b, None), - "Usage: 25%\n\nCredits: $100\nAuto topup: disabled" - ); - // A disabled rule renders the same. - assert_eq!( - format_usage_summary(&b, Some(&topup(false, Some(2000), Some(10000)))), - "Usage: 25%\n\nCredits: $100\nAuto topup: disabled" - ); - } - - #[test] - fn summary_autotopup_enabled_without_max_omits_max() { - let b = CreditBalance { - prepaid_balance_cents: Some(10000), - ..bal(25.0) - }; - assert_eq!( - format_usage_summary(&b, Some(&topup(true, Some(2000), None))), - "Usage: 25%\n\nCredits: $100\nAuto topup: $20" - ); - } - - #[test] - fn summary_autotopup_enabled_with_max_renders_all() { - let b = CreditBalance { - period_end_display: Some("June 14, 16:00".into()), - prepaid_balance_cents: Some(10000), - ..bal(25.0) - }; - assert_eq!( - format_usage_summary(&b, Some(&topup(true, Some(2000), Some(10000)))), - "Usage: 25%\nNext reset: June 14, 16:00\n\nCredits: $100\nAuto topup: $20\nMax monthly topup: $100" - ); - } - - #[test] - fn summary_formats_fractional_dollars() { - let b = CreditBalance { - prepaid_balance_cents: Some(1250), - ..bal(25.0) - }; - assert_eq!( - format_usage_summary(&b, Some(&topup(true, Some(550), None))), - "Usage: 25%\n\nCredits: $12.50\nAuto topup: $5.50" - ); - } - - #[test] - fn summary_abs_negative_billing_amounts() { - // Billing returns credit / top-up amounts as negative cents; the - // summary must render them as positive USD (matching the web). - let b = CreditBalance { - prepaid_balance_cents: Some(-500), - ..bal(100.0) - }; - assert_eq!( - format_usage_summary(&b, Some(&topup(true, Some(-500), Some(-1000)))), - "Usage: 100%\n\nCredits: $5\nAuto topup: $5\nMax monthly topup: $10" - ); - } - - #[test] - fn summary_pay_as_you_go_enabled_renders_used_of_limit() { - let b = CreditBalance { - pay_as_you_go: true, - on_demand_used_cents: Some(355), - on_demand_cap_cents: Some(5000), - period_type: Some("USAGE_PERIOD_TYPE_MONTHLY".into()), - period_end_display: Some("June 30, 16:00".into()), - ..bal(91.0) - }; - assert_eq!( - format_usage_summary(&b, None), - "Monthly limit: 91%\nNext reset: June 30, 16:00\n\nPay-as-you-go: $3.55 used of $50.00 limit" - ); - } - - #[test] - fn summary_pay_as_you_go_disabled_omits_line() { - let b = CreditBalance { - pay_as_you_go: false, - period_type: Some("USAGE_PERIOD_TYPE_MONTHLY".into()), - period_end_display: Some("June 30, 16:00".into()), - ..bal(91.0) - }; - assert_eq!( - format_usage_summary(&b, None), - "Monthly limit: 91%\nNext reset: June 30, 16:00" - ); - } - - // ── usage_label / period type ──────────────────────────────────── - - fn bal_period(pct: f64, period_type: &str) -> CreditBalance { - CreditBalance { - period_type: Some(period_type.to_string()), - ..bal(pct) - } - } - - #[test] - fn usage_label_from_period_type() { - assert_eq!( - bal_period(0.0, "USAGE_PERIOD_TYPE_WEEKLY").usage_label(), - "Weekly limit" - ); - assert_eq!( - bal_period(0.0, "USAGE_PERIOD_TYPE_MONTHLY").usage_label(), - "Monthly limit" - ); - // Unknown / unspecified / absent → falls back to "Usage". - assert_eq!( - bal_period(0.0, "USAGE_PERIOD_TYPE_UNSPECIFIED").usage_label(), - "Usage" - ); - assert_eq!(bal(0.0).usage_label(), "Usage"); - } - - #[test] - fn summary_uses_period_label() { - let weekly = bal_period(25.0, "USAGE_PERIOD_TYPE_WEEKLY"); - assert_eq!(format_usage_summary(&weekly, None), "Weekly limit: 25%"); - let monthly = bal_period(25.0, "USAGE_PERIOD_TYPE_MONTHLY"); - assert_eq!(format_usage_summary(&monthly, None), "Monthly limit: 25%"); - } - - #[test] - fn warning_uses_period_label() { - let weekly = bal_period(92.0, "USAGE_PERIOD_TYPE_WEEKLY"); - assert_eq!( - usage_warning(&weekly, None, true), - Some(("Weekly limit left: 8%".to_string(), false)) - ); - } - - #[test] - fn summary_floors_usage_percent() { - // Match the backend SpendingLimiter (`as u8` truncation): 99.994% must - // render as 99%, not round up to 100%. - let almost = bal_period(99.994, "USAGE_PERIOD_TYPE_WEEKLY"); - assert_eq!(format_usage_summary(&almost, None), "Weekly limit: 99%"); - // A true 100% still shows 100%. - let full = bal_period(100.0, "USAGE_PERIOD_TYPE_WEEKLY"); - assert_eq!(format_usage_summary(&full, None), "Weekly limit: 100%"); - } - - #[test] - fn warning_percent_left_is_floor_complement() { - // 99.994% used → floored to 99% → "1% left" (not "0% left"), so the - // warning and the floored summary always sum to 100. - let almost = bal_period(99.994, "USAGE_PERIOD_TYPE_WEEKLY"); - assert_eq!( - usage_warning(&almost, None, true), - Some(("Weekly limit left: 1%".to_string(), true)) - ); - // A true 100% (no credits) → "0% left". - let full = bal_period(100.0, "USAGE_PERIOD_TYPE_WEEKLY"); - assert_eq!( - usage_warning(&full, None, true), - Some(("Weekly limit left: 0%".to_string(), true)) - ); - } - - // ── usage_warning (prompt info row) ────────────────────────────── - - #[test] - fn warning_usage_model_thresholds() { - assert_eq!(usage_warning(&bal(50.0), None, true), None); - assert_eq!( - usage_warning(&bal(92.0), None, true), - Some(("Usage left: 8%".to_string(), false)) - ); - assert_eq!( - usage_warning(&bal(97.0), None, true), - Some(("Usage left: 3%".to_string(), true)) - ); - } - - #[test] - fn warning_hidden_for_team_users() { - assert_eq!(usage_warning(&bal(99.0), None, false), None); - let credits = CreditBalance { - prepaid_balance_cents: Some(100), - ..bal(0.0) - }; - assert_eq!(usage_warning(&credits, None, false), None); - } - - #[test] - fn warning_credits_unknown_topup_is_suppressed() { - // At 100% usage with prepaid credits, but the rule isn't known yet - // (None) — never warn; it resolves on the next billing fetch. - let b = CreditBalance { - prepaid_balance_cents: Some(100), - ..bal(100.0) - }; - assert_eq!(usage_warning(&b, None, true), None); - } - - #[test] - fn warning_credits_suppressed_below_full_usage() { - // Low credits + no auto top-up, but the included allowance still has - // room (usage < 100%) → no warning (credits aren't being spent yet). - let disabled = topup(false, None, None); - let low = CreditBalance { - prepaid_balance_cents: Some(453), - ..bal(0.0) - }; - assert_eq!(usage_warning(&low, Some(&disabled), true), None); - // Same balance once the allowance is exhausted → warn. - let exhausted = CreditBalance { - prepaid_balance_cents: Some(453), - ..bal(100.0) - }; - assert_eq!( - usage_warning(&exhausted, Some(&disabled), true), - Some(("Credits left: $4.53".to_string(), true)) - ); - } - - #[test] - fn warning_credits_no_topup_low_shows_dollars() { - // "No auto top-up" is a known, disabled rule (not an unresolved None). - let b = CreditBalance { - prepaid_balance_cents: Some(453), - ..bal(100.0) - }; - let disabled = topup(false, None, None); - assert_eq!( - usage_warning(&b, Some(&disabled), true), - Some(("Credits left: $4.53".to_string(), true)) - ); - } - - #[test] - fn warning_credits_no_topup_above_threshold_silent() { - let disabled = topup(false, None, None); - let b = CreditBalance { - prepaid_balance_cents: Some(1500), - ..bal(100.0) - }; - assert_eq!(usage_warning(&b, Some(&disabled), true), None); - // Exactly $10 is still "low". - let at_ten = CreditBalance { - prepaid_balance_cents: Some(1000), - ..bal(100.0) - }; - assert_eq!( - usage_warning(&at_ten, Some(&disabled), true), - Some(("Credits left: $10".to_string(), true)) - ); - } - - #[test] - fn warning_credits_topup_no_max_never_warns() { - let b = CreditBalance { - prepaid_balance_cents: Some(1), - ..bal(100.0) - }; - assert_eq!( - usage_warning(&b, Some(&topup(true, Some(2000), None)), true), - None - ); - } - - #[test] - fn warning_credits_topup_with_max_below_topup_amount() { - // $15 balance, $20 top-up amount, $100 max → below one top-up → warn. - let b = CreditBalance { - prepaid_balance_cents: Some(1500), - ..bal(100.0) - }; - assert_eq!( - usage_warning(&b, Some(&topup(true, Some(2000), Some(10000))), true), - Some(("Credits left: $15".to_string(), true)) - ); - let plenty = CreditBalance { - prepaid_balance_cents: Some(2500), - ..bal(100.0) - }; - assert_eq!( - usage_warning(&plenty, Some(&topup(true, Some(2000), Some(10000))), true), - None - ); - } - - #[test] - fn warning_credits_handles_negative_cents() { - let b = CreditBalance { - prepaid_balance_cents: Some(-453), - ..bal(100.0) - }; - assert_eq!( - usage_warning(&b, Some(&topup(true, Some(-2000), Some(-10000))), true), - Some(("Credits left: $4.53".to_string(), true)) - ); - } - - #[test] - fn warning_credits_take_precedence_over_usage() { - // A credits user below 100% usage gets no warning at all (no usage-% - // warning, and credits aren't being spent yet) — unlike a non-credits - // user, who would see "Usage left: 1%" at 99%. - let b = CreditBalance { - prepaid_balance_cents: Some(5000), - ..bal(99.0) - }; - assert_eq!( - usage_warning(&b, Some(&topup(false, None, None)), true), - None - ); - // Zero prepaid falls back to the usage model. - let zero = CreditBalance { - prepaid_balance_cents: Some(0), - ..bal(99.0) - }; - assert_eq!( - usage_warning(&zero, None, true), - Some(("Usage left: 1%".to_string(), true)) - ); - } - - // ── usage_warning: pay-as-you-go (monthly on-demand) ───────────── - - fn pay_as_you_go(usage_pct: f64, cap_cents: i64, used_cents: i64) -> CreditBalance { - CreditBalance { - pay_as_you_go: true, - on_demand_cap_cents: Some(cap_cents), - on_demand_used_cents: Some(used_cents), - period_type: Some("USAGE_PERIOD_TYPE_MONTHLY".into()), - ..bal(usage_pct) - } - } - - #[test] - fn warning_pay_as_you_go_low_dollars_shows_remaining() { - // $50 cap, $42 used → $8 left → grey (above $5). - let grey = pay_as_you_go(100.0, 5000, 4200); - assert_eq!( - usage_warning(&grey, None, true), - Some(("Pay-as-you-go limit left: $8".to_string(), false)) - ); - // $50 cap, $46 used → $4 left → critical (yellow). - let yellow = pay_as_you_go(100.0, 5000, 4600); - assert_eq!( - usage_warning(&yellow, None, true), - Some(("Pay-as-you-go limit left: $4".to_string(), true)) - ); - } - - #[test] - fn warning_pay_as_you_go_boundaries() { - // Exactly $10 left → show, grey. - let at_ten = pay_as_you_go(100.0, 5000, 4000); - assert_eq!( - usage_warning(&at_ten, None, true), - Some(("Pay-as-you-go limit left: $10".to_string(), false)) - ); - // Exactly $5 left → critical (yellow). - let at_five = pay_as_you_go(100.0, 5000, 4500); - assert_eq!( - usage_warning(&at_five, None, true), - Some(("Pay-as-you-go limit left: $5".to_string(), true)) - ); - } - - #[test] - fn warning_pay_as_you_go_above_threshold_silent() { - // $20 left (> $10) → no warning. - let b = pay_as_you_go(100.0, 5000, 3000); - assert_eq!(usage_warning(&b, None, true), None); - } - - #[test] - fn warning_pay_as_you_go_suppressed_below_full_usage() { - // Pay-as-you-go users get NO percentage warning before the included - // allowance is exhausted, even with low on-demand room remaining. - let b = pay_as_you_go(95.0, 5000, 4800); - assert_eq!(usage_warning(&b, None, true), None); - } - - #[test] - fn warning_pay_as_you_go_fractional_dollars() { - // $50 cap, $46.50 used → $3.50 left → critical, fractional formatting. - let b = pay_as_you_go(100.0, 5000, 4650); - assert_eq!( - usage_warning(&b, None, true), - Some(("Pay-as-you-go limit left: $3.50".to_string(), true)) - ); - } - - #[test] - fn test_credit_bar_line_shows_percentage() { - let theme = Theme::default(); - let line = credit_bar_line(&bal(24.0), false, &theme); - let text: String = line.spans.iter().map(|s| s.content.as_ref()).collect(); - assert_eq!(text, "Credits used: 24%"); - } - - #[test] - fn test_color_thresholds() { - let theme = Theme::default(); - - let low = credit_bar_line(&bal(50.0), false, &theme); - assert_eq!(low.spans[0].style.fg, Some(theme.accent_success)); - - let high = credit_bar_line(&bal(85.0), false, &theme); - assert_eq!(high.spans[0].style.fg, Some(theme.warning)); - - let over = credit_bar_line(&bal(100.0), false, &theme); - assert_eq!(over.spans[0].style.fg, Some(theme.accent_error)); - } - - #[test] - fn test_zero_percent() { - let theme = Theme::default(); - let line = credit_bar_line(&bal(0.0), false, &theme); - let text: String = line.spans.iter().map(|s| s.content.as_ref()).collect(); - assert_eq!(text, "Credits used: 0%"); - assert_eq!(line.spans[0].style.fg, Some(theme.accent_success)); - } - - #[test] - fn test_boundary_at_80_percent() { - let theme = Theme::default(); - // Exactly 80% should be warning (yellow). - let at_80 = credit_bar_line(&bal(80.0), false, &theme); - assert_eq!(at_80.spans[0].style.fg, Some(theme.warning)); - - // Just below 80% should be success (green). - let below_80 = credit_bar_line(&bal(79.9), false, &theme); - assert_eq!(below_80.spans[0].style.fg, Some(theme.accent_success)); - } - - #[test] - fn test_boundary_at_100_percent() { - let theme = Theme::default(); - // Exactly 100% should be error (red). - let at_100 = credit_bar_line(&bal(100.0), false, &theme); - assert_eq!(at_100.spans[0].style.fg, Some(theme.accent_error)); - - // Just below 100% should be warning (yellow). - let below_100 = credit_bar_line(&bal(99.9), false, &theme); - assert_eq!(below_100.spans[0].style.fg, Some(theme.warning)); - } - - #[test] - fn test_over_100_percent() { - let theme = Theme::default(); - let line = credit_bar_line(&bal(150.0), false, &theme); - let text: String = line.spans.iter().map(|s| s.content.as_ref()).collect(); - assert_eq!(text, "Credits used: 150%"); - assert_eq!(line.spans[0].style.fg, Some(theme.accent_error)); - } - - #[test] - fn test_fractional_percentage_rounds_display() { - let theme = Theme::default(); - let line = credit_bar_line(&bal(33.7), false, &theme); - let text: String = line.spans.iter().map(|s| s.content.as_ref()).collect(); - assert_eq!(text, "Credits used: 34%"); - } - - #[test] - fn test_credit_balance_with_on_demand_fields() { - let balance = CreditBalance { - effective_usage_pct: 25.0, - period_end_display: Some("Jun 1, 00:00".into()), - pay_as_you_go: true, - on_demand_cap_cents: Some(2000), - on_demand_used_cents: Some(500), - ..bal(50.0) - }; - let theme = Theme::default(); - // The credit bar uses usage_pct (not effective_usage_pct). - let line = credit_bar_line(&balance, false, &theme); - let text: String = line.spans.iter().map(|s| s.content.as_ref()).collect(); - assert_eq!(text, "Credits used: 50%"); - } - - #[test] - fn gateway_chat_suppresses_credit_bar_and_usage_warning() { - let theme = Theme::default(); - let b = bal(90.0); - assert!(credit_bar_line_for_session(&b, false, &theme, true).is_none()); - assert!(usage_warning_for_session(&b, None, true, true).is_none()); - // Build path still renders. - assert!(credit_bar_line_for_session(&b, false, &theme, false).is_some()); - } -} diff --git a/crates/codegen/kigi-tui/src/views/dashboard/peek.rs b/crates/codegen/kigi-tui/src/views/dashboard/peek.rs index 401604c..50ccdd1 100644 --- a/crates/codegen/kigi-tui/src/views/dashboard/peek.rs +++ b/crates/codegen/kigi-tui/src/views/dashboard/peek.rs @@ -520,8 +520,6 @@ fn paint_peek_config_badge( model_name: &model_label, flags: &flags, multiline, - usage_warning: None, - usage_warning_critical: false, }; // Bottom border row, inside the corners — the same content rect the // chat prompt and dispatch box use for their info line. @@ -1111,7 +1109,6 @@ pub fn extract_last_response_type(agent: &AgentView) -> String { RenderBlock::BgTask(_) => return "Task".to_string(), RenderBlock::Btw(_) => return "Btw".to_string(), RenderBlock::ContextInfo(_) => return "Context".to_string(), - RenderBlock::CreditLimit(_) => return "Credit limit".to_string(), // The user's latest input marks the turn boundary — there's // no agent response after it yet. RenderBlock::UserPrompt(_) => break, @@ -1310,7 +1307,6 @@ fn block_short_text(block: &crate::scrollback::block::RenderBlock) -> Option Some("(subagent)".to_string()), RenderBlock::Btw(_) => Some("(btw)".to_string()), RenderBlock::ContextInfo(_) => Some("(context info)".to_string()), - RenderBlock::CreditLimit(_) => Some("(credit limit)".to_string()), RenderBlock::Stub(_) => None, } } diff --git a/crates/codegen/kigi-tui/src/views/dashboard/render.rs b/crates/codegen/kigi-tui/src/views/dashboard/render.rs index b3f2cd8..df2d9fd 100644 --- a/crates/codegen/kigi-tui/src/views/dashboard/render.rs +++ b/crates/codegen/kigi-tui/src/views/dashboard/render.rs @@ -2596,8 +2596,6 @@ fn paint_dispatch_config_badge( model_name: &model_label, flags: &flags, multiline: state.multiline_mode, - usage_warning: None, - usage_warning_critical: false, }; // Bottom border row, inside the corners — the same content rect the chat // prompt uses for its info line. diff --git a/crates/codegen/kigi-tui/src/views/dashboard/row.rs b/crates/codegen/kigi-tui/src/views/dashboard/row.rs index 7dcdbec..e8a1018 100644 --- a/crates/codegen/kigi-tui/src/views/dashboard/row.rs +++ b/crates/codegen/kigi-tui/src/views/dashboard/row.rs @@ -1721,8 +1721,6 @@ mod tests { restore_degree: None, rate_limited: false, model_incompatible: false, - credit_limit_blocked: false, - free_usage_blocked: false, available_commands: Vec::new(), available_commands_generation: 0, available_tools: None, diff --git a/crates/codegen/kigi-tui/src/views/mod.rs b/crates/codegen/kigi-tui/src/views/mod.rs index dd59d11..bdab026 100644 --- a/crates/codegen/kigi-tui/src/views/mod.rs +++ b/crates/codegen/kigi-tui/src/views/mod.rs @@ -6,7 +6,6 @@ pub mod block_viewer; pub mod btw_overlay; pub mod completion_dropdown; pub mod context_bar; -pub mod credit_bar; pub mod dashboard; pub mod debug_style; pub mod extensions_modal; diff --git a/crates/codegen/kigi-tui/src/views/modal.rs b/crates/codegen/kigi-tui/src/views/modal.rs index 8518a93..d51c2b8 100644 --- a/crates/codegen/kigi-tui/src/views/modal.rs +++ b/crates/codegen/kigi-tui/src/views/modal.rs @@ -363,11 +363,8 @@ pub enum PaletteCommand { OpenAgentsModal, } /// Build the default set of palette entries with section grouping. -/// -/// `sharing_enabled` controls whether the `/share` entry is included. -/// Pass `true` to preserve the default behavior (show `/share`). -pub fn default_palette_entries(sharing_enabled: bool) -> Vec { - let mut entries = vec![ +pub fn default_palette_entries() -> Vec { + let entries = vec![ PaletteEntry { label: "Session".into(), shortcut: String::new(), @@ -398,11 +395,6 @@ pub fn default_palette_entries(sharing_enabled: bool) -> Vec { shortcut: "/resume".into(), command: PaletteCommand::SlashCommand("/resume".into()), }, - PaletteEntry { - label: "Share Session".into(), - shortcut: "/share".into(), - command: PaletteCommand::SlashCommand("/share".into()), - }, PaletteEntry { label: "Rename Session".into(), shortcut: "/rename ".into(), @@ -536,19 +528,12 @@ pub fn default_palette_entries(sharing_enabled: bool) -> Vec { command: PaletteCommand::Quit, }, ]; - if !sharing_enabled { - entries.retain(|e| { - !matches!( - & e.command, PaletteCommand::SlashCommand(s) if s.trim() == "/share" - ) - }); - } entries } #[allow(clippy::collapsible_if)] /// Filter palette entries for search, preserving section headers when any item in the section matches. -pub fn filter_palette_entries(query: &str, sharing_enabled: bool) -> Vec { - let all = default_palette_entries(sharing_enabled); +pub fn filter_palette_entries(query: &str) -> Vec { + let all = default_palette_entries(); let query_lower = query.to_lowercase(); if query_lower.is_empty() { return all; @@ -1233,26 +1218,11 @@ mod doc_viewer_scroll_tests { } } #[cfg(test)] -mod palette_sharing_tests { +mod palette_tests { use super::*; - fn has_share(entries: &[PaletteEntry]) -> bool { - entries.iter().any(|e| { - matches!( - & e.command, PaletteCommand::SlashCommand(s) if s.trim() == "/share" - ) - }) - } - #[test] - fn default_palette_includes_share_when_enabled() { - let entries = default_palette_entries(true); - assert!( - has_share(&entries), - "/share should be present when sharing_enabled=true" - ); - } #[test] fn default_palette_includes_dashboard() { - let entries = default_palette_entries(true); + let entries = default_palette_entries(); let has_dashboard = entries.iter().any(|e| { matches!( & e.command, PaletteCommand::SlashCommand(s) if s.trim() == @@ -1270,38 +1240,9 @@ mod palette_sharing_tests { ); } #[test] - fn default_palette_omits_share_when_disabled() { - let entries = default_palette_entries(false); - assert!( - !has_share(&entries), - "/share must not appear in palette when sharing_enabled=false" - ); - } - #[test] - fn filter_palette_omits_share_when_disabled() { - let entries = filter_palette_entries("", false); - assert!( - !has_share(&entries), - "/share must not appear in unfiltered palette when sharing_enabled=false" - ); - let entries = filter_palette_entries("share", false); - assert!( - !has_share(&entries), - "/share must not appear when filtering for 'share' with sharing_enabled=false" - ); - } - #[test] - fn filter_palette_includes_share_when_enabled_and_matched() { - let entries = filter_palette_entries("share", true); - assert!( - has_share(&entries), - "/share should match a 'share' query when sharing_enabled=true" - ); - } - #[test] fn palette_tools_section_routes_each_tab_to_itself() { use crate::views::extensions_modal::ExtensionsTab; - let entries = default_palette_entries(true); + let entries = default_palette_entries(); for (label, expected) in [ ("Hooks", ExtensionsTab::Hooks), ("Plugins", ExtensionsTab::Plugins), diff --git a/crates/codegen/kigi-tui/src/views/prompt_widget/mod.rs b/crates/codegen/kigi-tui/src/views/prompt_widget/mod.rs index 543ab13..59d5c92 100644 --- a/crates/codegen/kigi-tui/src/views/prompt_widget/mod.rs +++ b/crates/codegen/kigi-tui/src/views/prompt_widget/mod.rs @@ -286,11 +286,6 @@ pub struct PromptInfo<'a> { pub flags: &'a [PromptFlag<'a>], /// Whether multiline mode is active (shown right-aligned). pub multiline: bool, - /// Optional usage warning displayed right-aligned (e.g. "5% usage left"). - pub usage_warning: Option<&'a str>, - /// When true the warning uses the yellow warning color (<=5% left); - /// when false it uses dim grey text (5-10% left). - pub usage_warning_critical: bool, } /// Result of rendering the prompt. @@ -3198,16 +3193,6 @@ impl PromptWidget { // bottom-border fill — giving 1 cell of visual padding on each side. let pad_style = Style::default().bg(bg); let mut left_spans = vec![Span::styled(" ", pad_style)]; - if let Some(warning) = info.usage_warning { - let fg = if info.usage_warning_critical { - theme.warning - } else { - sep_fg - }; - let warning_style = Style::default().fg(fg).bg(bg); - left_spans.push(Span::styled(warning.to_owned(), warning_style)); - left_spans.push(Span::styled(" · ", sep_style)); - } left_spans.push(Span::styled(info.model_name, model_style)); for flag in info.flags { left_spans.push(Span::styled(" · ", sep_style)); diff --git a/crates/codegen/kigi-tui/src/views/question_view.rs b/crates/codegen/kigi-tui/src/views/question_view.rs index 3c84acd..dc1465c 100644 --- a/crates/codegen/kigi-tui/src/views/question_view.rs +++ b/crates/codegen/kigi-tui/src/views/question_view.rs @@ -105,14 +105,6 @@ pub enum LocalQuestionKind { /// On submit, the selected option index is translated into an /// [`crate::app::actions::Action::NewSessionAnswered`]. NewSession, - /// Modal shown when the user hits the credit/rate limit (403). - /// Options map to upsell URLs: upgrade tier or enable on-demand. - CreditLimitUpsell, - /// SuperGrok upsell modal: the free-usage paywall (429 + - /// `subscription:free-usage-exhausted`) or a tier-restricted slash - /// command invocation. Upgrade options carry their URL in the option - /// `id`. - FreeUsageUpsell, /// Modal shown when the shell rejects a model switch due to agent /// type incompatibility. Carries the target model + effort so the /// answer handler can create a new session with it. diff --git a/crates/codegen/kigi-tui/src/views/settings_modal.rs b/crates/codegen/kigi-tui/src/views/settings_modal.rs index a6cf2a7..ff52d65 100644 --- a/crates/codegen/kigi-tui/src/views/settings_modal.rs +++ b/crates/codegen/kigi-tui/src/views/settings_modal.rs @@ -755,11 +755,6 @@ fn action_for_enum_commit(key: SettingKey, choice: &'static str) -> Option None, }, - "coding_data_sharing" => match choice { - "opt-in" => Some(Action::SetCodingDataSharing { opted_in: true }), - "opt-out" => Some(Action::SetCodingDataSharing { opted_in: false }), - _ => None, - }, "plan_mode" => match choice { "on" => Some(Action::SetPlanMode(crate::app::actions::PlanModeKind::On)), "off" => Some(Action::SetPlanMode(crate::app::actions::PlanModeKind::Off)), @@ -5566,9 +5561,9 @@ mod tests { /// The default registry contains Appearance settings /// (3 bools + 3 enums + 1 int = 7 entries), the Editor entry /// `multiline_mode`, the Agent entries `permission_mode` and - /// `plan_mode`, the Privacy entry `coding_data_sharing`, the - /// Models entry `default_model`, and the Advanced entries - /// `show_tips` and `auto_update`. `default_reasoning_effort` and + /// `plan_mode`, the Models entry `default_model`, and the + /// Advanced entries `show_tips` and `auto_update`. + /// `default_reasoning_effort` and /// `auto_compact_threshold_percent` are not exposed in the modal. #[test] fn rows_contain_categories_and_settings_through_pr_14() { @@ -5591,7 +5586,6 @@ mod tests { &SettingCategory::Mouse, &SettingCategory::Editor, &SettingCategory::Agent, - &SettingCategory::Privacy, &SettingCategory::Models, // The Session category has no registered settings, so its // header is not emitted. @@ -5672,8 +5666,6 @@ mod tests { "toolset.ask_user_question.timeout_enabled", // PAGER-owned plan_mode (Agent category). "plan_mode", - // SHELL-owned coding_data_sharing (Privacy category). - "coding_data_sharing", // SHELL-owned default_model (Models category). "default_model", // Models category. `default_reasoning_effort`, @@ -8095,7 +8087,7 @@ mod tests { fn picker_visual_smoke_debug() { let entries = vec![SettingMeta { key: "wrap_enum", - category: SettingCategory::Privacy, + category: SettingCategory::Advanced, owner: SettingOwner::Shared, label: "Coding data sharing", description: "Controls whether SpaceXAI may retain and train on coding data.", @@ -8157,7 +8149,7 @@ mod tests { fn picker_long_description_wraps_to_multiple_lines() { let entries = vec![SettingMeta { key: "wrap_enum", - category: SettingCategory::Privacy, + category: SettingCategory::Advanced, owner: SettingOwner::Shared, label: "Coding data sharing", description: "Controls whether SpaceXAI may retain and train on coding data.", @@ -8442,7 +8434,7 @@ mod tests { // Reuse the wrap fixture: long descriptions on both choices. let entries = vec![SettingMeta { key: "wrap_enum", - category: SettingCategory::Privacy, + category: SettingCategory::Advanced, owner: SettingOwner::Shared, label: "Coding data sharing", description: "Controls whether SpaceXAI may retain coding data.", @@ -9501,7 +9493,7 @@ mod tests { fn synthetic_enum_chevron_meta() -> SettingMeta { SettingMeta { key: "test-enum-with-chevron", - category: SettingCategory::Privacy, + category: SettingCategory::Advanced, owner: SettingOwner::Shared, label: "Coding data sharing", description: "Enum row that opens a picker — chevron suffix applies.", @@ -9665,9 +9657,9 @@ mod tests { /// Two-line rows expand `state.row_rects` to span BOTH lines so /// mouse clicks on either line trigger the same default action. /// - /// `coding_data_sharing`: label 19 + value "Opt out" 7 + chevron - /// 2 + chrome 4 = 32 cells one-line. We render at width=28 so - /// the row drops to two lines. + /// `default_selected_permission`: label 27 + value "Always allow + /// on all sessions" 28 + chevron 2 + chrome 4 = 61 cells + /// one-line. We render at width=40 so the row drops to two lines. #[test] fn two_line_row_hit_rect_spans_both_lines() { let mut s = make_state(); @@ -9675,15 +9667,15 @@ mod tests { .rows .iter() .position( - |r| matches!(r, RowEntry::Setting { key, .. } if *key == "coding_data_sharing"), + |r| matches!(r, RowEntry::Setting { key, .. } if *key == "default_selected_permission"), ) - .expect("coding_data_sharing must be registered"); - // Render at a narrow width so coding_data_sharing forces a - // two-line layout. + .expect("default_selected_permission must be registered"); + // Render at a narrow width so default_selected_permission + // forces a two-line layout. let area = Rect { x: 0, y: 0, - width: 28, + width: 40, height: 60, }; let mut buf = Buffer::empty(area); @@ -9700,7 +9692,7 @@ mod tests { // Synthesize a click on line 2 of the row. The mouse handler // should fire the default action (open the enum picker for - // coding_data_sharing). + // default_selected_permission). s.list_area = area; let click_y = rect.y + 1; // Click somewhere in the middle of line 2. @@ -9736,22 +9728,22 @@ mod tests { #[test] fn two_line_row_with_expansion_renders_three_segments() { let mut s = make_state(); - // Coding data sharing's label + value (with chevron) won't - // fit on a 28-col line, forcing two-line layout. + // Default selected permission's label + value (with chevron) + // won't fit on a 40-col line, forcing two-line layout. let row_idx = s .rows .iter() .position( - |r| matches!(r, RowEntry::Setting { key, .. } if *key == "coding_data_sharing"), + |r| matches!(r, RowEntry::Setting { key, .. } if *key == "default_selected_permission"), ) - .expect("coding_data_sharing must be registered"); + .expect("default_selected_permission must be registered"); s.selected = row_idx; - s.expanded_keys.insert("coding_data_sharing"); + s.expanded_keys.insert("default_selected_permission"); let area = Rect { x: 0, y: 0, - width: 28, + width: 40, height: 60, }; let mut buf = Buffer::empty(area); @@ -9767,18 +9759,16 @@ mod tests { // The row label is on line 1. let label_line = buf_row_text(&buf, rect.y, area.x, area.width); assert!( - label_line.contains("Coding data sharing"), + label_line.contains("Default selected permission"), "line 1 must contain the row label: {label_line:?}" ); - // The value (display: "Opt out" or similar) is on line 2. + // The value is on line 2. It comes from the canonical → + // display mapping: `UiConfig::default()` resolves to the + // `always_allow_all_sessions` canonical, whose registered + // display is "Always allow on all sessions". let value_line = buf_row_text(&buf, rect.y + 1, area.x, area.width); - // Value comes from displaying the canonical → display mapping, - // which uses the synthetic enum's "Third Option" canonical of - // "opt-out". The display fallback returns the canonical when - // the lookup misses — registry has the real `CodingDataSharing` - // choices, so display should be "Opt out". assert!( - value_line.contains("Opt") || value_line.contains("opt") || value_line.contains("out"), + value_line.contains("Always allow"), "line 2 must contain the value text: {value_line:?}" ); // The expanded description renders on line 3 and below. diff --git a/crates/codegen/kigi-tui/src/views/welcome/mod.rs b/crates/codegen/kigi-tui/src/views/welcome/mod.rs index e71e634..c889043 100644 --- a/crates/codegen/kigi-tui/src/views/welcome/mod.rs +++ b/crates/codegen/kigi-tui/src/views/welcome/mod.rs @@ -100,9 +100,6 @@ pub struct WelcomeRenderResult { /// Hit-test rect for the "show full URL" fallback link. pub auth_fallback_rect: Option, /// Hit-test rect for the "[Refresh]" button on the paywall tier line. - pub refresh_rect: Option, - /// Hit-test rect for the gate URL link (click to open in browser). - pub gate_url_rect: Option, /// Whether a "Changelog" menu action was rendered (above Quit), so the /// input handler can map the extra menu row to the release-notes action /// once markdown is available. @@ -334,12 +331,12 @@ impl WelcomeLayout { } /// Controls what the version badge renders. -pub(super) enum VersionBadgeMode<'a> { - /// Full badge: team | tier | api_key | **Grok Build** VERSION+channel **Beta** (right-aligned). - Full { subscription_tier: Option<&'a str> }, - /// Hero footer: team | api_key | Grok Build Beta [channel] (right-aligned, gray). +pub(super) enum VersionBadgeMode { + /// Full badge: team | tier | api_key | **Kigi** VERSION+channel (right-aligned). + Full, + /// Hero footer: team | api_key | Kigi [channel] (right-aligned, gray). HeroFooter, - /// Hero inline: **Grok Build Beta** VERSION (left-aligned). + /// Hero inline: **Kigi** VERSION (left-aligned). HeroInline, } @@ -350,7 +347,7 @@ pub(super) fn render_version_badge( team_name: Option<&str>, h_margin: u16, is_api_key_auth: bool, - mode: VersionBadgeMode<'_>, + mode: VersionBadgeMode, ) { let version_area = Rect { width: version_rect.width.saturating_sub(h_margin), @@ -362,27 +359,16 @@ pub(super) fn render_version_badge( ); let mut spans = Vec::new(); - let (show_team, show_tier, show_api_key, align) = match &mode { - VersionBadgeMode::Full { .. } => (true, true, true, Alignment::Right), - VersionBadgeMode::HeroFooter => (true, false, true, Alignment::Right), - VersionBadgeMode::HeroInline => (false, false, false, Alignment::Left), + let (show_team, show_api_key, align) = match &mode { + VersionBadgeMode::Full => (true, true, Alignment::Right), + VersionBadgeMode::HeroFooter => (true, true, Alignment::Right), + VersionBadgeMode::HeroInline => (false, false, Alignment::Left), }; if show_team && let Some(team) = team_name { spans.push(Span::styled(team, Style::default().fg(theme.gray))); spans.push(sep.clone()); } - if show_tier - && let VersionBadgeMode::Full { - subscription_tier: Some(tier), - } = &mode - { - spans.push(Span::styled( - format!("Tier: {tier}"), - Style::default().fg(theme.gray), - )); - spans.push(sep.clone()); - } if show_api_key && is_api_key_auth { spans.push(Span::styled( "Logged in with API key", @@ -393,9 +379,9 @@ pub(super) fn render_version_badge( let channel = kigi_update::channel_label(); match &mode { - VersionBadgeMode::Full { .. } => { + VersionBadgeMode::Full => { spans.push(Span::styled( - "Grok Build ", + "Kigi ", Style::default() .fg(theme.text_primary) .add_modifier(Modifier::BOLD), @@ -404,16 +390,10 @@ pub(super) fn render_version_badge( format!("{}{}", kigi_version::VERSION, channel), Style::default().fg(theme.gray), )); - spans.push(Span::styled( - " Beta", - Style::default() - .fg(theme.text_primary) - .add_modifier(Modifier::BOLD), - )); } VersionBadgeMode::HeroFooter => { let channel_display = if channel.is_empty() { - "Beta" + "Kigi" } else { channel.trim() }; @@ -424,7 +404,7 @@ pub(super) fn render_version_badge( } VersionBadgeMode::HeroInline => { spans.push(Span::styled( - "Grok Build Beta ", + "Kigi ", Style::default() .fg(theme.text_primary) .add_modifier(Modifier::BOLD), @@ -520,9 +500,7 @@ fn render_prompt_and_version( team_name, h_margin, is_api_key_auth, - VersionBadgeMode::Full { - subscription_tier: None, - }, + VersionBadgeMode::Full, ); } else { render_version_badge( @@ -554,11 +532,8 @@ pub struct WelcomeRenderParams<'a> { pub model_name: &'a str, pub flags: &'a [PromptFlag<'a>], pub selected: Option, - pub team_name: Option<&'a str>, - pub has_access: bool, pub has_claude_import: bool, pub mouse_pos: Option<(u16, u16)>, - pub is_zdr_blocked: bool, pub session_picker: Option<&'a [SessionPickerEntry]>, pub session_picker_loading: bool, pub compact: bool, @@ -575,8 +550,6 @@ pub struct WelcomeRenderParams<'a> { /// [`crate::views::session_picker::effective_filter_query`]). pub session_picker_entries_query: Option<&'a str>, pub welcome_tick: u64, - pub gate: Option<&'a kigi_shell::auth::GateInfo>, - pub subscription_tier: Option<&'a str>, pub session_picker_grouped: bool, /// Source filter (local/remote/all) for the session picker. pub session_picker_source_filter: crate::views::session_picker::SourceFilter, @@ -586,12 +559,6 @@ pub struct WelcomeRenderParams<'a> { /// Live working directory (tracks `Effect::SetWorkingDir`), used to pin /// the current repo's session group to the top of the picker. pub cwd: &'a std::path::Path, - /// App-level credit balance for showing the usage warning on the welcome screen. - pub credit_balance: Option<&'a crate::views::credit_bar::CreditBalance>, - /// Auto top-up rule paired with `credit_balance` for the welcome warning. - pub auto_topup: Option<&'a crate::views::credit_bar::AutoTopupInfo>, - /// Whether /usage is visible (false for team users — suppresses the warning). - pub usage_visible: bool, /// Cached changelog bullets for the welcome screen (up to 3). pub changelog_bullets: &'a [String], /// Whether full release notes markdown is available (controls the CTA hint). @@ -635,7 +602,7 @@ pub fn render_welcome( let mut result = match params.auth_state { AuthState::Pending { error } => { - let label = params.login_label.unwrap_or("grok.com"); + let label = params.login_label.unwrap_or("kimi.com"); let login_text = format!("Login with {}", label); let menu = [("l", login_text.as_str()), ("q", "Quit")]; let msg = error.as_deref().map(|e| (e, theme.accent_error)); @@ -643,8 +610,6 @@ pub fn render_welcome( model_name: params.model_name, flags: params.flags, multiline: false, - usage_warning: None, - usage_warning_critical: false, }; let (menu_rects, post_flush_escapes) = render_welcome_blocked( content_area, @@ -665,8 +630,6 @@ pub fn render_welcome( import_banner_rect: None, auth_url_rect: None, auth_fallback_rect: None, - refresh_rect: None, - gate_url_rect: None, changelog_action_present: false, changelog_cta_rect: None, } @@ -693,49 +656,15 @@ pub fn render_welcome( import_banner_rect: None, auth_url_rect: url_rect, auth_fallback_rect: fallback_rect, - refresh_rect: None, - gate_url_rect: None, - changelog_action_present: false, - changelog_cta_rect: None, - } - } - AuthState::Done if params.is_zdr_blocked => { - let menu = [("l", "Switch account"), ("q", "Quit")]; - let (menu_rects, post_flush_escapes) = render_welcome_blocked( - content_area, - buf, - Some(( - "Grok Build is not yet available for this account.", - theme.gray_bright, - )), - &menu, - params.selected, - None, - h_margin, - params.compact, - ); - WelcomeRenderResult { - cursor_pos: None, - post_flush_escapes, - menu_rects, - prompt_rect: None, - session_picker_hit_areas: None, - import_banner_rect: None, - auth_url_rect: None, - auth_fallback_rect: None, - refresh_rect: None, - gate_url_rect: None, changelog_action_present: false, changelog_cta_rect: None, } } // Folder-trust question: shown after auth, before any session is // created, when the cwd has untrusted repo-local config. Mirrors the - // Pending login screen. Skipped under ZDR/access gates (the ZDR arm - // above and the !has_access arm below) since those already block - // sessions. The `if let` destructure makes the `Pending`-only render - // structurally exhaustive (no `unreachable!`). - AuthState::Done if params.has_access => { + // Pending login screen. The `if let` destructure makes the + // `Pending`-only render structurally exhaustive (no `unreachable!`). + AuthState::Done => { if let TrustState::Pending { workspace } = params.trust_state { render_welcome_trust( content_area, @@ -758,15 +687,6 @@ pub fn render_welcome( ) } } - AuthState::Done => render_welcome_done( - content_area, - buf, - &theme, - params, - prompt, - session_picker_state, - h_margin, - ), }; if result.post_flush_escapes.is_none() { result.post_flush_escapes = crate::terminal::overlay::clear().map(Into::into); @@ -851,16 +771,14 @@ fn render_welcome_blocked( None, h_margin, false, - VersionBadgeMode::Full { - subscription_tier: None, - }, + VersionBadgeMode::Full, ); (menu_rects, post_flush_escapes) } /// Render the folder-trust question. Mirrors [`render_welcome_blocked`]'s /// stacked layout (logo + message + menu + version badge), but the message is a -/// multi-line block showing the workspace path and the warning that Grok Build +/// multi-line block showing the workspace path and the warning that Kigi /// may run or modify contents in this directory (a security risk). The y/N /// answer is handled by the welcome input interceptor, so this only paints; /// `menu_rects` are returned for parity with the other welcome arms. @@ -889,7 +807,7 @@ fn render_welcome_trust( // Two lines so the warning never clips at narrow / compact widths // (a single ~78-char line would truncate "...posing security risks"). Line::from(Span::styled( - "Grok Build may run or modify contents in this directory,", + "Kigi may run or modify contents in this directory,", Style::default().fg(theme.gray), )) .alignment(Alignment::Center), @@ -926,9 +844,7 @@ fn render_welcome_trust( None, h_margin, false, - VersionBadgeMode::Full { - subscription_tier: None, - }, + VersionBadgeMode::Full, ); // Only `menu_rects` are meaningful here; the rest are absent (no prompt, @@ -1527,16 +1443,7 @@ fn render_welcome_done( // normal welcome layout. let welcome_compact = show_picker; - let cta = p - .gate - .and_then(|g| g.label.as_deref()) - .unwrap_or("Upgrade Subscription"); let in_vscode_family = welcome_in_vscode_family(); - let (key_g, key_l, key_q) = ( - "ctrl+g", - "ctrl+l", - if in_vscode_family { "ctrl+d" } else { "ctrl+q" }, - ); // Heights that don't depend on the menu — computed first so the menu // builder can probe the layout to decide whether to add a Changelog row. @@ -1561,21 +1468,17 @@ fn render_welcome_done( } else { 0 }; - let changelog_height = if p.has_access && !show_picker && !p.changelog_bullets.is_empty() { + let changelog_height = if !show_picker && !p.changelog_bullets.is_empty() { 2 + p.changelog_bullets.len() as u16 } else { 0 }; // Changelog is reachable via this menu row (ctrl+l). Show from the first // frame so the menu doesn't shift while the CDN fetch completes. - let show_changelog_action = p.has_access && !show_picker; + let show_changelog_action = !show_picker; - let gate_menu; let owned_menu; - let menu_items: &[(&str, &str)] = if !p.has_access { - gate_menu = [(key_g, cta), (key_l, "Logout"), (key_q, "Quit")]; - &gate_menu - } else { + let menu_items: &[(&str, &str)] = { let (key_w, key_s, key_q, key_i_with_x) = ( "ctrl+w", "ctrl+s", @@ -1723,116 +1626,8 @@ fn render_welcome_done( // Skip the prompt input when picker is visible to save space; // shortcuts are rendered inside the picker content area. - let mut refresh_hit_rect: Option = None; - let mut gate_url_hit_rect: Option = None; let (cursor_pos, post_flush_escapes) = if show_picker { (None, None) - } else if !p.has_access { - // Show CTA message and version instead of the prompt. - let [_, centered, _] = Layout::horizontal([ - Constraint::Min(0), - Constraint::Length(content_area.width), - Constraint::Min(0), - ]) - .flex(Flex::Center) - .areas(layout.prompt); - // Show the user's current tier + clickable refresh button above the gate message. - let tier_label = p.subscription_tier.unwrap_or("Free"); - let tier_prefix = format!("Tier: {tier_label} "); - let refresh_text = "[Refresh]"; - let total_width = tier_prefix.len() + refresh_text.len(); - let tier_line = Line::from(vec![ - Span::styled("Tier: ", Style::default().fg(theme.gray)), - Span::styled( - tier_label, - Style::default() - .fg(theme.gray_bright) - .add_modifier(Modifier::BOLD), - ), - Span::styled(" ", Style::default()), - Span::styled( - refresh_text, - Style::default() - .fg(theme.accent_user) - .add_modifier(Modifier::UNDERLINED), - ), - ]) - .alignment(Alignment::Center); - let tier_area = Rect { - height: 1, - ..centered - }; - Paragraph::new(tier_line).render(tier_area, buf); - - // Compute the click rect for "[Refresh]" within the centered line. - let line_start_x = tier_area.x + tier_area.width.saturating_sub(total_width as u16) / 2; - refresh_hit_rect = Some(Rect { - x: line_start_x + tier_prefix.len() as u16, - y: tier_area.y, - width: refresh_text.len() as u16, - height: 1, - }); - - let gate_text = p - .gate - .map(|g| g.message.as_str()) - .unwrap_or("SuperGrok subscription required"); - let msg = Line::from(Span::styled( - gate_text, - Style::default().fg(theme.gray_bright), - )) - .alignment(Alignment::Center); - Paragraph::new(msg).render( - Rect { - y: centered.y + 1, - height: 1, - ..centered - }, - buf, - ); - - if centered.height > 2 { - let url_area = Rect { - y: centered.y + 2, - height: 1, - ..centered - }; - let gate_link = p - .gate - .and_then(|g| g.url.as_deref()) - .unwrap_or("https://grok.com/supergrok?referrer=grok-build"); - let url = Line::from(Span::styled( - gate_link, - Style::default() - .fg(theme.accent_user) - .add_modifier(Modifier::UNDERLINED), - )) - .alignment(Alignment::Center); - Paragraph::new(url).render(url_area, buf); - - // Compute click rect for the gate URL text (centered within url_area). - let link_width = gate_link.len() as u16; - let link_x = url_area.x + url_area.width.saturating_sub(link_width) / 2; - gate_url_hit_rect = Some(Rect { - x: link_x, - y: url_area.y, - width: link_width.min(url_area.width), - height: 1, - }); - } - - render_version_badge( - layout.version, - buf, - theme, - p.team_name, - h_margin, - p.is_api_key_auth, - VersionBadgeMode::Full { - subscription_tier: p.subscription_tier, - }, - ); - (None, None) } else { // When a background update is available, show the update // notification in the tip area instead of the random tip. @@ -1913,19 +1708,10 @@ fn render_welcome_done( .render(tip_inset, buf); } - let warning = p.credit_balance.and_then(|bal| { - crate::views::credit_bar::usage_warning(bal, p.auto_topup, p.usage_visible) - }); - let (usage_warning_text, usage_warning_critical) = match warning { - Some((text, critical)) => (Some(text), critical), - None => (None, false), - }; let usage_info = PromptInfo { model_name: p.model_name, flags: p.flags, multiline: false, - usage_warning: usage_warning_text.as_deref(), - usage_warning_critical, }; render_prompt_and_version( @@ -1942,7 +1728,7 @@ fn render_welcome_done( } else { p.tip }, - p.team_name, + None, h_margin, p.compact, p.pending_hint, @@ -1955,7 +1741,7 @@ fn render_welcome_done( cursor_pos, post_flush_escapes, menu_rects, - prompt_rect: if show_picker || !p.has_access { + prompt_rect: if show_picker { None } else { Some(layout.prompt) @@ -1964,8 +1750,6 @@ fn render_welcome_done( import_banner_rect, auth_url_rect: None, auth_fallback_rect: None, - refresh_rect: refresh_hit_rect, - gate_url_rect: gate_url_hit_rect, changelog_action_present: show_changelog_action, changelog_cta_rect, } @@ -2377,11 +2161,8 @@ mod tests { model_name: "test", flags: &[], selected: None, - team_name: None, - has_access: true, has_claude_import: false, mouse_pos: None, - is_zdr_blocked: false, session_picker, session_picker_loading: false, compact: false, @@ -2394,15 +2175,10 @@ mod tests { session_picker_content_loading: false, session_picker_entries_query: None, welcome_tick: 0, - gate: None, - subscription_tier: None, session_picker_grouped: false, session_picker_source_filter: crate::views::session_picker::SourceFilter::All, chat_mode: false, cwd: std::path::Path::new("/repo"), - credit_balance: None, - auto_topup: None, - usage_visible: true, changelog_bullets: &[], changelog_has_full_notes: false, } @@ -3156,24 +2932,33 @@ mod tests { #[test] fn extract_user_code_parses_verification_url() { + // The live Kimi device flow returns this exact URL shape. assert_eq!( - extract_user_code("https://accounts.x.ai/oauth2/device?user_code=ABCD-EFGH"), + extract_user_code("https://www.kimi.com/code/authorize_device?user_code=ABCD-EFGH"), Some("ABCD-EFGH"), ); // Trailing params after the code are ignored. assert_eq!( - extract_user_code("https://x.ai/oauth2/device?user_code=WXYZ-1234&foo=bar"), + extract_user_code( + "https://www.kimi.com/code/authorize_device?user_code=WXYZ-1234&foo=bar" + ), Some("WXYZ-1234"), ); // A param whose name merely ends in `user_code` must not be matched. assert_eq!( - extract_user_code("https://x.ai/d?foo_user_code=BAD&user_code=GOOD"), + extract_user_code("https://example.com/d?foo_user_code=BAD&user_code=GOOD"), Some("GOOD"), ); // No code param, empty code, and unexpected characters all yield None. - assert_eq!(extract_user_code("https://x.ai/oauth2/device"), None); - assert_eq!(extract_user_code("https://x.ai/d?user_code="), None); - assert_eq!(extract_user_code("https://x.ai/d?user_code=AB%20CD"), None); + assert_eq!( + extract_user_code("https://www.kimi.com/code/authorize_device"), + None + ); + assert_eq!(extract_user_code("https://example.com/d?user_code="), None); + assert_eq!( + extract_user_code("https://example.com/d?user_code=AB%20CD"), + None + ); } #[test] @@ -3181,7 +2966,7 @@ mod tests { let area = Rect::new(0, 0, 80, 40); let mut buf = Buffer::empty(area); let theme = Theme::current(); - let url = "https://accounts.x.ai/oauth2/device?user_code=ABCD-EFGH"; + let url = "https://www.kimi.com/code/authorize_device?user_code=ABCD-EFGH"; let (copy_rect, fallback_rect) = render_welcome_authenticating( area, @@ -3235,7 +3020,7 @@ mod tests { let area = Rect::new(0, 0, 80, 40); let mut buf = Buffer::empty(area); let theme = Theme::current(); - let url = "https://accounts.x.ai/oauth2/device?user_code=WXYZ-1234"; + let url = "https://www.kimi.com/code/authorize_device?user_code=WXYZ-1234"; render_welcome_authenticating( area, @@ -3261,7 +3046,7 @@ mod tests { let area = Rect::new(0, 0, 80, 40); let mut buf = Buffer::empty(area); let theme = Theme::current(); - let url = "https://accounts.x.ai/oauth2/device?user_code=WXYZ-1234"; + let url = "https://www.kimi.com/code/authorize_device?user_code=WXYZ-1234"; render_welcome_authenticating( area, @@ -3298,7 +3083,7 @@ mod tests { let theme = Theme::current(); // 40-col terminal; URL longer than one row must wrap at the exact // screen edge with no leading spaces so copy-paste stays intact. - let url = "https://accounts.x.ai/oauth2/device?user_code=WXYZ-1234&extra=0123456789"; + let url = "https://www.kimi.com/code/authorize_device?user_code=WXYZ-1234&extra=0123456789"; render_welcome_authenticating( area, @@ -3337,7 +3122,7 @@ mod tests { let area = Rect::new(0, 0, 80, 40); let mut buf = Buffer::empty(area); let theme = Theme::current(); - let url = "https://accounts.x.ai/oauth2/authorize?client_id=grok"; + let url = "https://example.com/oauth2/authorize?client_id=kigi"; let (copy_rect, fallback_rect) = render_welcome_authenticating( area, diff --git a/crates/codegen/kigi-tui/src/views/welcome/prompt.rs b/crates/codegen/kigi-tui/src/views/welcome/prompt.rs index db2a2a7..4fb179d 100644 --- a/crates/codegen/kigi-tui/src/views/welcome/prompt.rs +++ b/crates/codegen/kigi-tui/src/views/welcome/prompt.rs @@ -79,8 +79,6 @@ mod tests { model_name: "test", flags: &[], multiline: false, - usage_warning: None, - usage_warning_critical: false, }; let (_, post_flush) = render_prompt( diff --git a/crates/codegen/kigi-tui/tests/pty_e2e/managed_policy_gate_refusal_reaches_real_terminal.rs b/crates/codegen/kigi-tui/tests/pty_e2e/managed_policy_gate_refusal_reaches_real_terminal.rs index 3cf35d8..d7a1a8d 100644 --- a/crates/codegen/kigi-tui/tests/pty_e2e/managed_policy_gate_refusal_reaches_real_terminal.rs +++ b/crates/codegen/kigi-tui/tests/pty_e2e/managed_policy_gate_refusal_reaches_real_terminal.rs @@ -15,7 +15,7 @@ async fn managed_policy_gate_refusal_reaches_real_terminal() { "[endpoints]\n\ deployment_key = \"KEY-AAA\"\n\ managed_config_url = \"http://127.0.0.1:1/deployment/config\"\n\ - cli_chat_proxy_base_url = \"http://127.0.0.1:1\"\n", + coding_api_base_url = \"http://127.0.0.1:1\"\n", ) .expect("write config.toml"); std::fs::write( diff --git a/crates/codegen/kigi-tui/tests/settings_e2e.rs b/crates/codegen/kigi-tui/tests/settings_e2e.rs index c479942..6a5cee0 100644 --- a/crates/codegen/kigi-tui/tests/settings_e2e.rs +++ b/crates/codegen/kigi-tui/tests/settings_e2e.rs @@ -50,7 +50,6 @@ const ALL_SETTINGS_EXERCISED: &[&str] = &[ "scroll_lines", "invert_scroll", "display_refresh_auto_cadence", - "coding_data_sharing", "default_selected_permission", "plan_mode", "show_tips", @@ -1608,7 +1607,6 @@ fn registry_kind_membership_through_pr_14() { vec![ "auto_dark_theme", "auto_light_theme", - "coding_data_sharing", "default_selected_permission", "hunk_tracker_mode", "keep_text_selection", @@ -1675,7 +1673,6 @@ fn enum_settings_membership_through_pr_14() { vec![ "auto_dark_theme", "auto_light_theme", - "coding_data_sharing", "default_selected_permission", "hunk_tracker_mode", "keep_text_selection", @@ -1736,7 +1733,6 @@ fn defaults_round_trip_through_registry() { "scroll_lines" => SettingValue::Int(3), "invert_scroll" => SettingValue::Bool(false), "display_refresh_auto_cadence" => SettingValue::Bool(false), - "coding_data_sharing" => SettingValue::Enum("opt-in"), "default_selected_permission" => SettingValue::Enum("always_allow_all_sessions"), "hunk_tracker_mode" => SettingValue::Enum("agent_only"), "plan_mode" => SettingValue::Enum("off"), @@ -4380,382 +4376,6 @@ fn pr8_default_model_and_max_thoughts_width_defaults_roundtrip() { ); } -// --------------------------------------------------------------------------- -// coding_data_sharing (Privacy Enum, no preview — async ACP) -// --------------------------------------------------------------------------- - -/// `coding_data_sharing` lives under `Privacy`. -#[test] -fn pr9_coding_data_sharing_renders_under_privacy_category() { - let reg = SettingsRegistry::defaults(); - let meta = reg - .find("coding_data_sharing") - .expect("coding_data_sharing must be registered"); - assert_eq!( - meta.category, - SettingCategory::Privacy, - "coding_data_sharing must live under Privacy" - ); - assert_eq!( - meta.owner, - SettingOwner::Shell, - "coding_data_sharing is SHELL-owned (auth-metadata-backed, persists via ACP)" - ); -} - -/// `coding_data_sharing` must be `supports_preview: false` (async ACP). -#[test] -fn pr9_coding_data_sharing_does_not_support_preview() { - let reg = SettingsRegistry::defaults(); - let meta = reg - .find("coding_data_sharing") - .expect("coding_data_sharing must be registered"); - match &meta.kind { - SettingKind::Enum { - supports_preview, .. - } => { - assert!( - !supports_preview, - "coding_data_sharing MUST be supports_preview: false — every preview \ - would fire an async ACP round-trip OR commit-on-every-nav, both \ - unacceptable", - ); - } - other => panic!("expected Enum kind for coding_data_sharing, got {other:?}"), - } -} - -/// Reads from pager snapshot; inverts `_opt_out` bool. -#[test] -fn pr9_current_value_for_reads_pager_snapshot_inverts_opt_out() { - use kigi_tui::settings::current_value_for; - - let ui = UiConfig::default(); - - let opted_in_snap = PagerLocalSnapshot { - coding_data_sharing_opt_out: false, - ..PagerLocalSnapshot::default() - }; - let opted_out_snap = PagerLocalSnapshot { - coding_data_sharing_opt_out: true, - ..PagerLocalSnapshot::default() - }; - - assert_eq!( - current_value_for("coding_data_sharing", &ui, &opted_in_snap), - Some(SettingValue::Enum("opt-in")), - "opt_out=false → canonical 'opt-in' (user IS sharing data)", - ); - assert_eq!( - current_value_for("coding_data_sharing", &ui, &opted_out_snap), - Some(SettingValue::Enum("opt-out")), - "opt_out=true → canonical 'opt-out' (user opted OUT of sharing)", - ); -} - -/// Enter opens picker seeded to current state. -#[test] -fn pr9_enter_on_coding_data_sharing_row_enters_picking_enum() { - let mut s = make_state(); - navigate_to(&mut s, "coding_data_sharing"); - let outcome = handle_settings_key(&mut s, &press(KeyCode::Enter)); - assert!( - matches!(outcome, SettingsKeyOutcome::Changed), - "Enter on coding_data_sharing row must transition to PickingEnum, got {outcome:?}" - ); - match &s.mode { - SettingsModalMode::PickingEnum { - key, - original_value, - .. - } => { - assert_eq!(*key, "coding_data_sharing"); - assert_eq!( - original_value, - &SettingValue::Enum("opt-in"), - "default snapshot opt_out=false → original 'opt-in'" - ); - } - other => panic!("expected PickingEnum mode, got {other:?}"), - } -} - -/// Nav in picker must NOT dispatch preview (async ACP). -#[test] -fn pr9_coding_data_sharing_picker_nav_does_not_dispatch_preview() { - for nav_key in &[ - KeyCode::Down, - KeyCode::Char('j'), - KeyCode::Up, - KeyCode::Char('k'), - ] { - let mut s = make_state(); - navigate_to(&mut s, "coding_data_sharing"); - let _ = handle_settings_key(&mut s, &press(KeyCode::Enter)); - assert!(matches!(s.mode, SettingsModalMode::PickingEnum { .. })); - - if matches!(nav_key, KeyCode::Up | KeyCode::Char('k')) { - let _ = handle_settings_key(&mut s, &press(KeyCode::Down)); - } - - let outcome = handle_settings_key(&mut s, &press(*nav_key)); - assert!( - matches!(outcome, SettingsKeyOutcome::Changed), - "Nav key {nav_key:?} in coding_data_sharing picker MUST NOT dispatch a preview \ - Action — that would fire a network round-trip per keystroke. Got {outcome:?}", - ); - assert!(matches!(s.mode, SettingsModalMode::PickingEnum { .. })); - } -} - -/// Enter commits `SetCodingDataSharing { opted_in }` (opt-in→true). -#[test] -fn pr9_coding_data_sharing_picker_enter_dispatches_set_commit() { - let reg = SettingsRegistry::defaults(); - let meta = reg.find("coding_data_sharing").unwrap(); - let (default_canonical, choices) = match &meta.kind { - SettingKind::Enum { - default, choices, .. - } => (*default, *choices), - _ => panic!("coding_data_sharing must be Enum"), - }; - // Resolve "the other" canonical from the registry rather than - // hardcoding — robust against future catalog additions. - let other_canonical = choices - .iter() - .map(|c| c.canonical) - .find(|c| *c != default_canonical) - .expect("coding_data_sharing must have ≥2 choices"); - let expected_opted_in = match other_canonical { - "opt-in" => true, - "opt-out" => false, - _ => panic!("unexpected canonical: {other_canonical:?}"), - }; - - let mut s = make_state(); - navigate_to(&mut s, "coding_data_sharing"); - let _ = handle_settings_key(&mut s, &press(KeyCode::Enter)); - // Nav to the OTHER choice. - let _ = handle_settings_key(&mut s, &press(KeyCode::Down)); - // Enter → commit. - let outcome = handle_settings_key(&mut s, &press(KeyCode::Enter)); - match outcome { - SettingsKeyOutcome::Action(Action::SetCodingDataSharing { opted_in }) => { - assert_eq!( - opted_in, expected_opted_in, - "Enter must commit `{other_canonical}` → SetCodingDataSharing(opted_in={expected_opted_in})" - ); - } - other => panic!("expected Action::SetCodingDataSharing commit, got {other:?}"), - } - assert!( - matches!(s.mode, SettingsModalMode::Browse), - "Enter commit must return to Browse" - ); -} - -/// Esc in non-preview picker returns to Browse without Action. -#[test] -fn pr9_coding_data_sharing_picker_esc_does_not_dispatch_action() { - let mut s = make_state(); - navigate_to(&mut s, "coding_data_sharing"); - let _ = handle_settings_key(&mut s, &press(KeyCode::Enter)); - let _ = handle_settings_key(&mut s, &press(KeyCode::Down)); - - let outcome = handle_settings_key(&mut s, &press(KeyCode::Esc)); - assert!( - matches!(outcome, SettingsKeyOutcome::Changed), - "Esc on non-preview Enum picker must NOT emit an Action — \ - doing so would fire an ACP round-trip on every Esc. Got {outcome:?}" - ); - assert!( - matches!(s.mode, SettingsModalMode::Browse), - "Esc must return to Browse" - ); -} - -/// Picker seeds at "opt-out" when `coding_data_sharing_opt_out: true`. -#[test] -fn pr9_picker_seeds_choices_idx_from_pager_snapshot_opt_out_true() { - let snapshot = PagerLocalSnapshot { - coding_data_sharing_opt_out: true, - ..PagerLocalSnapshot::default() - }; - let mut s = SettingsModalState::new( - Arc::new(SettingsRegistry::defaults()), - UiConfig::default(), - snapshot, - ); - navigate_to(&mut s, "coding_data_sharing"); - let _ = handle_settings_key(&mut s, &press(KeyCode::Enter)); - let reg = SettingsRegistry::defaults(); - let opt_out_idx = match ®.find("coding_data_sharing").unwrap().kind { - SettingKind::Enum { choices, .. } => choices - .iter() - .position(|c| c.canonical == "opt-out") - .expect("coding_data_sharing must have 'opt-out' choice"), - _ => panic!("coding_data_sharing must be Enum"), - }; - match s.mode { - SettingsModalMode::PickingEnum { - choices_idx, - ref original_value, - .. - } => { - assert_eq!( - choices_idx, opt_out_idx, - "picker must seed at the 'opt-out' index when snapshot says opt_out=true" - ); - assert_eq!( - original_value, - &SettingValue::Enum("opt-out"), - "original_value must match the live snapshot" - ); - } - ref other => panic!("expected PickingEnum mode, got {other:?}"), - } -} - -/// Exactly 2 canonical choices: {opt-in, opt-out}. -#[test] -fn pr9_coding_data_sharing_choices_use_canonical_strings() { - let reg = SettingsRegistry::defaults(); - let meta = reg.find("coding_data_sharing").unwrap(); - let canonicals: Vec<&str> = match &meta.kind { - SettingKind::Enum { choices, .. } => choices.iter().map(|c| c.canonical).collect(), - _ => panic!("coding_data_sharing must be Enum"), - }; - assert_eq!( - canonicals.len(), - 2, - "coding_data_sharing catalog must be exactly {{opt-in, opt-out}} — adding a \ - choice requires updating the action_for_enum_commit arm in \ - views/settings_modal.rs AND the action_for_reset arm in dispatch.rs", - ); - assert!( - canonicals.contains(&"opt-in"), - "coding_data_sharing must include 'opt-in' canonical" - ); - assert!( - canonicals.contains(&"opt-out"), - "coding_data_sharing must include 'opt-out' canonical" - ); -} - -/// Search "privacy" finds exactly `coding_data_sharing`. -#[test] -fn pr9_search_privacy_matches_coding_data_sharing() { - let reg = SettingsRegistry::defaults(); - let hits = reg.search("privacy"); - // The category label "Privacy" appears as a header but is not - // part of `search()`'s haystack (search ignores categories); - // matches come from the meta's keywords + label + description. - let hit_keys: Vec<&str> = hits.iter().map(|m| m.key).collect(); - assert_eq!( - hits.len(), - 1, - "search('privacy') must return EXACTLY one result (coding_data_sharing). \ - Found {} results: {hit_keys:?}. \ - If this fails because another setting added 'privacy' to its keywords/label/\ - description, decide: (a) is 'privacy' a real keyword for that setting? If yes, \ - loosen this assertion to a presence-only check `hit_keys.contains(&\"coding_data_sharing\")`. \ - (b) If no, remove 'privacy' from the other setting's haystack — search relevance \ - is more important than tag promiscuity.", - hits.len(), - ); - assert_eq!( - hits[0].key, "coding_data_sharing", - "search('privacy') unique result must be coding_data_sharing" - ); -} - -// --------------------------------------------------------------------------- -// Mouse path tests for coding_data_sharing -// --------------------------------------------------------------------------- - -/// First click on unselected row only selects. -#[test] -fn pr9_mouse_click_on_unselected_coding_data_sharing_row_only_selects() { - let mut s = make_state(); - synth_rects(&mut s); - let row_y = row_idx_for(&s, "coding_data_sharing") as u16; - - let outcome = handle_settings_mouse( - &mut s, - MouseEventKind::Down(crossterm::event::MouseButton::Left), - 10, - row_y, - ); - assert!( - matches!(outcome, SettingsKeyOutcome::Changed), - "first body-click on unselected coding_data_sharing row should only select, got: {outcome:?}", - ); - assert_eq!(s.selected, row_y as usize); - assert!(matches!(s.mode, SettingsModalMode::Browse)); -} - -/// Second click on selected row opens picker. -#[test] -fn pr9_mouse_click_on_selected_coding_data_sharing_row_opens_picker() { - let mut s = make_state(); - synth_rects(&mut s); - let row_y = row_idx_for(&s, "coding_data_sharing") as u16; - - // First click: select. - let _ = handle_settings_mouse( - &mut s, - MouseEventKind::Down(crossterm::event::MouseButton::Left), - 10, - row_y, - ); - assert_eq!(s.selected, row_y as usize); - - // Second click on the focused row: open the picker. - let outcome = handle_settings_mouse( - &mut s, - MouseEventKind::Down(crossterm::event::MouseButton::Left), - 10, - row_y, - ); - assert!( - matches!(outcome, SettingsKeyOutcome::Changed), - "second click on focused Enum row must open picker, got: {outcome:?}", - ); - match &s.mode { - SettingsModalMode::PickingEnum { key, .. } => { - assert_eq!(*key, "coding_data_sharing"); - } - _ => panic!("second click on focused coding_data_sharing row must enter PickingEnum"), - } -} - -/// Value-column click opens picker in one click. -#[test] -fn pr9_mouse_click_on_coding_data_sharing_indicator_opens_picker_in_one_click() { - let mut s = make_state(); - synth_rects(&mut s); - let row_y = row_idx_for(&s, "coding_data_sharing") as u16; - - let outcome = handle_settings_mouse( - &mut s, - MouseEventKind::Down(crossterm::event::MouseButton::Left), - 72, - row_y, - ); - assert!( - matches!(outcome, SettingsKeyOutcome::Changed), - "value click must open picker in one click, got: {outcome:?}", - ); - match &s.mode { - SettingsModalMode::PickingEnum { key, .. } => { - assert_eq!(*key, "coding_data_sharing"); - } - _ => { - panic!("value click on coding_data_sharing must enter PickingEnum") - } - } -} - // --------------------------------------------------------------------------- // default_selected_permission (Agent Enum, no preview — SHELL-owned, persists) // --------------------------------------------------------------------------- @@ -5024,53 +4644,6 @@ fn default_selected_permission_mouse_click_on_indicator_opens_picker_in_one_clic } } -/// The `/privacy` slash command's argument parser -/// is case-insensitive and supports a deliberately-pared-down list of -/// unambiguous-semantic aliases. The unit-level coverage lives in the -/// slash command module; this e2e test pins the integration contract -/// (the parser is reachable from the slash command and produces the -/// expected `Action`). -/// -/// Ambiguous aliases -/// (`on/off/true/false/enable/disable`) were DROPPED because they -/// could be read either as "turn on privacy" (=opt-out) or "turn on -/// sharing" (=opt-in). For a privacy-critical setting we err on the -/// side of explicit, unambiguous arguments. The test below verifies -/// both the accept list AND the reject list. -#[test] -fn pr9_privacy_slash_command_parses_aliases() { - use kigi_tui::slash::commands::privacy::parse_privacy_arg; - - // Canonical names. - assert_eq!(parse_privacy_arg("opt-in"), Some(true)); - assert_eq!(parse_privacy_arg("opt-out"), Some(false)); - - // Case-insensitive (sample). - assert_eq!(parse_privacy_arg("Opt-In"), Some(true)); - assert_eq!(parse_privacy_arg("OPT-OUT"), Some(false)); - - // Unambiguous-semantic aliases (pruned list). - assert_eq!(parse_privacy_arg("in"), Some(true)); - assert_eq!(parse_privacy_arg("out"), Some(false)); - assert_eq!(parse_privacy_arg("share"), Some(true)); - assert_eq!(parse_privacy_arg("private"), Some(false)); - - // Ambiguous aliases MUST be rejected. `/privacy on` - // could be read as "turn on privacy" (=opt-out, the OPPOSITE of - // what an earlier mapping returned). For a privacy - // setting, ambiguity = silent data-exfiltration risk. - for ambiguous in &["on", "off", "true", "false", "enable", "disable"] { - assert_eq!( - parse_privacy_arg(ambiguous), - None, - "ambiguous alias `{ambiguous}` MUST be rejected (PR 9 R1, Security Issue 10)", - ); - } - - // Unknown. - assert_eq!(parse_privacy_arg("maybe"), None); -} - // --------------------------------------------------------------------------- // `plan_mode` (Agent-category Enum, PAGER-owned + ACP-mediated, // supports_preview: false) diff --git a/crates/codegen/kigi-workspace/src/handle.rs b/crates/codegen/kigi-workspace/src/handle.rs index 34e21c1..7293053 100644 --- a/crates/codegen/kigi-workspace/src/handle.rs +++ b/crates/codegen/kigi-workspace/src/handle.rs @@ -3171,8 +3171,7 @@ pub async fn connect_local_workspace( workspace_home.display() )) })?; - let api_base_url = std::env::var("KIGI_CLI_CHAT_PROXY_BASE_URL") - .unwrap_or_else(|_| kigi_env::coding_api_base_url()); + let api_base_url = kigi_env::coding_api_base_url(); let mut factory = WorkspaceSessionContextFactory::with_auth(auth.clone(), api_base_url.clone()); if crate::session::tool_config::tool_state_enabled() { factory = factory.with_tool_state_home(workspace_home.clone()); diff --git a/crates/codegen/kigi-workspace/src/permission/types.rs b/crates/codegen/kigi-workspace/src/permission/types.rs index 493ff55..a81a086 100644 --- a/crates/codegen/kigi-workspace/src/permission/types.rs +++ b/crates/codegen/kigi-workspace/src/permission/types.rs @@ -76,7 +76,12 @@ pub struct PermissionEvent { pub enum ClientType { /// Generic client - show simple permission options with full command text #[default] - #[serde(rename = "generic", alias = "grok-shell", alias = "grok_shell")] + #[serde( + rename = "generic", + alias = "grok-shell", + alias = "grok_shell", + alias = "kigi" + )] Generic, /// Grok TUI client - show fancy options with interactive bash term selection #[serde(rename = "grok-tui", alias = "grok_tui")] @@ -106,15 +111,18 @@ pub enum ClientType { Desktop, } impl ClientType { - /// Product token for the `User-Agent` header (e.g. `grok-pager`). + /// Product token for the `User-Agent` header. The first-party clients + /// (the bundled TUI/pager and the headless runner) all report `kigi`, so + /// their User-Agent collapses to `kigi/{version}` (PRD F3); the remaining + /// labels identify external ACP clients. pub fn user_agent_label(&self) -> &'static str { match self { - Self::Generic => "grok-shell", + Self::Generic => "kigi", Self::GrokTUI => "grok-tui", Self::GrokWeb => "grok-web", Self::Nebula => "nebula", Self::Extension => "grok-code-extension", - Self::GrokPager => "grok-pager", + Self::GrokPager => "kigi", Self::Desktop => "grok-desktop", } } diff --git a/prod/mc/cli-chat-proxy-types/Cargo.toml b/prod/mc/cli-chat-proxy-types/Cargo.toml deleted file mode 100644 index 78e0af3..0000000 --- a/prod/mc/cli-chat-proxy-types/Cargo.toml +++ /dev/null @@ -1,20 +0,0 @@ -[package] -license = "Apache-2.0" -name = "prod-mc-cli-chat-proxy-types" -version.workspace = true -edition.workspace = true -description = "Lightweight request/response types for cli-chat-proxy API" - -[features] -default-bazel = [] - -[dependencies] -chrono = { workspace = true, features = ["serde"] } -serde = { workspace = true } -serde_json.workspace = true -# The canonical requirements-TOML fail_closed parse, shared by the signer -# (cli-chat-proxy) and the client (kigi-config) so they can't drift. -toml.workspace = true - -[lints] -workspace = true diff --git a/prod/mc/cli-chat-proxy-types/src/client_metrics_types.rs b/prod/mc/cli-chat-proxy-types/src/client_metrics_types.rs deleted file mode 100644 index 4d56051..0000000 --- a/prod/mc/cli-chat-proxy-types/src/client_metrics_types.rs +++ /dev/null @@ -1,55 +0,0 @@ -use serde::{Deserialize, Serialize}; - -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct ClientMetric { - pub metric: String, - pub value: f64, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub timestamp: Option>, - // Dedup key: server-side / downstream may use to drop replays. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub idempotency_key: Option, -} - -impl ClientMetric { - pub fn new(metric: impl Into, value: f64) -> Self { - Self { - metric: metric.into(), - value, - timestamp: None, - idempotency_key: None, - } - } - - pub fn with_timestamp(mut self, ts: chrono::DateTime) -> Self { - self.timestamp = Some(ts); - self - } - - pub fn with_idempotency_key(mut self, key: impl Into) -> Self { - self.idempotency_key = Some(key.into()); - self - } -} - -#[derive(Clone, Debug, Default, Serialize, Deserialize)] -pub struct ClientMetricsBatch { - pub events: Vec, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub process_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub session_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub client_version: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub client_type: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub os: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub arch: Option, -} - -#[derive(Clone, Debug, Default, Serialize, Deserialize)] -pub struct ClientMetricsResponse { - pub accepted: usize, -} diff --git a/prod/mc/cli-chat-proxy-types/src/deployment_config_types.rs b/prod/mc/cli-chat-proxy-types/src/deployment_config_types.rs deleted file mode 100644 index eb52aaa..0000000 --- a/prod/mc/cli-chat-proxy-types/src/deployment_config_types.rs +++ /dev/null @@ -1,105 +0,0 @@ -//! Signed deployment-config envelope: the wire contract between the -//! cli-chat-proxy signer and the client verifier. Shared so a field rename -//! breaks at compile time on both sides instead of silently failing verification. - -use serde::{Deserialize, Serialize}; - -/// The payload format version the server currently signs. Bump when the payload -/// gains semantics (e.g. an anti-replay counter or a key-fingerprint binding) so -/// verifiers can distinguish generations; `0` means a pre-versioned payload. -pub const SIGNED_PAYLOAD_VERSION: u32 = 1; - -/// The exact bytes the server signs: the served policy, the principal it is -/// bound to, and an expiry. Serialized once on the server and shipped verbatim -/// as `signed_payload`, so the client verifies the received bytes directly -/// instead of re-canonicalizing (no cross-language serialization drift). -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct SignedPayload { - /// Payload format version ([`SIGNED_PAYLOAD_VERSION`]); `default` 0 so - /// pre-versioned sidecars parse and verify unchanged. - #[serde(default)] - pub version: u32, - #[serde(default)] - pub deployment_id: Option, - #[serde(default)] - pub team_id: Option, - #[serde(default)] - pub managed_config: Option, - #[serde(default)] - pub requirements: Option, - /// Strict (fail-closed) opt-in, carried in the SIGNED bytes so a local actor can't - /// flip enforcement. `default` false so an older/unsigned payload stays lenient. - #[serde(default)] - pub fail_closed: bool, - /// Unix seconds after which the signature is no longer trusted. - pub expires_at: u64, - /// Identifies the signing key, so a rotation can be distinguished. - pub key_id: String, -} - -/// One signed envelope carried alongside the legacy policy fields in the -/// deployment-config response (additive: old clients ignore it). Also the -/// shape the client persists as its on-disk signature sidecar. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct SignatureEnvelope { - /// The exact JSON string that was signed (a serialized [`SignedPayload`]). - pub signed_payload: String, - /// Base64 (standard) Ed25519 signature over `signed_payload`'s UTF-8 bytes. - pub signature: String, - /// Untrusted (outside the signed bytes): a hint for picking among multiple - /// envelopes, never for selecting the verifying key — only the signed - /// payload's `key_id` is authoritative. - #[serde(default)] - pub key_id: String, -} - -/// Unix seconds now (saturating to 0 on a pre-epoch clock). -pub fn now_unix() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0) -} - -/// The `requirements.toml` opt-in key for strict (fail-closed) enforcement. -pub const FAIL_CLOSED_KEY: &str = "fail_closed"; - -/// Read the `fail_closed` opt-in from a requirements-TOML string — THE canonical parse, -/// shared by the cli-chat-proxy signer and the client so the two sides can't drift. -/// Invalid TOML or a non-bool value → `false`. -pub fn fail_closed_flag_from_str(requirements: &str) -> bool { - toml::from_str::(requirements) - .ok() - .and_then(|v| v.get(FAIL_CLOSED_KEY).and_then(toml::Value::as_bool)) - .unwrap_or(false) -} - -#[cfg(test)] -mod tests { - use super::*; - - /// The version field round-trips, and a pre-versioned payload (no `version` - /// key) defaults to 0 — old sidecars keep parsing. - #[test] - fn signed_payload_version_round_trips_and_defaults() { - let versioned = SignedPayload { - version: SIGNED_PAYLOAD_VERSION, - deployment_id: None, - team_id: Some("team-007".into()), - managed_config: None, - requirements: None, - fail_closed: false, - expires_at: 4_000_000_000, - key_id: "v1".into(), - }; - let json = serde_json::to_string(&versioned).unwrap(); - assert_eq!( - serde_json::from_str::(&json).unwrap(), - versioned - ); - - let legacy: SignedPayload = - serde_json::from_str(r#"{"expires_at": 1, "key_id": "v1"}"#).unwrap(); - assert_eq!(legacy.version, 0, "pre-versioned payloads default to 0"); - } -} diff --git a/prod/mc/cli-chat-proxy-types/src/feedback_types.rs b/prod/mc/cli-chat-proxy-types/src/feedback_types.rs deleted file mode 100644 index a063ce5..0000000 --- a/prod/mc/cli-chat-proxy-types/src/feedback_types.rs +++ /dev/null @@ -1,2024 +0,0 @@ -//! Feedback API request and response types. -//! -//! These types support the feedback collection system for Grok sessions. -//! The agent (kigi-shell) uses heuristics to determine when to request feedback, -//! and clients submit feedback through these types to the feedback backend. - -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; - -// ============================================================================ -// Enums -// ============================================================================ - -/// Type of client submitting feedback. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ClientType { - /// Terminal/CLI agent - #[default] - Agent, - /// Terminal UI - Tui, - /// Web interface - Web, - /// IDE extension (VS Code, JetBrains, etc.) - Extension, - /// Remote workspace / hosted agent client (wire value `nebula`). - Nebula, - /// Desktop (Electron app) - Desktop, -} - -impl std::fmt::Display for ClientType { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - ClientType::Agent => write!(f, "agent"), - ClientType::Tui => write!(f, "tui"), - ClientType::Web => write!(f, "web"), - ClientType::Extension => write!(f, "extension"), - ClientType::Nebula => write!(f, "nebula"), - ClientType::Desktop => write!(f, "desktop"), - } - } -} - -/// Type of feedback being submitted. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum FeedbackType { - /// Numeric rating only - #[default] - Rating, - /// Free-form text only - Text, - /// Both rating and text - RatingWithText, - /// Model preference comparison - ModelPreference, - /// Bug report - BugReport, - /// Feature request - FeatureRequest, -} - -impl std::fmt::Display for FeedbackType { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - FeedbackType::Rating => write!(f, "rating"), - FeedbackType::Text => write!(f, "text"), - FeedbackType::RatingWithText => write!(f, "rating_with_text"), - FeedbackType::ModelPreference => write!(f, "model_preference"), - FeedbackType::BugReport => write!(f, "bug_report"), - FeedbackType::FeatureRequest => write!(f, "feature_request"), - } - } -} - -/// Type of rating scale used. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum RatingType { - /// Thumbs up/down (-1, 0, 1) - Thumbs, - /// Star rating (1-5) - Stars, - /// Net Promoter Score (0-10) - Nps, -} - -impl std::fmt::Display for RatingType { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - RatingType::Thumbs => write!(f, "thumbs"), - RatingType::Stars => write!(f, "stars"), - RatingType::Nps => write!(f, "nps"), - } - } -} - -/// Strength of preference in a comparison. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum PreferenceStrength { - /// Strongly prefer one over the other - Strong, - /// Slightly prefer one over the other - Slight, - /// No preference, both equal - Tie, - /// Both responses are bad - BothBad, - /// Both responses are good - BothGood, -} - -impl std::fmt::Display for PreferenceStrength { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - PreferenceStrength::Strong => write!(f, "strong"), - PreferenceStrength::Slight => write!(f, "slight"), - PreferenceStrength::Tie => write!(f, "tie"), - PreferenceStrength::BothBad => write!(f, "both_bad"), - PreferenceStrength::BothGood => write!(f, "both_good"), - } - } -} - -/// Status of a feedback request. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum FeedbackRequestStatus { - /// Request created, not yet delivered to client - Pending, - /// Request delivered to client, awaiting response - Delivered, - /// User submitted feedback - Completed, - /// Request expired before user responded - Expired, - /// User dismissed the request without responding - Dismissed, -} - -impl std::fmt::Display for FeedbackRequestStatus { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - FeedbackRequestStatus::Pending => write!(f, "pending"), - FeedbackRequestStatus::Delivered => write!(f, "delivered"), - FeedbackRequestStatus::Completed => write!(f, "completed"), - FeedbackRequestStatus::Expired => write!(f, "expired"), - FeedbackRequestStatus::Dismissed => write!(f, "dismissed"), - } - } -} - -/// Context type for what the feedback is about. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum ContextType { - /// Feedback about a specific message - Message, - /// Feedback about the overall session/conversation - Session, - /// Feedback about a specific feature - Feature, - /// Feedback about tool usage - ToolUse, - /// General feedback - General, -} - -impl std::fmt::Display for ContextType { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - ContextType::Message => write!(f, "message"), - ContextType::Session => write!(f, "session"), - ContextType::Feature => write!(f, "feature"), - ContextType::ToolUse => write!(f, "tool_use"), - ContextType::General => write!(f, "general"), - } - } -} - -/// Type of feedback mode requested. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum FeedbackMode { - /// Thumbs up/down - Thumbs, - /// Star rating (1-5) - Stars, - /// Free-form text - Text, - /// Thumbs up/down with optional text comment - ThumbsText, - /// Star rating with optional text comment - StarsText, - /// Model comparison - Comparison, - /// Multi-question survey - Survey, - /// Net Promoter Score (0-10) - Nps, - /// NPS with optional text comment - NpsText, -} - -impl std::fmt::Display for FeedbackMode { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - FeedbackMode::Thumbs => write!(f, "thumbs"), - FeedbackMode::Stars => write!(f, "stars"), - FeedbackMode::Text => write!(f, "text"), - FeedbackMode::ThumbsText => write!(f, "thumbs_text"), - FeedbackMode::StarsText => write!(f, "stars_text"), - FeedbackMode::Comparison => write!(f, "comparison"), - FeedbackMode::Survey => write!(f, "survey"), - FeedbackMode::Nps => write!(f, "nps"), - FeedbackMode::NpsText => write!(f, "nps_text"), - } - } -} - -// ============================================================================ -// Request Types -// ============================================================================ - -/// Allowed `feedback_type` + value-field combinations. Construct submissions -/// via [`FeedbackSubmission::with_content`]. -#[derive(Debug, Clone)] -#[non_exhaustive] -pub enum FeedbackContent { - Rating { - rating_type: RatingType, - rating_value: i32, - }, - Text(String), - RatingWithText { - rating_type: RatingType, - rating_value: i32, - text: String, - }, -} - -impl FeedbackContent { - fn apply_to(self, s: &mut FeedbackSubmission) { - s.rating_type = None; - s.rating_value = None; - s.feedback_text = None; - match self { - Self::Rating { - rating_type, - rating_value, - } => { - s.feedback_type = FeedbackType::Rating; - s.rating_type = Some(rating_type); - s.rating_value = Some(rating_value); - } - Self::Text(text) => { - s.feedback_type = FeedbackType::Text; - s.feedback_text = Some(text); - } - Self::RatingWithText { - rating_type, - rating_value, - text, - } => { - s.feedback_type = FeedbackType::RatingWithText; - s.rating_type = Some(rating_type); - s.rating_value = Some(rating_value); - s.feedback_text = Some(text); - } - } - } -} - -/// Request body for POST /v1/feedback. Construct via -/// [`FeedbackSubmission::with_content`]; the `Default` impl exists for -/// builder-style construction and test fixtures and does not produce a valid -/// submission on its own (empty `session_id`). -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct FeedbackSubmission { - /// Session ID this feedback is for - pub session_id: String, - - /// User ID (optional, extracted from auth if not provided) - #[serde(skip_serializing_if = "Option::is_none")] - pub user_id: Option, - - /// Type of client submitting feedback - pub client_type: ClientType, - - /// Type of feedback being submitted - pub feedback_type: FeedbackType, - - /// Turn number within the session (optional) - #[serde(skip_serializing_if = "Option::is_none")] - pub turn_number: Option, - - /// Rating type (if applicable) - #[serde(skip_serializing_if = "Option::is_none")] - pub rating_type: Option, - - /// Rating value (interpretation depends on rating_type) - /// - thumbs: -1 (down), 0 (neutral), 1 (up) - /// - stars: 1-5 - /// - nps: 0-10 - #[serde(skip_serializing_if = "Option::is_none")] - pub rating_value: Option, - - /// Free-form feedback text - #[serde(skip_serializing_if = "Option::is_none")] - pub feedback_text: Option, - - /// Feedback categories (e.g., ["accuracy", "speed", "helpfulness"]) - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub feedback_categories: Vec, - - /// Message ID this feedback is about (if context_type is message) - #[serde(skip_serializing_if = "Option::is_none")] - pub message_id: Option, - - /// Model ID used for the response being rated - #[serde(skip_serializing_if = "Option::is_none")] - pub model_id: Option, - - /// Server-resolved model ID from the actual chat completion response. - #[serde(skip_serializing_if = "Option::is_none")] - pub resolved_model_id: Option, - - /// Checkpoint fingerprint from the inference provider (`system_fingerprint`). - #[serde( - default, - skip_serializing_if = "Option::is_none", - deserialize_with = "crate::serde_helpers::empty_string_as_none" - )] - pub model_fingerprint: Option, - - /// Context type for the feedback - #[serde(skip_serializing_if = "Option::is_none")] - pub context_type: Option, - - /// Feature name (if context_type is feature) - #[serde(skip_serializing_if = "Option::is_none")] - pub feature_name: Option, - - /// Tool name (if context_type is tool_use) - #[serde(skip_serializing_if = "Option::is_none")] - pub tool_name: Option, - - // === Experiment / Comparison Fields === - /// Experiment ID (e.g. for routing experiments) - #[serde(skip_serializing_if = "Option::is_none")] - pub experiment_id: Option, - - /// Comparison ID (e.g. for routing experiments) - #[serde(skip_serializing_if = "Option::is_none")] - pub comparison_id: Option, - - /// Preferred model ID in comparison - #[serde(skip_serializing_if = "Option::is_none")] - pub preferred_model_id: Option, - - /// Strength of preference - #[serde(skip_serializing_if = "Option::is_none")] - pub preference_strength: Option, - - /// Reasons for preference (e.g., ["more_accurate", "faster", "better_code"]) - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub preference_reasons: Vec, - - // === Request Link === - /// Feedback request ID (links to feedback_requests table if responding to a request) - #[serde(skip_serializing_if = "Option::is_none")] - pub request_id: Option, - - // === Client Metadata === - /// Client version - #[serde(skip_serializing_if = "Option::is_none")] - pub client_version: Option, - - /// Shell (kigi-shell) version - #[serde(skip_serializing_if = "Option::is_none")] - pub shell_version: Option, - - /// Extension host (for extension client type) - #[serde(skip_serializing_if = "Option::is_none")] - pub extension_host: Option, - - /// Additional metadata as JSON - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, - - // === Feedback Context === - // Persisted server-side alongside the feedback record. - /// Last user message at feedback time. - #[serde( - default, - skip_serializing_if = "Option::is_none", - alias = "last_user_turn" - )] - pub last_user_message: Option, - - /// Last assistant response at feedback time. - #[serde( - default, - skip_serializing_if = "Option::is_none", - alias = "last_assistant_turn" - )] - pub last_assistant_message: Option, - - /// Per-tool call counts for the rated turn. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub tool_outcomes: Vec, - - /// Session working directory. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub session_cwd: Option, - - /// Number of compactions in the session. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub compaction_count: Option, - - /// Context window usage percentage (0–100). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub context_window_usage: Option, - - /// Raw context tokens used at feedback time. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub context_tokens_used: Option, - - /// Raw model context window token limit at feedback time. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub context_window_tokens: Option, - - // === Terminal Context === - /// Terminal environment snapshot at feedback time (brand, multiplexer, SSH, etc.). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub terminal_info: Option, - - /// Backend URL linking this feedback to its server-side session log. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub unified_log_url: Option, -} - -impl FeedbackSubmission { - /// Construct from typed content; set optional fields after. - pub fn with_content( - session_id: String, - client_type: ClientType, - content: FeedbackContent, - ) -> Self { - let mut s = Self { - session_id, - client_type, - ..Default::default() - }; - content.apply_to(&mut s); - s - } - - /// Remove session context and model metadata, preserving only the - /// user's rating/text and essential identifiers (session_id, client_type). - pub fn strip_metadata(&mut self) { - self.model_id = None; - self.resolved_model_id = None; - self.turn_number = None; - self.last_user_message = None; - self.last_assistant_message = None; - self.tool_outcomes = vec![]; - self.session_cwd = None; - self.compaction_count = None; - self.context_window_usage = None; - self.context_tokens_used = None; - self.context_window_tokens = None; - self.metadata = None; - self.terminal_info = None; - } - - /// Merge a JSON object into `metadata`, inserting if absent. - pub fn merge_metadata(&mut self, extra: serde_json::Value) { - match &mut self.metadata { - Some(existing) if existing.is_object() => { - if let (Some(dst), Some(src)) = (existing.as_object_mut(), extra.as_object()) { - for (k, v) in src { - dst.insert(k.clone(), v.clone()); - } - } - } - _ => { - self.metadata = Some(extra); - } - } - } -} - -/// Snapshot of the user's terminal environment at feedback time. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct FeedbackTerminalInfo { - /// Terminal emulator brand (e.g. "Ghostty", "iTerm2", "Unknown"). - pub brand: String, - /// Multiplexer wrapping the session (e.g. "tmux", "Zellij", "None detected"). - pub multiplexer: String, - /// Whether the session is over SSH. - pub is_ssh: bool, - /// Whether Byobu is wrapping the session. - pub is_byobu: bool, - /// Raw `TERM` environment variable value. - pub term_var: String, - /// tmux server version if inside tmux, otherwise "n/a". - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tmux_version: Option, - /// Hyperlink (OSC 8) support level (e.g. "native", "hostile_parser"). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub hyperlink_osc8_support: Option, - /// Active clipboard legs, e.g. "native+osc52" or "native+tmux+osc52". - #[serde(default, skip_serializing_if = "Option::is_none")] - pub clipboard_route: Option, - /// Native clipboard tool: "pbcopy", "wl-copy", "xclip", "xsel", "arboard". - #[serde(default, skip_serializing_if = "Option::is_none")] - pub clipboard_native_tool: Option, - /// Display server: "wayland", "x11", "quartz", "win32", "unknown". - #[serde(default, skip_serializing_if = "Option::is_none")] - pub display_server: Option, -} - -/// Per-tool call/failure counts for a single tool in a turn. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct FeedbackToolOutcome { - pub tool_name: String, - pub calls: u32, - pub failures: u32, -} - -/// Response from submitting feedback. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct FeedbackResponse { - /// Unique ID of the created feedback entry - pub feedback_id: String, - /// Timestamp when feedback was recorded - pub created_at: DateTime, -} - -/// Request body for completing a feedback request (with feedback data). -/// POST /v1/feedback/requests/{request_id}/complete -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -#[allow(dead_code)] -pub struct FeedbackRequestCompleteRequest { - /// The feedback submission data - #[serde(flatten)] - pub feedback: FeedbackSubmission, -} - -/// Response from completing or dismissing a feedback request. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct FeedbackRequestUpdateResponse { - /// The request ID that was updated - pub request_id: String, - /// New status of the request - pub status: FeedbackRequestStatus, - /// Feedback ID (only present if completed with feedback) - #[serde(skip_serializing_if = "Option::is_none")] - pub feedback_id: Option, - /// Timestamp of the update - pub updated_at: DateTime, -} - -/// Request body for creating a new feedback request. -/// POST /v1/feedback/requests -/// -/// When the agent decides to request feedback (based on heuristics), it should -/// call this endpoint to create a feedback record before sending the -/// FeedbackRequest notification to the client. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct CreateFeedbackRequestInput { - /// Unique request ID (UUID v7 generated by client) - pub request_id: String, - - /// Session ID this request is for - pub session_id: String, - - /// Type of client making the request - pub client_type: ClientType, - - /// What kind of feedback to collect - pub feedback_mode: FeedbackMode, - - /// Custom prompt to show user - #[serde(skip_serializing_if = "Option::is_none")] - pub feedback_prompt: Option, - - /// Priority (1-10, higher = more important) - #[serde(default = "default_priority")] - pub priority: i32, - - /// Which heuristic triggered this request (e.g., "tier1_engagement") - pub trigger_type: String, - - /// Human-readable reason for the request - #[serde(skip_serializing_if = "Option::is_none")] - pub trigger_reason: Option, - - /// Context type - #[serde(skip_serializing_if = "Option::is_none")] - pub context_type: Option, - - /// Specific message IDs to rate (if applicable) - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub context_message_ids: Vec, - - /// When the request expires (optional, defaults to 24 hours) - #[serde(skip_serializing_if = "Option::is_none")] - pub expires_at: Option>, - - /// Experiment ID (if applicable) - #[serde(skip_serializing_if = "Option::is_none")] - pub experiment_id: Option, - - /// Trigger condition details (JSON with signal values at trigger time) - #[serde(skip_serializing_if = "Option::is_none")] - pub trigger_condition: Option, - - /// Per-turn prompt/request ID from the agent session (req_id). - /// Distinct from request_id which is the feedback request's own UUID. - #[serde(skip_serializing_if = "Option::is_none")] - pub prompt_id: Option, -} - -fn default_priority() -> i32 { - 5 -} - -/// Response from creating a feedback request. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct CreateFeedbackRequestResponse { - /// The request ID that was created - pub request_id: String, - /// Timestamp when the request was created - pub created_at: DateTime, -} - -/// Request body for recording a session event. -/// POST /v1/sessions/{session_id}/events -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SessionEventRequest { - /// Type of event - pub event_type: SessionEventType, - /// Event-specific data - #[serde(skip_serializing_if = "Option::is_none")] - pub event_data: Option, - /// Timestamp of the event (defaults to now if not provided) - #[serde(skip_serializing_if = "Option::is_none")] - pub timestamp: Option>, -} - -/// Types of session events that can be recorded. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum SessionEventType { - /// User cancelled a response generation - Cancellation, - /// An error occurred - Error, - /// Context was compacted - Compaction, - /// A tool failed - ToolFailure, - /// User requested regeneration - Regeneration, - /// Session ended - SessionEnd, -} - -impl std::fmt::Display for SessionEventType { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - SessionEventType::Cancellation => write!(f, "cancellation"), - SessionEventType::Error => write!(f, "error"), - SessionEventType::Compaction => write!(f, "compaction"), - SessionEventType::ToolFailure => write!(f, "tool_failure"), - SessionEventType::Regeneration => write!(f, "regeneration"), - SessionEventType::SessionEnd => write!(f, "session_end"), - } - } -} - -/// Response from recording a session event. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SessionEventResponse { - /// Whether the event was recorded successfully - pub success: bool, - /// Timestamp when event was recorded - pub recorded_at: DateTime, -} - -/// Request body for updating session signals. -/// POST /v1/sessions/{session_id}/signals -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SessionSignalsUpdate { - /// Client type - pub client_type: ClientType, - - /// Total turns in the session - #[serde(skip_serializing_if = "Option::is_none")] - pub total_turns: Option, - - /// Number of user messages - #[serde(skip_serializing_if = "Option::is_none")] - pub user_message_count: Option, - - /// Number of assistant messages - #[serde(skip_serializing_if = "Option::is_none")] - pub assistant_message_count: Option, - - /// Number of cancellations - #[serde(skip_serializing_if = "Option::is_none")] - pub cancellation_count: Option, - - /// Number of consecutive cancellations - #[serde(skip_serializing_if = "Option::is_none")] - pub consecutive_cancellations: Option, - - /// Number of errors - #[serde(skip_serializing_if = "Option::is_none")] - pub error_count: Option, - - /// Number of tool failures - #[serde(skip_serializing_if = "Option::is_none")] - pub tool_failure_count: Option, - - /// Total number of tool calls executed (successful + failed) - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tool_call_count: Option, - - /// Number of compactions - #[serde(skip_serializing_if = "Option::is_none")] - pub compaction_count: Option, - - /// Number of regenerations - #[serde(skip_serializing_if = "Option::is_none")] - pub regeneration_count: Option, - - /// Number of edit-and-retry actions (user rewinds and submits a different prompt) - #[serde(default, skip_serializing_if = "Option::is_none")] - pub edit_and_retry_count: Option, - - /// Number of positive ratings (thumbs-up / stars >= 4) - #[serde(default, skip_serializing_if = "Option::is_none")] - pub positive_ratings: Option, - - /// Number of negative ratings (thumbs-down / stars <= 2) - #[serde(default, skip_serializing_if = "Option::is_none")] - pub negative_ratings: Option, - - /// Number of long pauses between turns (idle > 60s) - #[serde(default, skip_serializing_if = "Option::is_none")] - pub long_pauses_count: Option, - - /// Session duration in seconds - #[serde(skip_serializing_if = "Option::is_none")] - pub session_duration_seconds: Option, - - /// Tools used in the session - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub tools_used: Vec, - - /// Models used in the session - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub models_used: Vec, - - /// Primary model ID - #[serde(skip_serializing_if = "Option::is_none")] - pub primary_model_id: Option, - - // === Latency Metrics === - /// Average time to first token in milliseconds - #[serde(skip_serializing_if = "Option::is_none")] - pub avg_time_to_first_token_ms: Option, - - /// Average total response time in milliseconds - #[serde(skip_serializing_if = "Option::is_none")] - pub avg_response_time_ms: Option, - - /// Minimum time to first token in milliseconds - #[serde(skip_serializing_if = "Option::is_none")] - pub min_time_to_first_token_ms: Option, - - /// Maximum time to first token in milliseconds - #[serde(skip_serializing_if = "Option::is_none")] - pub max_time_to_first_token_ms: Option, - - /// Number of latency samples - #[serde(skip_serializing_if = "Option::is_none")] - pub latency_sample_count: Option, - - // === Inter-Token Latency (ITL) Metrics === - /// Most recent response's ITL p50 in milliseconds - #[serde(skip_serializing_if = "Option::is_none")] - pub last_itl_p50_ms: Option, - /// Most recent response's ITL p99 in milliseconds - #[serde(skip_serializing_if = "Option::is_none")] - pub last_itl_p99_ms: Option, - /// Worst-case (max across all responses) ITL max in milliseconds - #[serde(skip_serializing_if = "Option::is_none")] - pub worst_itl_max_ms: Option, - /// Weighted running average of per-response ITL mean in milliseconds - #[serde(skip_serializing_if = "Option::is_none")] - pub avg_itl_mean_ms: Option, - /// Total content chunks received across all responses - #[serde(skip_serializing_if = "Option::is_none")] - pub total_chunk_count: Option, - /// Number of responses measured for ITL - #[serde(skip_serializing_if = "Option::is_none")] - pub itl_sample_count: Option, - - // === LOC Attribution === - #[serde(default, skip_serializing_if = "Option::is_none")] - pub agent_lines_added: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub agent_lines_removed: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub agent_lines_added_reverted: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub agent_lines_removed_reverted: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub human_lines_added: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub human_lines_removed: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub human_lines_added_reverted: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub human_lines_removed_reverted: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub agent_files_touched: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub human_files_touched: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub total_files_touched: Option, - - // === Inference Idle Timeout Tracing === - /// Number of inference idle timeout events in this session. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub inference_idle_timeouts: Option, - /// Configured idle timeout threshold (seconds) — set once at session start. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub inference_idle_timeout_configured_secs: Option, - - // === Doom Loop Detection Tracing === - /// Number of doom loop warnings (model warned but continued). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub doom_loop_warnings: Option, - /// Number of doom loop terminations (turn force-stopped). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub doom_loop_terminations: Option, - /// Configured repeat threshold — set once at session start. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub doom_loop_threshold: Option, - /// Configured read-only repeat threshold — set once at session start. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub doom_loop_ro_threshold: Option, - /// Whether server-side doom-loop recovery resampled at least once. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub doom_loop_recovery_fired: Option, - /// Number of doom-loop recovery resamples in this session. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub doom_loop_recovery_attempts: Option, - /// Completed responses accepted with confident doom-loop signals after - /// the resample budget was spent. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub doom_loop_recovery_accepted_after_budget: Option, - /// Tightest (lowest-threshold) raw trigger label observed, e.g. - /// `tail_repetition:4@thinking`. Labels only — never content. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub doom_loop_recovery_top_trigger: Option, - /// Stream chunks consumed by doomed attempts at their abort points, - /// summed across resamples. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub doom_loop_recovery_aborted_chunks: Option, - - // === GCS Upload Queue Tracing === - /// Total items enqueued for background upload. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gcs_queue_enqueued: Option, - /// Successful background uploads. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gcs_queue_uploaded: Option, - /// Items that exhausted retry budget (superset of expired). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gcs_queue_failed: Option, - /// Enqueue failures that fell back to inline upload. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gcs_queue_fallbacks: Option, - /// Circuit breaker activations. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gcs_queue_circuit_breaker_trips: Option, - /// Current queue depth (snapshot gauge). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gcs_queue_pending: Option, - /// Current disk usage of queue temp dir in bytes (snapshot gauge). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gcs_queue_pending_bytes: Option, - /// Orphaned temp files cleaned up at startup. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gcs_queue_orphans_cleaned: Option, - - /// Additional metadata - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} - -/// Session signals data (for GET response). -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SessionSignals { - /// Session ID - pub session_id: String, - - /// User ID (if known) - #[serde(skip_serializing_if = "Option::is_none")] - pub user_id: Option, - - /// Client type - pub client_type: ClientType, - - /// When the session started - pub session_start_at: DateTime, - - /// Last activity timestamp - pub last_activity_at: DateTime, - - /// Session duration in seconds - #[serde(skip_serializing_if = "Option::is_none")] - pub session_duration_seconds: Option, - - /// Total turns - pub total_turns: i64, - - /// User message count - pub user_message_count: i64, - - /// Assistant message count - pub assistant_message_count: i64, - - /// Cancellation count - pub cancellation_count: i64, - - /// Consecutive cancellations - pub consecutive_cancellations: i64, - - /// Error count - pub error_count: i64, - - /// Tool failure count - pub tool_failure_count: i64, - - /// Compaction count - pub compaction_count: i64, - - /// Regeneration count - pub regeneration_count: i64, - - /// Tools used - pub tools_used: Vec, - - /// Models used - pub models_used: Vec, - - /// Primary model ID - #[serde(skip_serializing_if = "Option::is_none")] - pub primary_model_id: Option, - - // === Latency Metrics === - /// Average time to first token in milliseconds - #[serde(default)] - pub avg_time_to_first_token_ms: i64, - - /// Average total response time in milliseconds - #[serde(default)] - pub avg_response_time_ms: i64, - - /// Minimum time to first token in milliseconds - #[serde(default)] - pub min_time_to_first_token_ms: i64, - - /// Maximum time to first token in milliseconds - #[serde(default)] - pub max_time_to_first_token_ms: i64, - - /// Number of latency samples - #[serde(default)] - pub latency_sample_count: i64, - - // === Inter-Token Latency (ITL) Metrics === - /// Most recent response's ITL p50 in milliseconds (nullable: None = not yet reported) - #[serde(default)] - pub last_itl_p50_ms: Option, - /// Most recent response's ITL p99 in milliseconds (nullable: None = not yet reported) - #[serde(default)] - pub last_itl_p99_ms: Option, - /// Worst-case (max across all responses) ITL max in milliseconds - #[serde(default)] - pub worst_itl_max_ms: i64, - /// Weighted running average of per-response ITL mean in milliseconds - #[serde(default)] - pub avg_itl_mean_ms: i64, - /// Total content chunks received across all responses - #[serde(default)] - pub total_chunk_count: i64, - /// Number of responses measured for ITL - #[serde(default)] - pub itl_sample_count: i64, - - /// Feedback requests sent - pub feedback_requests_sent: i64, - - /// Feedback requests completed - pub feedback_requests_completed: i64, - - /// Last feedback request timestamp - #[serde(skip_serializing_if = "Option::is_none")] - pub last_feedback_request_at: Option>, - - /// Record created at - pub created_at: DateTime, - - /// Record updated at - pub updated_at: DateTime, - - /// Additional metadata - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} - -/// Response for session signals update. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SessionSignalsUpdateResponse { - /// Session ID - pub session_id: String, - /// Timestamp when signals were updated - pub updated_at: DateTime, -} - -/// Pending feedback request (returned by GET /v1/feedback/requests). -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct FeedbackRequest { - /// Unique request ID - pub request_id: String, - - /// Session ID - pub session_id: String, - - /// What kind of feedback to collect - pub feedback_mode: FeedbackMode, - - /// Custom prompt to show user - #[serde(skip_serializing_if = "Option::is_none")] - pub feedback_prompt: Option, - - /// Priority (1-10, higher = more important) - pub priority: i32, - - /// Which heuristic triggered this request - pub trigger_type: String, - - /// Human-readable reason for the request - #[serde(skip_serializing_if = "Option::is_none")] - pub trigger_reason: Option, - - /// Context type - #[serde(skip_serializing_if = "Option::is_none")] - pub context_type: Option, - - /// Specific message IDs to rate - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub context_message_ids: Vec, - - /// Current status - pub status: FeedbackRequestStatus, - - /// When the request was created - pub created_at: DateTime, - - /// When the request expires - #[serde(skip_serializing_if = "Option::is_none")] - pub expires_at: Option>, - - /// Experiment ID (if applicable) - #[serde(skip_serializing_if = "Option::is_none")] - pub experiment_id: Option, - - /// Comparison ID (if applicable) - #[serde(skip_serializing_if = "Option::is_none")] - pub comparison_id: Option, -} - -/// Query parameters for GET /v1/feedback/requests. -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase")] -#[allow(dead_code)] -pub struct FeedbackRequestsQuery { - /// Session ID to filter by - pub session_id: String, - /// Status filter (defaults to "pending") - #[serde(default = "default_pending_status")] - pub status: FeedbackRequestStatus, -} - -#[allow(dead_code)] -fn default_pending_status() -> FeedbackRequestStatus { - FeedbackRequestStatus::Pending -} - -// ============================================================================ -// Feedback Heuristics Configuration -// ============================================================================ - -/// Configuration for a single feedback tier. -/// -/// Each tier has specific thresholds and conditions that must be met -/// for feedback to be requested at that tier. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub struct TierConfig { - /// Whether this tier is enabled - pub enabled: bool, - /// Sample rate (0.0 to 1.0, e.g., 0.0005 = 0.05%) - pub sample_rate: f64, - /// Minimum turns required to trigger - pub min_turns: i64, - /// Minimum tool calls required (Tier 1 & 2) - #[serde(default)] - pub min_tool_calls: i64, - /// Minimum compactions required (Tier 1 & 2) - #[serde(default)] - pub min_compactions: i64, - /// Minimum errors required (Tier 2 only) - #[serde(default)] - pub min_errors: i64, - /// Whether cancellations disqualify this tier (Tier 1) - #[serde(default)] - pub no_cancellations: bool, - /// Whether cancellation is required (Tier 3) - #[serde(default)] - pub requires_cancellation: bool, - /// Whether revert is required (Tier 3) - #[serde(default)] - pub requires_revert: bool, - /// Whether at least one of cancellation/revert is required (Tier 3) - #[serde(default)] - pub requires_recovery: bool, - /// Feedback mode to use when this tier triggers - pub feedback_mode: FeedbackMode, - /// Whether feedback requests from this tier are dismissible (non-intrusive) - #[serde(default = "default_true")] - pub dismissible: bool, - /// Prompt text shown to users when this tier's feedback is requested - #[serde(default)] - pub prompt: String, - /// Max times this tier can trigger per session (0 = unlimited) - #[serde(default = "default_one")] - pub max_triggers: i32, -} - -impl Default for TierConfig { - fn default() -> Self { - Self { - enabled: true, - sample_rate: 0.0005, - min_turns: 10, - min_tool_calls: 5, - min_compactions: 2, - min_errors: 0, - no_cancellations: false, - requires_cancellation: false, - requires_revert: false, - requires_recovery: false, - feedback_mode: FeedbackMode::Thumbs, - dismissible: true, - prompt: String::new(), - max_triggers: 1, - } - } -} - -/// Configuration for feedback heuristics loaded from remote config. -/// -/// This struct maps directly to the server's feedback-heuristics config storage. -/// It controls when and how feedback is requested from users. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub struct FeedbackHeuristicsConfig { - /// Unique configuration identifier - pub config_id: String, - /// Configuration version (monotonically increasing) - pub config_version: i64, - - // === Global Settings === - /// Master enable/disable switch for all feedback collection - pub enabled: bool, - /// Minimum seconds between feedback requests (cooldown period) - #[serde(default = "default_cooldown_seconds")] - pub cooldown_seconds: i64, - /// Maximum feedback requests per session - #[serde(default = "default_max_requests")] - pub max_requests_per_session: i64, - - // === Tier 1: Standard Engagement === - /// Whether Tier 1 is enabled - #[serde(default = "default_true")] - pub tier1_enabled: bool, - /// Sample rate for Tier 1 (0.0-1.0) - #[serde(default = "default_tier1_sample_rate")] - pub tier1_sample_rate: f64, - /// Minimum turns for Tier 1 - #[serde(default = "default_tier1_min_turns")] - pub tier1_min_turns: i64, - /// Minimum tool calls for Tier 1 - #[serde(default = "default_tier1_min_tool_calls")] - pub tier1_min_tool_calls: i64, - /// Minimum compactions for Tier 1 - #[serde(default = "default_tier1_min_compactions")] - pub tier1_min_compactions: i64, - /// Whether Tier 1 requires no cancellations - #[serde(default = "default_true")] - pub tier1_no_cancellations: bool, - /// Feedback mode for Tier 1 - #[serde(default = "default_feedback_mode_thumbs")] - pub tier1_feedback_mode: String, - /// Whether Tier 1 feedback requests are dismissible - #[serde(default = "default_true")] - pub tier1_dismissible: bool, - /// Prompt text shown to users when Tier 1 feedback is requested - #[serde(default = "default_tier1_prompt")] - pub tier1_prompt: String, - /// Max times Tier 1 can trigger per session (0 = unlimited) - #[serde(default = "default_one")] - pub tier1_max_triggers: i32, - - // === Tier 2: Complex Session with Recovery === - /// Whether Tier 2 is enabled - #[serde(default = "default_true")] - pub tier2_enabled: bool, - /// Sample rate for Tier 2 (0.0-1.0) - #[serde(default = "default_tier2_sample_rate")] - pub tier2_sample_rate: f64, - /// Minimum turns for Tier 2 - #[serde(default = "default_tier2_min_turns")] - pub tier2_min_turns: i64, - /// Minimum tool calls for Tier 2 - #[serde(default = "default_tier2_min_tool_calls")] - pub tier2_min_tool_calls: i64, - /// Minimum compactions for Tier 2 - #[serde(default = "default_tier2_min_compactions")] - pub tier2_min_compactions: i64, - /// Minimum errors for Tier 2 - #[serde(default = "default_tier2_min_errors")] - pub tier2_min_errors: i64, - /// Feedback mode for Tier 2 - #[serde(default = "default_feedback_mode_thumbs_text")] - pub tier2_feedback_mode: String, - /// Whether Tier 2 feedback requests are dismissible - #[serde(default = "default_true")] - pub tier2_dismissible: bool, - /// Prompt text shown to users when Tier 2 feedback is requested - #[serde(default = "default_tier2_prompt")] - pub tier2_prompt: String, - /// Max times Tier 2 can trigger per session (0 = unlimited) - #[serde(default = "default_one")] - pub tier2_max_triggers: i32, - - // === Tier 3: Recovery from Friction === - /// Whether Tier 3 is enabled - #[serde(default = "default_true")] - pub tier3_enabled: bool, - /// Sample rate for Tier 3 (0.0-1.0) - #[serde(default = "default_tier3_sample_rate")] - pub tier3_sample_rate: f64, - /// Minimum turns for Tier 3 - #[serde(default = "default_tier3_min_turns")] - pub tier3_min_turns: i64, - /// Whether Tier 3 requires at least one cancellation - #[serde(default)] - pub tier3_requires_cancellation: bool, - /// Whether Tier 3 requires at least one revert - #[serde(default)] - pub tier3_requires_revert: bool, - /// Whether Tier 3 requires recovery (cancellation OR revert) - #[serde(default = "default_true")] - pub tier3_requires_recovery: bool, - /// Feedback mode for Tier 3 - #[serde(default = "default_feedback_mode_stars_text")] - pub tier3_feedback_mode: String, - /// Whether Tier 3 feedback requests are dismissible - #[serde(default = "default_true")] - pub tier3_dismissible: bool, - /// Prompt text shown to users when Tier 3 feedback is requested - #[serde(default = "default_tier3_prompt")] - pub tier3_prompt: String, - /// Max times Tier 3 can trigger per session (0 = unlimited) - #[serde(default = "default_one")] - pub tier3_max_triggers: i32, - - // === Metadata === - /// When this config was created - #[serde(skip_serializing_if = "Option::is_none")] - pub created_at: Option>, - /// When this config was last updated - #[serde(skip_serializing_if = "Option::is_none")] - pub updated_at: Option>, - /// Who created this config - #[serde(skip_serializing_if = "Option::is_none")] - pub created_by: Option, - /// Human-readable description - #[serde(skip_serializing_if = "Option::is_none")] - pub description: Option, - - // === Lifecycle === - /// When this config becomes effective - #[serde(skip_serializing_if = "Option::is_none")] - pub effective_from: Option>, - /// When this config expires - #[serde(skip_serializing_if = "Option::is_none")] - pub effective_until: Option>, - /// Whether this config is active - #[serde(default = "default_true")] - pub is_active: bool, - /// User cohorts this config targets (e.g., ["all"], ["beta"]) - #[serde(default = "default_all_cohorts")] - pub target_user_cohorts: Vec, - /// Higher priority configs are preferred when multiple match - #[serde(default)] - pub priority: i32, -} - -impl Default for FeedbackHeuristicsConfig { - fn default() -> Self { - Self { - config_id: "default".to_string(), - config_version: 1, - enabled: true, - cooldown_seconds: 300, - max_requests_per_session: 3, - // Tier 1 - tier1_enabled: true, - tier1_sample_rate: 0.0005, - tier1_min_turns: 10, - tier1_min_tool_calls: 5, - tier1_min_compactions: 2, - tier1_no_cancellations: true, - tier1_feedback_mode: "thumbs".to_string(), - tier1_dismissible: true, - tier1_prompt: default_tier1_prompt(), - tier1_max_triggers: 1, - // Tier 2 - tier2_enabled: true, - tier2_sample_rate: 0.0002, - tier2_min_turns: 15, - tier2_min_tool_calls: 10, - tier2_min_compactions: 3, - tier2_min_errors: 1, - tier2_feedback_mode: "thumbs_text".to_string(), - tier2_dismissible: true, - tier2_prompt: default_tier2_prompt(), - tier2_max_triggers: 1, - // Tier 3 - tier3_enabled: true, - tier3_sample_rate: 0.0001, - tier3_min_turns: 20, - tier3_requires_cancellation: false, - tier3_requires_revert: false, - tier3_requires_recovery: true, - tier3_feedback_mode: "stars_text".to_string(), - tier3_dismissible: true, - tier3_prompt: default_tier3_prompt(), - tier3_max_triggers: 1, - // Metadata - created_at: None, - updated_at: None, - created_by: None, - description: None, - // Lifecycle - effective_from: None, - effective_until: None, - is_active: true, - target_user_cohorts: default_all_cohorts(), - priority: 0, - } - } -} - -impl FeedbackHeuristicsConfig { - /// Get the Tier 1 configuration as a TierConfig. - pub fn tier1_config(&self) -> TierConfig { - TierConfig { - enabled: self.tier1_enabled, - sample_rate: self.tier1_sample_rate, - min_turns: self.tier1_min_turns, - min_tool_calls: self.tier1_min_tool_calls, - min_compactions: self.tier1_min_compactions, - min_errors: 0, - no_cancellations: self.tier1_no_cancellations, - requires_cancellation: false, - requires_revert: false, - requires_recovery: false, - feedback_mode: parse_feedback_mode_str(&self.tier1_feedback_mode), - dismissible: self.tier1_dismissible, - prompt: self.tier1_prompt.clone(), - max_triggers: self.tier1_max_triggers, - } - } - - /// Get the Tier 2 configuration as a TierConfig. - pub fn tier2_config(&self) -> TierConfig { - TierConfig { - enabled: self.tier2_enabled, - sample_rate: self.tier2_sample_rate, - min_turns: self.tier2_min_turns, - min_tool_calls: self.tier2_min_tool_calls, - min_compactions: self.tier2_min_compactions, - min_errors: self.tier2_min_errors, - no_cancellations: false, - requires_cancellation: false, - requires_revert: false, - requires_recovery: false, - feedback_mode: parse_feedback_mode_str(&self.tier2_feedback_mode), - dismissible: self.tier2_dismissible, - prompt: self.tier2_prompt.clone(), - max_triggers: self.tier2_max_triggers, - } - } - - /// Get the Tier 3 configuration as a TierConfig. - pub fn tier3_config(&self) -> TierConfig { - TierConfig { - enabled: self.tier3_enabled, - sample_rate: self.tier3_sample_rate, - min_turns: self.tier3_min_turns, - min_tool_calls: 0, - min_compactions: 0, - min_errors: 0, - no_cancellations: false, - requires_cancellation: self.tier3_requires_cancellation, - requires_revert: self.tier3_requires_revert, - requires_recovery: self.tier3_requires_recovery, - feedback_mode: parse_feedback_mode_str(&self.tier3_feedback_mode), - dismissible: self.tier3_dismissible, - prompt: self.tier3_prompt.clone(), - max_triggers: self.tier3_max_triggers, - } - } -} - -// Default value functions for serde -fn default_true() -> bool { - true -} -fn default_cooldown_seconds() -> i64 { - 300 -} -fn default_max_requests() -> i64 { - 3 -} -fn default_tier1_sample_rate() -> f64 { - 0.0005 -} -fn default_tier1_min_turns() -> i64 { - 10 -} -fn default_tier1_min_tool_calls() -> i64 { - 5 -} -fn default_tier1_min_compactions() -> i64 { - 2 -} -fn default_tier2_sample_rate() -> f64 { - 0.0002 -} -fn default_tier2_min_turns() -> i64 { - 15 -} -fn default_tier2_min_tool_calls() -> i64 { - 10 -} -fn default_tier2_min_compactions() -> i64 { - 3 -} -fn default_tier2_min_errors() -> i64 { - 1 -} -fn default_tier3_sample_rate() -> f64 { - 0.0001 -} -fn default_tier3_min_turns() -> i64 { - 20 -} -fn default_feedback_mode_thumbs() -> String { - "thumbs".to_string() -} -fn default_feedback_mode_thumbs_text() -> String { - "thumbs_text".to_string() -} -fn default_feedback_mode_stars_text() -> String { - "stars_text".to_string() -} -fn default_tier1_prompt() -> String { - "You've been using Grok Code productively! Would you mind sharing quick feedback?".to_string() -} -fn default_tier2_prompt() -> String { - "You've worked through a complex session. Your feedback would help us improve.".to_string() -} -fn default_tier3_prompt() -> String { - "Thanks for sticking with us through that session. Got a moment to share feedback?".to_string() -} -fn default_all_cohorts() -> Vec { - vec!["all".to_string()] -} -fn default_one() -> i32 { - 1 -} - -/// Parse a feedback mode string to FeedbackMode enum. -pub fn parse_feedback_mode_str(s: &str) -> FeedbackMode { - match s { - "thumbs" => FeedbackMode::Thumbs, - "stars" => FeedbackMode::Stars, - "text" => FeedbackMode::Text, - "thumbs_text" => FeedbackMode::ThumbsText, - "stars_text" => FeedbackMode::StarsText, - "comparison" => FeedbackMode::Comparison, - "survey" => FeedbackMode::Survey, - "nps" => FeedbackMode::Nps, - "nps_text" => FeedbackMode::NpsText, - _ => FeedbackMode::Thumbs, - } -} - -// ============================================================================ -// Session Turn Deltas — per-turn time-series data for regression tracking -// ============================================================================ - -/// Per-turn delta sent at the end of every turn. -/// -/// Contains the *change* in session counters since the previous turn end, -/// plus absolute values for contextual fields (latency, model, experiment info). -/// This produces a time-series row per turn that can be used for regression -/// detection and session-level analytics. -/// -/// POST /v1/sessions/{session_id}/turn-deltas -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SessionTurnDelta { - pub client_type: ClientType, - pub turn_number: i64, - // Delta counters - pub delta_tool_calls: i64, - pub delta_tool_failures: i64, - pub delta_errors: i64, - pub delta_cancellations: i64, - pub delta_regenerations: i64, - pub delta_compactions: i64, - pub delta_edit_and_retries: i64, - pub delta_positive_ratings: i64, - pub delta_negative_ratings: i64, - pub delta_assistant_messages: i64, - pub delta_long_pauses: i64, - pub delta_successful_tool_uses: i64, - // Turn-level snapshot values - pub consecutive_cancellations: i64, - // Turn-level absolute values - #[serde(skip_serializing_if = "Option::is_none")] - pub time_to_first_token_ms: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub total_response_time_ms: Option, - /// Inter-token latency p50 for this turn's response (ms) - #[serde(skip_serializing_if = "Option::is_none")] - pub itl_p50_ms: Option, - /// Inter-token latency p99 for this turn's response (ms) - #[serde(skip_serializing_if = "Option::is_none")] - pub itl_p99_ms: Option, - /// Inter-token latency max for this turn's response (ms) - #[serde(skip_serializing_if = "Option::is_none")] - pub itl_max_ms: Option, - /// Inter-token latency mean for this turn's response (ms) - #[serde(skip_serializing_if = "Option::is_none")] - pub itl_mean_ms: Option, - pub context_window_usage: i64, - #[serde(skip_serializing_if = "Option::is_none")] - pub model_id: Option, - /// Whole-turn wall-clock duration (prompt→final response), ms. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub turn_duration_ms: Option, - /// Terminal outcome: "completed" | "cancelled" | "error". - #[serde(default, skip_serializing_if = "Option::is_none")] - pub turn_outcome: Option, - /// Served model fingerprint (provider `system_fingerprint`). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub model_fingerprint: Option, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub tools_used_this_turn: Vec, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub error_types_this_turn: Vec, - /// Per-tool success/failure breakdown for this turn, JSON-serialized. - /// Each entry: `{"toolName":"bash","successes":2,"failures":1}`. - #[serde(default, skip_serializing_if = "String::is_empty")] - pub tool_outcomes: String, - // Cumulative totals - pub cumulative_tool_calls: i64, - pub cumulative_errors: i64, - pub session_duration_seconds: i64, - /// Cumulative total tokens across all compactions - #[serde(default)] - pub total_tokens_before_compaction: i64, - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, - /// Prompt/request ID for the turn (if available). - #[serde(skip_serializing_if = "Option::is_none")] - pub request_id: Option, - /// Session start time — used for analytics ingestion. If not provided, defaults - /// to the server receipt time. - #[serde(skip_serializing_if = "Option::is_none")] - pub session_start_at: Option>, - /// Number of feedback requests sent this session (cumulative). - #[serde(default)] - pub feedback_requests_sent: i64, - /// Wall-clock timestamp of the last feedback request sent this session. - #[serde(skip_serializing_if = "Option::is_none")] - pub last_feedback_request_at: Option>, - /// Number of response (completion - reasoning) tokens generated this turn. - /// `None` when not reported (old client or no inference this turn). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub response_tokens: Option, - /// Number of thinking (reasoning) tokens generated this turn. - /// `None` when not reported (old client or no inference this turn). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub thinking_tokens: Option, - - // === LOC Attribution Deltas === - #[serde(default)] - pub delta_agent_lines_added: i64, - #[serde(default)] - pub delta_agent_lines_removed: i64, - #[serde(default)] - pub delta_agent_lines_added_reverted: i64, - #[serde(default)] - pub delta_agent_lines_removed_reverted: i64, - #[serde(default)] - pub delta_human_lines_added: i64, - #[serde(default)] - pub delta_human_lines_removed: i64, - #[serde(default)] - pub delta_human_lines_added_reverted: i64, - #[serde(default)] - pub delta_human_lines_removed_reverted: i64, - #[serde(default)] - pub delta_agent_files_touched: i64, - #[serde(default)] - pub delta_human_files_touched: i64, - #[serde(default)] - pub delta_total_files_touched: i64, - /// Whether LOC tracking was enabled on the client for this session. - /// When `false`, all `delta_*` LOC fields are meaningless zeros (the - /// hunk tracker was never spawned). When `true`, zeros indicate the - /// tracker was active but no code changed during this turn. - /// Defaults to `false` for backwards-compat with old clients. - #[serde(default)] - pub loc_tracking_enabled: bool, -} - -/// Response for POST /v1/sessions/{session_id}/turn-deltas. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SessionTurnDeltaResponse { - pub session_id: String, - pub turn_number: i64, - pub recorded_at: DateTime, -} - -#[cfg(test)] -mod tests { - use super::*; - - /// Verify backward compatibility: JSON from old agents (no ITL fields) - /// deserializes cleanly, and new agents' ITL fields are also parsed. - #[test] - fn session_signals_update_itl_backward_compat() { - // Old agent: no ITL fields at all - let json_no_itl = r#"{"clientType": "agent"}"#; - let update: SessionSignalsUpdate = serde_json::from_str(json_no_itl).unwrap(); - assert_eq!(update.last_itl_p50_ms, None); - assert_eq!(update.last_itl_p99_ms, None); - assert_eq!(update.worst_itl_max_ms, None); - assert_eq!(update.avg_itl_mean_ms, None); - assert_eq!(update.total_chunk_count, None); - assert_eq!(update.itl_sample_count, None); - - // New agent: ITL fields present - let json_with_itl = r#"{"clientType": "agent", "lastItlP50Ms": 45, "worstItlMaxMs": 500}"#; - let update: SessionSignalsUpdate = serde_json::from_str(json_with_itl).unwrap(); - assert_eq!(update.last_itl_p50_ms, Some(45)); - assert_eq!(update.worst_itl_max_ms, Some(500)); - // Fields not provided should remain None - assert_eq!(update.last_itl_p99_ms, None); - - // Explicit null - let json_null_itl = r#"{"clientType": "agent", "lastItlP50Ms": null}"#; - let update: SessionSignalsUpdate = serde_json::from_str(json_null_itl).unwrap(); - assert_eq!(update.last_itl_p50_ms, None); - } - - /// Verify backward compatibility for the per-turn latency fields on - /// `SessionTurnDelta`: JSON from old agents (no turn-duration / outcome / - /// fingerprint) deserializes with those fields as `None`, and new agents' - /// values parse correctly. Required (non-`Option`) counters must be present - /// because they have no serde default. - #[test] - fn session_turn_delta_latency_backward_compat() { - // Old agent: required counters present, no new turn-latency fields. - let json_old = r#"{ - "clientType": "agent", - "turnNumber": 1, - "deltaToolCalls": 0, - "deltaToolFailures": 0, - "deltaErrors": 0, - "deltaCancellations": 0, - "deltaRegenerations": 0, - "deltaCompactions": 0, - "deltaEditAndRetries": 0, - "deltaPositiveRatings": 0, - "deltaNegativeRatings": 0, - "deltaAssistantMessages": 0, - "deltaLongPauses": 0, - "deltaSuccessfulToolUses": 0, - "consecutiveCancellations": 0, - "contextWindowUsage": 0, - "cumulativeToolCalls": 0, - "cumulativeErrors": 0, - "sessionDurationSeconds": 0 - }"#; - let delta_old: SessionTurnDelta = serde_json::from_str(json_old).unwrap(); - assert_eq!(delta_old.turn_duration_ms, None); - assert_eq!(delta_old.turn_outcome, None); - assert_eq!(delta_old.model_fingerprint, None); - - // Re-serializing an old delta omits the new fields. - let reserialized = serde_json::to_string(&delta_old).unwrap(); - assert!(!reserialized.contains("turnDurationMs")); - assert!(!reserialized.contains("turnOutcome")); - assert!(!reserialized.contains("modelFingerprint")); - - // New agent: turn-latency fields present and parsed. - let json_new = r#"{ - "clientType": "agent", - "turnNumber": 2, - "deltaToolCalls": 0, - "deltaToolFailures": 0, - "deltaErrors": 0, - "deltaCancellations": 0, - "deltaRegenerations": 0, - "deltaCompactions": 0, - "deltaEditAndRetries": 0, - "deltaPositiveRatings": 0, - "deltaNegativeRatings": 0, - "deltaAssistantMessages": 0, - "deltaLongPauses": 0, - "deltaSuccessfulToolUses": 0, - "consecutiveCancellations": 0, - "contextWindowUsage": 0, - "cumulativeToolCalls": 0, - "cumulativeErrors": 0, - "sessionDurationSeconds": 0, - "turnDurationMs": 4200, - "turnOutcome": "completed", - "modelFingerprint": "fp_abc123" - }"#; - let delta_new: SessionTurnDelta = serde_json::from_str(json_new).unwrap(); - assert_eq!(delta_new.turn_duration_ms, Some(4200)); - assert_eq!(delta_new.turn_outcome.as_deref(), Some("completed")); - assert_eq!(delta_new.model_fingerprint.as_deref(), Some("fp_abc123")); - } - - /// Verify backward compatibility: JSON without tracing fields (old clients) - /// deserializes cleanly, and new clients' tracing fields are also parsed. - #[test] - fn session_signals_update_tracing_backward_compat() { - // Old client: no tracing fields — all default to None - let json_old = r#"{"clientType": "agent"}"#; - let update: SessionSignalsUpdate = serde_json::from_str(json_old).unwrap(); - assert_eq!(update.inference_idle_timeouts, None); - assert_eq!(update.inference_idle_timeout_configured_secs, None); - assert_eq!(update.doom_loop_warnings, None); - assert_eq!(update.doom_loop_terminations, None); - assert_eq!(update.doom_loop_threshold, None); - assert_eq!(update.doom_loop_ro_threshold, None); - assert_eq!(update.doom_loop_recovery_fired, None); - assert_eq!(update.doom_loop_recovery_attempts, None); - assert_eq!(update.doom_loop_recovery_accepted_after_budget, None); - assert_eq!(update.doom_loop_recovery_top_trigger, None); - assert_eq!(update.doom_loop_recovery_aborted_chunks, None); - assert_eq!(update.gcs_queue_enqueued, None); - assert_eq!(update.gcs_queue_uploaded, None); - assert_eq!(update.gcs_queue_failed, None); - assert_eq!(update.gcs_queue_fallbacks, None); - assert_eq!(update.gcs_queue_circuit_breaker_trips, None); - assert_eq!(update.gcs_queue_pending, None); - assert_eq!(update.gcs_queue_pending_bytes, None); - assert_eq!(update.gcs_queue_orphans_cleaned, None); - - // New client: tracing fields present - let json_new = r#"{ - "clientType": "agent", - "inferenceIdleTimeouts": 2, - "inferenceIdleTimeoutConfiguredSecs": 300, - "doomLoopWarnings": 1, - "doomLoopTerminations": 0, - "doomLoopThreshold": 4, - "doomLoopRoThreshold": 8, - "doomLoopRecoveryFired": true, - "doomLoopRecoveryAttempts": 2, - "doomLoopRecoveryAcceptedAfterBudget": 1, - "doomLoopRecoveryTopTrigger": "tail_repetition:4@thinking", - "doomLoopRecoveryAbortedChunks": 421, - "gcsQueueEnqueued": 50, - "gcsQueueUploaded": 48, - "gcsQueueFailed": 1, - "gcsQueueFallbacks": 1, - "gcsQueueCircuitBreakerTrips": 0, - "gcsQueuePending": 3, - "gcsQueuePendingBytes": 1048576, - "gcsQueueOrphansCleaned": 2 - }"#; - let update: SessionSignalsUpdate = serde_json::from_str(json_new).unwrap(); - assert_eq!(update.inference_idle_timeouts, Some(2)); - assert_eq!(update.inference_idle_timeout_configured_secs, Some(300)); - assert_eq!(update.doom_loop_warnings, Some(1)); - assert_eq!(update.doom_loop_terminations, Some(0)); - assert_eq!(update.doom_loop_threshold, Some(4)); - assert_eq!(update.doom_loop_ro_threshold, Some(8)); - assert_eq!(update.doom_loop_recovery_fired, Some(true)); - assert_eq!(update.doom_loop_recovery_attempts, Some(2)); - assert_eq!(update.doom_loop_recovery_accepted_after_budget, Some(1)); - assert_eq!( - update.doom_loop_recovery_top_trigger.as_deref(), - Some("tail_repetition:4@thinking") - ); - assert_eq!(update.doom_loop_recovery_aborted_chunks, Some(421)); - assert_eq!(update.gcs_queue_enqueued, Some(50)); - assert_eq!(update.gcs_queue_uploaded, Some(48)); - assert_eq!(update.gcs_queue_failed, Some(1)); - assert_eq!(update.gcs_queue_fallbacks, Some(1)); - assert_eq!(update.gcs_queue_circuit_breaker_trips, Some(0)); - assert_eq!(update.gcs_queue_pending, Some(3)); - assert_eq!(update.gcs_queue_pending_bytes, Some(1048576)); - assert_eq!(update.gcs_queue_orphans_cleaned, Some(2)); - - // Round-trip: serialize → deserialize preserves values - let serialized = serde_json::to_string(&update).unwrap(); - let round_tripped: SessionSignalsUpdate = serde_json::from_str(&serialized).unwrap(); - assert_eq!(round_tripped.inference_idle_timeouts, Some(2)); - assert_eq!(round_tripped.gcs_queue_uploaded, Some(48)); - assert_eq!(round_tripped.doom_loop_warnings, Some(1)); - } - - /// Verify backward compatibility: JSON from old agents (no token fields) - /// deserializes cleanly as None, and new agents' token fields are parsed. - #[test] - fn session_turn_delta_token_backward_compat() { - // Old agent: no token fields at all — should deserialize as None - let json_no_tokens = r#"{ - "clientType": "agent", - "turnNumber": 1, - "deltaToolCalls": 0, - "deltaToolFailures": 0, - "deltaErrors": 0, - "deltaCancellations": 0, - "deltaRegenerations": 0, - "deltaCompactions": 0, - "deltaEditAndRetries": 0, - "deltaPositiveRatings": 0, - "deltaNegativeRatings": 0, - "deltaAssistantMessages": 1, - "deltaLongPauses": 0, - "deltaSuccessfulToolUses": 0, - "consecutiveCancellations": 0, - "contextWindowUsage": 50, - "cumulativeToolCalls": 0, - "cumulativeErrors": 0, - "sessionDurationSeconds": 10 - }"#; - let delta: SessionTurnDelta = serde_json::from_str(json_no_tokens).unwrap(); - assert_eq!(delta.response_tokens, None); - assert_eq!(delta.thinking_tokens, None); - - // New agent: token fields present - let json_with_tokens = r#"{ - "clientType": "agent", - "turnNumber": 1, - "deltaToolCalls": 0, - "deltaToolFailures": 0, - "deltaErrors": 0, - "deltaCancellations": 0, - "deltaRegenerations": 0, - "deltaCompactions": 0, - "deltaEditAndRetries": 0, - "deltaPositiveRatings": 0, - "deltaNegativeRatings": 0, - "deltaAssistantMessages": 1, - "deltaLongPauses": 0, - "deltaSuccessfulToolUses": 0, - "consecutiveCancellations": 0, - "contextWindowUsage": 50, - "cumulativeToolCalls": 0, - "cumulativeErrors": 0, - "sessionDurationSeconds": 10, - "responseTokens": 512, - "thinkingTokens": 1024 - }"#; - let delta: SessionTurnDelta = serde_json::from_str(json_with_tokens).unwrap(); - assert_eq!(delta.response_tokens, Some(512)); - assert_eq!(delta.thinking_tokens, Some(1024)); - } - - /// Verify backward compatibility: JSON from old agents (no `locTrackingEnabled`) - /// defaults to `false`, and new agents that send `true` are parsed correctly. - #[test] - fn session_turn_delta_loc_tracking_backward_compat() { - // Old agent: no locTrackingEnabled field — should default to false - let json_no_loc = r#"{ - "clientType": "agent", - "turnNumber": 1, - "deltaToolCalls": 0, - "deltaToolFailures": 0, - "deltaErrors": 0, - "deltaCancellations": 0, - "deltaRegenerations": 0, - "deltaCompactions": 0, - "deltaEditAndRetries": 0, - "deltaPositiveRatings": 0, - "deltaNegativeRatings": 0, - "deltaAssistantMessages": 1, - "deltaLongPauses": 0, - "deltaSuccessfulToolUses": 0, - "consecutiveCancellations": 0, - "contextWindowUsage": 50, - "cumulativeToolCalls": 0, - "cumulativeErrors": 0, - "sessionDurationSeconds": 10 - }"#; - let delta: SessionTurnDelta = serde_json::from_str(json_no_loc).unwrap(); - assert!( - !delta.loc_tracking_enabled, - "old clients should default to false" - ); - // LOC fields should be 0 (serde default) - assert_eq!(delta.delta_agent_lines_added, 0); - - // New agent with LOC tracking enabled - let json_with_loc = r#"{ - "clientType": "agent", - "turnNumber": 2, - "deltaToolCalls": 0, - "deltaToolFailures": 0, - "deltaErrors": 0, - "deltaCancellations": 0, - "deltaRegenerations": 0, - "deltaCompactions": 0, - "deltaEditAndRetries": 0, - "deltaPositiveRatings": 0, - "deltaNegativeRatings": 0, - "deltaAssistantMessages": 1, - "deltaLongPauses": 0, - "deltaSuccessfulToolUses": 0, - "consecutiveCancellations": 0, - "contextWindowUsage": 50, - "cumulativeToolCalls": 0, - "cumulativeErrors": 0, - "sessionDurationSeconds": 20, - "locTrackingEnabled": true, - "deltaAgentLinesAdded": 10 - }"#; - let delta: SessionTurnDelta = serde_json::from_str(json_with_loc).unwrap(); - assert!(delta.loc_tracking_enabled); - assert_eq!(delta.delta_agent_lines_added, 10); - } - - /// Old agents that predate cohort targeting send no cohort fields. - /// Verify they deserialize to defaults (target_user_cohorts=["all"], priority=0) - /// and that new agents' cohort fields are parsed correctly. - #[test] - fn feedback_heuristics_config_cohort_backward_compat() { - // Old config: no cohort fields → defaults to ["all"] and priority 0 - let json_no_cohorts = r#"{"config_id":"v1","config_version":1,"enabled":true}"#; - let config: FeedbackHeuristicsConfig = serde_json::from_str(json_no_cohorts).unwrap(); - assert_eq!(config.target_user_cohorts, vec!["all"]); - assert_eq!(config.priority, 0); - - // New config: cohort fields present - let json_with_cohorts = r#"{ - "config_id":"v2","config_version":2,"enabled":true, - "target_user_cohorts":["beta"],"priority":10 - }"#; - let config: FeedbackHeuristicsConfig = serde_json::from_str(json_with_cohorts).unwrap(); - assert_eq!(config.target_user_cohorts, vec!["beta"]); - assert_eq!(config.priority, 10); - - // Round-trip: serialize → deserialize preserves cohort values - let serialized = serde_json::to_string(&config).unwrap(); - let round_tripped: FeedbackHeuristicsConfig = serde_json::from_str(&serialized).unwrap(); - assert_eq!(round_tripped.target_user_cohorts, vec!["beta"]); - assert_eq!(round_tripped.priority, 10); - } -} diff --git a/prod/mc/cli-chat-proxy-types/src/lib.rs b/prod/mc/cli-chat-proxy-types/src/lib.rs deleted file mode 100644 index a7e4e61..0000000 --- a/prod/mc/cli-chat-proxy-types/src/lib.rs +++ /dev/null @@ -1,23 +0,0 @@ -//! Lightweight request/response types for the cli-chat-proxy sandbox API. -//! -//! This crate contains only the API types with minimal dependencies (just serde), -//! suitable for use by clients that don't need the full cli-chat-proxy crate. - -pub mod client_metrics_types; -pub mod deployment_config_types; -pub mod feedback_types; -pub mod metadata_types; -mod sandbox_types; -pub mod serde_helpers; -pub mod session_types; -pub mod storage_types; -pub mod subagent_bundle; - -pub use client_metrics_types::*; -pub use deployment_config_types::*; -pub use feedback_types::*; -pub use metadata_types::*; -pub use sandbox_types::*; -pub use session_types::*; -pub use storage_types::*; -pub use subagent_bundle::*; diff --git a/prod/mc/cli-chat-proxy-types/src/metadata_types.rs b/prod/mc/cli-chat-proxy-types/src/metadata_types.rs deleted file mode 100644 index ccf3a1a..0000000 --- a/prod/mc/cli-chat-proxy-types/src/metadata_types.rs +++ /dev/null @@ -1,219 +0,0 @@ -//! Prompt metadata types shared between the CLI client and the cli-chat-proxy server. -//! -//! The CLI client serializes `PromptMetadata` and uploads it as `metadata.json` to GCS -//! via the `/v1/storage` endpoint. The server deserializes it to inject authenticated -//! user identity fields (`user_id`, `user_email`) before forwarding to GCS. -use serde::{Deserialize, Serialize}; -/// Schema version for the GCS metadata format. -/// Increment this when making breaking changes to PromptMetadata structure. -/// v1.2: Added `signals` and `turn_delta` fields to turn_result.json. -/// v1.3: Added `user_query` field to metadata.json. -/// v1.4: Renamed `user_query` to `prompt` (required), renamed `prompt` to `full_prompt` (optional). -/// v1.5: Added `prompt_has_image` field. -/// v1.6: Added `prompt_was_truncated` flag. -/// v1.7: Added `truncated_prompt_local_path`: local disk path embedded in truncated message for search-replace against GCS path. -/// v1.8: Added A/B fork provenance: `ab_root_session_id`, `ab_root_turn_number`, `ab_comparison_id`, `ab_experiment_type`, `ab_experiment_name`. -/// v1.9: Added `cwd` field (current working directory). -/// v1.10: `ab_root_turn_number` now uses the monotonic GCS trace counter -/// (same as `turn_number` and GCS paths) instead of the signal-based prompt count. -/// v1.11: Added `auto_model_hash` for auto-mode model assignment. -/// v1.12: Removed `auto_model_hash` (auto-mode feature was removed). -/// v1.13: Removed `ab_*` fields after the A/B experimentation feature was discontinued. -/// v1.14: Added `prompt_verbatim` field. -/// v1.15: Added `agent_type` field. -/// v1.16: Added `team_id` field (OAuth team identity). -/// v1.17: Added `input_tokens`, `cached_input_tokens`, `output_tokens` to -/// TurnResultMetadata for per-component token attribution. -/// v1.18: Added `shell_version`: the grok-shell agent binary version, distinct -/// from `client_version` (the UI client's version). They coincide for the -/// TUI but differ for embedding clients like grok-desktop. -/// v1.19: Added `workspace_type`: classifies the working directory as "git", -/// "project" (non-git project dir), or "non_project" (system/temp/home). -/// v1.20: Added `sandbox`: resolved OS sandbox profile and whether enforcement is active. -/// v1.21: Added an optional session-metadata field. -/// v1.22: Added `reasoning_effort`: the reasoning effort the turn was sampled -/// with (e.g. "low"/"medium"/"high"/"xhigh"). Omitted when the session -/// has no configured effort. -/// v1.23: Removed `prompt`, `full_prompt`, and `truncated_prompt_local_path` -/// from metadata.json (prompt content is no longer uploaded in metadata). -pub const GCS_SCHEMA_VERSION: &str = "v1.23"; -/// OS-level sandbox state for a trace turn (local `kigi-sandbox`, not cloud sandbox). -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct LocalSandboxTelemetry { - /// Resolved profile at process startup (e.g. "off", "workspace", "strict"). - pub profile: String, - /// Whether kernel-level enforcement is active for this process. - pub applied: bool, -} -/// Metadata about a prompt turn, uploaded as JSON for tracing/debugging. -/// -/// Path format: `{session_id}/turn_{N}/metadata.json` -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct PromptMetadata { - /// Schema version for this metadata format - pub schema_version: String, - /// Session id (UUIDv7) for this trace - pub session_id: String, - /// Monotonic turn number within the session - pub turn_number: u64, - /// Request id for this prompt (uuid v4 we generate per prompt) - pub request_id: String, - /// Timestamp at the start of prompt handling (UTC RFC3339) - pub turn_started_at: String, - /// Git repo root (if the session cwd is inside a git repository). - #[serde(skip_serializing_if = "Option::is_none")] - pub repo_root: Option, - /// Git remote URL (origin) for the repository. - #[serde(skip_serializing_if = "Option::is_none")] - pub remote_url: Option, - /// How workspace files were collected: "git", "project", or "non_project". - #[serde(default, skip_serializing_if = "Option::is_none")] - pub workspace_type: Option, - /// User ID from authentication - pub user_id: Option, - /// User email from authentication (may be None) - pub user_email: Option, - /// Team ID from OAuth authentication (may be None for personal accounts) - #[serde(default, skip_serializing_if = "Option::is_none")] - pub team_id: Option, - /// Client source identifier. - /// Pulled from InitializeRequest.meta (prefers clientSource, falls back to clientType, then clientIdentifier). - #[serde(skip_serializing_if = "Option::is_none")] - pub client_source: Option, - /// Client (TUI) version string, e.g., "0.1.70 (c28a985a1f1)" - /// This is sent by the TUI in InitializeRequest.meta.clientVersion - #[serde(skip_serializing_if = "Option::is_none")] - pub client_version: Option, - /// The model being used for this session - pub model: String, - /// Reasoning effort the turn was sampled with (e.g. "low", "medium", - /// "high", "xhigh"). Omitted when the session has no configured effort - /// (the model then uses its server-side default). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub reasoning_effort: Option, - /// Experiment ID when the model was overridden via experiment routing. Currently unused. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub experiment_id: Option, - /// Host OS where the agent is running (e.g., "macos", "linux") - pub host_os: String, - /// Host architecture (e.g., "x86_64", "aarch64") - pub host_arch: String, - /// Whether the user's prompt contains at least one image attachment. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub prompt_has_image: Option, - /// Whether the prompt was truncated. When `Some(true)`, the full text is at `full_prompt.txt`. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub prompt_was_truncated: Option, - /// Whether the prompt was sent in verbatim mode (skipping `` wrapping). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub prompt_verbatim: Option, - /// Current working directory of the session. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cwd: Option, - /// The agent type / harness name for this session (e.g. "grok-build", "codex"). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub agent_type: Option, - /// Version of the grok-shell agent binary that handled this turn - /// (`kigi_version::VERSION`). Self-reported by the agent, so it reflects - /// the binary actually running. Distinct from `client_version`, which is the - /// UI client's version — for the TUI these coincide, but for embedding clients - /// like grok-desktop the bundled shell differs from the app version. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub shell_version: Option, - /// Resolved OS sandbox profile and whether enforcement is active. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub sandbox: Option, -} -#[cfg(test)] -mod tests { - use super::*; - /// Minimal JSON matching the pre-prompt-content schema fields. - fn minimal_json() -> &'static str { - r#"{ - "schema_version": "v1.23", - "session_id": "abc", - "turn_number": 1, - "request_id": "req-1", - "turn_started_at": "2025-01-01T00:00:00Z", - "user_id": null, - "user_email": null, - "model": "grok-3", - "host_os": "linux", - "host_arch": "x86_64" - }"# - } - #[test] - fn missing_fields_deserialize_to_none_not_false() { - let meta: PromptMetadata = serde_json::from_str(minimal_json()).unwrap(); - assert_eq!(meta.prompt_has_image, None); - assert_eq!(meta.prompt_was_truncated, None); - assert_eq!(meta.cwd, None); - assert_eq!(meta.team_id, None); - } - #[test] - fn explicit_false_deserializes_to_some_false() { - let json = r#"{ - "schema_version": "v1.23", - "session_id": "abc", - "turn_number": 1, - "request_id": "req-1", - "turn_started_at": "2025-01-01T00:00:00Z", - "user_id": null, - "user_email": null, - "model": "grok-3", - "host_os": "linux", - "host_arch": "x86_64", - "prompt_has_image": false, - "prompt_was_truncated": false - }"#; - let meta: PromptMetadata = serde_json::from_str(json).unwrap(); - assert_eq!(meta.prompt_has_image, Some(false)); - assert_eq!(meta.prompt_was_truncated, Some(false)); - } - #[test] - fn none_fields_are_omitted_from_serialization() { - let meta: PromptMetadata = serde_json::from_str(minimal_json()).unwrap(); - let serialized = serde_json::to_string(&meta).unwrap(); - assert!(!serialized.contains("prompt_has_image")); - assert!(!serialized.contains("prompt_was_truncated")); - assert!(!serialized.contains("cwd")); - assert!(!serialized.contains("team_id")); - assert!(!serialized.contains("\"prompt\"")); - assert!(!serialized.contains("full_prompt")); - assert!(!serialized.contains("truncated_prompt_local_path")); - } - #[test] - fn some_fields_are_included_in_serialization() { - let mut meta: PromptMetadata = serde_json::from_str(minimal_json()).unwrap(); - meta.prompt_has_image = Some(false); - meta.prompt_was_truncated = Some(true); - let serialized = serde_json::to_string(&meta).unwrap(); - assert!(serialized.contains("\"prompt_has_image\":false")); - assert!(serialized.contains("\"prompt_was_truncated\":true")); - } - #[test] - fn cwd_round_trips() { - let mut meta: PromptMetadata = serde_json::from_str(minimal_json()).unwrap(); - meta.cwd = Some("/root/code/xai".into()); - let json = serde_json::to_string(&meta).unwrap(); - let deserialized: PromptMetadata = serde_json::from_str(&json).unwrap(); - assert_eq!(deserialized.cwd.as_deref(), Some("/root/code/xai")); - } - #[test] - fn sandbox_round_trips() { - let mut meta: PromptMetadata = serde_json::from_str(minimal_json()).unwrap(); - meta.sandbox = Some(LocalSandboxTelemetry { - profile: "strict".into(), - applied: true, - }); - let json = serde_json::to_string(&meta).unwrap(); - let deserialized: PromptMetadata = serde_json::from_str(&json).unwrap(); - assert_eq!( - deserialized.sandbox, - Some(LocalSandboxTelemetry { - profile: "strict".into(), - applied: true, - }) - ); - } -} diff --git a/prod/mc/cli-chat-proxy-types/src/sandbox_types.rs b/prod/mc/cli-chat-proxy-types/src/sandbox_types.rs deleted file mode 100644 index 54fbaae..0000000 --- a/prod/mc/cli-chat-proxy-types/src/sandbox_types.rs +++ /dev/null @@ -1,715 +0,0 @@ -//! Sandbox API request and response types. -//! -//! These types are shared between the server and clients that use the sandbox API. -//! All types use `camelCase` serialization to match the proto3 canonical JSON encoding -//! used on the wire. - -use std::collections::HashMap; - -use serde::{Deserialize, Serialize}; - -/// Request body for forking a sandbox session. -/// POST /v1/sandbox/sessions/fork -#[derive(Debug, Clone, Deserialize, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxForkRequest { - /// The source sandbox ID to fork from - pub source_sandbox_id: String, - /// Number of copies to create (defaults to 1) - #[serde(default)] - pub copies: Option, - /// Snapshot bucket to use. - /// - /// SECURITY (CWE-284): This field is accepted for backwards compatibility - /// but MUST NOT be forwarded to backend services. The server always uses the - /// configured default bucket and enforces this server-side. - #[serde(default)] - pub snapshot_bucket: Option, -} - -#[cfg(test)] -mod tests { - use super::*; - - /// Verify that a user-supplied snapshotBucket is deserialized but - /// the handler is expected to ignore it. This test documents the security - /// invariant: snapshot_bucket from user input must never control GCS access. - #[test] - fn test_fork_request_snapshot_bucket_is_ignored_by_convention() { - // User sends a malicious bucket name - let json = r#"{ - "sourceSandboxId": "session-123", - "copies": 2, - "snapshotBucket": "attacker-controlled-bucket" - }"#; - - let req: SandboxForkRequest = serde_json::from_str(json).unwrap(); - assert_eq!(req.source_sandbox_id, "session-123"); - assert_eq!(req.copies, Some(2)); - // Field is deserialized for backwards compat, but the handler MUST NOT use it. - assert_eq!( - req.snapshot_bucket, - Some("attacker-controlled-bucket".to_string()) - ); - } - - /// Verify fork request works without snapshot_bucket (the expected path). - #[test] - fn test_fork_request_without_snapshot_bucket() { - let json = r#"{"sourceSandboxId": "session-456"}"#; - - let req: SandboxForkRequest = serde_json::from_str(json).unwrap(); - assert_eq!(req.source_sandbox_id, "session-456"); - assert_eq!(req.copies, None); - assert_eq!(req.snapshot_bucket, None); - } - - // ==================================================================== - // SandboxMode enum serde - // ==================================================================== - - #[test] - fn test_sandbox_mode_serializes_as_proto3_string() { - assert_eq!( - serde_json::to_string(&SandboxMode::Agent).unwrap(), - r#""SANDBOX_MODE_AGENT""# - ); - assert_eq!( - serde_json::to_string(&SandboxMode::WorkspaceServer).unwrap(), - r#""SANDBOX_MODE_WORKSPACE_SERVER""# - ); - assert_eq!( - serde_json::to_string(&SandboxMode::Bare).unwrap(), - r#""SANDBOX_MODE_BARE""# - ); - assert_eq!( - serde_json::to_string(&SandboxMode::Invalid).unwrap(), - r#""SANDBOX_MODE_INVALID""# - ); - } - - #[test] - fn test_sandbox_mode_roundtrip() { - for mode in [ - SandboxMode::Invalid, - SandboxMode::Agent, - SandboxMode::WorkspaceServer, - SandboxMode::Bare, - ] { - let json = serde_json::to_string(&mode).unwrap(); - let back: SandboxMode = serde_json::from_str(&json).unwrap(); - assert_eq!(back, mode); - } - } - - #[test] - fn test_sandbox_mode_default_is_invalid() { - assert_eq!(SandboxMode::default(), SandboxMode::Invalid); - } - - // ==================================================================== - // SandboxStartResponse deserialization from realistic proto3 JSON - // ==================================================================== - - #[test] - fn test_start_response_from_proto3_json() { - // Realistic JSON using proto3 canonical JSON encoding. - // uint64 values like memoryLimitBytes are encoded as strings. - let json = r#"{ - "sandboxId": "sb-abc123", - "sessionId": "sess-xyz789", - "websocketUrl": "wss://sandbox.example.com/ws", - "environment": { - "environment": { - "environmentId": "env-001", - "name": "test-env", - "repository": "org/repo", - "requestedMemoryBytes": "17179869184", - "requestedCpus": 4, - "cachingEnabled": true, - "preinstalledPackages": {"python": "3.11"} - }, - "environmentVariables": [ - {"key": "FOO", "value": "bar"} - ], - "secrets": [], - "userRole": "ROLE_OWNER" - }, - "directUrls": {"6013": "http://direct.example.com:6013"}, - "cloudflareUrls": {"443": "https://cf.example.com"}, - "mode": "SANDBOX_MODE_AGENT" - }"#; - - let resp: SandboxStartResponse = serde_json::from_str(json).unwrap(); - assert_eq!(resp.sandbox_id, "sb-abc123"); - assert_eq!(resp.session_id, "sess-xyz789"); - assert_eq!(resp.websocket_url, "wss://sandbox.example.com/ws"); - assert_eq!(resp.mode, Some(SandboxMode::Agent)); - - // Verify direct_urls / cloudflare_urls maps - assert_eq!( - resp.direct_urls.get("6013").map(|s| s.as_str()), - Some("http://direct.example.com:6013") - ); - assert_eq!( - resp.cloudflare_urls.get("443").map(|s| s.as_str()), - Some("https://cf.example.com") - ); - - // Verify nested environment - let env_meta = resp.environment.as_ref().unwrap(); - let env = env_meta.environment.as_ref().unwrap(); - assert_eq!(env.environment_id.as_deref(), Some("env-001")); - assert_eq!(env.name.as_deref(), Some("test-env")); - assert_eq!(env.requested_memory_bytes.as_deref(), Some("17179869184")); - assert_eq!(env.requested_cpus, Some(4)); - assert_eq!(env.caching_enabled, Some(true)); - assert_eq!( - env.preinstalled_packages.get("python").map(|s| s.as_str()), - Some("3.11") - ); - - // Verify environment variables - assert_eq!(env_meta.environment_variables.len(), 1); - assert_eq!( - env_meta.environment_variables[0].key.as_deref(), - Some("FOO") - ); - assert_eq!(env_meta.user_role.as_deref(), Some("ROLE_OWNER")); - } - - /// Verify SandboxStartResponse handles missing optional fields gracefully. - #[test] - fn test_start_response_minimal_json() { - let json = r#"{ - "sandboxId": "sb-min", - "sessionId": "sess-min", - "websocketUrl": "wss://example.com" - }"#; - - let resp: SandboxStartResponse = serde_json::from_str(json).unwrap(); - assert_eq!(resp.sandbox_id, "sb-min"); - assert!(resp.environment.is_none()); - assert!(resp.direct_urls.is_empty()); - assert!(resp.cloudflare_urls.is_empty()); - assert!(resp.mode.is_none()); - } - - // ==================================================================== - // SandboxEnvironmentResponse roundtrip - // ==================================================================== - - #[test] - fn test_environment_response_roundtrip() { - let resp = SandboxEnvironmentResponse { - environment: Some(SandboxEnvironmentWithMetadata { - environment: Some(SandboxEnvironment { - environment_id: Some("env-rt".into()), - name: Some("roundtrip".into()), - caching_enabled: Some(false), - preinstalled_packages: HashMap::from([("node".into(), "20".into())]), - ..Default::default() - }), - environment_variables: vec![SandboxEnvironmentVariable { - key: Some("KEY".into()), - value: Some("VAL".into()), - }], - secrets: vec![], - user_role: Some("ROLE_EDITOR".into()), - }), - }; - - let json = serde_json::to_string(&resp).unwrap(); - let back: SandboxEnvironmentResponse = serde_json::from_str(&json).unwrap(); - - let env = back - .environment - .as_ref() - .unwrap() - .environment - .as_ref() - .unwrap(); - assert_eq!(env.environment_id.as_deref(), Some("env-rt")); - assert_eq!(env.name.as_deref(), Some("roundtrip")); - assert_eq!( - env.preinstalled_packages.get("node").map(|s| s.as_str()), - Some("20") - ); - } -} - -/// Information about a single forked session. -#[derive(Debug, Clone, Deserialize, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxForkedSession { - /// The provider sandbox ID - pub sandbox_id: String, - /// WebSocket URL for connecting to the sandbox - pub websocket_url: String, - /// JWT token for authenticating the WebSocket connection - pub jwt_token: String, -} - -/// Response from forking a sandbox session. -#[derive(Debug, Clone, Deserialize, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxForkResponse { - /// List of created sandbox IDs - pub sandbox_ids: Vec, - /// Detailed information about each forked session - pub sessions: Vec, -} - -/// Request body/query for terminating a sandbox session. -/// DELETE /v1/sandbox/sessions/{sandbox_id} -#[derive(Debug, Clone, Deserialize, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxTerminateRequest { - /// Environment ID (defaults to "universal") - #[serde(default)] - pub environment_id: Option, -} - -// ============================================================================ -// Session Lifecycle Types -// ============================================================================ - -/// Sandbox operating mode. -/// -/// Proto3 enum serialized as its string name on the wire -/// (e.g. `"SANDBOX_MODE_AGENT"`). -#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] -pub enum SandboxMode { - #[default] - #[serde(rename = "SANDBOX_MODE_INVALID")] - Invalid, - #[serde(rename = "SANDBOX_MODE_AGENT")] - Agent, - #[serde(rename = "SANDBOX_MODE_WORKSPACE_SERVER")] - WorkspaceServer, - #[serde(rename = "SANDBOX_MODE_BARE")] - Bare, -} - -/// Request body for starting a sandbox session (non-TUI). -/// POST /v1/sandbox/sessions/start -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxStartRequest { - /// Environment ID to use (defaults to "universal"). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub environment_id: Option, - /// Optional session ID to resume or associate. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub session_id: Option, - /// Repository to clone (e.g. "owner/repo" or full git URL). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub repository: Option, - /// Branch to checkout. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub branch: Option, - /// Memory limit in bytes. Proto3 uint64, serialized as a JSON string. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub memory_limit_bytes: Option, - /// Number of CPUs. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub cpus: Option, - /// Session timeout in seconds. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub session_timeout_seconds: Option, - /// Additional environment variables. - #[serde(default, skip_serializing_if = "HashMap::is_empty")] - pub env_vars: HashMap, - /// Disk size in bytes. Proto3 uint64, serialized as a JSON string. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub disk_bytes: Option, - /// Number of GPUs. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gpus: Option, - /// GPU type (e.g. "A100", "H100"). - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gpu_type: Option, - /// Sandbox operating mode. - pub mode: SandboxMode, -} - -/// Response from starting a sandbox session. -/// Returned by POST /v1/sandbox/sessions/start. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxStartResponse { - /// Provider sandbox ID. - #[serde(default)] - pub sandbox_id: String, - /// Session ID for persistence and reconnection. - #[serde(default)] - pub session_id: String, - /// WebSocket URL for connecting to the sandbox. - #[serde(default)] - pub websocket_url: String, - /// Environment configuration returned by the sandbox service. - #[serde(default)] - pub environment: Option, - /// Port-to-URL mapping for direct access. - #[serde(default)] - pub direct_urls: HashMap, - /// Port-to-URL mapping for Cloudflare-proxied access. - #[serde(default)] - pub cloudflare_urls: HashMap, - /// Which mode was actually started. - #[serde(default)] - pub mode: Option, -} - -/// Response from getting sandbox session status. -/// GET /v1/sandbox/sessions/{id}/status -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxStatusResponse { - /// Status string (e.g. "STARTING", "SETUP", "READY", "ERROR"). - #[serde(default)] - pub status: String, - /// Human-readable status message. - #[serde(default)] - pub message: String, - /// Additional metadata (e.g. repository size). - #[serde(default)] - pub metadata: HashMap, - /// ISO 8601 timestamp. - #[serde(default)] - pub timestamp: Option, -} - -/// Exit codes for sandbox log commands. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxLogsExitCodes { - /// Exit code of environment variables echo command. - #[serde(default)] - pub env: Option, - /// Exit code of direct mode logs. - #[serde(default)] - pub direct_mode: Option, - /// Exit code of git fetch logs. - #[serde(default)] - pub fetch: Option, -} - -/// Response from getting sandbox session logs. -/// GET /v1/sandbox/sessions/{id}/logs -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxLogsResponse { - /// Combined environment variables echo stdout/stderr. - #[serde(default)] - pub env_vars: String, - /// Combined direct mode logs stdout/stderr. - #[serde(default)] - pub direct_mode_logs: String, - /// Combined fetch/clone logs stdout/stderr. - #[serde(default)] - pub fetch_logs: String, - /// Exit codes for each command. - #[serde(default)] - pub exit_codes: Option, -} - -/// Response from hibernating a sandbox session. -/// POST /v1/sandbox/sessions/{id}/hibernate -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxHibernateResponse { - /// GCS path where the snapshot was stored. - #[serde(default)] - pub snapshot_path: String, -} - -/// Request body for restoring a hibernated sandbox session. -/// POST /v1/sandbox/sessions/{id}/restore -/// -/// The `session_id` is provided as a path parameter, not in the body. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxRestoreRequest { - /// Server key for the restored session's direct-mode agent. - pub server_key: String, -} - -/// Response from restoring a hibernated sandbox session. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxRestoreResponse { - /// Provider sandbox ID of the newly created restored sandbox. - #[serde(default)] - pub sandbox_id: String, - /// GCS path of the snapshot that was restored. - #[serde(default)] - pub snapshot_path: String, - /// WebSocket URL for the restored session. - #[serde(default)] - pub websocket_url: String, -} - -// ============================================================================ -// Environment Types -// ============================================================================ - -/// A sandbox environment configuration. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxEnvironment { - #[serde(default)] - pub environment_id: Option, - #[serde(default)] - pub user_id: Option, - #[serde(default)] - pub team_id: Option, - #[serde(default)] - pub name: Option, - #[serde(default)] - pub description: Option, - #[serde(default)] - pub repository: Option, - #[serde(default)] - pub default_branch: Option, - #[serde(default)] - pub workspace_directory: Option, - #[serde(default)] - pub container_image: Option, - #[serde(default)] - pub setup_script: Option, - #[serde(default)] - pub maintenance_script: Option, - #[serde(default)] - pub caching_enabled: Option, - #[serde(default)] - pub internet_enabled: Option, - #[serde(default)] - pub domain_allowlist_preset: Option, - #[serde(default)] - pub additional_domains: Option, - #[serde(default)] - pub allowed_http_methods: Option, - #[serde(default)] - pub preinstalled_packages: HashMap, - /// ISO 8601 timestamp. - #[serde(default)] - pub create_time: Option, - /// ISO 8601 timestamp. - #[serde(default)] - pub modify_time: Option, - #[serde(default)] - pub cached_commit_sha: Option, - #[serde(default)] - pub provider_id: Option, - #[serde(default)] - pub requested_cpus: Option, - /// Proto3 uint64, serialized as a JSON string. - #[serde(default)] - pub requested_memory_bytes: Option, - /// Proto3 uint64, serialized as a JSON string. - #[serde(default)] - pub requested_disk_bytes: Option, - #[serde(default)] - pub requested_gpus: Option, -} - -/// An environment variable key-value pair. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxEnvironmentVariable { - #[serde(default)] - pub key: Option, - #[serde(default)] - pub value: Option, -} - -/// A secret input key-value pair for environment creation/update. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxSecretInput { - #[serde(default)] - pub key: Option, - #[serde(default)] - pub value: Option, -} - -/// A sandbox environment with its associated metadata. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxEnvironmentWithMetadata { - /// The environment configuration. - #[serde(default)] - pub environment: Option, - /// Non-secret environment variables. - #[serde(default)] - pub environment_variables: Vec, - /// Secret environment variables (values may be redacted). - #[serde(default)] - pub secrets: Vec, - /// The requesting user's role for this environment (proto enum as string). - #[serde(default)] - pub user_role: Option, -} - -/// Query parameters for listing sandbox environments. -/// Used with GET /v1/sandbox/environments. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxListEnvironmentsRequest { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub page: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub page_size: Option, -} - -/// Response from listing sandbox environments. -/// GET /v1/sandbox/environments -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxListEnvironmentsResponse { - #[serde(default)] - pub environments: Vec, - #[serde(default)] - pub page: Option, - #[serde(default)] - pub page_size: Option, - #[serde(default)] - pub has_more: Option, -} - -/// Request body for creating a sandbox environment. -/// POST /v1/sandbox/environments -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxCreateEnvironmentRequest { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub name: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub description: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub repository: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub default_branch: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub workspace_directory: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub container_image: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub setup_script: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub maintenance_script: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub caching_enabled: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub internet_enabled: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub domain_allowlist_preset: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub additional_domains: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub allowed_http_methods: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub environment_variables: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub secrets: Option>, - #[serde(default, skip_serializing_if = "HashMap::is_empty")] - pub preinstalled_packages: HashMap, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub provider_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub requested_cpus: Option, - /// Proto3 uint64, serialized as a JSON string. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub requested_memory_bytes: Option, - /// Proto3 uint64, serialized as a JSON string. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub requested_disk_bytes: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub requested_gpus: Option, -} - -/// Response wrapping a single environment with metadata. -/// -/// Shared by the create, get, and update environment endpoints since they all -/// return the same shape: `{ "environment": SandboxEnvironmentWithMetadata }`. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxEnvironmentResponse { - #[serde(default)] - pub environment: Option, -} - -/// Request body for updating a sandbox environment. -/// PUT /v1/sandbox/environments/{environment_id} -/// -/// The `environment_id` is provided as a path parameter, not in the body. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxUpdateEnvironmentRequest { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub name: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub description: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub repository: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub default_branch: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub workspace_directory: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub container_image: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub setup_script: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub maintenance_script: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub caching_enabled: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub internet_enabled: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub domain_allowlist_preset: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub additional_domains: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub allowed_http_methods: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub environment_variables: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub secrets: Option>, - #[serde(default, skip_serializing_if = "HashMap::is_empty")] - pub preinstalled_packages: HashMap, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub provider_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub requested_cpus: Option, - /// Proto3 uint64, serialized as a JSON string. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub requested_memory_bytes: Option, - /// Proto3 uint64, serialized as a JSON string. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub requested_disk_bytes: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub requested_gpus: Option, -} - -/// A preinstalled package available for sandbox environments. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxPreinstalledPackage { - #[serde(default)] - pub name: Option, - #[serde(default)] - pub versions: Vec, - #[serde(default)] - pub default_version: Option, -} - -/// Response from listing preinstalled packages. -/// GET /v1/sandbox/environments/preinstalled-packages -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SandboxListPreinstalledPackagesResponse { - #[serde(default)] - pub packages: Vec, -} diff --git a/prod/mc/cli-chat-proxy-types/src/serde_helpers.rs b/prod/mc/cli-chat-proxy-types/src/serde_helpers.rs deleted file mode 100644 index 7462464..0000000 --- a/prod/mc/cli-chat-proxy-types/src/serde_helpers.rs +++ /dev/null @@ -1,9 +0,0 @@ -use serde::{Deserialize, Deserializer}; - -pub fn empty_string_as_none<'de, D>(deserializer: D) -> Result, D::Error> -where - D: Deserializer<'de>, -{ - let opt = Option::::deserialize(deserializer)?; - Ok(opt.filter(|s| !s.is_empty())) -} diff --git a/prod/mc/cli-chat-proxy-types/src/session_types.rs b/prod/mc/cli-chat-proxy-types/src/session_types.rs deleted file mode 100644 index 51e062c..0000000 --- a/prod/mc/cli-chat-proxy-types/src/session_types.rs +++ /dev/null @@ -1,116 +0,0 @@ -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct RegisterSessionRequest { - pub session_id: String, - pub cwd: String, - /// Ignored; server derives this from `session_id`. Kept for wire-compat. - #[serde(default)] - pub gcs_trace_prefix: Option, - #[serde(default)] - pub model_id: Option, - #[serde(default)] - pub repo_remote_url: Option, - #[serde(default)] - pub repo_branch: Option, - #[serde(default)] - pub repo_head_at_start: Option, - /// Ignored; server uses its own bucket constant. Kept for wire-compat. - #[serde(default)] - pub gcs_bucket: Option, - #[serde(default)] - pub hostname: Option, - #[serde(default)] - pub parent_session_id: Option, - /// Opaque per-machine device id (`deviceId` on the wire). Sent by the CLI - /// at register; optional for backward-compat with older clients. - #[serde(default)] - pub device_id: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdateSessionRequest { - #[serde(default)] - pub summary: Option, - #[serde(default)] - pub first_prompt: Option, - #[serde(default)] - pub last_turn_number: Option, - #[serde(default)] - pub repo_head_at_end: Option, - /// Latest turn whose restore artifacts are confirmed durable. - /// `None` = leave unchanged. Written separately from `last_turn_number` - /// once session-state upload is confirmed. - #[serde(default)] - pub restorable_turn_number: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SearchSessionsQuery { - #[serde(default)] - pub query: Option, - #[serde(default)] - pub status: Option, - #[serde(default)] - pub cwd: Option, - #[serde(default = "default_limit")] - pub limit: i64, -} - -fn default_limit() -> i64 { - 20 -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct SessionReplicaResponse { - pub session_id: String, - pub summary: String, - pub first_prompt: Option, - pub model_id: Option, - pub created_at: DateTime, - pub updated_at: DateTime, - pub ended_at: Option>, - pub last_turn_number: i32, - /// See `UpdateSessionRequest.restorable_turn_number`. Optional in the wire - /// type so newer CLI builds can parse responses from older servers gracefully. - pub restorable_turn_number: Option, - pub cwd: String, - pub repo_remote_url: Option, - pub repo_branch: Option, - pub repo_head_at_start: Option, - pub repo_head_at_end: Option, - pub gcs_trace_prefix: String, - pub gcs_bucket: String, - pub hostname: Option, - pub parent_session_id: Option, - pub status: String, - pub last_active_at: Option>, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct SearchSessionsResponse { - pub sessions: Vec, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct DownloadSessionQuery { - pub file: String, - #[serde(default)] - pub turn: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct DownloadSessionResponse { - pub download_url: String, - pub expires_in_seconds: u64, - pub file: String, - pub turn: i32, -} diff --git a/prod/mc/cli-chat-proxy-types/src/storage_types.rs b/prod/mc/cli-chat-proxy-types/src/storage_types.rs deleted file mode 100644 index a94e574..0000000 --- a/prod/mc/cli-chat-proxy-types/src/storage_types.rs +++ /dev/null @@ -1,226 +0,0 @@ -//! Signed upload URL types shared between cli-chat-proxy (server) and grok-shell (client). - -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Deserialize, Serialize)] -pub struct BatchExistsRequest { - pub paths: Vec, -} - -#[derive(Debug, Deserialize, Serialize)] -pub struct BatchExistsResponse { - pub exists: Vec, - pub missing: Vec, -} - -/// Response from the signed upload URL endpoint. -/// `POST /v1/storage/signed-upload-url` -/// -/// The client uses the returned `signed_url` to PUT the object directly to GCS, -/// completely bypassing the proxy for the data transfer. This avoids nginx / -/// Cloudflare body-size limits that would otherwise cause 413 errors on large -/// payloads (e.g. session share data). -#[derive(Debug, Clone, Deserialize, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct SignedUploadUrlResponse { - /// Pre-signed GCS PUT URL. Upload the object body here with a simple PUT. - pub signed_url: String, - /// GCS bucket where the object will be stored. - pub bucket: String, - /// Object path within the bucket. - pub path: String, - /// Content-Type that was baked into the signed URL. - /// The PUT request **must** use this exact Content-Type header. - pub content_type: String, - /// Validity window in seconds. - pub expires_in_secs: u64, -} - -#[derive(Debug, Serialize, Deserialize)] -pub struct BatchUploadResult { - pub path: String, - pub status: BatchUploadStatus, - #[serde(skip_serializing_if = "Option::is_none")] - pub bucket: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub size: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub generation: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub error: Option, -} - -#[derive(Debug, PartialEq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum BatchUploadStatus { - Ok, - Error, - Skipped, -} - -#[derive(Debug, Serialize, Deserialize)] -pub struct BatchUploadResponse { - pub results: Vec, -} - -/// JSON request body for `POST /v1/storage/batch_upload_json`. -/// -/// Each file's content is base64-encoded. The request is typically sent with -/// `Content-Encoding: zstd` so the JSON body is compressed on the wire. -#[derive(Debug, Deserialize, Serialize)] -pub struct BatchUploadRequest { - pub files: Vec, -} - -/// A single file entry in a [`BatchUploadRequest`]. -/// -/// All three fields are required on the wire. The server treats an empty -/// `content_type` as `"application/octet-stream"`, but the field itself -/// must be present in the JSON object. -#[derive(Debug, Deserialize, Serialize)] -pub struct BatchUploadFile { - /// GCS destination path. - pub path: String, - /// MIME type of the file content. Required on the wire; the server - /// defaults empty values to `"application/octet-stream"`. - pub content_type: String, - /// Base64-encoded file content (standard alphabet, with padding). - pub data: String, -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn batch_upload_status_serde_round_trip() { - for (variant, expected_json) in [ - (BatchUploadStatus::Ok, "\"ok\""), - (BatchUploadStatus::Error, "\"error\""), - (BatchUploadStatus::Skipped, "\"skipped\""), - ] { - let json = serde_json::to_string(&variant).unwrap(); - assert_eq!(json, expected_json); - let deserialized: BatchUploadStatus = serde_json::from_str(&json).unwrap(); - assert_eq!(deserialized, variant); - } - } - - #[test] - fn batch_upload_response_serializes_ok_result_with_metadata() { - let resp = BatchUploadResponse { - results: vec![BatchUploadResult { - path: "data/file.txt".to_string(), - status: BatchUploadStatus::Ok, - bucket: Some("my-bucket".to_string()), - size: Some(1024), - generation: Some(42), - error: None, - }], - }; - let json: serde_json::Value = serde_json::to_value(&resp).unwrap(); - let result = &json["results"][0]; - assert_eq!(result["path"], "data/file.txt"); - assert_eq!(result["status"], "ok"); - assert_eq!(result["bucket"], "my-bucket"); - assert_eq!(result["size"], 1024); - assert_eq!(result["generation"], 42); - assert!(result.get("error").is_none(), "None fields must be omitted"); - } - - #[test] - fn batch_upload_response_serializes_error_result_without_metadata() { - let resp = BatchUploadResponse { - results: vec![BatchUploadResult { - path: "data/fail.txt".to_string(), - status: BatchUploadStatus::Error, - bucket: None, - size: None, - generation: None, - error: Some("upload failed".to_string()), - }], - }; - let json: serde_json::Value = serde_json::to_value(&resp).unwrap(); - let result = &json["results"][0]; - assert_eq!(result["status"], "error"); - assert_eq!(result["error"], "upload failed"); - assert!(result.get("bucket").is_none()); - assert!(result.get("size").is_none()); - } - - #[test] - fn batch_upload_response_round_trips_mixed_results() { - let original = BatchUploadResponse { - results: vec![ - BatchUploadResult { - path: "ok.bin".to_string(), - status: BatchUploadStatus::Ok, - bucket: Some("b".to_string()), - size: Some(100), - generation: Some(1), - error: None, - }, - BatchUploadResult { - path: "err.bin".to_string(), - status: BatchUploadStatus::Error, - bucket: None, - size: None, - generation: None, - error: Some("boom".to_string()), - }, - BatchUploadResult { - path: "skip.bin".to_string(), - status: BatchUploadStatus::Skipped, - bucket: Some("b".to_string()), - size: Some(200), - generation: Some(5), - error: None, - }, - ], - }; - let json = serde_json::to_string(&original).unwrap(); - let deserialized: BatchUploadResponse = serde_json::from_str(&json).unwrap(); - assert_eq!(deserialized.results.len(), 3); - assert_eq!(deserialized.results[0].status, BatchUploadStatus::Ok); - assert_eq!(deserialized.results[1].status, BatchUploadStatus::Error); - assert_eq!(deserialized.results[1].error.as_deref(), Some("boom")); - assert_eq!(deserialized.results[2].status, BatchUploadStatus::Skipped); - assert_eq!(deserialized.results[2].size, Some(200)); - } - - #[test] - fn batch_upload_request_serializes_round_trip() { - let req = BatchUploadRequest { - files: vec![ - BatchUploadFile { - path: "a.txt".to_string(), - content_type: "text/plain".to_string(), - data: "SGVsbG8=".to_string(), // "Hello" in base64 - }, - BatchUploadFile { - path: "b.bin".to_string(), - content_type: "application/octet-stream".to_string(), - data: "AAEC/w==".to_string(), - }, - ], - }; - let json = serde_json::to_string(&req).unwrap(); - let deserialized: BatchUploadRequest = serde_json::from_str(&json).unwrap(); - assert_eq!(deserialized.files.len(), 2); - assert_eq!(deserialized.files[0].path, "a.txt"); - assert_eq!(deserialized.files[0].data, "SGVsbG8="); - assert_eq!(deserialized.files[1].path, "b.bin"); - assert_eq!( - deserialized.files[1].content_type, - "application/octet-stream" - ); - } - - #[test] - fn batch_upload_request_empty_files_round_trip() { - let req = BatchUploadRequest { files: vec![] }; - let json = serde_json::to_string(&req).unwrap(); - let parsed: BatchUploadRequest = serde_json::from_str(&json).unwrap(); - assert!(parsed.files.is_empty()); - } -} diff --git a/prod/mc/cli-chat-proxy-types/src/subagent_bundle.rs b/prod/mc/cli-chat-proxy-types/src/subagent_bundle.rs deleted file mode 100644 index 7bdc1d4..0000000 --- a/prod/mc/cli-chat-proxy-types/src/subagent_bundle.rs +++ /dev/null @@ -1,101 +0,0 @@ -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; - -/// Shared bundle payload for subagent persona, role, and agent definitions. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SubagentBundle { - pub version: String, - pub personas: HashMap, - pub roles: HashMap, - pub agents: HashMap, - #[serde(default)] - pub skills: HashMap, -} - -impl SubagentBundle { - pub fn empty(version: impl Into) -> Self { - Self { - version: version.into(), - personas: HashMap::new(), - roles: HashMap::new(), - agents: HashMap::new(), - skills: HashMap::new(), - } - } -} - -#[cfg(test)] -mod tests { - use super::SubagentBundle; - use std::collections::HashMap; - - #[test] - fn serializes_expected_shape() { - let bundle = SubagentBundle { - version: "bundle-v1".to_owned(), - personas: HashMap::from([("researcher".to_owned(), "persona body".to_owned())]), - roles: HashMap::from([("reviewer".to_owned(), "role body".to_owned())]), - agents: HashMap::from([("default".to_owned(), "agent body".to_owned())]), - skills: HashMap::from([("commit".to_owned(), "skill body".to_owned())]), - }; - - let actual = serde_json::to_value(bundle).unwrap(); - let expected = serde_json::json!({ - "version": "bundle-v1", - "personas": { - "researcher": "persona body" - }, - "roles": { - "reviewer": "role body" - }, - "agents": { - "default": "agent body" - }, - "skills": { - "commit": "skill body" - } - }); - - assert_eq!(expected, actual); - } - - #[test] - fn deserializes_without_skills_field() { - let json = serde_json::json!({ - "version": "bundle-v1", - "personas": {}, - "roles": {}, - "agents": {} - }); - - let bundle: SubagentBundle = serde_json::from_value(json).unwrap(); - assert_eq!(bundle.version, "bundle-v1"); - assert!(bundle.skills.is_empty()); - assert!(SubagentBundle::empty("v1").skills.is_empty()); - } - - #[test] - fn round_trips_with_skills() { - let bundle = SubagentBundle { - version: "v2".to_owned(), - personas: HashMap::new(), - roles: HashMap::new(), - agents: HashMap::new(), - skills: HashMap::from([ - ( - "commit".to_owned(), - "---\nname: commit\n---\n# Commit".to_owned(), - ), - ( - "review".to_owned(), - "---\nname: review\n---\n# Review".to_owned(), - ), - ]), - }; - - let json = serde_json::to_string(&bundle).unwrap(); - let deserialized: SubagentBundle = serde_json::from_str(&json).unwrap(); - assert_eq!(bundle, deserialized); - } -}